trevon/yams
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2
yams/plugins/onnx/onnx_model_pool.cpp
trvon fe07469620
 feat(plugins): prefer a system ONNX Runtime, fall back to a bundled copy 
The ONNX plugin dlopen()ed ONNX Runtime from directories next to the
plugin first and only then from a few system prefixes, with no notion of
which copy it got or why, and the shipped packages contained no copy at
all, so on most hosts the plugin could only report "runtime not found".

The runtime resolution moves into plugins/ort_runtime/, built as a
hidden-visibility static library both ONNX-using plugins link (their
export maps still expose only the plugin C ABI). The policy is split out
as pure code (ort_runtime_resolution.{h,cpp}) so it is testable without a
runtime. Candidates are tried in this order and the first compatible one
wins:

  1. YAMS_ONNX_RUNTIME_LIB / YAMS_ONNX_RUNTIME_DIR (existing pin, kept),
     then the plugin config key runtime_library ([plugins.onnx] /
     [plugins.glint], a file or directory)
  2. a system copy: the loader search by soname, then the usual library
     directories (multiarch dir, /usr/lib64, /usr/lib, /usr/local/lib,
     /opt/onnxruntime, Homebrew) for versioned names such as Debian's
     libonnxruntime.so.1.23
  3. the private copy in <libdir>/yams/onnxruntime/
  4. legacy locations next to the plugin (build trees)

