Maven/internal/router/onnxruntime.go

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)
}