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:
2026-08-13 03:03:25 +04:00
parent 8015fdbb79
commit 28c2ffb84f
13 changed files with 448 additions and 11 deletions
+17
View File
@@ -210,3 +210,20 @@ pops only the active top before resuming the flow underneath.
passes (162.310s for mavend; every package green). Focused structural
possession, repair-pointer, nested-stack, and repaired-clarify tests pass under
the race detector.
### ONNX test/runtime lifecycle
V-716 found that each embedder constructor tried to initialize ONNX Runtime,
while `Close` destroyed only its session. In one package process the first
model-aware test ran and later tests converted “already initialized” into a
green skip. The router now owns the process-global environment through
reference-counted leases held by each embedder and routing-head session; the
last owned lease performs cleanup, and close is idempotent.
The router and mavend test packages hold a lease across their model gates.
`make eval-router` additionally requires proof that both named aggregate gates
actually executed. In one process the classifier baseline scored 72/96 and the
routing heads 93/96; destination was 11/33 and 25/33 respectively, and ecosystem
reach remained 28/30. The lifecycle reacquire test, focused race suite, full
aggregate command, and portable no-runtime packages all pass. Measurement:
`docs/evals/2026-08-13-onnx-runtime-lifecycle.md`.
+6 -2
View File
@@ -228,11 +228,15 @@ t:
# Verbose so the report tables land in the terminal. MAVEN_ONNX_LIB points the
# prod-representative baseline at the vendored runtime; override it or set it
# empty to run only the deterministic hash ratchet. This is the measurement
# Vikunja #319 compares before #320 flips the route decider.
# Vikunja #319 compares before #320 flips the route decider. With a non-empty
# runtime path the package must prove that at least two model gates executed;
# a constructor skip after the first process-global initialization is a failure.
MAVEN_ONNX_LIB ?= $(shell pwd)/deps/onnxruntime-linux-x64-1.26.0/lib/libonnxruntime.so
eval-router:
MAVEN_ONNX_LIB="$(MAVEN_ONNX_LIB)" $(GO) test -v -count=1 ./internal/router/eval/
MAVEN_ONNX_LIB="$(MAVEN_ONNX_LIB)" \
MAVEN_ONNX_REQUIRED_GATES="$(if $(MAVEN_ONNX_LIB),2,0)" \
$(GO) test -v -count=1 ./internal/router/eval/
# eval-reach — score the held-out ecosystem reach fixture (internal/router/eval,
# ru_ecosystem_v1.json). Answers "does a real Russian utterance actually arrive
+36
View File
@@ -0,0 +1,36 @@
package main
import (
"fmt"
"os"
"testing"
"github.com/kami/maven/internal/router"
)
// TestMain holds one ONNX Runtime lease across the model-aware topic and
// personal-boundary gates. Each test still owns and closes its model session;
// the process-global environment is released only after the final test.
func TestMain(m *testing.M) {
var lease *router.ONNXRuntimeLease
lib := os.Getenv("MAVEN_ONNX_LIB")
if lib != "" {
if _, err := os.Stat(lib); err == nil {
var acquireErr error
lease, acquireErr = router.AcquireONNXRuntime(lib)
if acquireErr != nil {
fmt.Fprintf(os.Stderr, "initialize shared ONNX test runtime: %v\n", acquireErr)
os.Exit(2)
}
}
}
code := m.Run()
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)
}
@@ -0,0 +1,65 @@
# ONNX aggregate gate lifecycle — 2026-08-13
Vikunja: V-716.
## Finding
`router.NewONNXEmbedder` initialized `onnxruntime_go` unconditionally. Its
`Close` method destroyed the model session but left the package-global runtime
environment alive. In one Go process the first `TestONNX*` therefore ran and
every later constructor returned `The onnxruntime has already been initialized`.
Those tests converted the constructor error into `t.Skip`, so the aggregate
package still reported `PASS`.
Reproduction before the fix:
```text
TestONNXBaseline PASS classifier+onnx 72/96
TestONNXRoutingHeads SKIP onnx: init environment: already initialized
package PASS
```
## Contract
The router now owns ONNX Runtime through reference-counted leases. Every model
session holds one lease; routing heads hold their own, so daemon shutdown order
cannot unload the library underneath a live graph. The last lease destroys an
environment Maven initialized. Package `TestMain` holds one extra lease across
all model-aware tests in `internal/router/eval` and `cmd/mavend`.
The aggregate router command additionally requires two named gates to record
execution only after all of their model dependencies loaded. A second-test
skip can no longer satisfy the command.
## Measurement
Runtime and model:
```text
deps/onnxruntime-linux-x64-1.26.0/lib/libonnxruntime.so.1.26.0
models/embedder/multilingual-e5-small/model_quantized.onnx
models/embedder/router-heads/router_heads.onnx
```
Focused aggregate, one process:
```text
TestONNXBaseline PASS 72/96 full (75.0%), destination 11/33
TestONNXRoutingHeads PASS 93/96 full (96.9%), destination 25/33
ONNX aggregate proof: verified 2 required model gates:
[TestONNXBaseline TestONNXRoutingHeads]
```
The separate lifecycle test released the last lease and successfully acquired
the runtime again. A second same-process package run also executed all three
selected mavend gates without a skip:
```text
TestONNXPersonalBoundary PASS 29/29
TestONNXPersonalBoundaryFourFold PASS 99/104
TestONNXTopics PASS 43/43
```
Finally, `make eval-router` completed the full package. In addition to the two
required proof gates, `TestONNXClaimConfidenceDistribution` and
`TestReachBaselineONNX` executed; ecosystem reach remained 28/30.
+1
View File
@@ -110,6 +110,7 @@ A pair in `docs/routing.md` went stale unnoticed. Its source predated the
| measurement | state |
| --- | --- |
| [Both ONNX routing gates execute in one process](2026-08-13-onnx-runtime-lifecycle.md) | live |
| [Where the resident model's 7.9GB of RSS goes](2026-08-03-llama-prompt-cache.md) | live |
| [Does one sqlite connection make reads queue? No](2026-08-07-store-connection-cap.md) | live |
+17
View File
@@ -301,6 +301,23 @@ make eval-recall
A large miss against 72.7% means the deploy differs from the bench harness.
`make eval-router` is also the aggregate ONNX lifecycle gate. When the default
runtime path is non-empty it requires both `TestONNXBaseline` and
`TestONNXRoutingHeads` to load their models and execute in the same Go process.
The final line names both verified gates. A later test that skips because the
first one consumed the process-global runtime now makes the command fail rather
than leaving a green package result. To exercise the proof directly:
```sh
MAVEN_ONNX_LIB="$PWD/deps/onnxruntime-linux-x64-1.26.0/lib/libonnxruntime.so" \
MAVEN_ONNX_REQUIRED_GATES=2 \
go test -v -count=1 \
-run '^(TestONNXBaseline|TestONNXRoutingHeads)$' ./internal/router/eval/
```
Set `MAVEN_ONNX_LIB=` explicitly when the intended measurement is the portable
hash floor; that also disables the aggregate ONNX requirement.
**Run on 02-08-2026 @ af9d213. The deploy matches the bench.** `eval-models`
scored 56 of 77: 72.7% full, 77.9% intent-only, 2 false clarifies and 1 missed.
That is the recorded figure to the decimal, and calendar sat at 2 of 2, so the
+1
View File
@@ -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 {
+1
View File
@@ -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
}
+20 -3
View File
@@ -7,6 +7,7 @@ import (
"math"
"os"
"path/filepath"
"sync"
ort "github.com/yalue/onnxruntime_go"
)
@@ -54,9 +55,12 @@ const (
type RouterHeads struct {
tokenizer *unigramTokenizer
session *ort.DynamicSession[int64, float32]
runtime *ONNXRuntimeLease
intents []Intent
sources []Source
threshold float64
closeOnce sync.Once
closeErr error
}
// headsMeta — router_heads.json, written beside the weights by the exporter.
@@ -70,8 +74,9 @@ type headsMeta struct {
// NewRouterHeads loads the graph and its label order. modelPath points at the
// .onnx; the external weights and router_heads.json sit beside it.
//
// It assumes the ONNX environment is already initialised, because the embedder
// does that at startup and the runtime allows it once.
// It shares the process-global ONNX environment with the embedder. The runtime
// lease is independent so shutdown order cannot unload the library while this
// graph's session is still alive.
func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
metaPath := filepath.Join(filepath.Dir(modelPath), "router_heads.json")
raw, err := os.ReadFile(metaPath)
@@ -106,18 +111,24 @@ func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
if err != nil {
return nil, fmt.Errorf("heads: tokenizer: %w", err)
}
runtime, err := AcquireONNXRuntime("")
if err != nil {
return nil, fmt.Errorf("heads: %w", err)
}
session, err := ort.NewDynamicSession[int64, float32](
modelPath,
[]string{"input_ids", "attention_mask"},
[]string{"intent", "source", "slots", "clarify"},
)
if err != nil {
_ = runtime.Close()
return nil, fmt.Errorf("heads: create session: %w", err)
}
return &RouterHeads{
tokenizer: tok,
session: session,
runtime: runtime,
intents: intents,
sources: sources,
threshold: headsThreshold,
@@ -128,8 +139,14 @@ func (h *RouterHeads) Close() error {
if h == nil {
return nil
}
h.closeOnce.Do(func() {
if h.session != nil {
h.session.Destroy()
return nil
h.session = nil
}
h.closeErr = h.runtime.Close()
})
return h.closeErr
}
// headsResult — one forward pass, read back.
+20 -4
View File
@@ -7,6 +7,7 @@ import (
"math"
"os"
"strings"
"sync"
ort "github.com/yalue/onnxruntime_go"
"golang.org/x/text/unicode/norm"
@@ -33,17 +34,21 @@ const (
type onnxEmbedder struct {
tokenizer *unigramTokenizer
session *ort.DynamicSession[int64, float32]
runtime *ONNXRuntimeLease
id string
closeOnce sync.Once
closeErr error
}
func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, error) {
ort.SetSharedLibraryPath(libPath)
if err := ort.InitializeEnvironment(); err != nil {
return nil, fmt.Errorf("onnx: init environment: %w", err)
runtime, err := AcquireONNXRuntime(libPath)
if err != nil {
return nil, err
}
tok, err := newUnigramTokenizer(tokenizerPath)
if err != nil {
_ = runtime.Close()
return nil, fmt.Errorf("tokenizer: %w", err)
}
@@ -53,12 +58,14 @@ func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, e
[]string{"last_hidden_state"},
)
if err != nil {
_ = runtime.Close()
return nil, fmt.Errorf("onnx: create session: %w", err)
}
return &onnxEmbedder{
tokenizer: tok,
session: session,
runtime: runtime,
id: modelIDFromPath(modelPath),
}, nil
}
@@ -148,8 +155,17 @@ func (e *onnxEmbedder) embed(ctx context.Context, text string) ([]float32, error
}
func (e *onnxEmbedder) Close() error {
e.session.Destroy()
if e == nil {
return nil
}
e.closeOnce.Do(func() {
if e.session != nil {
e.session.Destroy()
e.session = nil
}
e.closeErr = e.runtime.Close()
})
return e.closeErr
}
func meanPool(hidden []float32, mask []int64, seqLen, dim int) []float32 {
+124
View File
@@ -0,0 +1,124 @@
package router
import (
"fmt"
"path/filepath"
"sync"
ort "github.com/yalue/onnxruntime_go"
)
// ONNXRuntimeLease keeps the process-global ONNX Runtime environment alive.
//
// onnxruntime_go exposes the environment as package-global state: it may be
// initialized only once, and destroying it while any session is alive is
// unsafe. A lease makes that ownership explicit. Model wrappers hold one for
// the lifetime of their session; TestMain may hold an additional lease so a
// package's model-aware gates share one environment instead of repeatedly
// loading and unloading the shared library.
type ONNXRuntimeLease struct {
once sync.Once
err error
}
var sharedONNXRuntime = struct {
sync.Mutex
refs int
libPath string
managed bool
}{}
// AcquireONNXRuntime acquires the process-global ONNX Runtime environment.
// The first caller initializes it; later callers share it. A non-empty library
// path must agree with the path that initialized the active environment.
//
// The external-environment branch is intentional. It keeps this package safe
// when an executable initialized onnxruntime_go directly before constructing
// a router model; in that case the lease never destroys state it does not own.
func AcquireONNXRuntime(libPath string) (*ONNXRuntimeLease, error) {
libPath = cleanONNXLibraryPath(libPath)
sharedONNXRuntime.Lock()
defer sharedONNXRuntime.Unlock()
if sharedONNXRuntime.refs > 0 {
if err := compatibleONNXLibraryPath(sharedONNXRuntime.libPath, libPath); err != nil {
return nil, err
}
sharedONNXRuntime.refs++
return &ONNXRuntimeLease{}, nil
}
if ort.IsInitialized() {
// Some other component owns the already-live environment. Attach to it,
// but never unload its shared library when our last lease closes.
sharedONNXRuntime.refs = 1
sharedONNXRuntime.libPath = libPath
sharedONNXRuntime.managed = false
return &ONNXRuntimeLease{}, nil
}
if libPath != "" {
ort.SetSharedLibraryPath(libPath)
}
if err := ort.InitializeEnvironment(); err != nil {
return nil, fmt.Errorf("onnx: init environment: %w", err)
}
sharedONNXRuntime.refs = 1
sharedONNXRuntime.libPath = libPath
sharedONNXRuntime.managed = true
return &ONNXRuntimeLease{}, nil
}
// Close releases one environment lease. The last lease destroys an
// environment initialized by this package, after every model session that
// held a lease has already been destroyed.
func (l *ONNXRuntimeLease) Close() error {
if l == nil {
return nil
}
l.once.Do(func() {
l.err = releaseONNXRuntime()
})
return l.err
}
func releaseONNXRuntime() error {
sharedONNXRuntime.Lock()
defer sharedONNXRuntime.Unlock()
if sharedONNXRuntime.refs == 0 {
return fmt.Errorf("onnx: release environment without a lease")
}
sharedONNXRuntime.refs--
if sharedONNXRuntime.refs != 0 {
return nil
}
managed := sharedONNXRuntime.managed
sharedONNXRuntime.libPath = ""
sharedONNXRuntime.managed = false
if !managed {
return nil
}
if err := ort.DestroyEnvironment(); err != nil {
return fmt.Errorf("onnx: destroy environment: %w", err)
}
return nil
}
func cleanONNXLibraryPath(path string) string {
if path == "" {
return ""
}
return filepath.Clean(path)
}
func compatibleONNXLibraryPath(active, requested string) error {
// Empty means "use the environment that is already active". RouterHeads
// intentionally has no second copy of the embedder's library-path config.
if active == "" || requested == "" || active == requested {
return nil
}
return fmt.Errorf("onnx: runtime already uses shared library %q, cannot also use %q", active, requested)
}
+39
View File
@@ -0,0 +1,39 @@
package router
import (
"os"
"testing"
)
// TestONNXRuntimeCanBeReacquired pins the process-global lifecycle itself.
// The model gates exercise live sessions; this smaller check proves that the
// last lease unloads a managed environment and that a later gate can acquire
// it again. It is opt-in for the same reason as every model measurement.
func TestONNXRuntimeCanBeReacquired(t *testing.T) {
lib := os.Getenv("MAVEN_ONNX_LIB")
if lib == "" {
t.Skip("MAVEN_ONNX_LIB unset — see AGENTS.md § Embedder model for intent routing")
}
if _, err := os.Stat(lib); err != nil {
t.Skipf("missing %s: %v", lib, err)
}
first, err := AcquireONNXRuntime(lib)
if err != nil {
t.Fatalf("first acquire: %v", err)
}
if err := first.Close(); err != nil {
t.Fatalf("first close: %v", err)
}
if err := first.Close(); err != nil {
t.Fatalf("idempotent close: %v", err)
}
second, err := AcquireONNXRuntime(lib)
if err != nil {
t.Fatalf("reacquire after last close: %v", err)
}
if err := second.Close(); err != nil {
t.Fatalf("second close: %v", err)
}
}