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
+21 -4
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.session.Destroy()
return nil
h.closeOnce.Do(func() {
if h.session != nil {
h.session.Destroy()
h.session = nil
}
h.closeErr = h.runtime.Close()
})
return h.closeErr
}
// headsResult — one forward pass, read back.