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:
+17
@@ -210,3 +210,20 @@ pops only the active top before resuming the flow underneath.
|
|||||||
passes (162.310s for mavend; every package green). Focused structural
|
passes (162.310s for mavend; every package green). Focused structural
|
||||||
possession, repair-pointer, nested-stack, and repaired-clarify tests pass under
|
possession, repair-pointer, nested-stack, and repaired-clarify tests pass under
|
||||||
the race detector.
|
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`.
|
||||||
|
|||||||
@@ -228,11 +228,15 @@ t:
|
|||||||
# Verbose so the report tables land in the terminal. MAVEN_ONNX_LIB points the
|
# 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
|
# prod-representative baseline at the vendored runtime; override it or set it
|
||||||
# empty to run only the deterministic hash ratchet. This is the measurement
|
# 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
|
MAVEN_ONNX_LIB ?= $(shell pwd)/deps/onnxruntime-linux-x64-1.26.0/lib/libonnxruntime.so
|
||||||
|
|
||||||
eval-router:
|
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,
|
# eval-reach — score the held-out ecosystem reach fixture (internal/router/eval,
|
||||||
# ru_ecosystem_v1.json). Answers "does a real Russian utterance actually arrive
|
# ru_ecosystem_v1.json). Answers "does a real Russian utterance actually arrive
|
||||||
|
|||||||
@@ -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.
|
||||||
@@ -110,6 +110,7 @@ A pair in `docs/routing.md` went stale unnoticed. Its source predated the
|
|||||||
|
|
||||||
| measurement | state |
|
| 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 |
|
| [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 |
|
| [Does one sqlite connection make reads queue? No](2026-08-07-store-connection-cap.md) | live |
|
||||||
|
|
||||||
|
|||||||
+17
@@ -301,6 +301,23 @@ make eval-recall
|
|||||||
|
|
||||||
A large miss against 72.7% means the deploy differs from the bench harness.
|
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`
|
**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.
|
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
|
That is the recorded figure to the decimal, and calendar sat at 2 of 2, so the
|
||||||
|
|||||||
@@ -197,6 +197,7 @@ func TestONNXBaseline(t *testing.T) {
|
|||||||
t.Skipf("onnx embedder unavailable: %v", err)
|
t.Skipf("onnx embedder unavailable: %v", err)
|
||||||
}
|
}
|
||||||
defer emb.Close()
|
defer emb.Close()
|
||||||
|
recordONNXGateExecuted(t)
|
||||||
|
|
||||||
f, err := Load()
|
f, err := Load()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -54,6 +54,7 @@ func TestONNXRoutingHeads(t *testing.T) {
|
|||||||
t.Skipf("routing heads unavailable: %v", err)
|
t.Skipf("routing heads unavailable: %v", err)
|
||||||
}
|
}
|
||||||
defer h.Close()
|
defer h.Close()
|
||||||
|
recordONNXGateExecuted(t)
|
||||||
|
|
||||||
f, err := Load()
|
f, err := Load()
|
||||||
if err != nil {
|
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
|
||||||
|
}
|
||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"math"
|
"math"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"sync"
|
||||||
|
|
||||||
ort "github.com/yalue/onnxruntime_go"
|
ort "github.com/yalue/onnxruntime_go"
|
||||||
)
|
)
|
||||||
@@ -54,9 +55,12 @@ const (
|
|||||||
type RouterHeads struct {
|
type RouterHeads struct {
|
||||||
tokenizer *unigramTokenizer
|
tokenizer *unigramTokenizer
|
||||||
session *ort.DynamicSession[int64, float32]
|
session *ort.DynamicSession[int64, float32]
|
||||||
|
runtime *ONNXRuntimeLease
|
||||||
intents []Intent
|
intents []Intent
|
||||||
sources []Source
|
sources []Source
|
||||||
threshold float64
|
threshold float64
|
||||||
|
closeOnce sync.Once
|
||||||
|
closeErr error
|
||||||
}
|
}
|
||||||
|
|
||||||
// headsMeta — router_heads.json, written beside the weights by the exporter.
|
// 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
|
// NewRouterHeads loads the graph and its label order. modelPath points at the
|
||||||
// .onnx; the external weights and router_heads.json sit beside it.
|
// .onnx; the external weights and router_heads.json sit beside it.
|
||||||
//
|
//
|
||||||
// It assumes the ONNX environment is already initialised, because the embedder
|
// It shares the process-global ONNX environment with the embedder. The runtime
|
||||||
// does that at startup and the runtime allows it once.
|
// 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) {
|
func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
|
||||||
metaPath := filepath.Join(filepath.Dir(modelPath), "router_heads.json")
|
metaPath := filepath.Join(filepath.Dir(modelPath), "router_heads.json")
|
||||||
raw, err := os.ReadFile(metaPath)
|
raw, err := os.ReadFile(metaPath)
|
||||||
@@ -106,18 +111,24 @@ func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("heads: tokenizer: %w", err)
|
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](
|
session, err := ort.NewDynamicSession[int64, float32](
|
||||||
modelPath,
|
modelPath,
|
||||||
[]string{"input_ids", "attention_mask"},
|
[]string{"input_ids", "attention_mask"},
|
||||||
[]string{"intent", "source", "slots", "clarify"},
|
[]string{"intent", "source", "slots", "clarify"},
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
_ = runtime.Close()
|
||||||
return nil, fmt.Errorf("heads: create session: %w", err)
|
return nil, fmt.Errorf("heads: create session: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
return &RouterHeads{
|
return &RouterHeads{
|
||||||
tokenizer: tok,
|
tokenizer: tok,
|
||||||
session: session,
|
session: session,
|
||||||
|
runtime: runtime,
|
||||||
intents: intents,
|
intents: intents,
|
||||||
sources: sources,
|
sources: sources,
|
||||||
threshold: headsThreshold,
|
threshold: headsThreshold,
|
||||||
@@ -128,8 +139,14 @@ func (h *RouterHeads) Close() error {
|
|||||||
if h == nil {
|
if h == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
h.closeOnce.Do(func() {
|
||||||
|
if h.session != nil {
|
||||||
h.session.Destroy()
|
h.session.Destroy()
|
||||||
return nil
|
h.session = nil
|
||||||
|
}
|
||||||
|
h.closeErr = h.runtime.Close()
|
||||||
|
})
|
||||||
|
return h.closeErr
|
||||||
}
|
}
|
||||||
|
|
||||||
// headsResult — one forward pass, read back.
|
// headsResult — one forward pass, read back.
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"math"
|
"math"
|
||||||
"os"
|
"os"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
|
|
||||||
ort "github.com/yalue/onnxruntime_go"
|
ort "github.com/yalue/onnxruntime_go"
|
||||||
"golang.org/x/text/unicode/norm"
|
"golang.org/x/text/unicode/norm"
|
||||||
@@ -33,17 +34,21 @@ const (
|
|||||||
type onnxEmbedder struct {
|
type onnxEmbedder struct {
|
||||||
tokenizer *unigramTokenizer
|
tokenizer *unigramTokenizer
|
||||||
session *ort.DynamicSession[int64, float32]
|
session *ort.DynamicSession[int64, float32]
|
||||||
|
runtime *ONNXRuntimeLease
|
||||||
id string
|
id string
|
||||||
|
closeOnce sync.Once
|
||||||
|
closeErr error
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, error) {
|
func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, error) {
|
||||||
ort.SetSharedLibraryPath(libPath)
|
runtime, err := AcquireONNXRuntime(libPath)
|
||||||
if err := ort.InitializeEnvironment(); err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("onnx: init environment: %w", err)
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
tok, err := newUnigramTokenizer(tokenizerPath)
|
tok, err := newUnigramTokenizer(tokenizerPath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
_ = runtime.Close()
|
||||||
return nil, fmt.Errorf("tokenizer: %w", err)
|
return nil, fmt.Errorf("tokenizer: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -53,12 +58,14 @@ func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, e
|
|||||||
[]string{"last_hidden_state"},
|
[]string{"last_hidden_state"},
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
_ = runtime.Close()
|
||||||
return nil, fmt.Errorf("onnx: create session: %w", err)
|
return nil, fmt.Errorf("onnx: create session: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
return &onnxEmbedder{
|
return &onnxEmbedder{
|
||||||
tokenizer: tok,
|
tokenizer: tok,
|
||||||
session: session,
|
session: session,
|
||||||
|
runtime: runtime,
|
||||||
id: modelIDFromPath(modelPath),
|
id: modelIDFromPath(modelPath),
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
@@ -148,8 +155,17 @@ func (e *onnxEmbedder) embed(ctx context.Context, text string) ([]float32, error
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (e *onnxEmbedder) Close() error {
|
func (e *onnxEmbedder) Close() error {
|
||||||
e.session.Destroy()
|
if e == nil {
|
||||||
return 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 {
|
func meanPool(hidden []float32, mask []int64, seqLen, dim int) []float32 {
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user