Make aggregate ONNX gates execute for real
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.
This commit is contained in:
@@ -197,6 +197,7 @@ func TestONNXBaseline(t *testing.T) {
|
||||
t.Skipf("onnx embedder unavailable: %v", err)
|
||||
}
|
||||
defer emb.Close()
|
||||
recordONNXGateExecuted(t)
|
||||
|
||||
f, err := Load()
|
||||
if err != nil {
|
||||
|
||||
@@ -54,6 +54,7 @@ func TestONNXRoutingHeads(t *testing.T) {
|
||||
t.Skipf("routing heads unavailable: %v", err)
|
||||
}
|
||||
defer h.Close()
|
||||
recordONNXGateExecuted(t)
|
||||
|
||||
f, err := Load()
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user