A candidate is accepted only if it exports OrtGetApiBase, GetApi() returns
an API table for the ORT_API_VERSION the plugins are compiled against, and
it meets the 1.23 minimum; otherwise it is skipped with a reason. A soname
lookup that returns a copy another plugin already loaded from the bundled
directory is reported as that copy. The choice ("Using system|bundled ONNX
Runtime <ver> from <path> (<why>)") and every skipped candidate go to the
plugin log and, through the libyams_onnx_resource host sink (new
OnnxConcurrencyRegistry::emitHostLog), to the daemon log; plugin health
reports runtime_source. With no compatible copy the plugin reports "ONNX
Runtime unavailable: ..." instead of failing.

Tests: tests/unit/plugins/ort_runtime_resolution_catch2_test.cpp covers
the order, the override pin, configured directories, system-over-bundled
preference, skipping an incompatible system copy, the not-found result
and the already-loaded classification.

Signed-off-by: trvon <me@trevon.dev>
2026-10-01 22:42:15 -06:00

3447 lines
142 KiB
C++
Raw Permalink Blame History

#include <yams/common/test_utils.h>
#include <yams/daemon/components/TuneAdvisor.h>
#include <yams/daemon/resource/onnx_model_pool.h>
#include <yams/vector/tokenizer.h>
#include "onnx_gpu_provider.h"
#include <yams/daemon/resource/gpu_info.h>
#include <nlohmann/json.hpp>
#include <spdlog/spdlog.h>
#include <filesystem>
#include <mutex>
#include <sstream>
#include <unordered_set>
#include <boost/asio/post.hpp>
#include <boost/asio/thread_pool.hpp>
#include "../ort_runtime/ort_cxx_api_wrapper.h"
#include "../ort_runtime/ort_runtime_loader.h"
#ifdef _WIN32
#include <fcntl.h>
#include <io.h>
#include <windows.h>
#include <sys/stat.h>
#else
#include <fcntl.h>
#include <unistd.h>
#include <sys/file.h>
#endif
// GPU execution provider headers - conditionally included based on Conan build options
// Note: Prebuilt ONNX Runtime binaries from GitHub releases use the C API for providers.
// The provider-specific headers (cuda_provider_factory.h, etc.) are only available when
// building ONNX Runtime from source. For prebuilt binaries, we use
// OrtSessionOptionsAppendExecutionProvider_* functions from the C API which are declared in
// onnxruntime_c_api.h (included via onnxruntime_cxx_api.h).
// Global ONNX Runtime environment - single instance per process, intentionally leaked.
// Using leaking singleton pattern because ONNX Runtime's LoggingManager destructor
// tries to lock a mutex that may be destroyed during static destruction, causing
// "mutex lock failed: Invalid argument" crashes.
static std::recursive_mutex* g_onnx_init_mutex = new std::recursive_mutex();
static Ort::Env* g_onnx_env = nullptr; // Intentionally leaked to avoid static destruction crash
static bool g_onnx_env_initialized = false;
static Ort::Env& get_global_ort_env() {
std::lock_guard<std::recursive_mutex> lock(*g_onnx_init_mutex);
if (!g_onnx_env_initialized) {
spdlog::info("[ONNX] Initializing global Ort::Env (thread-safe, single instance)");
OrtThreadingOptions* threading_options = nullptr;
(void)Ort::GetApi().CreateThreadingOptions(&threading_options);
if (threading_options) {
unsigned hw_threads = std::thread::hardware_concurrency();
if (hw_threads == 0)
hw_threads = 4;
#ifdef _WIN32
int intra_threads = 1;
int inter_threads = 1;
spdlog::info(
"[ONNX] Windows: using single-threaded globalOrt::Env (EDEADLK workaround)");
#else
int intra_threads = static_cast<int>(std::max(1u, hw_threads / 2));
int inter_threads = 1;
spdlog::info("[ONNX] Using multi-threaded global Ort::Env (intra={}, inter={})",
intra_threads, inter_threads);
#endif
(void)Ort::GetApi().SetGlobalIntraOpNumThreads(threading_options, intra_threads);
(void)Ort::GetApi().SetGlobalInterOpNumThreads(threading_options, inter_threads);
(void)Ort::GetApi().SetGlobalSpinControl(threading_options, 0);
// Process-lifetime singleton: Ort::Env intentionally never freed — ORT
// requires the env for the process lifetime.
g_onnx_env = new Ort::Env(threading_options, ORT_LOGGING_LEVEL_WARNING, "YamsDaemon");
Ort::GetApi().ReleaseThreadingOptions(threading_options);
spdlog::info("[ONNX] Global Ort::Env initialized with custom threading options");
} else {
// Process-lifetime singleton: see comment above.
g_onnx_env = new Ort::Env(ORT_LOGGING_LEVEL_WARNING, "YamsDaemon");
spdlog::info("[ONNX] Global Ort::Env initialized (default threading)");
}
g_onnx_env_initialized = true;
}
return *g_onnx_env;
}
// Check if we should skip model loading for tests
#ifdef GTEST_API_
#include <yams/common/test_utils.h>
#endif
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <future>
#include <limits>
#include <thread>
#include <unordered_map>
namespace yams::daemon {
namespace fs = std::filesystem;
namespace {
bool envTruthy(const char* value) {
if (!value || !*value) {
return false;
}
std::string v(value);
std::transform(v.begin(), v.end(), v.begin(),
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
return !(v == "0" || v == "false" || v == "off" || v == "no");
}
bool allowGpuAutoUnload() {
return envTruthy(std::getenv("YAMS_ONNX_GPU_AUTO_UNLOAD"));
}
int envIntOr(const char* name, int fallback) {
if (const char* value = std::getenv(name)) {
try {
return std::stoi(value);
} catch (...) {
}
}
return fallback;
}
std::string lowerCopy(std::string value) {
std::transform(value.begin(), value.end(), value.begin(),
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
return value;
}
bool supportsCoreMLCpuFallbackModel(const std::string& modelName) {
const std::string lowerName = lowerCopy(modelName);
return lowerName.find("embeddinggemma") != std::string::npos ||
lowerName.find("gemma") != std::string::npos;
}
bool isGpuExecutionProviderName(std::string provider) {
provider = lowerCopy(std::move(provider));
return provider.find("cuda") != std::string::npos ||
provider.find("migraphx") != std::string::npos ||
provider.find("coreml") != std::string::npos ||
provider.find("dml") != std::string::npos ||
provider.find("directml") != std::string::npos;
}
GraphOptimizationLevel resolveGraphOptimizationLevel(bool cpuFallback) {
GraphOptimizationLevel optLevel = cpuFallback ? GraphOptimizationLevel::ORT_ENABLE_BASIC
: GraphOptimizationLevel::ORT_ENABLE_EXTENDED;
if (const char* opt = std::getenv("YAMS_ONNX_OPT_LEVEL")) {
std::string level = lowerCopy(opt);
if (level == "all") {
optLevel = GraphOptimizationLevel::ORT_ENABLE_ALL;
if (!cpuFallback) {
spdlog::info("[ONNX] Using ORT_ENABLE_ALL optimization (slow first load)");
}
} else if (level == "extended") {
optLevel = GraphOptimizationLevel::ORT_ENABLE_EXTENDED;
} else if (level == "basic") {
optLevel = GraphOptimizationLevel::ORT_ENABLE_BASIC;
}
}
return optLevel;
}
int resolveCpuFallbackIntraThreads(int configuredIntraThreads) {
if (const char* value = std::getenv("YAMS_ONNX_INTRA_OP_THREADS")) {
try {
const int parsed = std::stoi(value);
if (parsed >= 1 && parsed <= 64) {
return parsed;
}
} catch (...) {
}
}
const int fallback = configuredIntraThreads > 0 ? configuredIntraThreads : 2;
return std::clamp(fallback, 1, 2);
}
int resolveCpuFallbackInterThreads(int configuredInterThreads) {
if (const char* value = std::getenv("YAMS_ONNX_INTER_OP_THREADS")) {
try {
const int parsed = std::stoi(value);
if (parsed >= 1 && parsed <= 64) {
return parsed;
}
} catch (...) {
}
}
(void)configuredInterThreads;
return 1;
}
bool resolveCpuFallbackAllowSpinning() {
if (const char* value = std::getenv("YAMS_ONNX_ALLOW_SPINNING")) {
return envTruthy(value);
}
return false;
}
// Resolve path for caching ONNX Runtime optimized graphs on disk.
// Precedence: YAMS_DATA_DIR > XDG_DATA_HOME > ~/.local/share > model directory.
fs::path resolveOptimizedModelPath(const std::string& modelPath, GraphOptimizationLevel optLevel) {
std::string filename = fs::path(modelPath).stem().string() + "." +
std::to_string(static_cast<int>(optLevel)) + ".ort";
if (const char* dir = std::getenv("YAMS_DATA_DIR"); dir && *dir) {
return fs::path(dir) / "cache" / "onnx-opt" / filename;
}
if (const char* xdg = std::getenv("XDG_DATA_HOME"); xdg && *xdg) {
return fs::path(xdg) / "yams" / "cache" / "onnx-opt" / filename;
}
if (const char* home = std::getenv("HOME"); home && *home) {
return fs::path(home) / ".local" / "share" / "yams" / "cache" / "onnx-opt" / filename;
}
return fs::path(modelPath).parent_path() / filename;
}
bool gpuProcessLockEnabled() {
if (const char* v = std::getenv("YAMS_GPU_PROCESS_LOCK")) {
return envTruthy(v);
}
if (const char* v = std::getenv("YAMS_MIGRAPHX_PROCESS_LOCK")) {
return envTruthy(v);
}
return true;
}
std::filesystem::path resolveGpuProcessLockPath() {
if (const char* configured = std::getenv("YAMS_GPU_LOCK_PATH"); configured && *configured) {
std::filesystem::path path(configured);
if (path.extension() == ".lock") {
return path;
}
auto devEnv = std::getenv("YAMS_GPU_DEVICE_ID");
int deviceId = devEnv ? std::stoi(devEnv) : 0;
return path / ("yams-gpu-device-" + std::to_string(deviceId) + ".lock");
}
if (const char* configured = std::getenv("YAMS_MIGRAPHX_LOCK_PATH");
configured && *configured) {
std::filesystem::path path(configured);
if (path.extension() == ".lock") {
return path;
}
return path / "migraphx.lock"; // Fallback legacy behavior if needed, or just standard
}
int deviceId = envIntOr("YAMS_GPU_DEVICE_ID", -1);
if (deviceId < 0) {
deviceId = envIntOr("YAMS_MIGRAPHX_DEVICE_ID", 0);
}
return std::filesystem::temp_directory_path() /
("yams-gpu-device-" + std::to_string(deviceId) + ".lock");
}
class ScopedProcessFileLock {
public:
explicit ScopedProcessFileLock(const std::filesystem::path& path) : path_(path) {
std::error_code ec;
auto parent = path_.parent_path();
if (!parent.empty()) {
std::filesystem::create_directories(parent, ec);
}
#ifdef _WIN32
fd_ = _open(path_.string().c_str(), _O_CREAT | _O_RDWR | _O_BINARY, _S_IREAD | _S_IWRITE);
if (fd_ < 0) {
return;
}
hFile_ = reinterpret_cast<HANDLE>(_get_osfhandle(fd_));
if (hFile_ == INVALID_HANDLE_VALUE) {
_close(fd_);
fd_ = -1;
return;
}
OVERLAPPED overlapped = {0};
if (!LockFileEx(hFile_, LOCKFILE_EXCLUSIVE_LOCK, 0, 1, 0, &overlapped)) {
_close(fd_);
fd_ = -1;
hFile_ = INVALID_HANDLE_VALUE;
return;
}
#else
fd_ = ::open(path_.c_str(), O_CREAT | O_RDWR, 0644);
if (fd_ < 0) {
return;
}
if (::flock(fd_, LOCK_EX) != 0) {
::close(fd_);
fd_ = -1;
return;
}
#endif
}
ScopedProcessFileLock(const ScopedProcessFileLock&) = delete;
ScopedProcessFileLock& operator=(const ScopedProcessFileLock&) = delete;
bool locked() const { return fd_ >= 0; }
~ScopedProcessFileLock() {
if (fd_ < 0) {
return;
}
#ifdef _WIN32
if (hFile_ != INVALID_HANDLE_VALUE) {
OVERLAPPED overlapped = {0};
(void)UnlockFileEx(hFile_, 0, 1, 0, &overlapped);
}
_close(fd_);
#else
(void)::flock(fd_, LOCK_UN);
::close(fd_);
#endif
}
private:
std::filesystem::path path_;
int fd_{-1};
#ifdef _WIN32
HANDLE hFile_{INVALID_HANDLE_VALUE};
#endif
};
size_t& gpuProcessLockDepth() {
static thread_local size_t depth = 0;
return depth;
}
// Thread-local reusable buffers for ONNX tensor inputs.
// Avoids per-inference heap allocation for tokens, masks, and token_type_ids.
struct TensorBufferPool {
std::vector<int64_t> tokens;
std::vector<int64_t> masks;
std::vector<int64_t> token_type_ids;
};
TensorBufferPool& getThreadLocalBuffers() {
static thread_local TensorBufferPool pool;
return pool;
}
std::string providerBatchLimitKey(const std::string& modelName, const std::string& provider) {
std::string providerLower = provider;
std::transform(providerLower.begin(), providerLower.end(), providerLower.begin(),
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
return modelName + "|" + providerLower;
}
std::recursive_mutex& sharedBatchLimitMutex() {
static auto* mu = new std::recursive_mutex();
return *mu;
}
std::unordered_map<std::string, size_t>& sharedBatchLimitMap() {
static auto* caps = new std::unordered_map<std::string, size_t>();
return *caps;
}
size_t loadSharedBatchLimit(const std::string& modelName, const std::string& provider) {
const std::string key = providerBatchLimitKey(modelName, provider);
std::lock_guard<std::recursive_mutex> lock(sharedBatchLimitMutex());
const auto& caps = sharedBatchLimitMap();
const auto it = caps.find(key);
return (it == caps.end()) ? 0 : it->second;
}
size_t mergeSharedBatchLimit(const std::string& modelName, const std::string& provider,
size_t observedCap) {
if (observedCap == 0) {
return 0;
}
const std::string key = providerBatchLimitKey(modelName, provider);
std::lock_guard<std::recursive_mutex> lock(sharedBatchLimitMutex());
auto& caps = sharedBatchLimitMap();
const auto it = caps.find(key);
const size_t merged = (it == caps.end()) ? observedCap : std::min(it->second, observedCap);
caps[key] = merged;
return merged;
}
size_t resolveCpuFallbackBatchCap() {
constexpr size_t kSafeCpuFallbackBatchCap = 4;
const size_t configuredCap = TuneAdvisor::getEmbedDocCap();
if (configuredCap == 0) {
return kSafeCpuFallbackBatchCap;
}
return std::max<size_t>(1u, std::min(configuredCap, kSafeCpuFallbackBatchCap));
}
size_t applyCpuFallbackSafetyCap(const std::string& modelName, const std::string& provider,
std::atomic<size_t>& providerBatchLimit) {
const size_t safeCap = resolveCpuFallbackBatchCap();
const size_t mergedCap = mergeSharedBatchLimit(modelName, provider, safeCap);
providerBatchLimit.store(mergedCap, std::memory_order_relaxed);
const size_t currentCap = TuneAdvisor::getEmbedDocCap();
if (currentCap == 0 || currentCap > mergedCap) {
TuneAdvisor::setEmbedDocCap(mergedCap);
}
return mergedCap;
}
} // namespace
// ============================================================================
// Cross-instance CoreML denylist
// ============================================================================
// When a model fails CoreML inference (e.g., mask feature partition bug),
// record it here so other pool instances skip CoreML and go straight to CPU.
// This avoids wasting ~20s per instance on CoreML compilation + failed retries.
static std::mutex g_coreml_denylist_mu;
static std::unordered_set<std::string> g_coreml_denylist;
bool OnnxModelSession::isCoreMLDenylisted(const std::string& modelName) {
std::lock_guard<std::mutex> lk(g_coreml_denylist_mu);
return g_coreml_denylist.count(modelName) > 0;
}
void OnnxModelSession::denylistCoreML(const std::string& modelName) {
std::lock_guard<std::mutex> lk(g_coreml_denylist_mu);
g_coreml_denylist.insert(modelName);
}
void OnnxModelSession::clearCoreMLDenylist() {
std::lock_guard<std::mutex> lk(g_coreml_denylist_mu);
g_coreml_denylist.clear();
}
// ============================================================================
// OnnxModelSession Implementation
// ============================================================================
class OnnxModelSession::Impl {
public:
Impl(const std::string& modelPath, const std::string& modelName, const OnnxTextConfig& config)
: modelPath_(modelPath), modelName_(modelName), config_(config), preprocessor_(config) {
// Check if dynamic padding is disabled via environment variable
if (const char* dp = std::getenv("YAMS_ONNX_DYNAMIC_PADDING")) {
std::string val(dp);
std::transform(val.begin(), val.end(), val.begin(), ::tolower);
if (val == "0" || val == "false" || val == "no" || val == "off") {
dynamicPaddingEnabled_ = false;
}
}
#ifdef GTEST_API_
// Check if we should skip model loading in test mode
if (yams::test::shouldSkipModelLoading()) {
test_mode_ = true;
spdlog::debug("[ONNX] Test mode enabled - skipping actual model loading");
return;
}
#endif
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
// Lightweight, environment-driven mock mode for tests.
if (std::getenv("YAMS_USE_MOCK_PROVIDER") || std::getenv("YAMS_SKIP_MODEL_LOADING") ||
std::getenv("YAMS_TEST_MODE")) {
test_mode_ = true;
spdlog::warn("[ONNX] Mock provider mode enabled via env; skipping ONNX init");
return;
}
#endif
const auto& runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
if (!runtimeInfo.available) {
throw std::runtime_error(runtimeInfo.errorMessage.empty()
? std::string("ONNX Runtime unavailable")
: "ONNX Runtime unavailable: " + runtimeInfo.errorMessage);
}
memoryInfo_ = std::make_unique<Ort::MemoryInfo>(
Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault));
// Initialize ONNX Runtime environment
// Use a single global Ort::Env per process (best practice)
// The global env is protected by g_onnx_init_mutex during initialization.
spdlog::debug("[ONNX] Impl constructor: getting global Ort::Env for '{}'", modelName);
env_ = &get_global_ort_env();
spdlog::debug("[ONNX] Impl constructor: global Ort::Env obtained for '{}'", modelName);
// Configure session options
sessionOptions_ = std::make_unique<Ort::SessionOptions>();
// Allow env overrides for thread tuning
auto detect_threads = [](const char* name, int fallback) {
if (const char* s = std::getenv(name)) {
try {
int v = std::stoi(s);
if (v > 0 && v <= 64)
return v;
} catch (...) {
}
}
return fallback;
};
auto detect_bool = [](const char* name, bool fallback) {
if (const char* s = std::getenv(name)) {
std::string val(s);
std::transform(val.begin(), val.end(), val.begin(), ::tolower);
if (val == "0" || val == "false" || val == "no")
return false;
if (val == "1" || val == "true" || val == "yes")
return true;
}
return fallback;
};
#ifdef _WIN32
// Windows-specific: ALWAYS use minimal thread counts to avoid thread pool
// exhaustion and "resource deadlock would occur" errors. The Windows thread
// scheduler behaves differently than Linux. Do NOT use config.num_threads
// here as it defaults to hardware_concurrency() which causes issues.
int intra_default = 1; // Single thread for intra-op
int inter_default = 1; // Single thread for inter-op
bool allow_spinning_default = false; // No spinning on Windows
#else
int intra_default = 4;
int inter_default = 1;
bool allow_spinning_default = true;
#endif
// On Windows, ignore config.num_threads to enforce safe defaults
// On other platforms, respect the config
#ifdef _WIN32
int intra = detect_threads("YAMS_ONNX_INTRA_OP_THREADS", intra_default);
#else
int intra = config.num_threads > 0 ? config.num_threads : intra_default;
intra = detect_threads("YAMS_ONNX_INTRA_OP_THREADS", intra);
#endif
int inter = config.inter_op_threads > 0 ? config.inter_op_threads : inter_default;
inter = detect_threads("YAMS_ONNX_INTER_OP_THREADS", inter);
configuredIntraThreads_ = intra;
configuredInterThreads_ = inter;
sessionOptions_->SetIntraOpNumThreads(intra);
sessionOptions_->SetInterOpNumThreads(inter);
// Thread spinning control - spinning threads consume CPU cycles waiting for work
// which can cause contention issues on Windows with concurrent embedding requests
bool allow_spinning = detect_bool("YAMS_ONNX_ALLOW_SPINNING", allow_spinning_default);
sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning",
allow_spinning ? "1" : "0");
sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning",
allow_spinning ? "1" : "0");
spdlog::debug("[ONNX] SessionOptions threads: intra-op={} inter-op={} spinning={}", intra,
inter, allow_spinning);
#ifdef _WIN32
// On Windows, use sequential execution mode to avoid thread pool issues
// This ensures ONNX Runtime doesn't spawn additional threads internally
sessionOptions_->SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
spdlog::info("[ONNX] Windows: using sequential execution mode");
#endif
// Graph optimization level: EXTENDED enables operator fusion (MatMul+Add, GELU, etc.)
// which significantly benefits transformer embedding models.
// ORT_ENABLE_ALL performs additional expensive transformations (minutes for large models).
// Override with YAMS_ONNX_OPT_LEVEL=basic|extended|all
GraphOptimizationLevel optLevel = resolveGraphOptimizationLevel(false);
sessionOptions_->SetGraphOptimizationLevel(optLevel);
// Cache the optimized graph to disk so subsequent loads skip optimization passes.
if (detect_bool("YAMS_ONNX_GRAPH_CACHE", false)) {
auto optPath = resolveOptimizedModelPath(modelPath, optLevel);
std::error_code ec;
fs::create_directories(optPath.parent_path(), ec);
if (!ec) {
sessionOptions_->SetOptimizedModelFilePath(optPath.c_str());
spdlog::info("[ONNX] Optimized graph cache: {}", optPath.string());
} else {
spdlog::warn("[ONNX] Could not create graph cache dir {}: {}",
optPath.parent_path().string(), ec.message());
}
}
// Enable memory pattern optimization for faster inference
sessionOptions_->EnableMemPattern();
// Enable CPU memory arena for better memory reuse
sessionOptions_->EnableCpuMemArena();
if (config.enable_gpu) {
if (OnnxModelSession::isCoreMLDenylisted(modelName)) {
const int fallbackIntra = resolveCpuFallbackIntraThreads(configuredIntraThreads_);
const int fallbackInter = resolveCpuFallbackInterThreads(configuredInterThreads_);
const bool allowSpinning = resolveCpuFallbackAllowSpinning();
sessionOptions_->SetIntraOpNumThreads(fallbackIntra);
sessionOptions_->SetInterOpNumThreads(fallbackInter);
sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning",
allowSpinning ? "1" : "0");
sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning",
allowSpinning ? "1" : "0");
configuredIntraThreads_ = fallbackIntra;
configuredInterThreads_ = fallbackInter;
const size_t fallbackBatchCap =
applyCpuFallbackSafetyCap(modelName_, "cpu", providerBatchLimit_);
spdlog::info("[ONNX] Skipping CoreML for '{}' (previously failed; "
"CPU with intra_threads={} inter_threads={} batch_cap={} "
"spinning={})",
modelName, fallbackIntra, fallbackInter, fallbackBatchCap,
allowSpinning ? 1 : 0);
actualExecutionProvider_ = "cpu";
} else {
appendGpuExecutionProvider();
}
}
}
Result<void> loadModel() {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
if (test_mode_) {
// In test mode, pretend loading succeeded
// Derive embeddingDim_/maxSequenceLength_/pooling/normalize from nearby configs when
// available to better match real model characteristics.
try {
parseModelConfigHints();
} catch (...) {
}
// Heuristics based on common model families when config hints are absent
auto lname = modelName_;
std::transform(lname.begin(), lname.end(), lname.begin(),
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
// Dimension defaults
if (embeddingDim_ == 0) {
auto has = [&](const char* s) { return lname.find(s) != std::string::npos; };
// OpenAI TE3 sizes (mock support for external references)
if (has("text-embedding-3-large")) {
embeddingDim_ = 3072;
} else if (has("text-embedding-3-small")) {
embeddingDim_ = 1536;
} else if (has("e5")) {
if (has("large"))
embeddingDim_ = 1024;
else if (has("small"))
embeddingDim_ = 384;
else
embeddingDim_ = 768; // base/default
} else if (has("gte")) {
if (has("large"))
embeddingDim_ = 1024;
else if (has("small"))
embeddingDim_ = 384;
else
embeddingDim_ = 768; // base/default
} else if (has("bge")) {
if (has("-m3") || has("m3") || has("large"))
embeddingDim_ = 1024;
else if (has("small"))
embeddingDim_ = 384;
else
embeddingDim_ = 768; // base/default
} else if (has("nomic") || has("mpnet") || has("jina")) {
embeddingDim_ = 768;
} else if (has("minilm")) {
embeddingDim_ = 384;
} else if (has("colbert")) {
embeddingDim_ = 48;
}
}
// Sequence length defaults
if (maxSequenceLength_ == 0 || maxSequenceLength_ == 512) {
auto has = [&](const char* s) { return lname.find(s) != std::string::npos; };
if (has("nomic") || has("bge-m3") || has("text-embedding-3-")) {
maxSequenceLength_ = 8192;
} else if (has("minilm")) {
maxSequenceLength_ = 256;
} else if (has("mpnet") || has("bge") || has("e5") || has("gte") ||
has("colbert")) {
maxSequenceLength_ = 512;
}
}
// Respect explicit config cap
if (config_.max_sequence_length > 0 &&
maxSequenceLength_ > config_.max_sequence_length) {
maxSequenceLength_ = config_.max_sequence_length;
}
isLoaded_ = true;
return Result<void>();
}
auto providerLower = actualExecutionProvider_;
std::transform(providerLower.begin(), providerLower.end(), providerLower.begin(),
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
std::unique_lock<std::recursive_mutex> gpuLoadLock;
std::unique_ptr<ScopedProcessFileLock> gpuProcessLoadLock;
bool usesGpuProvider = false;
if (providerLower.find("migraphx") != std::string::npos ||
providerLower.find("cuda") != std::string::npos ||
providerLower.find("dml") != std::string::npos ||
providerLower.find("directml") != std::string::npos) {
usesGpuProvider = true;
}
if (usesGpuProvider) {
// Serialize GPU load paths to avoid non-deterministic GPU faults during startup.
static std::recursive_mutex* g_gpu_load_mutex = new std::recursive_mutex();
gpuLoadLock = std::unique_lock<std::recursive_mutex>(*g_gpu_load_mutex);
spdlog::debug("[ONNX] GPU load gate acquired for '{}'", modelName_);
if (gpuProcessLockEnabled()) {
const auto processLockPath = resolveGpuProcessLockPath();
gpuProcessLoadLock = std::make_unique<ScopedProcessFileLock>(processLockPath);
if (gpuProcessLoadLock->locked()) {
spdlog::debug("[ONNX] GPU process load lock acquired: {}",
processLockPath.string());
} else {
spdlog::warn("[ONNX] Failed to acquire GPU process load lock: {}",
processLockPath.string());
}
}
}
try {
spdlog::debug("[ONNX] Creating Ort::Session for model '{}' at {}", modelName_.c_str(),
modelPath_.c_str());
int retries = 3;
while (retries-- > 0) {
try {
// NOTE: Removed g_onnx_init_mutex lock here - it may have been conflicting
// with ONNX Runtime's internal locking. We rely on:
// 1. Single-flight pattern in loadModel() to prevent concurrent model loads
// 2. ONNX Runtime configured for single-threaded operation on Windows
spdlog::debug("[ONNX] Configuring session options for '{}'", modelName_);
// Clone session options; on Windows force single-thread to avoid deadlocks
auto options = sessionOptions_->Clone();
#ifdef _WIN32
// Windows: force single-threaded to avoid "resource deadlock would occur"
options.SetIntraOpNumThreads(1);
options.SetInterOpNumThreads(1);
#endif
// Non-Windows: use configured threads from sessionOptions_ (default intra=4)
applyCoreMLFreeDimensionOverrides(options);
// Create session directly - no async wrapper needed for local file operations
spdlog::debug("[ONNX] Creating Ort::Session for '{}' at '{}'", modelName_,
modelPath_);
session_ = std::make_unique<Ort::Session>(
*env_, std::filesystem::path(modelPath_).c_str(), options);
spdlog::info("[ONNX] Ort::Session created successfully for '{}'", modelName_);
// Success!
break;
} catch (const std::system_error& e) {
spdlog::warn("[ONNX] loadModel: system_error '{}' (code={}). Retries left: {}",
e.what(), e.code().value(), retries);
if (retries == 0)
throw;
std::this_thread::sleep_for(std::chrono::milliseconds(200));
} catch (const Ort::Exception& e) {
spdlog::warn("[ONNX] loadModel: Ort::Exception '{}'. Retries left: {}",
e.what(), retries);
if (retries == 0)
throw;
std::this_thread::sleep_for(std::chrono::milliseconds(200));
}
}
// Get input/output information
size_t numInputs = session_->GetInputCount();
size_t numOutputs = session_->GetOutputCount();
if (numInputs < 2) {
return Error{ErrorCode::InvalidData,
"Model must have at least 2 inputs (input_ids, attention_mask)"};
}
// Get input/output names
Ort::AllocatorWithDefaultOptions allocator;
for (size_t i = 0; i < numInputs; ++i) {
auto inputName = session_->GetInputNameAllocated(i, allocator);
inputNames_.push_back(inputName.get());
// Inspect first input shape to detect model sequence length when available
if (i == 0) {
try {
auto inInfo = session_->GetInputTypeInfo(i);
auto inTensor = inInfo.GetTensorTypeAndShapeInfo();
auto inShape = inTensor.GetShape();
// Expect shape like [batch, seq_len]; seq_len > 0 means fixed
if (inShape.size() >= 2) {
auto seq = inShape[1];
if (seq > 0) {
maxSequenceLength_ = static_cast<size_t>(seq);
hasFixedInputShape_ = true;
}
}
} catch (const std::exception& e) {
spdlog::debug("[ONNX] Failed to read input shape: {}", e.what());
}
}
}
for (size_t i = 0; i < numOutputs; ++i) {
auto outputName = session_->GetOutputNameAllocated(i, allocator);
outputNames_.push_back(outputName.get());
}
// Get output shape to determine embedding dimension
if (numOutputs > 0) {
auto outputInfo = session_->GetOutputTypeInfo(0);
auto tensorInfo = outputInfo.GetTensorTypeAndShapeInfo();
auto shape = tensorInfo.GetShape();
if (shape.size() >= 2 && shape.back() > 0) {
embeddingDim_ = static_cast<size_t>(shape.back());
}
}
// Try to parse nearby config (Sentence-Transformers) for pooling/max
// length/normalization
parseModelConfigHints();
// Load tokenizer.json from model directory (same search dirs as config.json)
try {
namespace fs = std::filesystem;
fs::path mp(modelPath_);
for (const auto& dir : {mp.parent_path(), mp.parent_path().has_parent_path()
? mp.parent_path().parent_path()
: fs::path{}}) {
if (dir.empty())
continue;
fs::path tok = dir / "tokenizer.json";
if (fs::exists(tok) && tokenizer_.load(tok.string())) {
spdlog::info("[ONNX] Loaded tokenizer from {} (vocab={})", tok.string(),
tokenizer_.vocabSize());
break;
}
}
if (!tokenizer_.isLoaded()) {
spdlog::warn(
"[ONNX] tokenizer.json not found; using TextPreprocessor fallback");
}
} catch (...) {
}
// If input shape and config did not give a positive seq_len, fall back to a heuristic.
if (maxSequenceLength_ == 0) {
if (modelName_.find("MiniLM") != std::string::npos) {
maxSequenceLength_ = 256;
} else if (modelName_.find("mpnet") != std::string::npos) {
maxSequenceLength_ = 512;
} else if (modelName_.find("nomic") != std::string::npos) {
// Nomic BERT supports very long contexts (e.g., 8192). Default to 2048 unless
// user config overrides to avoid excessive memory on defaults.
maxSequenceLength_ = 2048;
} else if (modelName_.find("embeddinggemma") != std::string::npos ||
modelName_.find("gemma") != std::string::npos) {
maxSequenceLength_ = 2048;
} else {
maxSequenceLength_ = 512;
}
}
// Respect configured max from EmbeddingConfig if provided
if (config_.max_sequence_length > 0) {
maxSequenceLength_ = std::min(maxSequenceLength_, config_.max_sequence_length);
}
if (gemmaCoreMLStaticBatchOneEnabled()) {
providerBatchLimit_.store(1, std::memory_order_relaxed);
spdlog::info("[ONNX] Enabled static batch_size=1 CoreML path for '{}'; larger "
"embedding batches will be chunked per item",
modelName_);
}
isLoaded_ = true;
// IoBinding benefits real GPU providers (pre-resolved memory placement);
// on CPU/CoreML-to-CPU it adds overhead vs direct Session::Run().
if (actualExecutionProvider_ == "cuda" || actualExecutionProvider_ == "migraphx" ||
actualExecutionProvider_ == "directml") {
initIoBinding();
}
spdlog::info(
"[ONNX] Session ready for '{}' (inputs={}, outputs={}, dim={}, max_seq_len={})",
modelName_.c_str(), numInputs, numOutputs, embeddingDim_, maxSequenceLength_);
return Result<void>();
} catch (const Ort::Exception& e) {
spdlog::error("[ONNX] Failed to load '{}' (Ort::Exception): {}", modelName_.c_str(),
e.what());
return Error{ErrorCode::InternalError,
std::string("Failed to load ONNX model: ") + e.what()};
} catch (const std::system_error& e) {
spdlog::error("[ONNX] Failed to load '{}' (std::system_error): {}", modelName_.c_str(),
e.what());
return Error{ErrorCode::InternalError,
std::string("Failed to load ONNX model (system error): ") + e.what()};
}
}
Result<std::vector<float>> generateEmbedding(const std::string& text) {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
if (test_mode_) {
// Representative mock: deterministic per-text embedding derived from tokens.
// - Uses same preprocessing (tokenize/truncate/pad + attention mask)
// - Pseudo-random per-token vector seeded by token id and position
// - Pooling respects configured pooling_mode_
// - Optional L2 normalization
const size_t seq_len = maxSequenceLength_ > 0 ? maxSequenceLength_ : 512;
auto toks = preprocessor_.tokenize(text);
toks = preprocessor_.truncateTokens(toks, seq_len);
toks = preprocessor_.padTokens(toks, seq_len);
auto mask = preprocessor_.generateAttentionMask(toks);
const size_t D = embeddingDim_ > 0 ? embeddingDim_ : 0;
auto prng_val = [](uint32_t seed) -> float {
// xorshift32 → map to [-1, 1]
uint32_t x = seed ? seed : 0x9e3779b9u;
x ^= x << 13;
x ^= x >> 17;
x ^= x << 5;
// Scale to [0,1)
float u = (x & 0x00FFFFFFu) / static_cast<float>(0x01000000u);
return 2.0f * u - 1.0f;
};
std::vector<float> emb(D, 0.0f);
if (pooling_mode_ == Pooling::CLS) {
const int32_t t = toks.empty() ? 0 : toks[0];
for (size_t d = 0; d < D; ++d) {
emb[d] = prng_val(static_cast<uint32_t>(t * 1315423911u + d * 2654435761u));
}
} else if (pooling_mode_ == Pooling::MAX) {
for (size_t i = 0; i < std::min(seq_len, toks.size()); ++i) {
if (i >= mask.size() || mask[i] <= 0)
continue;
const int32_t t = toks[i];
for (size_t d = 0; d < D; ++d) {
float v = prng_val(static_cast<uint32_t>(t * 2246822519u + i * 3266489917u +
d * 668265263u));
emb[d] = (i == 0) ? v : std::max(emb[d], v);
}
}
} else { // MEAN pooling
double denom = 0.0;
for (size_t i = 0; i < std::min(seq_len, toks.size()); ++i) {
float w = (i < mask.size() && mask[i] > 0) ? 1.0f : 0.0f;
if (w <= 0.0f)
continue;
denom += w;
const int32_t t = toks[i];
for (size_t d = 0; d < D; ++d) {
float v = prng_val(
static_cast<uint32_t>(t * 374761393u + i * 1664525u + d * 1013904223u));
emb[d] += v * w;
}
}
if (denom > 0.0) {
for (size_t d = 0; d < D; ++d)
emb[d] = static_cast<float>(emb[d] / denom);
}
}
if (normalize_) {
double n2 = 0.0;
for (auto v : emb) {
double dv = static_cast<double>(v);
n2 += dv * dv;
}
n2 = std::sqrt(n2);
if (n2 > 1e-8) {
for (auto& v : emb)
v = static_cast<float>(v / n2);
} else {
std::fill(emb.begin(), emb.end(), 0.0f);
}
}
return emb;
}
if (!isLoaded_) {
if (auto result = loadModel(); !result) {
return result.error();
}
}
// Try real ONNX path first (preferred)
if (auto r = withCoreMLFallback([&]() { return runOnnx(text); }); r) {
return r.value();
}
return Error{ErrorCode::InternalError,
"ONNX embedding inference failed and no fallback backend is available"};
}
// Public batch API that uses the optimized batched ONNX path
Result<std::vector<std::vector<float>>>
generateBatchEmbeddings(const std::vector<std::string>& texts) {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
if (test_mode_) {
std::vector<std::vector<float>> res;
res.reserve(texts.size());
for (const auto& t : texts) {
auto one = generateEmbedding(t);
if (!one)
return one.error();
res.emplace_back(std::move(one.value()));
}
return res;
}
if (!isLoaded_) {
if (auto result = loadModel(); !result) {
return result.error();
}
}
return withCoreMLFallback([&]() { return runOnnxBatch(texts); });
}
bool isValid() const {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
return test_mode_ ? isLoaded_ : (isLoaded_ && session_ != nullptr);
}
std::pair<int, int> getThreading() const {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
return {configuredIntraThreads_, configuredInterThreads_};
}
size_t getLearnedBatchLimit() const {
return providerBatchLimit_.load(std::memory_order_relaxed);
}
bool prefersExplicitSentenceEmbeddingOutput() const {
return modelName_.find("embeddinggemma") != std::string::npos ||
modelName_.find("gemma") != std::string::npos;
}
bool prefersStaticShapesOnCoreML() const {
return actualExecutionProvider_ == "coreml" &&
(modelName_.find("embeddinggemma") != std::string::npos ||
modelName_.find("gemma") != std::string::npos);
}
bool gemmaCoreMLStaticBatchOneEnabled() const { return prefersStaticShapesOnCoreML(); }
bool shouldApplyCoreMLFreeDimensionOverrides() const {
return gemmaCoreMLStaticBatchOneEnabled();
}
int64_t resolveCoreMLStaticSequenceLength() const {
if (config_.max_sequence_length > 0) {
return std::max<int64_t>(1, std::min<int64_t>(config_.max_sequence_length, 2048));
}
if (maxSequenceLength_ > 512) {
return static_cast<int64_t>(maxSequenceLength_);
}
return int64_t{2048};
}
void applyCoreMLFreeDimensionOverrides(Ort::SessionOptions& options) {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
if (!shouldApplyCoreMLFreeDimensionOverrides()) {
return;
}
const int64_t effectiveBatchSize = 1;
const int64_t sequenceLength = resolveCoreMLStaticSequenceLength();
const int64_t tripleSequenceLength =
sequenceLength > (std::numeric_limits<int64_t>::max() / 3) ? sequenceLength
: (sequenceLength * 3);
Ort::ThrowOnError(Ort::GetApi().AddFreeDimensionOverrideByName(options, "batch_size",
effectiveBatchSize));
Ort::ThrowOnError(Ort::GetApi().AddFreeDimensionOverrideByName(
options, "total_sequence_length", sequenceLength));
Ort::ThrowOnError(Ort::GetApi().AddFreeDimensionOverrideByName(
options, "3*total_sequence_length", tripleSequenceLength));
spdlog::info("[ONNX] Applied Gemma CoreML static-shape overrides for '{}': batch_size={} "
"total_sequence_length={} triple_total_sequence_length={}",
modelName_, effectiveBatchSize, sequenceLength, tripleSequenceLength);
}
bool supportsCpuFallbackForCurrentModel() const {
return supportsCoreMLCpuFallbackModel(modelName_);
}
bool isCoreMLMaskFeatureFailure(const std::string& message) const {
return message.find("_model_st_pool_0_mask_unsqueeze_output_0") != std::string::npos ||
message.find("required but not specified") != std::string::npos;
}
Result<void> reinitializeSessionOnCpu() {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
if (!supportsCpuFallbackForCurrentModel()) {
return Error{ErrorCode::InternalError, "CPU fallback is not enabled for this model"};
}
if (actualExecutionProvider_ == "cpu") {
return Result<void>();
}
try {
session_.reset();
sessionOptions_ = std::make_unique<Ort::SessionOptions>();
actualExecutionProvider_ = "cpu";
providerBatchLimit_.store(0, std::memory_order_relaxed);
const int hwThreads = static_cast<int>(std::thread::hardware_concurrency());
const int fallbackIntra = resolveCpuFallbackIntraThreads(configuredIntraThreads_);
const int fallbackInter = resolveCpuFallbackInterThreads(configuredInterThreads_);
sessionOptions_->SetIntraOpNumThreads(fallbackIntra);
sessionOptions_->SetInterOpNumThreads(fallbackInter);
configuredIntraThreads_ = fallbackIntra;
configuredInterThreads_ = fallbackInter;
const bool allowSpinning = resolveCpuFallbackAllowSpinning();
sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning",
allowSpinning ? "1" : "0");
sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning",
allowSpinning ? "1" : "0");
#ifdef _WIN32
sessionOptions_->SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
#endif
GraphOptimizationLevel optLevel = resolveGraphOptimizationLevel(true);
sessionOptions_->SetGraphOptimizationLevel(optLevel);
// Graph cache is on by default; disable with YAMS_ONNX_GRAPH_CACHE=0
if (!std::getenv("YAMS_ONNX_GRAPH_CACHE") ||
envTruthy(std::getenv("YAMS_ONNX_GRAPH_CACHE"))) {
auto optPath = resolveOptimizedModelPath(modelPath_, optLevel);
std::error_code ec;
fs::create_directories(optPath.parent_path(), ec);
if (!ec) {
sessionOptions_->SetOptimizedModelFilePath(optPath.c_str());
}
}
sessionOptions_->EnableMemPattern();
sessionOptions_->EnableCpuMemArena();
auto options = sessionOptions_->Clone();
#ifdef _WIN32
options.SetIntraOpNumThreads(1);
options.SetInterOpNumThreads(1);
#endif
session_ = std::make_unique<Ort::Session>(
*env_, std::filesystem::path(modelPath_).c_str(), options);
// CPU fallback — no IoBinding (adds overhead on CPU)
useIoBinding_ = false;
ioBinding_.reset();
OnnxModelSession::denylistCoreML(modelName_);
const size_t fallbackBatchCap = applyCpuFallbackSafetyCap(
modelName_, actualExecutionProvider_, providerBatchLimit_);
spdlog::warn("[ONNX] Reinitialized '{}' on CPU after CoreML inference failure "
"(denylisted for future instances, intra_threads={}, inter_threads={}, "
"batch_cap={}, spinning={}, hw={})",
modelName_, fallbackIntra, fallbackInter, fallbackBatchCap,
allowSpinning ? 1 : 0, hwThreads);
return Result<void>();
} catch (const Ort::Exception& e) {
return Error{ErrorCode::InternalError,
std::string("CPU fallback session init failed: ") + e.what()};
} catch (const std::exception& e) {
return Error{ErrorCode::InternalError,
std::string("CPU fallback session init failed: ") + e.what()};
}
}
template <typename Fn> auto withCoreMLFallback(Fn&& fn) -> decltype(fn()) {
auto result = fn();
if (result || actualExecutionProvider_ != "coreml" ||
!supportsCpuFallbackForCurrentModel()) {
return result;
}
const auto& message = result.error().message;
if (!isCoreMLMaskFeatureFailure(message)) {
return result;
}
auto fallback = reinitializeSessionOnCpu();
if (!fallback) {
return decltype(fn())(fallback.error());
}
return fn();
}
std::vector<const char*> resolveOutputNames() const {
if (outputNames_.empty()) {
static const char* default_out = "sentence_embedding";
return {default_out};
}
if (prefersExplicitSentenceEmbeddingOutput()) {
for (const auto& name : outputNames_) {
if (name == "sentence_embedding") {
return {name.c_str()};
}
}
}
std::vector<const char*> output_names_c;
output_names_c.reserve(outputNames_.size());
for (const auto& n : outputNames_) {
output_names_c.push_back(n.c_str());
}
return output_names_c;
}
Result<void> setThreading(int intraThreads, int interThreads) {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
auto valid = [](int value) { return value == -1 || (value >= 1 && value <= 64); };
if (!valid(intraThreads) || !valid(interThreads)) {
return Error{ErrorCode::InvalidArgument,
"Threading values must be -1 or in range [1, 64]"};
}
if (intraThreads > 0) {
config_.num_threads = intraThreads;
}
if (interThreads > 0) {
config_.inter_op_threads = interThreads;
}
configuredIntraThreads_ =
(config_.num_threads > 0) ? config_.num_threads : configuredIntraThreads_;
configuredInterThreads_ =
(config_.inter_op_threads > 0) ? config_.inter_op_threads : configuredInterThreads_;
if (sessionOptions_) {
sessionOptions_->SetIntraOpNumThreads(configuredIntraThreads_);
sessionOptions_->SetInterOpNumThreads(configuredInterThreads_);
}
// Force lazy reload on next inference with updated threading.
session_.reset();
isLoaded_ = false;
return Result<void>();
}
size_t getEmbeddingDim() const {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
return embeddingDim_;
}
size_t getMaxSequenceLength() const {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
return maxSequenceLength_;
}
std::string getExecutionProvider() const {
std::lock_guard<std::recursive_mutex> lock(stateMutex_);
return actualExecutionProvider_;
}
private:
void appendGpuExecutionProvider() {
std::string cacheDir;
if (!modelPath_.empty()) {
auto parentDir = fs::path(modelPath_).parent_path();
if (fs::exists(parentDir)) {
cacheDir = parentDir.string();
}
}
actualExecutionProvider_ = onnx_util::appendGpuProvider(*sessionOptions_, cacheDir);
}
struct GpuInferenceGuard {
std::unique_lock<std::recursive_mutex> threadLock;
std::unique_ptr<ScopedProcessFileLock> processLock;
bool depthIncremented{false};
GpuInferenceGuard() = default;
GpuInferenceGuard(const GpuInferenceGuard&) = delete;
GpuInferenceGuard& operator=(const GpuInferenceGuard&) = delete;
GpuInferenceGuard(GpuInferenceGuard&&) noexcept = default;
GpuInferenceGuard& operator=(GpuInferenceGuard&&) noexcept = default;
~GpuInferenceGuard() {
if (depthIncremented) {
auto& depth = gpuProcessLockDepth();
if (depth > 0) {
--depth;
}
}
}
};
GpuInferenceGuard acquireGpuInferenceLock() const {
GpuInferenceGuard guard;
std::string providerLower = actualExecutionProvider_;
std::transform(providerLower.begin(), providerLower.end(), providerLower.begin(),
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
bool usesGpuProvider = false;
if (providerLower.find("migraphx") != std::string::npos ||
providerLower.find("cuda") != std::string::npos ||
providerLower.find("dml") != std::string::npos ||
providerLower.find("directml") != std::string::npos) {
usesGpuProvider = true;
}
if (!usesGpuProvider) {
return guard;
}
static std::recursive_mutex* g_gpu_infer_mutex = new std::recursive_mutex();
guard.threadLock = std::unique_lock<std::recursive_mutex>(*g_gpu_infer_mutex);
if (!gpuProcessLockEnabled()) {
return guard;
}
auto& depth = gpuProcessLockDepth();
++depth;
guard.depthIncremented = true;
if (depth > 1) {
return guard;
}
const auto processLockPath = resolveGpuProcessLockPath();
guard.processLock = std::make_unique<ScopedProcessFileLock>(processLockPath);
if (!guard.processLock->locked()) {
spdlog::warn("[ONNX] Failed to acquire GPU process inference lock: {}",
processLockPath.string());
guard.processLock.reset();
}
return guard;
}
// Align a sequence length up to a SIMD-friendly multiple (AVX: 8 x int32, NEON: 4).
static size_t alignSeqLen(size_t len, size_t alignment = 8) {
return ((len + alignment - 1) / alignment) * alignment;
}
void initIoBinding() {
if (!session_)
return;
try {
ioBinding_ = std::make_unique<Ort::IoBinding>(*session_);
useIoBinding_ = true;
spdlog::debug("[ONNX] IoBinding created for '{}'", modelName_);
} catch (const std::exception& e) {
spdlog::warn("[ONNX] IoBinding creation failed for '{}': {}; using Run()", modelName_,
e.what());
useIoBinding_ = false;
}
}
// Run inference using IoBinding if available, otherwise fall back to Session::Run().
std::vector<Ort::Value> runSession(std::vector<Ort::Value>& inputs) {
if (useIoBinding_ && ioBinding_) {
ioBinding_->ClearBoundInputs();
ioBinding_->ClearBoundOutputs();
for (size_t i = 0; i < inputs.size() && i < inputNames_.size(); ++i) {
ioBinding_->BindInput(inputNames_[i].c_str(), inputs[i]);
}
for (const auto& name : resolveOutputNames()) {
ioBinding_->BindOutput(name, *memoryInfo_);
}
session_->Run(Ort::RunOptions{nullptr}, *ioBinding_);
return ioBinding_->GetOutputValues();
}
// Fallback: direct Run with name arrays
std::vector<const char*> in_names;
for (auto& n : inputNames_)
in_names.push_back(n.c_str());
auto out_names = resolveOutputNames();
return session_->Run(Ort::RunOptions{nullptr}, in_names.data(), inputs.data(),
inputs.size(), out_names.data(), out_names.size());
}
Result<std::vector<float>> runOnnx(const std::string& t) {
auto gpuInferenceLock = acquireGpuInferenceLock();
const size_t seq_len = maxSequenceLength_ > 0 ? maxSequenceLength_ : 512;
auto tokens = tokenizer_.isLoaded() ? tokenizer_.encode(t) : preprocessor_.tokenize(t);
tokens = preprocessor_.truncateTokens(tokens, seq_len);
// Dynamic padding: pad to aligned actual length instead of model max
size_t effective_seq_len = seq_len;
if (dynamicPaddingEnabled_ && !hasFixedInputShape_ && !prefersStaticShapesOnCoreML()) {
effective_seq_len = std::min(alignSeqLen(tokens.size()), seq_len);
if (effective_seq_len == 0)
effective_seq_len = 1;
spdlog::debug("[ONNX] Dynamic padding (single): actual={} effective={} (model_max={})",
tokens.size(), effective_seq_len, seq_len);
} else if (prefersStaticShapesOnCoreML()) {
spdlog::debug("[ONNX] Static padding forced for '{}' on CoreML: seq_len={}", modelName_,
effective_seq_len);
}
tokens = preprocessor_.padTokens(tokens, effective_seq_len);
auto attention_mask = preprocessor_.generateAttentionMask(tokens);
std::vector<int64_t> input_shape = {1, static_cast<int64_t>(effective_seq_len)};
const size_t input_tensor_size = effective_seq_len;
// Reuse thread-local buffers to avoid per-call heap allocations.
auto& bufs = getThreadLocalBuffers();
bufs.tokens.resize(effective_seq_len);
bufs.masks.resize(effective_seq_len);
if (tokenizer_.isLoaded()) {
for (size_t i = 0; i < effective_seq_len && i < tokens.size(); ++i)
bufs.tokens[i] = static_cast<int64_t>(tokens[i]);
for (size_t i = tokens.size(); i < effective_seq_len; ++i)
bufs.tokens[i] = 0;
} else {
const int64_t MAX_TOKEN_ID = 30527;
const int64_t UNK_TOKEN_ID = 100;
for (size_t i = 0; i < effective_seq_len && i < tokens.size(); ++i) {
auto v = tokens[i];
bufs.tokens[i] =
(v < 0 || v > MAX_TOKEN_ID) ? UNK_TOKEN_ID : static_cast<int64_t>(v);
}
for (size_t i = tokens.size(); i < effective_seq_len; ++i)
bufs.tokens[i] = UNK_TOKEN_ID;
}
for (size_t i = 0; i < effective_seq_len && i < attention_mask.size(); ++i)
bufs.masks[i] = static_cast<int64_t>(attention_mask[i]);
for (size_t i = attention_mask.size(); i < effective_seq_len; ++i)
bufs.masks[i] = 0;
std::vector<Ort::Value> input_tensors;
input_tensors.push_back(Ort::Value::CreateTensor<int64_t>(
*memoryInfo_, bufs.tokens.data(), input_tensor_size, input_shape.data(), 2));
input_tensors.push_back(Ort::Value::CreateTensor<int64_t>(
*memoryInfo_, bufs.masks.data(), input_tensor_size, input_shape.data(), 2));
if (inputNames_.size() >= 3) {
bufs.token_type_ids.assign(effective_seq_len, 0);
input_tensors.push_back(
Ort::Value::CreateTensor<int64_t>(*memoryInfo_, bufs.token_type_ids.data(),
input_tensor_size, input_shape.data(), 2));
}
auto outputs = runSession(input_tensors);
if (outputs.empty()) {
return Error{ErrorCode::InternalError, "ONNX session returned no outputs"};
}
auto& out = outputs[0];
auto out_shape = out.GetTensorTypeAndShapeInfo().GetShape();
const size_t hidden_dim = embeddingDim_;
if (out_shape.size() == 3) {
const size_t seq = static_cast<size_t>(out_shape[1]);
const size_t dim = static_cast<size_t>(out_shape[2]);
if (dim != hidden_dim) {
return Error{ErrorCode::InvalidData, "Output dim mismatch"};
}
const float* data = out.GetTensorData<float>();
std::vector<float> embedding(hidden_dim, 0.0f);
if (pooling_mode_ == Pooling::CLS) {
// CLS token (first token)
const size_t off = 0;
for (size_t d = 0; d < dim; ++d)
embedding[d] = data[off + d];
} else if (pooling_mode_ == Pooling::MAX) {
for (size_t i = 0; i < std::min(seq, attention_mask.size()); ++i) {
if (attention_mask[i] <= 0)
continue;
const size_t off = i * dim;
for (size_t d = 0; d < dim; ++d)
embedding[d] = std::max(embedding[d], data[off + d]);
}
} else { // MEAN
float total_attn = 0.0f;
for (size_t i = 0; i < std::min(seq, attention_mask.size()); ++i) {
const float w = static_cast<float>(attention_mask[i]);
if (w <= 0.0f)
continue;
total_attn += w;
const size_t off = i * dim;
for (size_t d = 0; d < dim; ++d)
embedding[d] += data[off + d] * w;
}
if (total_attn > 0.0f) {
for (auto& v : embedding)
v = static_cast<float>(v / total_attn);
}
}
if (normalize_) {
double norm = 0.0;
for (auto v : embedding) {
double dv = static_cast<double>(v);
norm += dv * dv;
}
norm = std::sqrt(norm);
if (norm > 1e-8) {
for (auto& v : embedding)
v = static_cast<float>(v / norm);
} else {
std::fill(embedding.begin(), embedding.end(), 0.0f);
}
}
return embedding;
}
if (out_shape.size() == 2) {
const size_t dim = static_cast<size_t>(out_shape[1]);
if (dim != hidden_dim) {
return Error{ErrorCode::InvalidData, "Output dim mismatch"};
}
const float* data = out.GetTensorData<float>();
std::vector<float> embedding(data, data + dim);
if (normalize_) {
double norm = 0.0;
for (auto v : embedding) {
double dv = static_cast<double>(v);
norm += dv * dv;
}
norm = std::sqrt(norm);
if (norm > 1e-8) {
for (auto& v : embedding)
v = static_cast<float>(v / norm);
} else {
std::fill(embedding.begin(), embedding.end(), 0.0f);
}
}
return embedding;
}
return Error{ErrorCode::InvalidData, "Unexpected output tensor rank"};
}
Result<std::vector<std::vector<float>>> runOnnxBatch(const std::vector<std::string>& texts) {
auto gpuInferenceLock = acquireGpuInferenceLock();
if (texts.empty())
return std::vector<std::vector<float>>{};
if (providerBatchLimit_.load(std::memory_order_relaxed) == 0) {
const size_t sharedCap = loadSharedBatchLimit(modelName_, actualExecutionProvider_);
if (sharedCap > 0) {
providerBatchLimit_.store(sharedCap, std::memory_order_relaxed);
const auto currentCap = TuneAdvisor::getEmbedDocCap();
if (currentCap == 0 || currentCap > sharedCap) {
TuneAdvisor::setEmbedDocCap(sharedCap);
}
}
}
const size_t requestedBatchSize = texts.size();
size_t learnedBatchCap = providerBatchLimit_.load(std::memory_order_relaxed);
if (gemmaCoreMLStaticBatchOneEnabled() && requestedBatchSize > 1 && learnedBatchCap == 0) {
providerBatchLimit_.store(1, std::memory_order_relaxed);
learnedBatchCap = 1;
}
if (learnedBatchCap > 0 && requestedBatchSize > learnedBatchCap) {
std::vector<std::vector<float>> merged;
merged.reserve(requestedBatchSize);
for (size_t offset = 0; offset < requestedBatchSize; offset += learnedBatchCap) {
const size_t count =
std::min(learnedBatchCap, static_cast<size_t>(requestedBatchSize - offset));
std::vector<std::string> chunk;
chunk.reserve(count);
for (size_t i = 0; i < count; ++i) {
chunk.push_back(texts[offset + i]);
}
auto chunkResult = runOnnxBatch(chunk);
if (!chunkResult) {
return chunkResult.error();
}
auto& vectors = chunkResult.value();
merged.insert(merged.end(), std::make_move_iterator(vectors.begin()),
std::make_move_iterator(vectors.end()));
}
return merged;
}
const size_t seq_len = maxSequenceLength_ > 0 ? maxSequenceLength_ : 512;
// Pass 1: tokenize + truncate (no padding yet)
std::vector<std::vector<int32_t>> token_seqs;
token_seqs.reserve(texts.size());
for (const auto& t : texts) {
auto tokens = tokenizer_.isLoaded() ? tokenizer_.encode(t) : preprocessor_.tokenize(t);
tokens = preprocessor_.truncateTokens(tokens, seq_len);
token_seqs.push_back(std::move(tokens));
}
// Determine effective padding length
size_t effective_seq_len = seq_len;
if (dynamicPaddingEnabled_ && !hasFixedInputShape_ && !prefersStaticShapesOnCoreML()) {
size_t max_actual = 0;
for (const auto& toks : token_seqs)
max_actual = std::max(max_actual, toks.size());
effective_seq_len = std::min(alignSeqLen(max_actual), seq_len);
if (effective_seq_len == 0)
effective_seq_len = 1;
spdlog::debug("[ONNX] Dynamic padding: B={} max_actual={} effective={} (model_max={})",
texts.size(), max_actual, effective_seq_len, seq_len);
} else if (prefersStaticShapesOnCoreML()) {
spdlog::debug("[ONNX] Static padding forced for '{}' batch on CoreML: B={} seq_len={}",
modelName_, texts.size(), effective_seq_len);
}
// Pass 2: pad to effective_seq_len + generate masks
std::vector<std::vector<int32_t>> masks;
masks.reserve(texts.size());
for (auto& toks : token_seqs) {
toks = preprocessor_.padTokens(toks, effective_seq_len);
masks.push_back(preprocessor_.generateAttentionMask(toks));
}
// Prepare batched tensors [B, S]
const size_t B = texts.size();
std::vector<int64_t> input_shape = {static_cast<int64_t>(B),
static_cast<int64_t>(effective_seq_len)};
const size_t tensor_size = B * effective_seq_len;
// Reuse thread-local buffers to avoid per-call heap allocations.
auto& bufs = getThreadLocalBuffers();
bufs.tokens.resize(tensor_size);
bufs.masks.resize(tensor_size);
const bool useRealTokenizer = tokenizer_.isLoaded();
const int64_t MAX_TOKEN_ID = 30527;
const int64_t UNK_TOKEN_ID = 100;
size_t offset = 0;
for (size_t i = 0; i < B; ++i) {
const auto& toks = token_seqs[i];
const auto& m = masks[i];
for (size_t j = 0; j < effective_seq_len && j < toks.size(); ++j) {
int64_t v = toks[j];
if (useRealTokenizer) {
bufs.tokens[offset] = v;
} else {
bufs.tokens[offset] = (v < 0 || v > MAX_TOKEN_ID) ? UNK_TOKEN_ID : v;
}
bufs.masks[offset] = static_cast<int64_t>(m[j]);
++offset;
}
for (size_t j = toks.size(); j < effective_seq_len; ++j) {
bufs.tokens[offset] = useRealTokenizer ? 0 : UNK_TOKEN_ID;
bufs.masks[offset] = 0;
++offset;
}
}
std::vector<Ort::Value> inputs;
inputs.push_back(Ort::Value::CreateTensor<int64_t>(*memoryInfo_, bufs.tokens.data(),
tensor_size, input_shape.data(), 2));
inputs.push_back(Ort::Value::CreateTensor<int64_t>(*memoryInfo_, bufs.masks.data(),
tensor_size, input_shape.data(), 2));
if (inputNames_.size() >= 3) {
bufs.token_type_ids.assign(tensor_size, 0);
inputs.push_back(Ort::Value::CreateTensor<int64_t>(
*memoryInfo_, bufs.token_type_ids.data(), tensor_size, input_shape.data(), 2));
}
// Wrap ONNX Run in try/catch to prevent exceptions from bubbling through ABI boundary
// On Windows, ORT can throw std::system_error("resource deadlock would occur") when
// thread pool resources are exhausted (PBI-1c1)
std::vector<Ort::Value> outputs;
int retries = 3;
while (retries-- > 0) {
try {
outputs = runSession(inputs);
break; // synchronous success
} catch (const Ort::Exception& e) {
if (retries == 0) {
fprintf(stderr, "[ONNX Plugin] Ort::Exception in Run: %s\n", e.what());
return Error{ErrorCode::InternalError,
std::string("ONNX runtime error: ") + e.what()};
}
std::this_thread::sleep_for(std::chrono::milliseconds(50));
} catch (const std::system_error& e) {
if (std::string(e.what()).find("resource deadlock") != std::string::npos ||
retries > 0) {
spdlog::warn("[ONNX Plugin] Deadlock/system error caught, retrying... ({})",
e.what());
std::this_thread::sleep_for(std::chrono::milliseconds(100));
if (retries == 0) {
fprintf(stderr, "[ONNX Plugin] std::system_error in Run (exhausted): %s\n",
e.what());
return Error{ErrorCode::InternalError,
std::string("System error during ONNX inference: ") +
e.what()};
}
continue;
}
fprintf(stderr, "[ONNX Plugin] std::system_error in Run: %s\n", e.what());
return Error{ErrorCode::InternalError,
std::string("System error during ONNX inference: ") + e.what()};
} catch (const std::exception& e) {
fprintf(stderr, "[ONNX Plugin] std::exception in Run: %s\n", e.what());
return Error{ErrorCode::InternalError,
std::string("Exception during ONNX inference: ") + e.what()};
}
}
if (outputs.empty())
return Error{ErrorCode::InternalError, "ONNX session returned no outputs"};
auto& out = outputs[0];
auto shape = out.GetTensorTypeAndShapeInfo().GetShape();
// Some execution providers can leave output dims unknown until the first Run.
// If we don't know the embedding dim yet, infer it from the first output tensor.
size_t hidden_dim = embeddingDim_;
std::vector<std::vector<float>> result;
if (shape.size() == 3) {
// [B, S, D]
const int64_t b0 = shape[0];
const int64_t s1 = shape[1];
const int64_t d2 = shape[2];
const size_t SS = (s1 > 0) ? static_cast<size_t>(s1) : maxSequenceLength_;
const size_t DD = (d2 > 0) ? static_cast<size_t>(d2) : 0;
if (hidden_dim == 0 && DD > 0) {
hidden_dim = DD;
embeddingDim_ = hidden_dim;
}
if (b0 > 0 && static_cast<size_t>(b0) < B && (hidden_dim == 0 || DD == hidden_dim)) {
const size_t observedCap = static_cast<size_t>(b0);
const size_t mergedCap =
mergeSharedBatchLimit(modelName_, actualExecutionProvider_, observedCap);
if (mergedCap > 0) {
providerBatchLimit_.store(mergedCap, std::memory_order_relaxed);
const auto currentCap = TuneAdvisor::getEmbedDocCap();
if (currentCap == 0 || currentCap > mergedCap) {
TuneAdvisor::setEmbedDocCap(mergedCap);
}
}
spdlog::warn(
"[ONNX] Provider returned smaller batch than requested: requested_B={} "
"output_B={} (learned_cap={}), processing partial result + chunking remainder",
B, observedCap, mergedCap);
std::vector<std::vector<float>> partial(observedCap,
std::vector<float>(hidden_dim, 0.0f));
const float* data = out.GetTensorData<float>();
for (size_t b = 0; b < observedCap; ++b) {
auto& emb = partial[b];
if (pooling_mode_ == Pooling::CLS) {
const size_t off = b * SS * DD;
for (size_t d = 0; d < DD; ++d)
emb[d] = data[off + d];
} else if (pooling_mode_ == Pooling::MAX) {
for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
if (masks[b][i] <= 0)
continue;
const size_t off = b * SS * DD + i * DD;
for (size_t d = 0; d < DD; ++d)
emb[d] = std::max(emb[d], data[off + d]);
}
} else {
float total = 0.0f;
for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
float w = static_cast<float>(masks[b][i]);
if (w <= 0.0f) {
continue;
}
total += w;
const size_t off = b * SS * DD + i * DD;
for (size_t d = 0; d < DD; ++d)
emb[d] += data[off + d] * w;
}
if (total > 0.0f) {
for (auto& v : emb)
v = static_cast<float>(v / total);
}
}
if (normalize_) {
double norm = 0.0;
for (auto v : emb) {
const double dv = static_cast<double>(v);
norm += dv * dv;
}
norm = std::sqrt(norm);
if (norm > 1e-8) {
for (auto& v : emb)
v = static_cast<float>(v / norm);
} else {
std::fill(emb.begin(), emb.end(), 0.0f);
}
}
}
std::vector<std::string> remainder;
remainder.reserve(B - observedCap);
for (size_t idx = observedCap; idx < B; ++idx) {
remainder.push_back(texts[idx]);
}
if (!remainder.empty()) {
auto rem = runOnnxBatch(remainder);
if (!rem) {
return rem.error();
}
auto& remVals = rem.value();
partial.insert(partial.end(), std::make_move_iterator(remVals.begin()),
std::make_move_iterator(remVals.end()));
}
return partial;
}
if (hidden_dim > 0 && d2 > 0 && DD != hidden_dim) {
std::ostringstream oss;
oss << "Output shape mismatch: expected [B,S,D] with B=" << B;
if (hidden_dim > 0) {
oss << " D=" << hidden_dim;
}
oss << ", got [" << b0 << "," << s1 << "," << d2 << "]";
return Error{ErrorCode::InvalidData, oss.str()};
}
if (b0 > 0 && static_cast<size_t>(b0) > B) {
spdlog::debug("[ONNX] Provider returned padded batch output: requested_B={} "
"output_B={} (using first B)",
B, b0);
}
const float* data = out.GetTensorData<float>();
result.resize(B, std::vector<float>(hidden_dim, 0.0f));
for (size_t b = 0; b < B; ++b) {
auto& emb = result[b];
if (pooling_mode_ == Pooling::CLS) {
const size_t off = b * SS * DD;
for (size_t d = 0; d < DD; ++d)
emb[d] = data[off + d];
} else if (pooling_mode_ == Pooling::MAX) {
for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
if (masks[b][i] <= 0)
continue;
const size_t off = b * SS * DD + i * DD;
for (size_t d = 0; d < DD; ++d)
emb[d] = std::max(emb[d], data[off + d]);
}
} else {
float total = 0.0f;
for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
float w = static_cast<float>(masks[b][i]);
if (w <= 0.0f) {
continue;
}
total += w;
const size_t off = b * SS * DD + i * DD;
for (size_t d = 0; d < DD; ++d)
emb[d] += data[off + d] * w;
}
if (total > 0.0f) {
for (auto& v : emb)
v = static_cast<float>(v / total);
}
}
if (normalize_) {
double norm = 0.0;
for (auto v : emb) {
double dv = static_cast<double>(v);
norm += dv * dv;
}
norm = std::sqrt(norm);
if (norm > 1e-8) {
for (auto& v : emb)
v = static_cast<float>(v / norm);
} else {
std::fill(emb.begin(), emb.end(), 0.0f);
}
}
}
return result;
}
if (shape.size() == 2) {
// [B, D]
const int64_t b0 = shape[0];
const int64_t d1 = shape[1];
const size_t DD = (d1 > 0) ? static_cast<size_t>(d1) : 0;
if (hidden_dim == 0 && DD > 0) {
hidden_dim = DD;
embeddingDim_ = hidden_dim;
}
if (b0 > 0 && static_cast<size_t>(b0) < B && (hidden_dim == 0 || DD == hidden_dim)) {
const size_t observedCap = static_cast<size_t>(b0);
const size_t mergedCap =
mergeSharedBatchLimit(modelName_, actualExecutionProvider_, observedCap);
if (mergedCap > 0) {
providerBatchLimit_.store(mergedCap, std::memory_order_relaxed);
const auto currentCap = TuneAdvisor::getEmbedDocCap();
if (currentCap == 0 || currentCap > mergedCap) {
TuneAdvisor::setEmbedDocCap(mergedCap);
}
}
spdlog::warn(
"[ONNX] Provider returned smaller batch than requested: requested_B={} "
"output_B={} (learned_cap={}), processing partial result + chunking remainder",
B, observedCap, mergedCap);
std::vector<std::vector<float>> partial;
partial.reserve(B);
const float* data = out.GetTensorData<float>();
for (size_t b = 0; b < observedCap; ++b) {
partial.emplace_back(data + b * DD, data + (b + 1) * DD);
if (normalize_) {
double norm = 0.0;
for (auto v : partial.back()) {
const double dv = static_cast<double>(v);
norm += dv * dv;
}
norm = std::sqrt(norm);
if (norm > 1e-8) {
for (auto& v : partial.back())
v = static_cast<float>(v / norm);
} else {
std::fill(partial.back().begin(), partial.back().end(), 0.0f);
}
}
}
std::vector<std::string> remainder;
remainder.reserve(B - observedCap);
for (size_t idx = observedCap; idx < B; ++idx) {
remainder.push_back(texts[idx]);
}
if (!remainder.empty()) {
auto rem = runOnnxBatch(remainder);
if (!rem) {
return rem.error();
}
auto& remVals = rem.value();
partial.insert(partial.end(), std::make_move_iterator(remVals.begin()),
std::make_move_iterator(remVals.end()));
}
return partial;
}
if (hidden_dim > 0 && d1 > 0 && DD != hidden_dim) {
std::ostringstream oss;
oss << "Output shape mismatch: expected [B,D] with B=" << B;
if (hidden_dim > 0) {
oss << " D=" << hidden_dim;
}
oss << ", got [" << b0 << "," << d1 << "]";
return Error{ErrorCode::InvalidData, oss.str()};
}
if (b0 > 0 && static_cast<size_t>(b0) > B) {
spdlog::debug("[ONNX] Provider returned padded batch output: requested_B={} "
"output_B={} (using first B)",
B, b0);
}
const float* data = out.GetTensorData<float>();
result.resize(B);
for (size_t b = 0; b < B; ++b) {
result[b].assign(data + b * DD, data + (b + 1) * DD);
if (normalize_) {
double norm = 0.0;
for (auto v : result[b]) {
double dv = static_cast<double>(v);
norm += dv * dv;
}
norm = std::sqrt(norm);
if (norm > 1e-8) {
for (auto& v : result[b])
v = static_cast<float>(v / norm);
} else {
std::fill(result[b].begin(), result[b].end(), 0.0f);
}
}
}
return result;
}
return Error{ErrorCode::InvalidData, "Unexpected output tensor rank (batch)"};
}
std::string modelPath_;
std::string modelName_;
OnnxTextConfig config_;
OnnxTextPreprocessor preprocessor_;
yams::vector::HuggingFaceTokenizer tokenizer_;
Ort::Env* env_ = nullptr; // shared global env
std::unique_ptr<Ort::SessionOptions> sessionOptions_;
std::unique_ptr<Ort::Session> session_;
std::string actualExecutionProvider_{"cpu"};
std::vector<std::string> inputNames_;
std::vector<std::string> outputNames_;
size_t embeddingDim_ = 0;
size_t maxSequenceLength_ = 512;
bool hasFixedInputShape_ = false; // true when ONNX model input dim[1] > 0
bool isLoaded_ = false;
bool test_mode_ = false;
bool dynamicPaddingEnabled_ = true; // controlled by YAMS_ONNX_DYNAMIC_PADDING
std::atomic<size_t> providerBatchLimit_{0};
int configuredIntraThreads_{1};
int configuredInterThreads_{1};
enum class Pooling { MEAN, CLS, MAX };
Pooling pooling_mode_ = Pooling::MEAN;
bool normalize_ = true;
// Cached MemoryInfo to avoid re-creating per inference call.
std::unique_ptr<Ort::MemoryInfo> memoryInfo_;
// I/O Binding for optimized inference dispatch (pre-resolved name lookups,
// better GPU memory placement). Falls back to Session::Run() if unavailable.
std::unique_ptr<Ort::IoBinding> ioBinding_;
bool useIoBinding_ = false;
mutable std::recursive_mutex stateMutex_;
void parseModelConfigHints() {
try {
namespace fs = std::filesystem;
fs::path mp(modelPath_);
std::vector<fs::path> candidates;
if (mp.has_parent_path()) {
candidates.push_back(mp.parent_path() / "sentence_bert_config.json");
candidates.push_back(mp.parent_path() / "config.json");
if (mp.parent_path().has_parent_path()) {
auto up = mp.parent_path().parent_path();
candidates.push_back(up / "sentence_bert_config.json");
candidates.push_back(up / "config.json");
}
}
for (const auto& c : candidates) {
if (!fs::exists(c))
continue;
std::ifstream in(c);
if (!in.good())
continue;
nlohmann::json j;
in >> j;
if (j.contains("max_seq_length") && j["max_seq_length"].is_number_integer()) {
maxSequenceLength_ = static_cast<size_t>(j["max_seq_length"].get<int>());
} else if (j.contains("model_max_length") &&
j["model_max_length"].is_number_integer()) {
maxSequenceLength_ = static_cast<size_t>(j["model_max_length"].get<int>());
}
// Common HF key for context length
if (j.contains("max_position_embeddings") &&
j["max_position_embeddings"].is_number_integer()) {
maxSequenceLength_ =
static_cast<size_t>(j["max_position_embeddings"].get<int>());
}
if (j.contains("pooling_mode_cls_token") &&
j["pooling_mode_cls_token"].is_boolean() &&
j["pooling_mode_cls_token"].get<bool>()) {
pooling_mode_ = Pooling::CLS;
}
if (j.contains("pooling_mode_max_tokens") &&
j["pooling_mode_max_tokens"].is_boolean() &&
j["pooling_mode_max_tokens"].get<bool>()) {
pooling_mode_ = Pooling::MAX;
}
if (j.contains("pooling_mode_mean_tokens") &&
j["pooling_mode_mean_tokens"].is_boolean()) {
if (j["pooling_mode_mean_tokens"].get<bool>())
pooling_mode_ = Pooling::MEAN;
}
if (j.contains("normalize_embeddings") && j["normalize_embeddings"].is_boolean()) {
normalize_ = j["normalize_embeddings"].get<bool>();
}
// Only use config hints for embedding dimension when we could not infer it from
// the ONNX graph output shape. The ONNX output tensor is authoritative.
if (embeddingDim_ == 0) {
if (j.contains("embedding_dimension") &&
j["embedding_dimension"].is_number_integer()) {
embeddingDim_ = static_cast<size_t>(j["embedding_dimension"].get<int>());
}
// Hidden size often present in HF config (standard models use hidden_size,
// Nomic models use n_embd)
if (j.contains("hidden_size") && j["hidden_size"].is_number_integer()) {
embeddingDim_ = static_cast<size_t>(j["hidden_size"].get<int>());
} else if (j.contains("n_embd") && j["n_embd"].is_number_integer()) {
embeddingDim_ = static_cast<size_t>(j["n_embd"].get<int>());
}
}
// Nomic-specific hints via architectures
try {
if (j.contains("architectures") && j["architectures"].is_array()) {
for (const auto& a : j["architectures"]) {
if (a.is_string()) {
auto s = a.get<std::string>();
if (s.find("Nomic") != std::string::npos) {
pooling_mode_ = Pooling::MEAN;
normalize_ = true;
// If very long context supported and not set, prefer 8192
if (maxSequenceLength_ == 0 || maxSequenceLength_ < 2048) {
maxSequenceLength_ = 8192;
}
break;
}
// EmbeddingGemma: mean pool + L2 normalize, 2048 context
if (s.find("Gemma") != std::string::npos) {
pooling_mode_ = Pooling::MEAN;
normalize_ = true;
if (maxSequenceLength_ == 0) {
maxSequenceLength_ = 2048;
}
spdlog::info("[ONNX] Detected Gemma architecture: "
"mean pooling, L2 normalize, max_seq={}",
maxSequenceLength_);
break;
}
}
}
}
} catch (...) {
}
break; // stop after first readable config
}
} catch (...) {
}
}
};
// ============================================================================
// OnnxModelSession Public Interface
// ============================================================================
OnnxModelSession::OnnxModelSession(const std::string& modelPath, const std::string& modelName,
const OnnxTextConfig& config)
: pImpl(std::make_unique<Impl>(modelPath, modelName, config)) {
info_.name = modelName;
info_.path = modelPath;
info_.loadTime = std::chrono::system_clock::now();
// Get file size for memory estimation
if (fs::exists(modelPath)) {
info_.memoryUsageBytes = fs::file_size(modelPath);
}
// Lazy loading: model will be loaded on first inference request
// This prevents slow session creation during pool initialization
spdlog::debug("[ONNX] Created session for '{}' - lazy loading (load on first inference)",
modelName);
// Dimension info will be populated on first inference or explicit warm-up
info_.embeddingDim = 0;
info_.maxSequenceLength = 0;
}
OnnxModelSession::~OnnxModelSession() = default;
Result<std::vector<float>> OnnxModelSession::generateEmbedding(const std::string& text) {
info_.requestCount++;
auto result = pImpl->generateEmbedding(text);
if (!result) {
info_.errorCount++;
return result;
}
// Lazy load dimension info on first successful inference
if (info_.embeddingDim == 0 && pImpl->isValid()) {
info_.embeddingDim = pImpl->getEmbeddingDim();
info_.maxSequenceLength = pImpl->getMaxSequenceLength();
spdlog::debug("[ONNX] Lazy loaded dimensions for '{}': dim={}, max_seq_len={}", info_.name,
info_.embeddingDim, info_.maxSequenceLength);
}
return result;
}
Result<std::vector<std::vector<float>>>
OnnxModelSession::generateBatchEmbeddings(const std::vector<std::string>& texts) {
info_.requestCount++;
auto result = pImpl->generateBatchEmbeddings(texts);
if (!result) {
info_.errorCount++;
return result;
}
// Lazy load dimension info on first successful inference
if (info_.embeddingDim == 0 && pImpl->isValid()) {
info_.embeddingDim = pImpl->getEmbeddingDim();
info_.maxSequenceLength = pImpl->getMaxSequenceLength();
spdlog::debug("[ONNX] Lazy loaded dimensions for '{}' (batch): dim={}, max_seq_len={}",
info_.name, info_.embeddingDim, info_.maxSequenceLength);
}
return result;
}
bool OnnxModelSession::isValid() const {
return pImpl->isValid();
}
std::string OnnxModelSession::getExecutionProvider() const {
return pImpl->getExecutionProvider();
}
size_t OnnxModelSession::getLearnedBatchLimit() const {
return pImpl->getLearnedBatchLimit();
}
Result<void> OnnxModelSession::setThreading(int intraThreads, int interThreads) {
return pImpl->setThreading(intraThreads, interThreads);
}
std::pair<int, int> OnnxModelSession::getThreading() const {
return pImpl->getThreading();
}
// ============================================================================
// OnnxModelPool Implementation
// ============================================================================
OnnxModelPool::OnnxModelPool(const ModelPoolConfig& config) : config_(config) {
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
// In mock mode, avoid initializing any ONNX Runtime state to prevent
// platform-specific crashes on hosts without compatible runtimes.
if (std::getenv("YAMS_USE_MOCK_PROVIDER") || std::getenv("YAMS_SKIP_MODEL_LOADING") ||
std::getenv("YAMS_TEST_MODE")) {
spdlog::warn("[ONNX] Mock provider mode enabled; skipping global ONNX environment init");
return;
}
#endif
if (config_.asyncLoading) {
size_t threads = config_.loadWorkerThreads > 0 ? config_.loadWorkerThreads : 1;
loadPool_ = std::make_unique<boost::asio::thread_pool>(threads);
spdlog::info("[ONNX] Async model loading enabled (threads={})", threads);
}
// NOTE: We do NOT create a separate Ort::Env here.
// Instead, OnnxModelSession::Impl uses the global get_global_ort_env() function.
// Creating multiple Ort::Env instances causes thread pool conflicts that lead to
// "resource deadlock would occur" errors (EDEADLK, code 36).
// The ortEnv_ and sessionOptions_ members are kept for API compatibility but not used.
try {
// Log that we're using the shared global environment
// NOTE: We do NOT create a separate Ort::Env here - OnnxModelSession::Impl
// uses the global get_global_ort_env() function to avoid thread pool conflicts.
spdlog::info(
"[ONNX] Pool: Using shared global Ort::Env (sessions use single-threaded execution)");
runtimeGpuEnabled_ = false;
if (config.enableGPU) {
const auto& gpu = yams::daemon::resource::detectGpu();
if (gpu.detected) {
runtimeGpuEnabled_ = true;
spdlog::info("[ONNX] GPU enabled: {} ({:.1f} GB VRAM, provider={})", gpu.name,
static_cast<double>(gpu.vramBytes) / (1024.0 * 1024.0 * 1024.0),
gpu.provider);
} else {
spdlog::info("[ONNX] GPU requested but no GPU detected (will use CPU)");
}
}
} catch (const std::exception& e) {
// Do not crash the daemon if runtime init fails; log and continue.
spdlog::warn("[ONNX] Runtime environment init failed: {} — continuing (mock may be used)",
e.what());
} catch (...) {
spdlog::warn("[ONNX] Runtime environment init failed with unknown error — continuing");
}
}
OnnxModelPool::~OnnxModelPool() {
try {
shutdown();
} catch (const std::exception& e) {
try {
spdlog::warn("[OnnxModelPool] destructor exception: {}", e.what());
} catch (...) {
}
} catch (...) {
try {
spdlog::warn("[OnnxModelPool] destructor unknown exception");
} catch (...) {
}
}
}
Result<void> OnnxModelPool::initialize() {
std::lock_guard<std::recursive_mutex> lock(mutex_);
if (initialized_.load(std::memory_order_acquire)) {
return Result<void>();
}
spdlog::info("Initializing ONNX model pool with max {} models", config_.maxLoadedModels);
const auto& runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
if (!runtimeInfo.available) {
return Error{ErrorCode::InternalError,
runtimeInfo.errorMessage.empty()
? std::string("ONNX Runtime unavailable")
: "ONNX Runtime unavailable: " + runtimeInfo.errorMessage};
}
// Eagerly initialize the global Ort::Env NOW, before any concurrent access.
// This prevents "resource deadlock would occur" errors that can happen when
// multiple threads try to initialize ONNX Runtime simultaneously.
try {
spdlog::info("[ONNX Pool] Eagerly initializing global Ort::Env...");
(void)get_global_ort_env();
spdlog::info("[ONNX Pool] Global Ort::Env ready");
} catch (const std::exception& e) {
spdlog::error("[ONNX Pool] Failed to initialize Ort::Env: {}", e.what());
return Error{ErrorCode::InternalError, std::string("ONNX Env init failed: ") + e.what()};
}
// Start background preloading if configured (non-blocking)
if (!config_.lazyLoading && !config_.preloadModels.empty()) {
// Launch preloading in background thread to avoid blocking
// Store the thread so we can join it during shutdown
preloadThread_ = std::thread([this]() {
try {
spdlog::info("Starting background model preloading");
// Check shutdown flag periodically during preload
for (const auto& modelName : config_.preloadModels) {
if (shutdown_.load(std::memory_order_acquire)) {
spdlog::info("Background preloading interrupted by shutdown");
return;
}
auto result = loadModel(modelName);
if (!result) {
spdlog::warn("Failed to preload model {}: {}", modelName,
result.error().message);
} else {
// Warm up the model to trigger lazy loading
// This ensures the first real inference request is fast
spdlog::info("Warming up preloaded model: {}", modelName);
auto warmupResult = warmupModel(modelName);
if (!warmupResult) {
spdlog::warn("Failed to warm up model {}: {}", modelName,
warmupResult.error().message);
}
}
}
spdlog::info("Background model preloading completed");
} catch (const std::exception& e) {
try {
spdlog::error("[ONNX] Background preload crashed: {}", e.what());
} catch (...) {
}
} catch (...) {
try {
spdlog::error("[ONNX] Background preload crashed with unknown exception");
} catch (...) {
}
}
});
}
initialized_.store(true, std::memory_order_release);
return Result<void>();
}
void OnnxModelPool::shutdown() {
if (shutdown_.exchange(true, std::memory_order_acq_rel)) {
return;
}
if (preloadThread_.joinable()) {
try {
spdlog::info("Waiting for background preload thread to finish...");
} catch (...) {
}
try {
preloadThread_.join();
try {
spdlog::info("Background preload thread joined");
} catch (...) {
}
} catch (const std::exception& e) {
try {
spdlog::warn("Failed to join preload thread: {}", e.what());
} catch (...) {
}
} catch (...) {
}
}
if (loadPool_) {
loadPool_->join();
loadPool_.reset();
}
try {
std::lock_guard<std::recursive_mutex> lock(mutex_);
try {
spdlog::info("Shutting down ONNX model pool");
} catch (...) {
}
loadingFutures_.clear();
for (auto& [name, entry] : models_) {
if (entry.pool) {
try {
entry.pool->shutdown();
} catch (const std::exception& e) {
try {
spdlog::warn("Failed to shutdown model pool '{}': {}", name, e.what());
} catch (...) {
}
} catch (...) {
}
}
}
models_.clear();
initialized_.store(false, std::memory_order_release);
} catch (const std::exception& e) {
try {
spdlog::warn("ONNX model pool shutdown exception: {}", e.what());
} catch (...) {
}
} catch (...) {
try {
spdlog::warn("ONNX model pool shutdown unknown exception");
} catch (...) {
}
}
}
Result<OnnxModelPool::ModelHandle> OnnxModelPool::acquireModel(const std::string& modelName,
std::chrono::milliseconds timeout) {
if (!initialized_.load(std::memory_order_acquire)) {
if (auto result = initialize(); !result) {
return result.error();
}
}
totalRequests_++;
// Check if model is already loaded
// IMPORTANT: Get pool pointer under lock, then release lock before calling acquire()
// to avoid deadlock with ResourcePool's internal locking (PBI-1c1 fix)
std::shared_ptr<ResourcePool<OnnxModelSession>> poolPtr;
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it != models_.end() && it->second.pool) {
updateAccessStats(modelName);
cacheHits_++;
poolPtr = it->second.pool;
}
}
// Acquire outside of lock to avoid deadlock
if (poolPtr) {
return poolPtr->acquire(timeout);
}
cacheMisses_++;
// Model not loaded, need to load it
if (config_.asyncLoading) {
if (auto result = loadModel(modelName, /*wait=*/false); !result) {
return result.error();
}
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = loadingFutures_.find(modelName);
if (it != loadingFutures_.end()) {
if (it->second.wait_for(std::chrono::seconds(0)) != std::future_status::ready) {
return Error{ErrorCode::OperationInProgress, "Model loading in background"};
}
const auto loadOk = static_cast<bool>(it->second.get());
std::optional<Error> loadError;
if (!loadOk) {
loadError = it->second.get().error();
}
loadingFutures_.erase(it);
if (loadError) {
return std::move(*loadError);
}
}
}
} else {
if (auto result = loadModel(modelName); !result) {
return result.error();
}
}
// Try again after loading - get pool pointer while holding lock
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it != models_.end() && it->second.pool) {
updateAccessStats(modelName);
poolPtr = it->second.pool;
}
}
// Acquire outside of lock
if (poolPtr) {
return poolPtr->acquire(timeout);
}
return Error{ErrorCode::NotFound, "Failed to load model: " + modelName};
}
Result<void> OnnxModelPool::loadModel(const std::string& modelName, bool wait) {
if (!config_.asyncLoading || wait) {
return loadModelSync(modelName);
}
return scheduleLoad(modelName);
}
Result<void> OnnxModelPool::scheduleLoad(const std::string& modelName) {
if (!loadPool_) {
return loadModelSync(modelName);
}
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it != models_.end() && it->second.pool) {
return Result<void>();
}
if (loadingFutures_.count(modelName) > 0) {
return Result<void>();
}
auto prom = std::make_shared<std::promise<Result<void>>>();
auto fut = prom->get_future().share();
loadingFutures_[modelName] = fut;
boost::asio::post(*loadPool_, [this, modelName, prom]() {
auto res = loadModelSync(modelName);
prom->set_value(res);
});
}
return Result<void>();
}
bool OnnxModelPool::isModelLoaded(const std::string& modelName) const {
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
return it != models_.end() && it->second.pool != nullptr;
}
std::vector<std::string> OnnxModelPool::getLoadedModels() const {
std::lock_guard<std::recursive_mutex> lock(mutex_);
std::vector<std::string> loaded;
for (const auto& [name, entry] : models_) {
if (entry.pool) {
loaded.push_back(name);
}
}
return loaded;
}
Result<void> OnnxModelPool::loadModelSync(const std::string& modelName) {
spdlog::info("[ONNX Plugin] loadModel() called for: {}", modelName);
// In test environments, avoid heavy model loading entirely
if (yams::test::shouldSkipModelLoading()) {
return Error{ErrorCode::NotFound, "Model loading skipped in test mode"};
}
// Single-flight pattern: only one thread loads a given model at a time.
// Other threads wait for the loading thread to finish.
{
std::unique_lock<std::recursive_mutex> lock(mutex_);
// Check if already loaded
auto it = models_.find(modelName);
if (it != models_.end() && it->second.pool) {
spdlog::info("[ONNX Plugin] Model already loaded: {}", modelName);
return Result<void>();
}
// Check if another thread is currently loading this model
if (loadingModels_.count(modelName) > 0) {
spdlog::info("[ONNX Plugin] Model '{}' is being loaded by another thread, waiting...",
modelName);
// Wait for the other thread to finish loading
loadingCv_.wait(lock, [this, &modelName]() {
// Wake up when either: model is loaded, or no longer in loading set
auto it = models_.find(modelName);
bool loaded = (it != models_.end() && it->second.pool);
bool notLoading = (loadingModels_.count(modelName) == 0);
return loaded || notLoading;
});
// Check if model was loaded successfully by the other thread
it = models_.find(modelName);
if (it != models_.end() && it->second.pool) {
spdlog::info("[ONNX Plugin] Model '{}' loaded by another thread", modelName);
return Result<void>();
}
// If not loaded and not loading, we'll try to load it ourselves
spdlog::warn("[ONNX Plugin] Model '{}' load failed by other thread, retrying",
modelName);
}
// Mark this model as being loaded by us
loadingModels_.insert(modelName);
}
// RAII guard to remove from loadingModels_ and notify waiters
struct LoadingGuard {
OnnxModelPool* pool;
std::string name;
~LoadingGuard() {
std::lock_guard<std::recursive_mutex> lock(pool->mutex_);
pool->loadingModels_.erase(name);
pool->loadingCv_.notify_all();
}
} guard{this, modelName};
// Resolve model path WITHOUT holding the lock
std::string modelPath = resolveModelPath(modelName);
spdlog::info("[ONNX Plugin] Resolved path: {}", modelPath);
// Direct filesystem check
spdlog::info("[ONNX Plugin] Checking file existence: {}", modelPath);
bool fileExists = false;
try {
fileExists = fs::exists(modelPath) && fs::is_regular_file(modelPath);
} catch (const std::exception& e) {
spdlog::warn("Error checking model file existence: {}", e.what());
}
if (!fileExists) {
spdlog::info("Model file not found: {} (searched for: {})", modelName, modelPath);
const bool offline = []() {
if (const char* s = std::getenv("YAMS_ONNX_OFFLINE")) {
std::string v(s);
std::transform(v.begin(), v.end(), v.begin(), ::tolower);
return v == "1" || v == "true" || v == "yes" || v == "on";
}
return false;
}();
if (offline) {
return Error{ErrorCode::NotFound,
"Model file not found in offline mode. Pre-place artifacts in cache or "
"~/.yams/models, or unset YAMS_ONNX_OFFLINE"};
}
spdlog::info("Download the model with: yams model --download {}", modelName);
return Error{ErrorCode::NotFound,
"Model file not found: " + modelPath +
". Download with: yams model --download " + modelName +
" or set YAMS_ALLOW_MODEL_DOWNLOAD=1 for automatic downloads"};
}
// Check file size
size_t modelSize = 0;
try {
modelSize = fs::file_size(modelPath);
} catch (const std::exception& e) {
spdlog::warn("Error getting model file size: {}", e.what());
modelSize = 100 * 1024 * 1024; // Assume 100MB if can't determine
}
// Configure pool for this model (no lock needed - local variables)
// Use TuneAdvisor for GPU-aware pool sizing:
// GPU mode: larger pool (GPU handles inference, more sessions for throughput)
// CPU mode: smaller pool (avoid CPU saturation from parallel inference)
const bool gpuEnabled = runtimeGpuEnabled_ && !OnnxModelSession::isCoreMLDenylisted(modelName);
const size_t embedConcurrencyCap =
std::max<size_t>(1u, std::min<size_t>(TuneAdvisor::postEmbedConcurrent(),
TuneAdvisor::getEmbedMaxConcurrency()));
const size_t dynamic_max = std::max<size_t>(
1u, std::min<size_t>(TuneAdvisor::onnxSessionsPerModel(gpuEnabled), embedConcurrencyCap));
const size_t dynamic_min = std::min(size_t{2}, dynamic_max);
spdlog::info("[ONNX Plugin] Configuring resource pool (gpu={}, pool_max={}, embed_cap={})",
gpuEnabled ? "enabled" : "disabled", dynamic_max, embedConcurrencyCap);
PoolConfig<OnnxModelSession> poolConfig;
poolConfig.minSize = dynamic_min;
poolConfig.maxSize = dynamic_max;
// Keep maxIdle = maxSize to prevent session churn. With maxIdle < maxSize,
// excess sessions would be destroyed on return and recreated on next acquire.
poolConfig.maxIdle = dynamic_max;
poolConfig.idleTimeout = config_.modelIdleTimeout;
// CRITICAL: Disable preCreateResources to avoid creating ONNX session during
// ResourcePool construction. Sessions will be created lazily on first acquire().
// This prevents EDEADLK errors from ONNX Runtime's internal thread pool.
poolConfig.preCreateResources = false;
// Create embedding config
OnnxTextConfig embConfig;
embConfig.model_path = modelPath;
embConfig.model_name = modelName;
embConfig.enable_gpu = gpuEnabled;
embConfig.num_threads = config_.numThreads;
// Leave max_sequence_length unset so session/model metadata can determine the native limit.
// The OnnxTextConfig default of 512 would otherwise silently clamp models like Gemma.
embConfig.max_sequence_length = 0;
// Create ResourcePool WITHOUT holding mutex_ - the factory lambda will be called
// later during acquire(), also without holding mutex_.
spdlog::debug("[ONNX Plugin] Creating ResourcePool for '{}'", modelName);
auto t0 = std::chrono::steady_clock::now();
std::shared_ptr<ResourcePool<OnnxModelSession>> pool;
try {
spdlog::debug("[ONNX Plugin] About to construct ResourcePool...");
pool = std::make_shared<ResourcePool<OnnxModelSession>>(
poolConfig,
[modelPath, modelName, embConfig](const std::string&) -> Result<ModelSessionPtr> {
try {
spdlog::debug("[ONNX Plugin] Factory: creating OnnxModelSession for '{}'",
modelName);
return std::make_shared<OnnxModelSession>(modelPath, modelName, embConfig);
} catch (const std::exception& e) {
spdlog::error("[ONNX Plugin] Factory: failed to create session: {}", e.what());
return Error{ErrorCode::InternalError,
std::string("Failed to create model session: ") + e.what()};
} catch (...) {
spdlog::error("[ONNX Plugin] Factory: unknown error creating session");
return Error{ErrorCode::InternalError, "Unknown error creating ONNX session"};
}
},
[](const OnnxModelSession& session) { return session.isValid(); });
spdlog::debug("[ONNX Plugin] ResourcePool constructed successfully");
} catch (const std::system_error& e) {
spdlog::error("[ONNX Plugin] ResourcePool construction threw system_error: {} (code={})",
e.what(), e.code().value());
return Error{ErrorCode::InternalError,
std::string("ResourcePool construction failed: ") + e.what()};
} catch (const std::exception& e) {
spdlog::error("[ONNX Plugin] ResourcePool construction threw exception: {}", e.what());
return Error{ErrorCode::InternalError,
std::string("ResourcePool construction failed: ") + e.what()};
}
// Now acquire lock briefly to update shared state
spdlog::debug("[ONNX Plugin] Acquiring mutex for model registration...");
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
// Check if we need to evict models
size_t loadedCount = 0;
for (const auto& [name, entry] : models_) {
if (entry.pool)
loadedCount++;
}
if (loadedCount >= config_.maxLoadedModels) {
evictLRU(1);
}
// Check memory constraints
if (!canLoadModel(modelSize)) {
return Error{ErrorCode::ResourceExhausted, "Insufficient memory to load model"};
}
// Create/update model entry
spdlog::info("[ONNX Plugin] Registering model entry for: {}", modelName);
ModelEntry& entry = models_[modelName];
entry.name = modelName;
entry.path = modelPath;
entry.lastAccess = std::chrono::steady_clock::now();
entry.pool = std::move(pool);
// Update cached model count for non-blocking status queries
loadedModelCount_.fetch_add(1, std::memory_order_relaxed);
}
auto t1 = std::chrono::steady_clock::now();
auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
spdlog::info("Registered ONNX model: {} ({}MB) load_ms={} path={}", modelName,
modelSize / (1024 * 1024), ms, modelPath);
spdlog::info("[ONNX Plugin] Model '{}' registered successfully (session created on first use)",
modelName);
return Result<void>();
}
Result<void> OnnxModelPool::unloadModel(const std::string& modelName) {
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it == models_.end()) {
return Error{ErrorCode::NotFound, "Model not found: " + modelName};
}
if (it->second.isHot) {
return Error{ErrorCode::InvalidOperation, "Cannot unload hot model: " + modelName};
}
if (it->second.pool) {
auto stats = it->second.pool->getStats();
if (stats.inUseResources > 0) {
return Error{ErrorCode::OperationInProgress,
"Cannot unload model with in-use sessions: " + modelName};
}
it->second.pool->shutdown();
it->second.pool.reset();
// Update cached model count for non-blocking status queries
loadedModelCount_.fetch_sub(1, std::memory_order_relaxed);
}
spdlog::info("Unloaded ONNX model: {}", modelName);
return Result<void>();
}
Result<void> OnnxModelPool::preloadModels() {
for (const auto& modelName : config_.preloadModels) {
spdlog::info("Preloading model: {}", modelName);
// Load model directly - no async wrapper needed
auto result = loadModel(modelName);
if (!result) {
spdlog::warn("Failed to preload model {}: {}", modelName, result.error().message);
} else {
// Warm up the model to trigger lazy loading
spdlog::info("Warming up preloaded model: {}", modelName);
auto warmupResult = warmupModel(modelName);
if (!warmupResult) {
spdlog::warn("Failed to warm up model {}: {}", modelName,
warmupResult.error().message);
}
// Mark as hot model
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it != models_.end()) {
it->second.isHot = true;
}
}
}
return Result<void>();
}
Result<void> OnnxModelPool::warmupModel(const std::string& modelName) {
spdlog::info("[ONNX Plugin] Warming up model: {}", modelName);
auto t0 = std::chrono::steady_clock::now();
// Acquire a session from the pool
auto handleResult = acquireModel(modelName, std::chrono::seconds(30));
if (!handleResult) {
spdlog::error("[ONNX Plugin] Failed to acquire session for warmup: {}",
handleResult.error().message);
return handleResult.error();
}
auto handle = std::move(handleResult.value());
// Do a batch inference to trigger lazy loading and warm the batch code path
std::vector<std::string> warmupTexts = {"warmup text one", "warmup text two",
"warmup text three", "warmup text four"};
auto embedResult = handle->generateBatchEmbeddings(warmupTexts);
if (!embedResult) {
spdlog::error("[ONNX Plugin] Warmup inference failed for {}: {}", modelName,
embedResult.error().message);
return embedResult.error();
}
auto t1 = std::chrono::steady_clock::now();
auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
spdlog::info("[ONNX Plugin] Model '{}' warmed up in {}ms (batch={}, embedding dim={})",
modelName, ms, warmupTexts.size(),
embedResult.value().empty() ? 0 : embedResult.value()[0].size());
// Auto-tune batch size if not explicitly set and auto-tuning not disabled
auto autotuneDisabled = []() {
if (const char* s = std::getenv("YAMS_EMBED_AUTOTUNE")) {
std::string v(s);
std::transform(v.begin(), v.end(), v.begin(), ::tolower);
return v == "0" || v == "false" || v == "no" || v == "off";
}
return false;
};
if (TuneAdvisor::getEmbedDocCap() == 0 && !autotuneDisabled()) {
spdlog::info("[ONNX Plugin] Auto-tuning batch size for '{}'...", modelName);
std::vector<std::string> probeTexts;
probeTexts.reserve(64);
for (uint32_t i = 0; i < 64; ++i)
probeTexts.push_back("probe text number " + std::to_string(i) + " for auto-tuning");
std::vector<uint32_t> candidates = {4, 8, 16, 32, 64};
double bestTps = 0.0;
uint32_t bestBatch = 16;
for (uint32_t bs : candidates) {
std::vector<std::string> batch(probeTexts.begin(),
probeTexts.begin() + static_cast<ptrdiff_t>(bs));
auto pt0 = std::chrono::steady_clock::now();
auto r = handle->generateBatchEmbeddings(batch);
auto elapsed = std::chrono::steady_clock::now() - pt0;
if (!r)
break;
double probeMs =
std::chrono::duration_cast<std::chrono::microseconds>(elapsed).count() / 1000.0;
double tps = static_cast<double>(bs) / (probeMs / 1000.0);
spdlog::debug("[ONNX Plugin] Autotune probe: batch={} {:.1f}ms ({:.1f} texts/sec)", bs,
probeMs, tps);
if (tps > bestTps) {
bestTps = tps;
bestBatch = bs;
}
}
TuneAdvisor::setEmbedDocCap(bestBatch);
spdlog::info("[ONNX Plugin] Auto-tuned batch size: {} ({:.1f} texts/sec)", bestBatch,
bestTps);
}
// Pre-warm additional sessions to avoid cold-start penalty on concurrent requests.
// GPU mode: up to 3 sessions (GPU provider init is expensive).
// CPU mode: up to 2 sessions (eliminates ~1788ms cold start for second request).
// The initial handle is still held; additional sessions are acquired and immediately
// released back to the pool so they are pre-initialized for future use.
{
const std::string activeProvider = handle->getExecutionProvider();
const bool providerGpuEnabled = isGpuExecutionProviderName(activeProvider);
const bool cpuFallbackPinned = !providerGpuEnabled && runtimeGpuEnabled_ &&
OnnxModelSession::isCoreMLDenylisted(modelName);
size_t targetWarm =
std::min(static_cast<size_t>(TuneAdvisor::onnxSessionsPerModel(providerGpuEnabled)),
providerGpuEnabled ? size_t{3} : size_t{2});
if (cpuFallbackPinned) {
targetWarm = std::min<size_t>(targetWarm, 1u);
}
// Keep warm session count aligned with current embedding throughput budget
// to avoid retaining idle ONNX sessions under low-concurrency settings.
const size_t embedConcurrencyCap =
std::max<size_t>(1u, std::min<size_t>(TuneAdvisor::postEmbedConcurrent(),
TuneAdvisor::getEmbedMaxConcurrency()));
targetWarm = std::min(targetWarm, embedConcurrencyCap);
// Optional explicit cap for test/benchmark scenarios.
if (const char* s = std::getenv("YAMS_ONNX_PREWARM_SESSIONS")) {
try {
size_t forced = static_cast<size_t>(std::stoull(s));
if (forced >= 1 && forced <= 8) {
targetWarm = std::min(targetWarm, forced);
}
} catch (...) {
}
}
for (size_t i = 1; i < targetWarm; ++i) {
auto extra = acquireModel(modelName, std::chrono::seconds(60));
if (!extra) {
spdlog::debug("[ONNX Plugin] Pre-warm: could not acquire session {} for {}", i + 1,
modelName);
break;
}
// Session returned to pool when extra goes out of scope
}
if (targetWarm > 1) {
spdlog::info("[ONNX Plugin] Pre-warmed {} sessions for '{}' (provider={} gpu={})",
targetWarm, modelName, activeProvider, providerGpuEnabled ? "yes" : "no");
}
}
return Result<void>();
}
Result<void> OnnxModelPool::setModelHot(const std::string& modelName, bool hot) {
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it == models_.end()) {
return Error{ErrorCode::NotFound, "Model not found: " + modelName};
}
it->second.isHot = hot;
spdlog::info("[ONNX Plugin] Model '{}' hot={} ", modelName, hot ? "true" : "false");
return Result<void>();
}
Result<void> OnnxModelPool::setModelThreading(const std::string& modelName, int intraThreads,
int interThreads, bool applyNow) {
auto valid = [](int value) { return value == -1 || (value >= 1 && value <= 64); };
if (modelName.empty() || !valid(intraThreads) || !valid(interThreads)) {
return Error{ErrorCode::InvalidArgument,
"setModelThreading requires model name and valid thread values"};
}
std::shared_ptr<ResourcePool<OnnxModelSession>> pool;
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
threadingOverrides_[modelName] = {intraThreads, interThreads};
auto it = models_.find(modelName);
if (it != models_.end()) {
pool = it->second.pool;
}
}
if (!applyNow || !pool) {
return Result<void>();
}
if (pool->getStats().inUseResources > 0) {
return Error{ErrorCode::ResourceBusy,
"Model is currently in use; threading override will apply on next reload"};
}
auto handle = pool->acquire(std::chrono::milliseconds(500));
if (!handle) {
return Error{ErrorCode::Timeout,
"Timed out acquiring model session to apply threading override"};
}
return handle.value()->setThreading(intraThreads, interThreads);
}
void OnnxModelPool::evictLRU(size_t numToEvict) {
if (runtimeGpuEnabled_ && !allowGpuAutoUnload()) {
spdlog::debug(
"Skipping LRU model eviction on GPU path (set YAMS_ONNX_GPU_AUTO_UNLOAD=1 to enable)");
return;
}
// Find least recently used models that are not hot
std::vector<std::pair<std::string, std::chrono::steady_clock::time_point>> candidates;
for (const auto& [name, entry] : models_) {
if (entry.pool && !entry.isHot) {
candidates.emplace_back(name, entry.lastAccess);
}
}
// Sort by last access time
std::sort(candidates.begin(), candidates.end(),
[](const auto& a, const auto& b) { return a.second < b.second; });
// Evict the oldest models
for (size_t i = 0; i < std::min(numToEvict, candidates.size()); ++i) {
const auto& modelName = candidates[i].first;
auto it = models_.find(modelName);
if (it != models_.end() && it->second.pool) {
auto stats = it->second.pool->getStats();
if (stats.inUseResources > 0) {
spdlog::debug("Skipping LRU eviction for model '{}' (in_use={})", modelName,
stats.inUseResources);
continue;
}
spdlog::info("Evicting model due to LRU: {}", modelName);
it->second.pool->shutdown();
it->second.pool.reset();
const size_t count = loadedModelCount_.load(std::memory_order_relaxed);
if (count > 0) {
loadedModelCount_.fetch_sub(1, std::memory_order_relaxed);
}
}
}
}
OnnxModelPool::PoolStats OnnxModelPool::getStats() const {
std::lock_guard<std::recursive_mutex> lock(mutex_);
PoolStats stats;
stats.totalRequests = totalRequests_.load();
stats.cacheHits = cacheHits_.load();
stats.cacheMisses = cacheMisses_.load();
stats.hitRate =
stats.totalRequests > 0 ? static_cast<double>(stats.cacheHits) / stats.totalRequests : 0.0;
stats.loadedModels = 0;
stats.totalMemoryBytes = 0;
for (const auto& [name, entry] : models_) {
if (entry.pool) {
stats.loadedModels++;
// Get model info (approximate from first session in pool)
// Note: This is simplified - in production, track actual memory usage
if (fs::exists(entry.path)) {
size_t modelSize = fs::file_size(entry.path);
stats.totalMemoryBytes += modelSize;
}
// Populate per-model stats for status reporting (copyable ModelInfo)
ModelInfo mi;
mi.name = name;
mi.path = entry.path;
mi.embeddingDim = 384; // default when unknown
mi.maxSequenceLength = 512; // default when unknown
try {
if (fs::exists(entry.path)) {
mi.memoryUsageBytes = fs::file_size(entry.path);
}
} catch (...) {
mi.memoryUsageBytes = 0;
}
mi.loadTime = std::chrono::system_clock::now();
mi.requestCount = 0;
mi.errorCount = 0;
stats.modelStats.emplace(name, std::move(mi));
}
}
return stats;
}
size_t OnnxModelPool::getMemoryUsage() const {
return getStats().totalMemoryBytes;
}
void OnnxModelPool::performMaintenance() {
if (runtimeGpuEnabled_ && !allowGpuAutoUnload()) {
return;
}
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto now = std::chrono::steady_clock::now();
for (auto& [name, entry] : models_) {
if (!entry.pool) {
continue;
}
auto stats = entry.pool->getStats();
if (stats.inUseResources > 0) {
continue;
}
if (entry.isHot) {
// Keep the model loaded, but still evict expired idle sessions from its pool.
entry.pool->evictExpired();
continue;
}
// Check if model has been idle too long
auto idleTime = now - entry.lastAccess;
if (idleTime > config_.modelIdleTimeout) {
spdlog::info("Unloading idle model: {}", name);
entry.pool->shutdown();
entry.pool.reset();
const size_t count = loadedModelCount_.load(std::memory_order_relaxed);
if (count > 0) {
loadedModelCount_.fetch_sub(1, std::memory_order_relaxed);
}
} else if (entry.pool) {
// Clean up expired resources in the pool
entry.pool->evictExpired();
}
}
}
#ifdef YAMS_TESTING
std::optional<OnnxModelPool::TestingModelPoolStats>
OnnxModelPool::testingModelPoolStats(const std::string& modelName) const {
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = models_.find(modelName);
if (it == models_.end()) {
return std::nullopt;
}
TestingModelPoolStats out;
out.loaded = (it->second.pool != nullptr);
out.isHot = it->second.isHot;
if (it->second.pool) {
auto stats = it->second.pool->getStats();
out.totalResources = stats.totalResources;
out.availableResources = stats.availableResources;
out.inUseResources = stats.inUseResources;
}
return out;
}
#endif
std::string OnnxModelPool::resolveModelPath(const std::string& modelName) const {
auto is_hf_like = [](const std::string& s) {
if (s.rfind("hf://", 0) == 0)
return true;
if (s.rfind("https://huggingface.co/", 0) == 0)
return true;
return (s.find('/') != std::string::npos) && (s.find('/') == s.rfind('/')) &&
s.find(".onnx") == std::string::npos && s.find('/') > 0;
};
auto to_repo_id = [](const std::string& s) {
if (s.rfind("hf://", 0) == 0)
return s.substr(5);
const std::string prefix = "https://huggingface.co/";
if (s.rfind(prefix, 0) == 0) {
auto rest = s.substr(prefix.size());
auto slash = rest.find('/');
if (slash != std::string::npos) {
auto slash2 = rest.find('/', slash + 1);
if (slash2 != std::string::npos)
return rest.substr(0, slash2);
}
return rest;
}
return s;
};
// Get home directory (cross-platform)
std::string homeDir;
if (const char* home = std::getenv("HOME")) {
homeDir = home;
} else if (const char* userprofile = std::getenv("USERPROFILE")) {
// Windows: USERPROFILE is the user's home directory
homeDir = userprofile;
}
// Windows-specific paths
std::string localAppData;
std::string appData;
#ifdef _WIN32
if (const char* lad = std::getenv("LOCALAPPDATA")) {
localAppData = lad;
}
if (const char* ad = std::getenv("APPDATA")) {
appData = ad;
}
#endif
// Build search paths with priority:
// 1) Treat modelName as a full path
// 2) Configured modelsRoot (if set), with common filename layouts
// 3) User models (~/.yams/models), local models/, system share models
std::vector<std::string> searchPaths;
searchPaths.push_back(modelName);
if (is_hf_like(modelName)) {
std::string repo = to_repo_id(modelName);
// Check for revision hint
std::string preferredRev;
{
std::lock_guard<std::recursive_mutex> lk(mutex_);
auto it = modelHints_.find(modelName);
if (it != modelHints_.end())
preferredRev = it->second.hfRevision;
}
auto cache_roots = [&]() {
std::vector<std::string> roots;
if (const char* env = std::getenv("YAMS_ONNX_HF_CACHE"))
roots.push_back(env);
if (const char* hf = std::getenv("HF_HOME")) {
roots.push_back(std::string(hf));
roots.push_back(std::string(hf) + "/hub");
}
if (!homeDir.empty()) {
roots.push_back(homeDir + "/.cache/huggingface");
roots.push_back(homeDir + "/.cache/huggingface/hub");
}
return roots;
}();
std::string repo_dir = "models--";
for (auto c : repo)
repo_dir += (c == '/' ? std::string("--") : std::string(1, c));
for (const auto& root : cache_roots) {
try {
fs::path base = fs::path(root) / repo_dir;
if (!fs::exists(base))
continue;
fs::path found;
if (!preferredRev.empty()) {
fs::path snapdir;
fs::path ref = base / "refs" / preferredRev;
if (fs::exists(ref)) {
try {
std::ifstream in(ref);
std::string s;
std::getline(in, s);
if (!s.empty())
snapdir = base / "snapshots" / s;
} catch (...) {
}
}
if (snapdir.empty())
snapdir = base / "snapshots" / preferredRev;
if (fs::exists(snapdir)) {
for (auto& p : fs::recursive_directory_iterator(snapdir)) {
if (!p.is_regular_file())
continue;
if (p.path().filename() == "model.onnx") {
found = p.path();
break;
}
}
}
}
if (found.empty()) {
for (auto& p : fs::recursive_directory_iterator(base)) {
if (!p.is_regular_file())
continue;
if (p.path().filename() == "model.onnx") {
found = p.path();
break;
}
}
}
if (!found.empty()) {
spdlog::debug("[ONNX] Resolved HF {} to {}", repo, found.string());
return found.string();
}
} catch (...) {
}
}
spdlog::debug("[ONNX] HF repo '{}' not found in cache; falling back to filesystem search",
repo);
}
// Expand configured modelsRoot if present (supports '~')
if (!config_.modelsRoot.empty()) {
std::string root = config_.modelsRoot;
if (!root.empty() && root[0] == '~') {
if (!homeDir.empty()) {
root = homeDir + root.substr(1);
}
}
// Common layouts under configured root
searchPaths.push_back(root + "/" + modelName + "/model.onnx");
searchPaths.push_back(root + "/" + modelName + "/" + modelName + ".onnx");
searchPaths.push_back(root + "/" + modelName + ".onnx");
}
// Default locations (XDG standard first, then fallbacks)
// Priority: XDG_DATA_HOME > ~/.local/share/yams > ~/.yams > relative > system
std::string xdg_data_home;
if (const char* xdg = std::getenv("XDG_DATA_HOME")) {
xdg_data_home = xdg;
} else if (!homeDir.empty()) {
xdg_data_home = homeDir + "/.local/share";
}
// XDG paths (highest priority)
if (!xdg_data_home.empty()) {
searchPaths.push_back(xdg_data_home + "/yams/models/" + modelName + "/model.onnx");
searchPaths.push_back(xdg_data_home + "/yams/models/" + modelName + "/" + modelName +
".onnx");
searchPaths.push_back(xdg_data_home + "/yams/models/" + modelName + ".onnx");
}
// Legacy ~/.yams/models (for backward compatibility)
if (!homeDir.empty()) {
searchPaths.push_back(homeDir + "/.yams/models/" + modelName + "/model.onnx");
searchPaths.push_back(homeDir + "/.yams/models/" + modelName + "/" + modelName + ".onnx");
searchPaths.push_back(homeDir + "/.yams/models/" + modelName + ".onnx");
}
#ifdef _WIN32
// Windows: Check LOCALAPPDATA/yams/models (primary Windows location)
if (!localAppData.empty()) {
searchPaths.push_back(localAppData + "\\yams\\models\\" + modelName + "\\model.onnx");
searchPaths.push_back(localAppData + "\\yams\\models\\" + modelName + "\\" + modelName +
".onnx");
searchPaths.push_back(localAppData + "\\yams\\models\\" + modelName + ".onnx");
}
// Windows: Check APPDATA/yams/models (roaming profile)
if (!appData.empty()) {
searchPaths.push_back(appData + "\\yams\\models\\" + modelName + "\\model.onnx");
searchPaths.push_back(appData + "\\yams\\models\\" + modelName + "\\" + modelName +
".onnx");
searchPaths.push_back(appData + "\\yams\\models\\" + modelName + ".onnx");
}
#endif
// Relative paths
searchPaths.push_back("models/" + modelName + "/model.onnx");
searchPaths.push_back("models/" + modelName + "/" + modelName + ".onnx");
searchPaths.push_back("models/" + modelName + ".onnx");
#ifdef _WIN32
// Windows: System paths
searchPaths.push_back("C:\\ProgramData\\yams\\models\\" + modelName + "\\model.onnx");
searchPaths.push_back("C:\\ProgramData\\yams\\models\\" + modelName + "\\" + modelName +
".onnx");
searchPaths.push_back("C:\\ProgramData\\yams\\models\\" + modelName + ".onnx");
#else
// Unix: System paths
searchPaths.push_back("/usr/local/share/yams/models/" + modelName + "/model.onnx");
searchPaths.push_back("/usr/local/share/yams/models/" + modelName + "/" + modelName + ".onnx");
searchPaths.push_back("/usr/local/share/yams/models/" + modelName + ".onnx");
#endif
// Log search paths on first attempt for debugging
spdlog::info("[ONNX Plugin] Resolving model path for '{}', searching {} locations", modelName,
searchPaths.size());
#ifdef _WIN32
spdlog::debug("[ONNX Plugin] Windows paths: LOCALAPPDATA={}, APPDATA={}, HOME={}", localAppData,
appData, homeDir);
#endif
// Direct filesystem checks - no async, no timeouts needed for local files
for (const auto& path : searchPaths) {
if (path.empty())
continue;
try {
if (fs::exists(path)) {
spdlog::info("[ONNX Plugin] Found model at: {}", path);
return path;
}
} catch (const std::exception& e) {
spdlog::debug("[ONNX Plugin] Error checking path {}: {}", path, e.what());
}
}
spdlog::warn("[ONNX Plugin] Model '{}' not found in any search path. Searched:", modelName);
for (const auto& path : searchPaths) {
if (!path.empty()) {
spdlog::warn("[ONNX Plugin] - {}", path);
}
}
// Default to models directory
return "models/" + modelName + ".onnx";
}
Result<OnnxModelPool::ModelSessionPtr>
OnnxModelPool::createModelSession(const std::string& modelName) {
std::string modelPath = resolveModelPath(modelName);
OnnxTextConfig config;
config.model_path = modelPath;
config.model_name = modelName;
config.enable_gpu = config_.enableGPU;
config.num_threads = config_.numThreads;
config.max_sequence_length = 0;
{
std::lock_guard<std::recursive_mutex> lock(mutex_);
auto it = threadingOverrides_.find(modelName);
if (it != threadingOverrides_.end()) {
config.num_threads = it->second.first;
config.inter_op_threads = it->second.second;
}
}
try {
auto session = std::make_shared<OnnxModelSession>(modelPath, modelName, config);
return session;
} catch (const std::exception& e) {
return Error{ErrorCode::InternalError,
std::string("Failed to create model session: ") + e.what()};
}
}
void OnnxModelPool::updateAccessStats(const std::string& modelName) {
auto it = models_.find(modelName);
if (it != models_.end()) {
it->second.lastAccess = std::chrono::steady_clock::now();
it->second.accessCount++;
}
}
bool OnnxModelPool::canLoadModel(size_t estimatedSize) const {
size_t currentUsage = getMemoryUsage();
size_t maxBytes = config_.maxMemoryGB * 1024 * 1024 * 1024;
return (currentUsage + estimatedSize) <= maxBytes;
}
void OnnxModelPool::setResolutionHints(const std::string& modelName, const ResolutionHints& hints) {
std::lock_guard<std::recursive_mutex> lock(mutex_);
modelHints_[modelName] = hints;
}
} // namespace yams::daemon
Reference in a new issue View git blame Copy permalink