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 -5
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()
return nil
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 {