trevon/yams
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2
yams/plugins/onnx/onnx_model_pool.cpp

3456 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_cxx_api_wrapper.h"
#include "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);
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 {
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
// Lightweight, environment-driven mock mode for CI and constrained hosts
// (works regardless of gtest being present in this translation unit).
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;
}
const auto& runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
if (!runtimeInfo.available) {
throw std::runtime_error(runtimeInfo.errorMessage.empty()
? std::string("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) {
// 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;
}
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>();
}
// Mark as startup mode - reduces pool sizes during initial loading
TuneAdvisor::setOnnxStartupMode(true);
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")
: 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);
}
}
}
// Clear startup mode - normal pool sizing now applies to new models
TuneAdvisor::setOnnxStartupMode(false);
spdlog::info("Background model preloading completed");
} catch (const std::exception& e) {
try {
spdlog::error("[ONNX] Background preload crashed: {}", e.what());
} catch (...) {
}
try {
TuneAdvisor::setOnnxStartupMode(false);
} catch (...) {
}
} catch (...) {
try {
spdlog::error("[ONNX] Background preload crashed with unknown exception");
} catch (...) {
}
try {
TuneAdvisor::setOnnxStartupMode(false);
} catch (...) {
}
}
});
} else {
// No preload configured - clear startup mode immediately
TuneAdvisor::setOnnxStartupMode(false);
}
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