claude
28c2ffb84f
Reference-count the process-global ONNX Runtime across embedder and routing-head sessions, make close idempotent, and require named proof that both aggregate routing gates executed rather than self-skipped (V-716). Owner explicitly requested direct commits to master.
100 lines
2.4 KiB
Go
100 lines
2.4 KiB
Go
package eval
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"sort"
|
|
"strconv"
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/kami/maven/internal/router"
|
|
)
|
|
|
|
const requiredONNXGatesEnv = "MAVEN_ONNX_REQUIRED_GATES"
|
|
|
|
var executedONNXGates struct {
|
|
sync.Mutex
|
|
names []string
|
|
}
|
|
|
|
// recordONNXGateExecuted is called only after a gate has loaded every model it
|
|
// needs. TestMain uses the names as proof that an aggregate command did not
|
|
// turn a second initialization error into a green SKIP.
|
|
func recordONNXGateExecuted(t *testing.T) {
|
|
t.Helper()
|
|
executedONNXGates.Lock()
|
|
executedONNXGates.names = append(executedONNXGates.names, t.Name())
|
|
executedONNXGates.Unlock()
|
|
}
|
|
|
|
func TestMain(m *testing.M) {
|
|
required, err := requiredONNXGateCount()
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
os.Exit(2)
|
|
}
|
|
|
|
lease, err := configuredONNXTestRuntime(required > 0)
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
os.Exit(2)
|
|
}
|
|
|
|
code := m.Run()
|
|
if required > 0 {
|
|
executedONNXGates.Lock()
|
|
names := append([]string(nil), executedONNXGates.names...)
|
|
executedONNXGates.Unlock()
|
|
sort.Strings(names)
|
|
fmt.Fprintf(os.Stderr, "ONNX aggregate proof: verified %d required model gates: %v\n", len(names), names)
|
|
if len(names) < required {
|
|
fmt.Fprintf(os.Stderr, "ONNX aggregate gate failed: executed %d, require at least %d\n", len(names), required)
|
|
code = 1
|
|
}
|
|
}
|
|
|
|
if lease != nil {
|
|
if err := lease.Close(); err != nil {
|
|
fmt.Fprintf(os.Stderr, "ONNX test runtime cleanup: %v\n", err)
|
|
code = 1
|
|
}
|
|
}
|
|
os.Exit(code)
|
|
}
|
|
|
|
func requiredONNXGateCount() (int, error) {
|
|
raw := os.Getenv(requiredONNXGatesEnv)
|
|
if raw == "" {
|
|
return 0, nil
|
|
}
|
|
n, err := strconv.Atoi(raw)
|
|
if err != nil || n < 0 {
|
|
return 0, fmt.Errorf("%s must be a non-negative integer, got %q", requiredONNXGatesEnv, raw)
|
|
}
|
|
return n, nil
|
|
}
|
|
|
|
func configuredONNXTestRuntime(required bool) (*router.ONNXRuntimeLease, error) {
|
|
lib := os.Getenv("MAVEN_ONNX_LIB")
|
|
if lib == "" {
|
|
if required {
|
|
return nil, fmt.Errorf("ONNX aggregate gate requires MAVEN_ONNX_LIB")
|
|
}
|
|
return nil, nil
|
|
}
|
|
if _, err := os.Stat(lib); err != nil {
|
|
if required {
|
|
return nil, fmt.Errorf("ONNX aggregate gate runtime %s: %w", lib, err)
|
|
}
|
|
// Preserve the ordinary portable test path: individual model tests
|
|
// report a skip when optional local model dependencies are absent.
|
|
return nil, nil
|
|
}
|
|
lease, err := router.AcquireONNXRuntime(lib)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("initialize shared ONNX test runtime: %w", err)
|
|
}
|
|
return lease, nil
|
|
}
|