fe07469620
The ONNX plugin dlopen()ed ONNX Runtime from directories next to the
plugin first and only then from a few system prefixes, with no notion of
which copy it got or why, and the shipped packages contained no copy at
all, so on most hosts the plugin could only report "runtime not found".
The runtime resolution moves into plugins/ort_runtime/, built as a
hidden-visibility static library both ONNX-using plugins link (their
export maps still expose only the plugin C ABI). The policy is split out
as pure code (ort_runtime_resolution.{h,cpp}) so it is testable without a
runtime. Candidates are tried in this order and the first compatible one
wins:
1. YAMS_ONNX_RUNTIME_LIB / YAMS_ONNX_RUNTIME_DIR (existing pin, kept),
then the plugin config key runtime_library ([plugins.onnx] /
[plugins.glint], a file or directory)
2. a system copy: the loader search by soname, then the usual library
directories (multiarch dir, /usr/lib64, /usr/lib, /usr/local/lib,
/opt/onnxruntime, Homebrew) for versioned names such as Debian's
libonnxruntime.so.1.23
3. the private copy in <libdir>/yams/onnxruntime/
4. legacy locations next to the plugin (build trees)
A candidate is accepted only if it exports OrtGetApiBase, GetApi() returns
an API table for the ORT_API_VERSION the plugins are compiled against, and
it meets the 1.23 minimum; otherwise it is skipped with a reason. A soname
lookup that returns a copy another plugin already loaded from the bundled
directory is reported as that copy. The choice ("Using system|bundled ONNX
Runtime <ver> from <path> (<why>)") and every skipped candidate go to the
plugin log and, through the libyams_onnx_resource host sink (new
OnnxConcurrencyRegistry::emitHostLog), to the daemon log; plugin health
reports runtime_source. With no compatible copy the plugin reports "ONNX
Runtime unavailable: ..." instead of failing.
Tests: tests/unit/plugins/ort_runtime_resolution_catch2_test.cpp covers
the order, the override pin, configured directories, system-over-bundled
preference, skipping an incompatible system copy, the not-found result
and the already-loaded classification.
Signed-off-by: trvon <me@trevon.dev>
2151 lines
93 KiB
C++
2151 lines
93 KiB
C++
// Replace stub with actual ONNX-backed implementation
|
|
#include "model_provider.h"
|
|
#include <nlohmann/json.hpp>
|
|
#include <spdlog/spdlog.h>
|
|
#include <algorithm>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <condition_variable>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <functional>
|
|
#include <future>
|
|
#include <optional>
|
|
#include <thread>
|
|
#include "../ort_runtime/ort_runtime_loader.h"
|
|
#include <yams/daemon/components/InternalEventBus.h>
|
|
#include <yams/daemon/components/TuneAdvisor.h>
|
|
#include <yams/daemon/resource/onnx_colbert_session.h>
|
|
#include <yams/daemon/resource/onnx_model_pool.h>
|
|
|
|
namespace {
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
// Global storage for plugin config JSON passed from daemon at init time
|
|
static std::string g_plugin_config_json;
|
|
|
|
static std::vector<int32_t> readColbertSkiplist(const std::string& modelsRoot,
|
|
const std::string& modelName) {
|
|
try {
|
|
fs::path file = fs::path(modelsRoot) / modelName / "skiplist.json";
|
|
if (!fs::exists(file))
|
|
return {};
|
|
std::ifstream in(file);
|
|
if (!in.good())
|
|
return {};
|
|
nlohmann::json j;
|
|
in >> j;
|
|
if (!j.is_array())
|
|
return {};
|
|
std::vector<int32_t> out;
|
|
out.reserve(j.size());
|
|
for (const auto& item : j) {
|
|
if (item.is_number_integer()) {
|
|
out.push_back(item.get<int32_t>());
|
|
}
|
|
}
|
|
return out;
|
|
} catch (...) {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
static bool isColbertModelName(const std::string& name) {
|
|
if (name.empty()) {
|
|
return false;
|
|
}
|
|
std::string lower = name;
|
|
for (auto& c : lower) {
|
|
c = static_cast<char>(std::tolower(c));
|
|
}
|
|
return lower.find("colbert") != std::string::npos;
|
|
}
|
|
|
|
struct ProviderCtx;
|
|
|
|
extern "C" yams_model_provider_v1* yams_onnx_get_model_provider();
|
|
|
|
static bool envTruthy(const char* value) {
|
|
if (!value || !*value) {
|
|
return false;
|
|
}
|
|
std::string v(value);
|
|
for (auto& c : v) {
|
|
c = static_cast<char>(std::tolower(c));
|
|
}
|
|
return v != "0" && v != "false" && v != "off" && v != "no";
|
|
}
|
|
|
|
static std::chrono::milliseconds modelAcquireTimeout(bool batchCall) {
|
|
// Batch calls can hold the model session for longer on large chunks/sub-batches.
|
|
// Keep a longer default to avoid turning transient contention into hard failures.
|
|
std::chrono::milliseconds timeout =
|
|
batchCall ? std::chrono::milliseconds(120000) : std::chrono::milliseconds(30000);
|
|
|
|
if (const char* global = std::getenv("YAMS_ONNX_MODEL_ACQUIRE_TIMEOUT_MS")) {
|
|
try {
|
|
const auto parsed = std::stoll(global);
|
|
if (parsed > 0) {
|
|
timeout = std::chrono::milliseconds(parsed);
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
|
|
if (batchCall) {
|
|
if (const char* batch = std::getenv("YAMS_ONNX_MODEL_ACQUIRE_TIMEOUT_BATCH_MS")) {
|
|
try {
|
|
const auto parsed = std::stoll(batch);
|
|
if (parsed > 0) {
|
|
timeout = std::chrono::milliseconds(parsed);
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
}
|
|
|
|
if (envTruthy(std::getenv("YAMS_ONNX_SHUTDOWN_IN_PROGRESS"))) {
|
|
timeout = std::min(timeout, std::chrono::milliseconds(1000));
|
|
}
|
|
|
|
constexpr auto kMinTimeout = std::chrono::milliseconds(1000);
|
|
constexpr auto kMaxTimeout = std::chrono::milliseconds(600000);
|
|
timeout = std::clamp(timeout, kMinTimeout, kMaxTimeout);
|
|
return timeout;
|
|
}
|
|
|
|
// Provide a C-callable function for threading control (not supported yet)
|
|
static yams_status_t onnx_set_threading(void* /*self*/, const char* /*model_id*/, int /*intra*/,
|
|
int /*inter*/);
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// BatchCoalescer — accumulates concurrent embedding requests into larger
|
|
// batches before running inference. Opt-in via YAMS_ONNX_BATCH_COALESCE=1.
|
|
// ---------------------------------------------------------------------------
|
|
class BatchCoalescer {
|
|
public:
|
|
struct Config {
|
|
bool enabled = false;
|
|
int windowMs = 5;
|
|
size_t maxBatch = 64;
|
|
};
|
|
|
|
explicit BatchCoalescer(Config cfg) : cfg_(cfg) {}
|
|
|
|
bool isEnabled() const { return cfg_.enabled; }
|
|
|
|
using InferFn = std::function<yams::Result<std::vector<std::vector<float>>>(
|
|
const std::string& /*modelId*/, const std::vector<std::string>& /*texts*/)>;
|
|
|
|
// Submit texts for embedding. Blocks until the coalesced batch completes
|
|
// and returns only this caller's slice of the result.
|
|
yams::Result<std::vector<std::vector<float>>> submit(const std::string& modelId,
|
|
const std::vector<std::string>& texts,
|
|
const InferFn& inferFn) {
|
|
if (!cfg_.enabled || texts.empty()) {
|
|
return inferFn(modelId, texts);
|
|
}
|
|
|
|
std::unique_lock<std::mutex> lk(mu_);
|
|
|
|
auto& batch = pending_[modelId];
|
|
if (!batch) {
|
|
batch = std::make_unique<PendingBatch>();
|
|
batch->modelId = modelId;
|
|
batch->deadline =
|
|
std::chrono::steady_clock::now() + std::chrono::milliseconds(cfg_.windowMs);
|
|
}
|
|
|
|
// Record this request's slice within the coalesced batch.
|
|
const size_t startIdx = batch->allTexts.size();
|
|
const size_t count = texts.size();
|
|
for (auto& t : texts)
|
|
batch->allTexts.push_back(std::move(t));
|
|
|
|
auto future = batch->addRequest(startIdx, count);
|
|
|
|
// Should we dispatch now?
|
|
const bool sizeTriggered = batch->allTexts.size() >= cfg_.maxBatch;
|
|
const bool timeTriggered = std::chrono::steady_clock::now() >= batch->deadline;
|
|
|
|
if (sizeTriggered || timeTriggered) {
|
|
// This thread becomes the dispatcher.
|
|
auto owned = std::move(batch);
|
|
pending_.erase(modelId);
|
|
lk.unlock();
|
|
dispatch(std::move(owned), inferFn);
|
|
} else {
|
|
// Wait for either the deadline or the dispatcher to signal.
|
|
cv_.wait_until(lk, batch->deadline, [&] {
|
|
// If our batch was consumed (pending_ no longer holds it), we're done.
|
|
auto it = pending_.find(modelId);
|
|
return it == pending_.end() || it->second.get() != batch.get();
|
|
});
|
|
|
|
// If we woke up because of timeout and the batch is still ours, dispatch it.
|
|
auto it = pending_.find(modelId);
|
|
if (it != pending_.end() && it->second.get() == batch.get()) {
|
|
auto owned = std::move(it->second);
|
|
pending_.erase(it);
|
|
lk.unlock();
|
|
dispatch(std::move(owned), inferFn);
|
|
} else {
|
|
lk.unlock();
|
|
}
|
|
}
|
|
|
|
return future.get();
|
|
}
|
|
|
|
// Drain all pending batches with errors (call on shutdown).
|
|
void drain() {
|
|
std::lock_guard<std::mutex> lk(mu_);
|
|
for (auto& [id, batch] : pending_) {
|
|
if (batch) {
|
|
batch->failAll(yams::Error{yams::ErrorCode::InternalError, "plugin shutting down"});
|
|
}
|
|
}
|
|
pending_.clear();
|
|
cv_.notify_all();
|
|
}
|
|
|
|
private:
|
|
struct PendingBatch {
|
|
struct Request {
|
|
size_t startIdx;
|
|
size_t count;
|
|
std::promise<yams::Result<std::vector<std::vector<float>>>> promise;
|
|
};
|
|
|
|
std::string modelId;
|
|
std::vector<std::string> allTexts;
|
|
std::vector<Request> requests;
|
|
std::chrono::steady_clock::time_point deadline;
|
|
|
|
std::future<yams::Result<std::vector<std::vector<float>>>> addRequest(size_t startIdx,
|
|
size_t count) {
|
|
requests.emplace_back();
|
|
auto& req = requests.back();
|
|
req.startIdx = startIdx;
|
|
req.count = count;
|
|
return req.promise.get_future();
|
|
}
|
|
|
|
void failAll(const yams::Error& err) {
|
|
for (auto& req : requests) {
|
|
try {
|
|
req.promise.set_value(yams::Result<std::vector<std::vector<float>>>(err));
|
|
} catch (...) {
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
void dispatch(std::unique_ptr<PendingBatch> batch, const InferFn& inferFn) {
|
|
auto result = inferFn(batch->modelId, batch->allTexts);
|
|
for (auto& req : batch->requests) {
|
|
try {
|
|
if (!result) {
|
|
req.promise.set_value(
|
|
yams::Result<std::vector<std::vector<float>>>(result.error()));
|
|
} else {
|
|
auto& mat = result.value();
|
|
if (req.startIdx + req.count <= mat.size()) {
|
|
std::vector<std::vector<float>> slice(
|
|
mat.begin() + static_cast<ptrdiff_t>(req.startIdx),
|
|
mat.begin() + static_cast<ptrdiff_t>(req.startIdx + req.count));
|
|
req.promise.set_value(std::move(slice));
|
|
} else {
|
|
req.promise.set_value(yams::Error{yams::ErrorCode::InternalError,
|
|
"batch coalescer: output size mismatch"});
|
|
}
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
cv_.notify_all();
|
|
}
|
|
|
|
Config cfg_;
|
|
std::mutex mu_;
|
|
std::condition_variable cv_;
|
|
std::unordered_map<std::string, std::unique_ptr<PendingBatch>> pending_;
|
|
};
|
|
|
|
struct ProviderCtx {
|
|
enum class State : uint8_t { Unloaded, Loading, Ready, Failed };
|
|
|
|
// Track model load failures for backoff and event signaling
|
|
struct FailureInfo {
|
|
std::chrono::steady_clock::time_point lastAttempt;
|
|
std::chrono::steady_clock::time_point lastFailure;
|
|
uint32_t consecutiveFailures = 0;
|
|
std::string lastError;
|
|
bool eventFired = false; // Whether we've fired ModelLoadFailedEvent
|
|
};
|
|
|
|
std::mutex mu;
|
|
yams_model_load_progress_cb progress_cb = nullptr;
|
|
void* progress_user = nullptr;
|
|
std::unique_ptr<yams::daemon::OnnxModelPool> pool;
|
|
std::unique_ptr<yams::daemon::OnnxColbertSession> colbert; // ColBERT session
|
|
std::unordered_map<std::string, State> model_states; // per-model FSM
|
|
std::unordered_map<std::string, FailureInfo> model_failures; // failure tracking
|
|
bool ready = false;
|
|
bool disabled = false;
|
|
std::atomic<bool> shutdownRequested{false};
|
|
bool gpuEnabled = false; // Tracks whether GPU acceleration is configured
|
|
std::string actualExecutionProvider{"cpu"};
|
|
std::string last_error;
|
|
std::string colbertModelPath; // Path to ColBERT model
|
|
std::size_t configuredMaxLoadedModels = 0;
|
|
std::size_t configuredHotPoolSize = 0;
|
|
yams::onnx_util::OrtRuntimeInfo runtimeInfo;
|
|
std::unique_ptr<BatchCoalescer> coalescer;
|
|
|
|
// Check if model is in cooldown period after failures
|
|
bool isInCooldown(const std::string& modelId) const {
|
|
auto it = model_failures.find(modelId);
|
|
if (it == model_failures.end() || it->second.consecutiveFailures == 0) {
|
|
return false;
|
|
}
|
|
// Exponential backoff: 2^min(failures, 6) seconds (max ~64 seconds)
|
|
auto backoffSeconds =
|
|
std::chrono::seconds(1 << std::min(it->second.consecutiveFailures, 6u));
|
|
auto elapsed = std::chrono::steady_clock::now() - it->second.lastFailure;
|
|
return elapsed < backoffSeconds;
|
|
}
|
|
|
|
// Record a model load failure
|
|
void recordFailure(const std::string& modelId, const std::string& error) {
|
|
auto& info = model_failures[modelId];
|
|
auto now = std::chrono::steady_clock::now();
|
|
info.lastAttempt = now;
|
|
info.lastFailure = now;
|
|
info.consecutiveFailures++;
|
|
info.lastError = error;
|
|
}
|
|
|
|
// Clear failure tracking on successful load
|
|
void clearFailure(const std::string& modelId) {
|
|
auto it = model_failures.find(modelId);
|
|
if (it != model_failures.end()) {
|
|
it->second.consecutiveFailures = 0;
|
|
it->second.lastError.clear();
|
|
it->second.eventFired = false;
|
|
}
|
|
}
|
|
|
|
ProviderCtx() {
|
|
// Allow disabling via env for diagnostics or platform constraints
|
|
if (const char* d = std::getenv("YAMS_ONNX_PLUGIN_DISABLE"); d && *d) {
|
|
disabled = true;
|
|
last_error = "disabled_by_env";
|
|
return;
|
|
}
|
|
|
|
runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
|
|
if (!runtimeInfo.available) {
|
|
last_error = runtimeInfo.reason.empty() ? "runtime_not_available" : runtimeInfo.reason;
|
|
spdlog::warn("[ONNX-Plugin] ONNX Runtime unavailable: {}",
|
|
runtimeInfo.errorMessage.empty() ? last_error : runtimeInfo.errorMessage);
|
|
return;
|
|
}
|
|
|
|
yams::daemon::ModelPoolConfig cfg;
|
|
// Defaults: prefer lazy loading to avoid blocking startup
|
|
// Models will be loaded on first use, not during plugin initialization
|
|
cfg.lazyLoading = true;
|
|
cfg.enableGPU = true; // runtime detection in appendGpuExecutionProvider()
|
|
const unsigned hwThreads = std::max(1u, std::thread::hardware_concurrency());
|
|
cfg.numThreads = static_cast<int>(std::clamp<unsigned>(hwThreads / 2, 1u, 4u));
|
|
cfg.loadWorkerThreads = std::max<size_t>(1u, static_cast<size_t>(cfg.numThreads / 2));
|
|
cfg.asyncLoading = false;
|
|
|
|
// Variables for config (both file-based and JSON-based)
|
|
std::string dataDir;
|
|
bool keepModelHot = true;
|
|
std::string preferredModel;
|
|
std::vector<std::string> preloadList;
|
|
|
|
// Read embeddings settings from config.toml (preferred_model, model_path, keep_model_hot)
|
|
try {
|
|
namespace fs = std::filesystem;
|
|
fs::path cfgPath;
|
|
if (const char* explicitConfig = std::getenv("YAMS_CONFIG");
|
|
explicitConfig && *explicitConfig) {
|
|
cfgPath = explicitConfig;
|
|
} else if (const char* xdg = std::getenv("XDG_CONFIG_HOME")) {
|
|
cfgPath = fs::path(xdg) / "yams" / "config.toml";
|
|
} else if (const char* home = std::getenv("HOME")) {
|
|
cfgPath = fs::path(home) / ".config" / "yams" / "config.toml";
|
|
}
|
|
|
|
// First, check for YAMS_STORAGE environment variable
|
|
if (const char* storage = std::getenv("YAMS_STORAGE")) {
|
|
dataDir = storage;
|
|
}
|
|
|
|
spdlog::debug("[ONNX-Plugin] Config path: {}", cfgPath.string());
|
|
if (!cfgPath.empty() && fs::exists(cfgPath)) {
|
|
spdlog::debug("[ONNX-Plugin] Config file exists, parsing...");
|
|
std::ifstream file(cfgPath);
|
|
std::string line;
|
|
std::string section;
|
|
auto trim = [](std::string& s) {
|
|
if (s.empty())
|
|
return;
|
|
s.erase(0, s.find_first_not_of(" \t"));
|
|
auto pos = s.find_last_not_of(" \t");
|
|
if (pos != std::string::npos)
|
|
s.erase(pos + 1);
|
|
};
|
|
auto apply_num_threads = [&](const std::string& value) {
|
|
try {
|
|
int v = std::stoi(value);
|
|
if (v > 0) {
|
|
cfg.numThreads = v;
|
|
}
|
|
} catch (...) {
|
|
}
|
|
};
|
|
auto apply_max_loaded_models = [&](const std::string& value) {
|
|
try {
|
|
std::size_t v = static_cast<std::size_t>(std::stoull(value));
|
|
if (v >= 1 && v <= 64) {
|
|
cfg.maxLoadedModels = v;
|
|
}
|
|
} catch (...) {
|
|
}
|
|
};
|
|
auto apply_hot_pool_size = [&](const std::string& value) {
|
|
try {
|
|
std::size_t v = static_cast<std::size_t>(std::stoull(value));
|
|
if (v >= 1 && v <= 64) {
|
|
cfg.hotPoolSize = v;
|
|
}
|
|
} catch (...) {
|
|
}
|
|
};
|
|
auto apply_async_loading = [&](const std::string& value) {
|
|
std::string v;
|
|
v.reserve(value.size());
|
|
for (unsigned char c : value)
|
|
v.push_back(static_cast<char>(std::tolower(c)));
|
|
if (v == "1" || v == "true" || v == "yes" || v == "on")
|
|
cfg.asyncLoading = true;
|
|
if (v == "0" || v == "false" || v == "no" || v == "off")
|
|
cfg.asyncLoading = false;
|
|
};
|
|
while (std::getline(file, line)) {
|
|
if (line.empty() || line[0] == '#')
|
|
continue;
|
|
if (line[0] == '[') {
|
|
auto end = line.find(']');
|
|
section = (end != std::string::npos) ? line.substr(1, end - 1) : "";
|
|
spdlog::debug("[ONNX-Plugin] Entering section: [{}]", section);
|
|
continue;
|
|
}
|
|
auto eq = line.find('=');
|
|
if (eq == std::string::npos)
|
|
continue;
|
|
std::string key = line.substr(0, eq);
|
|
std::string value = line.substr(eq + 1);
|
|
trim(key);
|
|
trim(value);
|
|
auto hash = value.find('#');
|
|
if (hash != std::string::npos) {
|
|
value = value.substr(0, hash);
|
|
trim(value);
|
|
}
|
|
if (value.size() >= 2 && value.front() == '"' && value.back() == '"')
|
|
value = value.substr(1, value.size() - 2);
|
|
|
|
// Read data directory from [storage].base_path if not set by env
|
|
if (section == "storage" && key == "base_path") {
|
|
spdlog::debug("[ONNX-Plugin] Found [storage].base_path = '{}'", value);
|
|
if (dataDir.empty()) {
|
|
dataDir = value;
|
|
}
|
|
}
|
|
|
|
if (section == "embeddings") {
|
|
if (key == "preferred_model" && preferredModel.empty())
|
|
preferredModel = std::move(value);
|
|
else if (key == "num_threads")
|
|
apply_num_threads(value);
|
|
else if (key == "keep_model_hot") {
|
|
std::string v;
|
|
v.reserve(value.size());
|
|
for (unsigned char c : value)
|
|
v.push_back(static_cast<char>(std::tolower(c)));
|
|
keepModelHot = !(v == "false" || v == "0" || v == "no" || v == "off");
|
|
}
|
|
}
|
|
if (section == "daemon.models") {
|
|
if (key == "num_threads")
|
|
apply_num_threads(value);
|
|
else if (key == "max_loaded_models")
|
|
apply_max_loaded_models(value);
|
|
else if (key == "hot_pool_size")
|
|
apply_hot_pool_size(value);
|
|
}
|
|
// New: preload list under [plugins.onnx]
|
|
if (section == "plugins.onnx" && key == "preload") {
|
|
// parse CSV style list: preload = "modelA,modelB"
|
|
std::string s = value;
|
|
for (char& ch : s) {
|
|
if (ch == '[' || ch == ']')
|
|
ch = ' ';
|
|
}
|
|
trim(s);
|
|
// allow ["a","b"] or a,b
|
|
if (!s.empty() && s.front() == '"' && s.back() == '"')
|
|
s = s.substr(1, s.size() - 2);
|
|
size_t start = 0;
|
|
while (start < s.size()) {
|
|
auto comma = s.find(',', start);
|
|
std::string item =
|
|
s.substr(start, comma == std::string::npos ? std::string::npos
|
|
: (comma - start));
|
|
trim(item);
|
|
if (!item.empty() && item.front() == '"' && item.back() == '"')
|
|
item = item.substr(1, item.size() - 2);
|
|
if (!item.empty())
|
|
preloadList.push_back(item);
|
|
if (comma == std::string::npos)
|
|
break;
|
|
start = comma + 1;
|
|
}
|
|
}
|
|
if (section == "plugins.onnx" && key == "max_loaded_models")
|
|
apply_max_loaded_models(value);
|
|
if (section == "plugins.onnx" && key == "hot_pool_size")
|
|
apply_hot_pool_size(value);
|
|
if (section == "plugins.onnx" && key == "async_load")
|
|
apply_async_loading(value);
|
|
// New: explicit models table entries [plugins.onnx.models.NAME]
|
|
if (section.rfind("plugins.onnx.models.", 0) == 0 && key == "task") {
|
|
// Section name encodes model_id; mark for preload as hot
|
|
std::string model_id =
|
|
section.substr(std::string("plugins.onnx.models.").size());
|
|
if (!model_id.empty())
|
|
preloadList.push_back(model_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Set modelsRoot to $dataDir/models if dataDir is available
|
|
spdlog::debug("[ONNX-Plugin] After config parsing: dataDir='{}'", dataDir);
|
|
if (!dataDir.empty()) {
|
|
// Expand ~ if present
|
|
if (dataDir[0] == '~') {
|
|
if (const char* home = std::getenv("HOME"))
|
|
dataDir = std::string(home) + dataDir.substr(1);
|
|
}
|
|
cfg.modelsRoot = dataDir + "/models";
|
|
spdlog::info("[ONNX-Plugin] Using models directory: {}", cfg.modelsRoot);
|
|
}
|
|
// Build preload list if any; eager load behavior is controlled later.
|
|
if (!preloadList.empty()) {
|
|
cfg.preloadModels = preloadList;
|
|
} else if (keepModelHot && !preferredModel.empty()) {
|
|
cfg.preloadModels = {preferredModel};
|
|
}
|
|
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
|
|
// Force no-preload in test/mock environments to avoid background threads
|
|
if (std::getenv("YAMS_SKIP_MODEL_LOADING") || std::getenv("YAMS_TEST_MODE") ||
|
|
std::getenv("YAMS_USE_MOCK_PROVIDER")) {
|
|
cfg.preloadModels.clear();
|
|
cfg.lazyLoading = true;
|
|
}
|
|
#endif
|
|
} catch (const std::exception&) {
|
|
// ignore config parse errors
|
|
}
|
|
|
|
// Env overrides for async loading
|
|
auto parse_bool_env = [](const char* name) -> std::optional<bool> {
|
|
if (const char* s = std::getenv(name)) {
|
|
std::string v(s);
|
|
for (auto& c : v)
|
|
c = static_cast<char>(std::tolower(c));
|
|
if (v == "1" || v == "true" || v == "yes" || v == "on")
|
|
return true;
|
|
if (v == "0" || v == "false" || v == "no" || v == "off")
|
|
return false;
|
|
}
|
|
return std::nullopt;
|
|
};
|
|
if (auto val = parse_bool_env("YAMS_ONNX_ASYNC_LOAD")) {
|
|
cfg.asyncLoading = *val;
|
|
}
|
|
if (parse_bool_env("YAMS_ONNX_FORCE_CPU").value_or(false)) {
|
|
cfg.enableGPU = false;
|
|
}
|
|
const bool eagerPreload = parse_bool_env("YAMS_ONNX_EAGER_PRELOAD").value_or(false);
|
|
cfg.lazyLoading = !eagerPreload;
|
|
if (const char* w = std::getenv("YAMS_ONNX_LOAD_WORKERS")) {
|
|
try {
|
|
size_t threads = static_cast<size_t>(std::stoul(w));
|
|
if (threads > 0 && threads <= 32) {
|
|
cfg.loadWorkerThreads = threads;
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
|
|
// Parse JSON config from plugin init (overrides file-based config)
|
|
if (!g_plugin_config_json.empty()) {
|
|
try {
|
|
auto j = nlohmann::json::parse(g_plugin_config_json, nullptr, false);
|
|
if (!j.is_discarded() && j.is_object()) {
|
|
spdlog::info("[ONNX-Plugin] Applying JSON config: {}", g_plugin_config_json);
|
|
// preferred_model
|
|
if (j.contains("preferred_model") && j["preferred_model"].is_string()) {
|
|
preferredModel = j["preferred_model"].get<std::string>();
|
|
spdlog::info("[ONNX-Plugin] JSON config: preferred_model='{}'",
|
|
preferredModel);
|
|
}
|
|
// preload (CSV string or single model)
|
|
if (j.contains("preload") && j["preload"].is_string()) {
|
|
std::string preloadStr = j["preload"].get<std::string>();
|
|
preloadList.clear();
|
|
// Parse CSV: "modelA,modelB" or single "modelA"
|
|
size_t start = 0;
|
|
while (start < preloadStr.size()) {
|
|
auto comma = preloadStr.find(',', start);
|
|
std::string item = preloadStr.substr(start, comma == std::string::npos
|
|
? std::string::npos
|
|
: (comma - start));
|
|
// Trim whitespace
|
|
auto trim = [](std::string& s) {
|
|
if (s.empty())
|
|
return;
|
|
s.erase(0, s.find_first_not_of(" \t"));
|
|
auto pos = s.find_last_not_of(" \t");
|
|
if (pos != std::string::npos)
|
|
s.erase(pos + 1);
|
|
};
|
|
trim(item);
|
|
if (!item.empty())
|
|
preloadList.push_back(item);
|
|
if (comma == std::string::npos)
|
|
break;
|
|
start = comma + 1;
|
|
}
|
|
spdlog::info("[ONNX-Plugin] JSON config: preload list has {} models",
|
|
preloadList.size());
|
|
}
|
|
// keep_model_hot
|
|
if (j.contains("keep_model_hot") && j["keep_model_hot"].is_boolean()) {
|
|
keepModelHot = j["keep_model_hot"].get<bool>();
|
|
spdlog::info("[ONNX-Plugin] JSON config: keep_model_hot={}", keepModelHot);
|
|
}
|
|
// num_threads
|
|
if (j.contains("num_threads") && j["num_threads"].is_number_integer()) {
|
|
int v = j["num_threads"].get<int>();
|
|
if (v > 0) {
|
|
cfg.numThreads = v;
|
|
spdlog::info("[ONNX-Plugin] JSON config: num_threads={}",
|
|
cfg.numThreads);
|
|
}
|
|
}
|
|
if (j.contains("max_loaded_models") &&
|
|
j["max_loaded_models"].is_number_integer()) {
|
|
auto v = static_cast<std::size_t>(j["max_loaded_models"].get<int>());
|
|
if (v >= 1 && v <= 64) {
|
|
cfg.maxLoadedModels = v;
|
|
spdlog::info("[ONNX-Plugin] JSON config: max_loaded_models={}",
|
|
cfg.maxLoadedModels);
|
|
}
|
|
}
|
|
if (j.contains("hot_pool_size") && j["hot_pool_size"].is_number_integer()) {
|
|
auto v = static_cast<std::size_t>(j["hot_pool_size"].get<int>());
|
|
if (v >= 1 && v <= 64) {
|
|
cfg.hotPoolSize = v;
|
|
spdlog::info("[ONNX-Plugin] JSON config: hot_pool_size={}",
|
|
cfg.hotPoolSize);
|
|
}
|
|
}
|
|
if (j.contains("async_load") && j["async_load"].is_boolean()) {
|
|
cfg.asyncLoading = j["async_load"].get<bool>();
|
|
spdlog::info("[ONNX-Plugin] JSON config: async_load={}", cfg.asyncLoading);
|
|
}
|
|
if (j.contains("load_worker_threads") &&
|
|
j["load_worker_threads"].is_number_integer()) {
|
|
auto v = static_cast<std::size_t>(j["load_worker_threads"].get<int>());
|
|
if (v >= 1 && v <= 32) {
|
|
cfg.loadWorkerThreads = v;
|
|
spdlog::info("[ONNX-Plugin] JSON config: load_worker_threads={}",
|
|
cfg.loadWorkerThreads);
|
|
}
|
|
}
|
|
// models_root - override the models directory
|
|
if (j.contains("models_root") && j["models_root"].is_string()) {
|
|
cfg.modelsRoot = j["models_root"].get<std::string>();
|
|
spdlog::info("[ONNX-Plugin] JSON config: models_root='{}'", cfg.modelsRoot);
|
|
}
|
|
}
|
|
} catch (const std::exception& e) {
|
|
spdlog::warn("[ONNX-Plugin] Failed to parse JSON config: {}", e.what());
|
|
}
|
|
// Re-apply preload settings from JSON config
|
|
if (!preloadList.empty()) {
|
|
cfg.preloadModels = preloadList;
|
|
} else if (keepModelHot && !preferredModel.empty()) {
|
|
cfg.preloadModels = {preferredModel};
|
|
}
|
|
}
|
|
|
|
cfg.maxLoadedModels = std::max<std::size_t>(1, cfg.maxLoadedModels);
|
|
cfg.hotPoolSize = std::clamp(cfg.hotPoolSize, std::size_t{1}, cfg.maxLoadedModels);
|
|
|
|
if (!cfg.preloadModels.empty()) {
|
|
const std::size_t preloadCap = std::min(cfg.hotPoolSize, cfg.maxLoadedModels);
|
|
if (cfg.preloadModels.size() > preloadCap) {
|
|
spdlog::warn("[ONNX-Plugin] Trimming preload list from {} to {} (hot_pool_size={}, "
|
|
"max_loaded_models={})",
|
|
cfg.preloadModels.size(), preloadCap, cfg.hotPoolSize,
|
|
cfg.maxLoadedModels);
|
|
cfg.preloadModels.resize(preloadCap);
|
|
}
|
|
if (cfg.lazyLoading) {
|
|
spdlog::info("[ONNX-Plugin] preload list configured ({} models), lazy loading "
|
|
"enabled; models will load on first use",
|
|
cfg.preloadModels.size());
|
|
}
|
|
}
|
|
|
|
configuredMaxLoadedModels = cfg.maxLoadedModels;
|
|
configuredHotPoolSize = cfg.hotPoolSize;
|
|
gpuEnabled = cfg.enableGPU;
|
|
spdlog::info("[ONNX-Plugin] Creating OnnxModelPool with modelsRoot={}, gpuEnabled={}, "
|
|
"async_load={}",
|
|
cfg.modelsRoot, gpuEnabled, cfg.asyncLoading);
|
|
pool = std::make_unique<yams::daemon::OnnxModelPool>(cfg);
|
|
spdlog::info("[ONNX-Plugin] OnnxModelPool created, calling initialize()...");
|
|
try {
|
|
auto res = pool->initialize();
|
|
if (res) {
|
|
ready = true;
|
|
last_error.clear();
|
|
spdlog::info("[ONNX-Plugin] Pool initialized successfully, ready=true");
|
|
|
|
// Best-effort: derive actual EP from a preloaded model (if any)
|
|
try {
|
|
if (!cfg.preloadModels.empty()) {
|
|
auto h = pool->acquireModel(cfg.preloadModels.front(),
|
|
std::chrono::milliseconds(500));
|
|
if (h) {
|
|
actualExecutionProvider = h.value()->getExecutionProvider();
|
|
}
|
|
}
|
|
} catch (...) {
|
|
}
|
|
} else {
|
|
ready = false;
|
|
last_error = res.error().message;
|
|
spdlog::warn("[ONNX-Plugin] Pool initialize failed: {}", last_error);
|
|
}
|
|
} catch (const std::exception& e) {
|
|
spdlog::error("[ONNX-Plugin] Pool initialize exception: {}", e.what());
|
|
ready = false;
|
|
last_error = e.what();
|
|
}
|
|
|
|
std::string colbertModelName = preferredModel;
|
|
if (colbertModelName.empty() && !preloadList.empty()) {
|
|
colbertModelName = preloadList.front();
|
|
}
|
|
if (isColbertModelName(colbertModelName)) {
|
|
colbertModelPath = cfg.modelsRoot + "/" + colbertModelName + "/model.onnx";
|
|
yams::daemon::ColbertConfig colbertCfg{};
|
|
colbertCfg.model_path = colbertModelPath;
|
|
colbertCfg.model_name = colbertModelName;
|
|
colbertCfg.enable_gpu = cfg.enableGPU;
|
|
colbertCfg.num_threads = static_cast<int>(cfg.numThreads);
|
|
colbertCfg.max_sequence_length = 512;
|
|
colbertCfg.token_dim = 48;
|
|
colbertCfg.query_marker_id = 50368;
|
|
colbertCfg.doc_marker_id = 50369;
|
|
colbertCfg.skiplist = readColbertSkiplist(cfg.modelsRoot, colbertModelName);
|
|
colbert = std::make_unique<yams::daemon::OnnxColbertSession>(
|
|
colbertModelPath, colbertModelName, colbertCfg);
|
|
}
|
|
// Initialize batch coalescer (off by default; enable with YAMS_ONNX_BATCH_COALESCE=1).
|
|
// Only useful for workloads with many concurrent small requests (1-4 texts).
|
|
// Callers that already send well-sized batches should NOT enable this — it
|
|
// adds latency (coalesce window) and serializes concurrent inferences.
|
|
{
|
|
BatchCoalescer::Config coalesceCfg;
|
|
coalesceCfg.enabled = envTruthy(std::getenv("YAMS_ONNX_BATCH_COALESCE"));
|
|
if (const char* w = std::getenv("YAMS_ONNX_COALESCE_WINDOW_MS")) {
|
|
try {
|
|
coalesceCfg.windowMs = std::stoi(w);
|
|
} catch (...) {
|
|
}
|
|
}
|
|
if (const char* m = std::getenv("YAMS_ONNX_COALESCE_MAX_BATCH")) {
|
|
try {
|
|
coalesceCfg.maxBatch = static_cast<size_t>(std::stoull(m));
|
|
} catch (...) {
|
|
}
|
|
}
|
|
coalescer = std::make_unique<BatchCoalescer>(coalesceCfg);
|
|
if (coalesceCfg.enabled) {
|
|
spdlog::info("[ONNX-Plugin] Batch coalescing enabled: window={}ms max_batch={}",
|
|
coalesceCfg.windowMs, coalesceCfg.maxBatch);
|
|
}
|
|
}
|
|
|
|
spdlog::info("[ONNX-Plugin] ProviderCtx init complete: ready={} pool={}", ready,
|
|
pool ? "valid" : "null");
|
|
}
|
|
};
|
|
|
|
static yams_status_t onnx_set_threading(void* self, const char* model_id, int intra, int inter) {
|
|
if (!self || !model_id || !*model_id) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
if (intra == 0 || inter == 0) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
if ((intra < -1 || intra > 64) || (inter < -1 || inter > 64)) {
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled) {
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
}
|
|
if (!c->pool) {
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
auto r = c->pool->setModelThreading(model_id, intra, inter, true);
|
|
if (!r) {
|
|
spdlog::warn("[ONNX Plugin] set_threading failed model='{}' intra={} inter={}: {}",
|
|
model_id, intra, inter, r.error().message);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
spdlog::info("[ONNX Plugin] set_threading applied model='{}' intra={} inter={}", model_id,
|
|
intra, inter);
|
|
return YAMS_OK;
|
|
}
|
|
|
|
static bool shutdownInProgress(const ProviderCtx* c) {
|
|
if (c && c->shutdownRequested.load(std::memory_order_acquire)) {
|
|
return true;
|
|
}
|
|
return envTruthy(std::getenv("YAMS_ONNX_SHUTDOWN_IN_PROGRESS"));
|
|
}
|
|
|
|
static ProviderCtx* refreshCtxIfStaleShutdown(ProviderCtx* c) {
|
|
if (!c) {
|
|
return nullptr;
|
|
}
|
|
|
|
const bool marker = envTruthy(std::getenv("YAMS_ONNX_SHUTDOWN_IN_PROGRESS"));
|
|
if (marker || !c->shutdownRequested.load(std::memory_order_acquire)) {
|
|
return c;
|
|
}
|
|
|
|
// Stale in-process lifecycle: provider self pointer still references a shutdown context
|
|
// from a prior daemon instance. Force provider revival via canonical accessor.
|
|
auto* provider = yams_onnx_get_model_provider();
|
|
if (provider && provider->self) {
|
|
spdlog::info("[ONNX Plugin] Revived stale provider context after prior shutdown");
|
|
return static_cast<ProviderCtx*>(provider->self);
|
|
}
|
|
|
|
return c;
|
|
}
|
|
|
|
// Helpers for progress emission
|
|
static void emit_progress(ProviderCtx* ctx, const char* model, int phase, const char* msg,
|
|
uint64_t cur = 0, uint64_t tot = 0) {
|
|
if (!ctx || !ctx->progress_cb)
|
|
return;
|
|
try {
|
|
ctx->progress_cb(ctx->progress_user, model, phase, cur, tot, msg);
|
|
} catch (...) {
|
|
}
|
|
}
|
|
|
|
struct ProviderSingleton {
|
|
ProviderCtx ctx;
|
|
yams_model_provider_v1 vtable;
|
|
|
|
std::atomic<bool> shutdownCalled{false};
|
|
std::atomic<bool> explicitShutdownCalled{false}; // Set by yams_onnx_shutdown_provider()
|
|
std::mutex lifecycleMu;
|
|
|
|
void ensureActive() {
|
|
std::lock_guard<std::mutex> lk(lifecycleMu);
|
|
if (!shutdownCalled.load(std::memory_order_acquire)) {
|
|
return;
|
|
}
|
|
|
|
spdlog::info("[ONNX Plugin] Reviving provider after explicit shutdown");
|
|
|
|
try {
|
|
ctx.~ProviderCtx();
|
|
} catch (...) {
|
|
}
|
|
new (&ctx) ProviderCtx();
|
|
vtable.self = &ctx;
|
|
|
|
shutdownCalled.store(false, std::memory_order_release);
|
|
explicitShutdownCalled.store(false, std::memory_order_release);
|
|
}
|
|
|
|
void shutdown(bool isExplicit = false) noexcept {
|
|
std::lock_guard<std::mutex> lk(lifecycleMu);
|
|
if (shutdownCalled.exchange(true, std::memory_order_acq_rel)) {
|
|
return;
|
|
}
|
|
|
|
ctx.shutdownRequested.store(true, std::memory_order_release);
|
|
ctx.ready = false;
|
|
|
|
// Drain any pending coalesced batches so waiting threads don't hang.
|
|
if (ctx.coalescer) {
|
|
ctx.coalescer->drain();
|
|
}
|
|
|
|
// If this is called from destructor during static destruction (not explicit),
|
|
// skip cleanup to avoid crashes from corrupted global state.
|
|
// At process exit, resources will be freed anyway.
|
|
if (!isExplicit && !explicitShutdownCalled.load()) {
|
|
return; // Skip cleanup during static destruction
|
|
}
|
|
|
|
if (ctx.pool) {
|
|
try {
|
|
ctx.pool->shutdown();
|
|
} catch (const std::exception& e) {
|
|
try {
|
|
spdlog::warn("[ONNX] pool shutdown exception: {}", e.what());
|
|
} catch (...) {
|
|
}
|
|
} catch (...) {
|
|
try {
|
|
spdlog::warn("[ONNX] pool shutdown unknown exception");
|
|
} catch (...) {
|
|
}
|
|
}
|
|
try {
|
|
ctx.pool.reset();
|
|
} catch (const std::exception& e) {
|
|
try {
|
|
spdlog::warn("[ONNX] pool reset exception: {}", e.what());
|
|
} catch (...) {
|
|
}
|
|
} catch (...) {
|
|
try {
|
|
spdlog::warn("[ONNX] pool reset unknown exception");
|
|
} catch (...) {
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
~ProviderSingleton() noexcept {
|
|
shutdown(false); // Called from destructor, not explicit
|
|
}
|
|
|
|
ProviderSingleton() {
|
|
vtable.abi_version = 4; // v1.4 (added evict_under_pressure for memory pressure)
|
|
vtable.self = &ctx;
|
|
|
|
// Use daemon's default logger - no separate file needed
|
|
spdlog::debug("[ONNX Plugin] Initialized (v1.2) - using daemon logger");
|
|
|
|
vtable.set_progress_callback = [](void* self, yams_model_load_progress_cb cb,
|
|
void* user) -> yams_status_t {
|
|
if (!self)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->progress_cb = cb;
|
|
c->progress_user = user;
|
|
return YAMS_OK;
|
|
};
|
|
|
|
vtable.load_model = [](void* self, const char* model_id, const char* /*model_path*/,
|
|
const char* options_json) -> yams_status_t {
|
|
if (!self || !model_id)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
// Wrap entire lambda body in try/catch to prevent exceptions from crossing ABI
|
|
// boundary. On Windows, ONNX Runtime can throw std::system_error("resource deadlock
|
|
// would occur") when thread pool resources are exhausted (PBI-1c1).
|
|
try {
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled)
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
if (shutdownInProgress(c)) {
|
|
spdlog::info("[ONNX Plugin] load_model skipped during shutdown for '{}'",
|
|
model_id);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
if (!c->ready || !c->pool)
|
|
return YAMS_ERR_INTERNAL;
|
|
// Parse per-load options (hf.revision, offline)
|
|
try {
|
|
yams::daemon::OnnxModelPool::ResolutionHints hints;
|
|
if (options_json && std::strlen(options_json) > 0) {
|
|
auto j = nlohmann::json::parse(options_json, nullptr, false);
|
|
if (!j.is_discarded()) {
|
|
if (j.contains("hf") && j["hf"].is_object()) {
|
|
auto& hf = j["hf"];
|
|
if (hf.contains("revision") && hf["revision"].is_string())
|
|
hints.hfRevision = hf["revision"].get<std::string>();
|
|
if (hf.contains("offline") && hf["offline"].is_boolean())
|
|
hints.offlineOnly = hf["offline"].get<bool>();
|
|
}
|
|
if (j.contains("offline") && j["offline"].is_boolean())
|
|
hints.offlineOnly = j["offline"].get<bool>();
|
|
}
|
|
}
|
|
if (!hints.hfRevision.empty() || hints.offlineOnly) {
|
|
c->pool->setResolutionHints(model_id, hints);
|
|
}
|
|
} catch (...) {
|
|
}
|
|
std::optional<bool> hotModel;
|
|
std::string modelIdStr(model_id);
|
|
if (options_json && *options_json) {
|
|
try {
|
|
auto j = nlohmann::json::parse(options_json);
|
|
if (j.contains("hot") && j["hot"].is_boolean()) {
|
|
hotModel = j["hot"].get<bool>();
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
|
|
// FSM: if already loading/ready, return OK
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
auto& st = c->model_states[modelIdStr];
|
|
if (st == ProviderCtx::State::Ready) {
|
|
if (hotModel.has_value()) {
|
|
(void)c->pool->setModelHot(modelIdStr, hotModel.value());
|
|
}
|
|
spdlog::debug("[ONNX Plugin] load_model: '{}' already ready", model_id);
|
|
return YAMS_OK;
|
|
}
|
|
if (st == ProviderCtx::State::Loading) {
|
|
spdlog::debug("[ONNX Plugin] load_model: '{}' already loading", model_id);
|
|
return YAMS_OK;
|
|
}
|
|
// Check if in cooldown from previous failures
|
|
if (c->isInCooldown(modelIdStr)) {
|
|
auto& info = c->model_failures[modelIdStr];
|
|
spdlog::warn(
|
|
"[ONNX Plugin] load_model: '{}' in cooldown ({} failures), skipping",
|
|
model_id, info.consecutiveFailures);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
st = ProviderCtx::State::Loading;
|
|
}
|
|
spdlog::info("[ONNX Plugin] load_model: starting load for '{}'", model_id);
|
|
emit_progress(c, model_id, YAMS_MODEL_PHASE_PROBE, "probe");
|
|
|
|
// Simply call loadModel synchronously - it returns immediately since the pool uses
|
|
// lazy loading
|
|
auto r = c->pool->loadModel(modelIdStr);
|
|
if (r) {
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->model_states[modelIdStr] = ProviderCtx::State::Ready;
|
|
c->clearFailure(modelIdStr);
|
|
}
|
|
spdlog::info("[ONNX Plugin] load_model: '{}' loaded successfully", model_id);
|
|
if (hotModel.has_value()) {
|
|
(void)c->pool->setModelHot(modelIdStr, hotModel.value());
|
|
}
|
|
emit_progress(c, model_id, YAMS_MODEL_PHASE_READY, "ready");
|
|
auto q =
|
|
yams::daemon::InternalEventBus::instance()
|
|
.get_or_create_channel<yams::daemon::InternalEventBus::ModelReadyEvent>(
|
|
"model.events", 256);
|
|
(void)q->try_push({modelIdStr});
|
|
} else {
|
|
std::string errMsg = r.error().message;
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->model_states[modelIdStr] = ProviderCtx::State::Failed;
|
|
c->recordFailure(modelIdStr, errMsg);
|
|
}
|
|
spdlog::error("[ONNX Plugin] load_model: '{}' failed: {}", model_id, errMsg);
|
|
emit_progress(c, model_id, YAMS_MODEL_PHASE_UNKNOWN, errMsg.c_str());
|
|
auto q = yams::daemon::InternalEventBus::instance()
|
|
.get_or_create_channel<
|
|
yams::daemon::InternalEventBus::ModelLoadFailedEvent>(
|
|
"model.events", 256);
|
|
(void)q->try_push({std::string(model_id), r.error().message});
|
|
}
|
|
return YAMS_OK;
|
|
} catch (const std::system_error& e) {
|
|
fprintf(stderr, "[ONNX Plugin] load_model std::system_error: %s\n", e.what());
|
|
spdlog::error("[ONNX Plugin] load_model std::system_error: {}", e.what());
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (const std::exception& e) {
|
|
fprintf(stderr, "[ONNX Plugin] load_model exception: %s\n", e.what());
|
|
spdlog::error("[ONNX Plugin] load_model exception: {}", e.what());
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (...) {
|
|
fprintf(stderr, "[ONNX Plugin] load_model unknown exception\n");
|
|
spdlog::error("[ONNX Plugin] load_model unknown exception");
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
};
|
|
|
|
vtable.unload_model = [](void* self, const char* model_id) -> yams_status_t {
|
|
if (!self || !model_id)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled)
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
if (!c->ready || !c->pool)
|
|
return YAMS_ERR_INTERNAL;
|
|
auto r = c->pool->unloadModel(model_id);
|
|
return r ? YAMS_OK : YAMS_ERR_NOT_FOUND;
|
|
};
|
|
|
|
vtable.is_model_loaded = [](void* self, const char* model_id,
|
|
bool* out_loaded) -> yams_status_t {
|
|
if (!self || !model_id || !out_loaded)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled) {
|
|
*out_loaded = false;
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
}
|
|
if (!c->ready || !c->pool) {
|
|
*out_loaded = false;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
bool loaded = c->pool->isModelLoaded(model_id);
|
|
// Consider FSM Ready as loaded (pool may be lazy)
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
auto it = c->model_states.find(model_id);
|
|
if (!loaded && it != c->model_states.end() &&
|
|
it->second == ProviderCtx::State::Ready)
|
|
loaded = true;
|
|
}
|
|
*out_loaded = loaded;
|
|
return YAMS_OK;
|
|
};
|
|
|
|
vtable.get_loaded_models = [](void* self, const char*** out_ids,
|
|
size_t* out_count) -> yams_status_t {
|
|
if (!self || !out_ids || !out_count)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled)
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
if (!c->ready || !c->pool) {
|
|
*out_ids = nullptr;
|
|
*out_count = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
auto v = c->pool->getLoadedModels();
|
|
const size_t n = v.size();
|
|
const char** ids = nullptr;
|
|
if (n > 0) {
|
|
ids = static_cast<const char**>(std::malloc(sizeof(char*) * n));
|
|
if (!ids)
|
|
return YAMS_ERR_INTERNAL;
|
|
size_t i = 0;
|
|
for (const auto& s : v) {
|
|
char* dup = static_cast<char*>(std::malloc(s.size() + 1));
|
|
if (!dup) {
|
|
for (size_t j = 0; j < i; ++j)
|
|
std::free(const_cast<char*>(ids[j]));
|
|
std::free(ids);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
std::memcpy(dup, s.c_str(), s.size() + 1);
|
|
ids[i++] = dup;
|
|
}
|
|
}
|
|
*out_ids = ids;
|
|
*out_count = n;
|
|
return YAMS_OK;
|
|
};
|
|
|
|
vtable.free_model_list = [](void* /*self*/, const char** ids, size_t count) {
|
|
if (!ids)
|
|
return;
|
|
for (size_t i = 0; i < count; ++i)
|
|
std::free(const_cast<char*>(ids[i]));
|
|
std::free(const_cast<char**>(ids));
|
|
};
|
|
|
|
vtable.generate_embedding = [](void* self, const char* model_id, const uint8_t* input,
|
|
size_t input_len, float** out_vec,
|
|
size_t* out_dim) -> yams_status_t {
|
|
if (!self || !model_id || !input || !out_vec || !out_dim)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
// Wrap entire lambda body in try/catch to prevent exceptions from crossing ABI
|
|
// boundary. On Windows, ONNX Runtime can throw std::system_error("resource deadlock
|
|
// would occur") when thread pool resources are exhausted (PBI-1c1).
|
|
try {
|
|
auto* c = refreshCtxIfStaleShutdown(static_cast<ProviderCtx*>(self));
|
|
if (c->disabled) {
|
|
spdlog::warn("[ONNX Plugin] generate_embedding: plugin disabled");
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
}
|
|
if (shutdownInProgress(c)) {
|
|
spdlog::info("[ONNX Plugin] generate_embedding: shutdown in progress");
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
if (!c->ready) {
|
|
spdlog::warn("[ONNX Plugin] generate_embedding: plugin not ready");
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
if (!c->pool) {
|
|
spdlog::warn("[ONNX Plugin] generate_embedding: pool is null");
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
std::string modelIdStr(model_id);
|
|
|
|
// Check if model is in cooldown after previous failures
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
if (c->isInCooldown(modelIdStr)) {
|
|
auto& info = c->model_failures[modelIdStr];
|
|
spdlog::debug("[ONNX Plugin] Model '{}' in cooldown ({} failures)",
|
|
model_id, info.consecutiveFailures);
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
}
|
|
|
|
std::string text(reinterpret_cast<const char*>(input), input_len);
|
|
auto h = c->pool->acquireModel(model_id, modelAcquireTimeout(false));
|
|
if (!h) {
|
|
const auto err = h.error();
|
|
const bool acquireTimeout = (err.code == yams::ErrorCode::Timeout);
|
|
std::string errMsg = h.error().message;
|
|
spdlog::warn("[ONNX Plugin] acquireModel failed for {}: {}", model_id, errMsg);
|
|
|
|
// Resource-pool acquisition timeout is usually transient contention, not a
|
|
// model load failure. Do not poison cooldown state for this path.
|
|
if (acquireTimeout) {
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
// Record failure and fire event if this is first failure
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->recordFailure(modelIdStr, errMsg);
|
|
auto& info = c->model_failures[modelIdStr];
|
|
c->model_states[modelIdStr] = ProviderCtx::State::Failed;
|
|
|
|
if (!info.eventFired) {
|
|
info.eventFired = true;
|
|
spdlog::error("[ONNX Plugin] Model '{}' load failed, firing event: {}",
|
|
model_id, errMsg);
|
|
try {
|
|
auto q =
|
|
yams::daemon::InternalEventBus::instance()
|
|
.get_or_create_channel<
|
|
yams::daemon::InternalEventBus::ModelLoadFailedEvent>(
|
|
"model.events", 256);
|
|
(void)q->try_push({modelIdStr, errMsg});
|
|
} catch (...) {
|
|
}
|
|
}
|
|
}
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
// Model acquired successfully - clear failure tracking
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->clearFailure(modelIdStr);
|
|
c->model_states[modelIdStr] = ProviderCtx::State::Ready;
|
|
}
|
|
|
|
auto& session = *h.value();
|
|
auto r =
|
|
const_cast<yams::daemon::OnnxModelSession&>(session).generateEmbedding(text);
|
|
if (!r) {
|
|
spdlog::warn("[ONNX Plugin] generateEmbedding failed for {}: {}", model_id,
|
|
r.error().message);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
auto& vec = r.value();
|
|
const size_t n = vec.size();
|
|
float* buf = static_cast<float*>(std::malloc(sizeof(float) * n));
|
|
if (!buf)
|
|
return YAMS_ERR_INTERNAL;
|
|
std::memcpy(buf, vec.data(), sizeof(float) * n);
|
|
*out_vec = buf;
|
|
*out_dim = n;
|
|
return YAMS_OK;
|
|
} catch (const std::system_error& e) {
|
|
fprintf(stderr, "[ONNX Plugin] generate_embedding std::system_error: %s\n",
|
|
e.what());
|
|
spdlog::error("[ONNX Plugin] generate_embedding std::system_error: {}", e.what());
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (const std::exception& e) {
|
|
fprintf(stderr, "[ONNX Plugin] generate_embedding exception: %s\n", e.what());
|
|
spdlog::error("[ONNX Plugin] generate_embedding exception: {}", e.what());
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (...) {
|
|
fprintf(stderr, "[ONNX Plugin] generate_embedding unknown exception\n");
|
|
spdlog::error("[ONNX Plugin] generate_embedding unknown exception");
|
|
*out_vec = nullptr;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
};
|
|
|
|
vtable.free_embedding = [](void* /*self*/, float* vec, size_t /*dim*/) {
|
|
if (vec)
|
|
std::free(vec);
|
|
};
|
|
|
|
vtable.generate_embedding_batch = [](void* self, const char* model_id,
|
|
const uint8_t* const* inputs, const size_t* input_lens,
|
|
size_t batch_size, float** out_vecs, size_t* out_batch,
|
|
size_t* out_dim) -> yams_status_t {
|
|
spdlog::debug("[ONNX Plugin] generate_embedding_batch called: model={} batch={}",
|
|
model_id ? model_id : "(null)", batch_size);
|
|
|
|
if (!self || !model_id || !inputs || !input_lens || !out_vecs || !out_batch ||
|
|
!out_dim) {
|
|
spdlog::error("[ONNX Plugin] Invalid arguments");
|
|
return YAMS_ERR_INVALID_ARG;
|
|
}
|
|
|
|
auto* c = refreshCtxIfStaleShutdown(static_cast<ProviderCtx*>(self));
|
|
int retries = shutdownInProgress(c) ? 1 : 3;
|
|
while (retries > 0) {
|
|
retries--;
|
|
try {
|
|
c = refreshCtxIfStaleShutdown(c);
|
|
const bool timingDebug = []() {
|
|
if (const char* s = std::getenv("YAMS_ONNX_DEBUG_BATCH_TIMINGS")) {
|
|
return std::string{s} == "1" || std::string{s} == "true" ||
|
|
std::string{s} == "yes";
|
|
}
|
|
return false;
|
|
}();
|
|
|
|
if (c->disabled) {
|
|
spdlog::warn("[ONNX Plugin] generate_embedding_batch: plugin disabled");
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
}
|
|
if (shutdownInProgress(c)) {
|
|
spdlog::info(
|
|
"[ONNX Plugin] generate_embedding_batch: shutdown in progress");
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
std::string modelIdStr(model_id);
|
|
const bool isColbert = isColbertModelName(modelIdStr);
|
|
if (!c->ready && !isColbert) {
|
|
spdlog::warn("[ONNX Plugin] generate_embedding_batch: plugin not ready");
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
if (!c->pool && !isColbert) {
|
|
spdlog::warn("[ONNX Plugin] generate_embedding_batch: pool is null");
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
// Check if model is in cooldown
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
if (c->isInCooldown(modelIdStr)) {
|
|
auto& info = c->model_failures[modelIdStr];
|
|
spdlog::warn("[ONNX Plugin] Model '{}' in cooldown ({} failures)",
|
|
model_id, info.consecutiveFailures);
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
}
|
|
|
|
std::vector<std::string> texts;
|
|
texts.reserve(batch_size);
|
|
for (size_t i = 0; i < batch_size; ++i) {
|
|
texts.emplace_back(reinterpret_cast<const char*>(inputs[i]), input_lens[i]);
|
|
}
|
|
|
|
// Batch coalescing: route small non-ColBERT requests through
|
|
// the coalescer to accumulate concurrent small batches.
|
|
// Only coalesce small requests (<=4 texts); callers that
|
|
// already send well-sized batches skip this to avoid
|
|
// serialization overhead and the coalesce window delay.
|
|
constexpr size_t kCoalesceMaxInputBatch = 4;
|
|
if (!isColbert && c->coalescer && c->coalescer->isEnabled() &&
|
|
batch_size <= kCoalesceMaxInputBatch) {
|
|
auto inferFn = [c](const std::string& mid,
|
|
const std::vector<std::string>& batch)
|
|
-> yams::Result<std::vector<std::vector<float>>> {
|
|
auto h2 = c->pool->acquireModel(mid.c_str(), modelAcquireTimeout(true));
|
|
if (!h2)
|
|
return h2.error();
|
|
return const_cast<yams::daemon::OnnxModelSession&>(*h2.value())
|
|
.generateBatchEmbeddings(batch);
|
|
};
|
|
auto r = c->coalescer->submit(modelIdStr, std::move(texts), inferFn);
|
|
if (!r) {
|
|
spdlog::error(
|
|
"[ONNX Plugin] coalesced generateBatchEmbeddings failed: {}",
|
|
r.error().message);
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
auto& mat = r.value();
|
|
if (mat.empty()) {
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_OK;
|
|
}
|
|
const size_t b = mat.size();
|
|
const size_t d = mat[0].size();
|
|
float* buf = static_cast<float*>(std::malloc(sizeof(float) * b * d));
|
|
if (!buf)
|
|
return YAMS_ERR_INTERNAL;
|
|
for (size_t i = 0; i < b; ++i) {
|
|
std::memcpy(buf + i * d, mat[i].data(), sizeof(float) * d);
|
|
}
|
|
*out_vecs = buf;
|
|
*out_batch = b;
|
|
*out_dim = d;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
spdlog::debug("[ONNX Plugin] acquiring model '{}'...", model_id);
|
|
const auto tAcquire = std::chrono::steady_clock::now();
|
|
auto h = c->pool->acquireModel(model_id, modelAcquireTimeout(true));
|
|
const auto acquireMs = std::chrono::duration_cast<std::chrono::milliseconds>(
|
|
std::chrono::steady_clock::now() - tAcquire)
|
|
.count();
|
|
if (timingDebug) {
|
|
spdlog::info("[ONNX Plugin] batch_timing acquire model='{}' batch={} ms={} "
|
|
"retries_left={}",
|
|
model_id, batch_size, acquireMs, retries);
|
|
}
|
|
if (!h) {
|
|
const auto err = h.error();
|
|
const bool acquireTimeout = (err.code == yams::ErrorCode::Timeout);
|
|
std::string errMsg = h.error().message;
|
|
spdlog::error("[ONNX Plugin] acquireModel failed for '{}': {}", model_id,
|
|
errMsg);
|
|
|
|
// Timeout while waiting for a session is transient contention. Retry
|
|
// without tripping model cooldown/failure bookkeeping.
|
|
if (acquireTimeout) {
|
|
if (retries > 0) {
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(250));
|
|
continue;
|
|
}
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
if (shutdownInProgress(c)) {
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->recordFailure(modelIdStr, errMsg);
|
|
auto& info = c->model_failures[modelIdStr];
|
|
c->model_states[modelIdStr] = ProviderCtx::State::Failed;
|
|
|
|
if (!info.eventFired) {
|
|
info.eventFired = true;
|
|
try {
|
|
auto q =
|
|
yams::daemon::InternalEventBus::instance()
|
|
.get_or_create_channel<yams::daemon::InternalEventBus::
|
|
ModelLoadFailedEvent>(
|
|
"model.events", 256);
|
|
(void)q->try_push({modelIdStr, errMsg});
|
|
} catch (...) {
|
|
}
|
|
}
|
|
}
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->clearFailure(modelIdStr);
|
|
c->model_states[modelIdStr] = ProviderCtx::State::Ready;
|
|
}
|
|
|
|
auto& session = *h.value();
|
|
const auto tInfer = std::chrono::steady_clock::now();
|
|
yams::Result<std::vector<std::vector<float>>> r;
|
|
if (isColbert) {
|
|
if (!c->colbert || !c->colbert->isValid()) {
|
|
r = yams::Error{yams::ErrorCode::NotSupported,
|
|
"ColBERT session not available for embeddings"};
|
|
} else {
|
|
r = c->colbert->encodeDocumentEmbeddings(texts);
|
|
}
|
|
} else {
|
|
r = const_cast<yams::daemon::OnnxModelSession&>(session)
|
|
.generateBatchEmbeddings(texts);
|
|
}
|
|
if (!r) {
|
|
spdlog::error("[ONNX Plugin] generateBatchEmbeddings failed: {}",
|
|
r.error().message);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
const auto inferMs = std::chrono::duration_cast<std::chrono::milliseconds>(
|
|
std::chrono::steady_clock::now() - tInfer)
|
|
.count();
|
|
auto& mat = r.value();
|
|
if (mat.empty()) {
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_OK;
|
|
}
|
|
const size_t b = mat.size();
|
|
const size_t d = mat[0].size();
|
|
float* buf = static_cast<float*>(std::malloc(sizeof(float) * b * d));
|
|
if (!buf)
|
|
return YAMS_ERR_INTERNAL;
|
|
for (size_t i = 0; i < b; ++i) {
|
|
std::memcpy(buf + i * d, mat[i].data(), sizeof(float) * d);
|
|
}
|
|
*out_vecs = buf;
|
|
*out_batch = b;
|
|
*out_dim = d;
|
|
if (timingDebug) {
|
|
spdlog::info(
|
|
"[ONNX Plugin] batch_timing infer model='{}' batch={} out_batch={} "
|
|
"dim={} acquire_ms={} infer_ms={}",
|
|
model_id, batch_size, b, d, acquireMs, inferMs);
|
|
}
|
|
return YAMS_OK;
|
|
|
|
} catch (const std::system_error& e) {
|
|
spdlog::error("[ONNX Plugin] system_error: {} (code={}). Retries: {}", e.what(),
|
|
e.code().value(), retries);
|
|
if (retries > 0) {
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(200));
|
|
continue;
|
|
}
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (const std::exception& e) {
|
|
spdlog::error("[ONNX Plugin] exception: {}", e.what());
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (...) {
|
|
spdlog::error("[ONNX Plugin] unknown exception");
|
|
*out_vecs = nullptr;
|
|
*out_batch = 0;
|
|
*out_dim = 0;
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
} // end while
|
|
|
|
return YAMS_ERR_INTERNAL;
|
|
};
|
|
|
|
vtable.free_embedding_batch = [](void* /*self*/, float* vecs, size_t /*batch*/,
|
|
size_t /*dim*/) {
|
|
if (vecs)
|
|
std::free(vecs);
|
|
};
|
|
|
|
// v1.2: get_embedding_dim (non-loading path)
|
|
vtable.get_embedding_dim = [](void* self, const char* model_id,
|
|
size_t* out_dim) -> yams_status_t {
|
|
if (!self || !model_id || !out_dim)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (!c->pool)
|
|
return YAMS_ERR_INTERNAL;
|
|
|
|
// 1) If a hot session already exists, read its dim without blocking
|
|
if (c->ready && c->pool->isModelLoaded(model_id)) {
|
|
auto h = c->pool->acquireModel(model_id, std::chrono::milliseconds(0));
|
|
if (h) {
|
|
*out_dim = h.value()->getEmbeddingDim();
|
|
if (*out_dim > 0)
|
|
return YAMS_OK;
|
|
}
|
|
}
|
|
|
|
// 2) Probe model directory metadata without loading the model
|
|
namespace fs = std::filesystem;
|
|
size_t dim = 0;
|
|
try {
|
|
auto read_dim = [&](const fs::path& p) -> size_t {
|
|
try {
|
|
if (!fs::exists(p))
|
|
return 0u;
|
|
std::ifstream in(p);
|
|
if (!in)
|
|
return 0u;
|
|
nlohmann::json j;
|
|
in >> j;
|
|
if (j.contains("embedding_dimension"))
|
|
return j.value("embedding_dimension", 0u);
|
|
if (j.contains("embedding_dim"))
|
|
return j.value("embedding_dim", 0u);
|
|
// Standard models use hidden_size, Nomic models use n_embd
|
|
if (j.contains("hidden_size"))
|
|
return j.value("hidden_size", 0u);
|
|
if (j.contains("n_embd"))
|
|
return j.value("n_embd", 0u);
|
|
} catch (...) {
|
|
}
|
|
return 0u;
|
|
};
|
|
|
|
// Candidate roots: XDG_DATA_HOME/yams/models, HOME/.local/share/yams/models
|
|
std::vector<fs::path> roots;
|
|
if (const char* xdg = std::getenv("XDG_DATA_HOME"))
|
|
roots.emplace_back(fs::path(xdg) / "yams" / "models");
|
|
if (const char* home = std::getenv("HOME"))
|
|
roots.emplace_back(fs::path(home) / ".local" / "share" / "yams" / "models");
|
|
|
|
for (const auto& r : roots) {
|
|
fs::path base = r / model_id;
|
|
for (const auto& fn :
|
|
{"sentence_bert_config.json", "config.json", "model.onnx.yams.meta.json",
|
|
"config.json.yams.meta.json",
|
|
"sentence_bert_config.json.yams.meta.json"}) {
|
|
dim = read_dim(base / fn);
|
|
if (dim > 0)
|
|
break;
|
|
}
|
|
if (dim > 0)
|
|
break;
|
|
}
|
|
} catch (...) {
|
|
}
|
|
if (dim > 0) {
|
|
*out_dim = dim;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
// 3) Heuristic from common embedding model names (no load)
|
|
std::string lm(model_id);
|
|
for (auto& ch : lm)
|
|
ch = static_cast<char>(std::tolower(ch));
|
|
if (lm.find("minilm") != std::string::npos)
|
|
dim = 384;
|
|
else if (lm.find("mpnet") != std::string::npos ||
|
|
lm.find("nomic") != std::string::npos || lm.find("jina") != std::string::npos)
|
|
dim = 768;
|
|
else if (lm.find("bge-m3") != std::string::npos ||
|
|
(lm.find("bge") != std::string::npos && lm.find("large") != std::string::npos))
|
|
dim = 1024;
|
|
else if (lm.find("bge") != std::string::npos && lm.find("small") != std::string::npos)
|
|
dim = 384;
|
|
else if (lm.find("colbert") != std::string::npos)
|
|
dim = 48;
|
|
else if (lm.find("e5") != std::string::npos && lm.find("large") != std::string::npos)
|
|
dim = 1024;
|
|
else if (lm.find("e5") != std::string::npos && lm.find("small") != std::string::npos)
|
|
dim = 384;
|
|
else if (lm.find("gte") != std::string::npos && lm.find("large") != std::string::npos)
|
|
dim = 1024;
|
|
else if (lm.find("gte") != std::string::npos && lm.find("small") != std::string::npos)
|
|
dim = 384;
|
|
else if (lm.find("text-embedding-3-large") != std::string::npos)
|
|
dim = 3072;
|
|
else if (lm.find("text-embedding-3-small") != std::string::npos)
|
|
dim = 1536;
|
|
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
|
|
else if (lm.find("test-model") != std::string::npos &&
|
|
(std::getenv("YAMS_TEST_MODE") || std::getenv("YAMS_USE_MOCK_PROVIDER")))
|
|
dim = 384;
|
|
#endif
|
|
|
|
if (dim > 0) {
|
|
*out_dim = dim;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
*out_dim = 0;
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
};
|
|
|
|
// v1.2: get_runtime_info_json
|
|
vtable.get_runtime_info_json = [](void* self, const char* model_id,
|
|
char** out_json) -> yams_status_t {
|
|
if (!self || !model_id || !out_json)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (!c->pool)
|
|
return YAMS_ERR_INTERNAL;
|
|
nlohmann::json j;
|
|
j["backend"] = "onnxruntime";
|
|
j["pipeline"] = "raw_ort";
|
|
j["model"] = model_id;
|
|
// Best-effort dimension and runtime hints
|
|
j["graph_optimization"] = "enabled";
|
|
|
|
// Best-effort: report ONNX Runtime build's available providers
|
|
auto providers = yams::onnx_util::OrtRuntimeLoader::instance().availableProviders();
|
|
if (!providers.empty()) {
|
|
auto arr = nlohmann::json::array();
|
|
for (const auto& p : providers)
|
|
arr.push_back(p);
|
|
j["available_providers"] = std::move(arr);
|
|
}
|
|
|
|
const auto& runtimeInfo = c->runtimeInfo;
|
|
if (!runtimeInfo.version.empty())
|
|
j["runtime_version"] = runtimeInfo.version;
|
|
if (!runtimeInfo.libraryPath.empty())
|
|
j["runtime_library"] = runtimeInfo.libraryPath;
|
|
if (!runtimeInfo.reason.empty())
|
|
j["runtime_status"] = runtimeInfo.reason;
|
|
if (!runtimeInfo.source.empty())
|
|
j["runtime_source"] = runtimeInfo.source;
|
|
|
|
// Report the actual execution provider used by a session (not a platform guess)
|
|
std::string ep = "cpu";
|
|
if (c->ready && c->pool) {
|
|
try {
|
|
auto h = c->pool->acquireModel(model_id, std::chrono::milliseconds(500));
|
|
if (h) {
|
|
auto* session = h.value().operator->();
|
|
if (!session) {
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
ep = session->getExecutionProvider();
|
|
auto [intraThreads, interThreads] = session->getThreading();
|
|
j["intra_threads"] = intraThreads;
|
|
j["inter_threads"] = interThreads;
|
|
j["learned_batch_cap"] = session->getLearnedBatchLimit();
|
|
{
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
c->actualExecutionProvider = ep;
|
|
}
|
|
}
|
|
} catch (...) {
|
|
}
|
|
}
|
|
if (ep == "cpu") {
|
|
std::lock_guard<std::mutex> lk(c->mu);
|
|
if (!c->actualExecutionProvider.empty())
|
|
ep = c->actualExecutionProvider;
|
|
}
|
|
j["execution_provider"] = ep;
|
|
size_t dim = 0;
|
|
if (c->ready && c->pool) {
|
|
auto h = c->pool->acquireModel(model_id, std::chrono::seconds(2));
|
|
if (h) {
|
|
auto* session = h.value().operator->();
|
|
if (!session) {
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
dim = session->getEmbeddingDim();
|
|
try {
|
|
const auto& info = session->getInfo();
|
|
if (!info.path.empty())
|
|
j["path"] = info.path;
|
|
} catch (...) {
|
|
}
|
|
}
|
|
}
|
|
j["dim"] = dim;
|
|
if (!j.contains("intra_threads")) {
|
|
j["intra_threads"] =
|
|
static_cast<int>(std::max(1u, std::thread::hardware_concurrency()));
|
|
}
|
|
if (!j.contains("inter_threads")) {
|
|
j["inter_threads"] = 1;
|
|
}
|
|
// Source hint from path when present
|
|
try {
|
|
if (j.contains("path") && j["path"].is_string()) {
|
|
std::string p = j["path"].get<std::string>();
|
|
if (p.find("/.cache/huggingface/") != std::string::npos)
|
|
j["source"] = "huggingface-cache";
|
|
else if (p.find(".yams/models/") != std::string::npos)
|
|
j["source"] = "yams-models";
|
|
else if (p.rfind("http://", 0) == 0 || p.rfind("https://", 0) == 0)
|
|
j["source"] = "remote";
|
|
else
|
|
j["source"] = "local";
|
|
}
|
|
} catch (...) {
|
|
}
|
|
std::string s;
|
|
try {
|
|
s = j.dump();
|
|
} catch (...) {
|
|
s = "{}";
|
|
}
|
|
char* buf = static_cast<char*>(std::malloc(s.size() + 1));
|
|
if (!buf)
|
|
return YAMS_ERR_INTERNAL;
|
|
std::memcpy(buf, s.c_str(), s.size() + 1);
|
|
*out_json = buf;
|
|
return YAMS_OK;
|
|
};
|
|
|
|
vtable.free_string = [](void* /*self*/, char* s) {
|
|
if (s)
|
|
std::free(s);
|
|
};
|
|
|
|
// Assign function pointer for set_threading
|
|
vtable.set_threading = &onnx_set_threading;
|
|
|
|
// v1.3: Cross-encoder reranking
|
|
vtable.score_documents = [](void* self, const char* /*reranker_model_id*/,
|
|
const char* query, const char* const* documents,
|
|
size_t doc_count, float** out_scores,
|
|
size_t* out_count) -> yams_status_t {
|
|
if (!self || !query || !documents || !out_scores || !out_count)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
if (doc_count == 0) {
|
|
*out_scores = nullptr;
|
|
*out_count = 0;
|
|
return YAMS_OK;
|
|
}
|
|
|
|
try {
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled)
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
|
|
if (c->colbert && c->colbert->isValid()) {
|
|
auto queryEmb = c->colbert->encodeQuery(query);
|
|
if (!queryEmb) {
|
|
spdlog::warn("[ONNX Plugin] ColBERT encode query failed: {}",
|
|
queryEmb.error().message);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
std::vector<std::string> docs;
|
|
docs.reserve(doc_count);
|
|
for (size_t i = 0; i < doc_count; ++i) {
|
|
docs.emplace_back(documents[i] ? documents[i] : "");
|
|
}
|
|
auto docEmbeddings = c->colbert->encodeDocuments(docs);
|
|
if (!docEmbeddings) {
|
|
spdlog::warn("[ONNX Plugin] ColBERT encode documents failed: {}",
|
|
docEmbeddings.error().message);
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
|
|
const auto& docVecs = docEmbeddings.value();
|
|
std::vector<float> scores;
|
|
scores.reserve(docVecs.size());
|
|
for (const auto& docEmb : docVecs) {
|
|
scores.push_back(c->colbert->computeMaxSim(queryEmb.value(), docEmb));
|
|
}
|
|
|
|
float* buf = static_cast<float*>(std::malloc(sizeof(float) * scores.size()));
|
|
if (!buf)
|
|
return YAMS_ERR_INTERNAL;
|
|
std::memcpy(buf, scores.data(), sizeof(float) * scores.size());
|
|
*out_scores = buf;
|
|
*out_count = scores.size();
|
|
return YAMS_OK;
|
|
}
|
|
|
|
spdlog::debug("[ONNX Plugin] Cross-encoder reranker removed; ColBERT is the only "
|
|
"ONNX reranker path");
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
|
|
} catch (const std::exception& e) {
|
|
spdlog::error("[ONNX Plugin] score_documents exception: {}", e.what());
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (...) {
|
|
spdlog::error("[ONNX Plugin] score_documents unknown exception");
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
};
|
|
|
|
vtable.free_scores = [](void* /*self*/, float* scores, size_t /*count*/) {
|
|
if (scores)
|
|
std::free(scores);
|
|
};
|
|
|
|
// v1.4: Memory pressure eviction
|
|
vtable.evict_under_pressure = [](void* self, double pressure_level, bool allow_hot_eviction,
|
|
size_t* out_evicted) -> yams_status_t {
|
|
if (!self || !out_evicted)
|
|
return YAMS_ERR_INVALID_ARG;
|
|
*out_evicted = 0;
|
|
|
|
try {
|
|
auto* c = static_cast<ProviderCtx*>(self);
|
|
if (c->disabled)
|
|
return YAMS_ERR_UNSUPPORTED;
|
|
if (!c->ready || !c->pool)
|
|
return YAMS_ERR_INTERNAL;
|
|
|
|
// Get profile-aware eviction thresholds from TuneAdvisor
|
|
const double warningThresh =
|
|
yams::daemon::TuneAdvisor::modelEvictWarningThreshold();
|
|
const double criticalThresh =
|
|
yams::daemon::TuneAdvisor::modelEvictCriticalThreshold();
|
|
const double emergencyThresh =
|
|
yams::daemon::TuneAdvisor::modelEvictEmergencyThreshold();
|
|
|
|
// Determine how many models to evict based on pressure level:
|
|
// - Emergency: evict all non-hot (or all if allow_hot_eviction)
|
|
// - Critical: evict up to 2 models
|
|
// - Warning: evict 1 model
|
|
size_t numToEvict = 0;
|
|
if (pressure_level >= emergencyThresh) {
|
|
numToEvict = allow_hot_eviction
|
|
? c->pool->getLoadedModelCount()
|
|
: std::max<size_t>(1, c->pool->getLoadedModelCount() / 2);
|
|
} else if (pressure_level >= criticalThresh) {
|
|
numToEvict = 2;
|
|
} else if (pressure_level >= warningThresh) {
|
|
numToEvict = 1;
|
|
}
|
|
|
|
if (numToEvict == 0) {
|
|
// pressure below warning threshold: do idle maintenance instead
|
|
c->pool->performMaintenance();
|
|
return YAMS_OK;
|
|
}
|
|
|
|
// Use existing evictLRU method
|
|
size_t countBefore = c->pool->getLoadedModelCount();
|
|
c->pool->evictLRU(numToEvict);
|
|
size_t countAfter = c->pool->getLoadedModelCount();
|
|
|
|
size_t evicted = (countBefore > countAfter) ? (countBefore - countAfter) : 0;
|
|
*out_evicted = evicted;
|
|
|
|
if (evicted > 0) {
|
|
spdlog::info("[ONNX Plugin] Evicted {} models under memory pressure ({:.0f}%)",
|
|
evicted, pressure_level * 100.0);
|
|
}
|
|
|
|
return YAMS_OK;
|
|
} catch (const std::exception& e) {
|
|
spdlog::error("[ONNX Plugin] evict_under_pressure exception: {}", e.what());
|
|
return YAMS_ERR_INTERNAL;
|
|
} catch (...) {
|
|
spdlog::error("[ONNX Plugin] evict_under_pressure unknown exception");
|
|
return YAMS_ERR_INTERNAL;
|
|
}
|
|
};
|
|
}
|
|
};
|
|
|
|
static ProviderSingleton& singleton() {
|
|
static ProviderSingleton s;
|
|
return s;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// Set the plugin config JSON (called from yams_plugin_init before provider is accessed)
|
|
extern "C" void yams_onnx_set_config_json(const char* json) {
|
|
if (json && *json) {
|
|
g_plugin_config_json = json;
|
|
spdlog::info("[ONNX Plugin] Config JSON set: {}", g_plugin_config_json);
|
|
// Optional pin for the ONNX Runtime copy ([plugins.onnx] runtime_library).
|
|
auto cfg = nlohmann::json::parse(g_plugin_config_json, nullptr, false);
|
|
if (cfg.is_object() && cfg.contains("runtime_library") &&
|
|
cfg["runtime_library"].is_string()) {
|
|
yams::onnx_util::OrtRuntimeLoader::instance().setConfiguredLibrary(
|
|
cfg["runtime_library"].get<std::string>());
|
|
}
|
|
}
|
|
}
|
|
|
|
extern "C" void yams_onnx_shutdown_provider() {
|
|
singleton().explicitShutdownCalled.store(true);
|
|
singleton().shutdown(true);
|
|
}
|
|
|
|
extern "C" yams_model_provider_v1* yams_onnx_get_model_provider() {
|
|
auto& s = singleton();
|
|
s.ensureActive();
|
|
return &s.vtable;
|
|
}
|
|
|
|
// Provide a small JSON health snapshot for the ABI host
|
|
extern "C" const char* yams_onnx_get_health_json_cstr() {
|
|
static thread_local std::string json;
|
|
auto& c = singleton().ctx;
|
|
// Build a minimal JSON reflecting readiness/disabled and last error
|
|
nlohmann::json j;
|
|
if (c.disabled) {
|
|
j["status"] = "disabled";
|
|
j["reason"] = c.last_error.empty() ? "disabled_by_env" : c.last_error;
|
|
} else if (!c.runtimeInfo.available) {
|
|
j["status"] = "unavailable";
|
|
j["reason"] = c.last_error.empty() ? "runtime_not_available" : c.last_error;
|
|
} else if (!c.pool) {
|
|
j["status"] = "unavailable";
|
|
j["reason"] = "no_pool";
|
|
} else if (!c.ready) {
|
|
j["status"] = "degraded";
|
|
j["reason"] = c.last_error.empty() ? "init_failed" : c.last_error;
|
|
} else {
|
|
j["status"] = "ok";
|
|
}
|
|
|
|
// Execution provider + available providers (best-effort)
|
|
{
|
|
std::string ep = "cpu";
|
|
{
|
|
std::lock_guard<std::mutex> lk(c.mu);
|
|
ep = c.actualExecutionProvider;
|
|
}
|
|
|
|
if (c.pool) {
|
|
try {
|
|
// Try a quick acquire of any ready model to get the real EP.
|
|
std::string candidateModelId;
|
|
{
|
|
std::lock_guard<std::mutex> lk(c.mu);
|
|
for (const auto& [modelId, state] : c.model_states) {
|
|
if (state == ProviderCtx::State::Ready) {
|
|
candidateModelId = modelId;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!candidateModelId.empty()) {
|
|
auto h = c.pool->acquireModel(candidateModelId, std::chrono::milliseconds(50));
|
|
if (h) {
|
|
ep = h.value()->getExecutionProvider();
|
|
std::lock_guard<std::mutex> lk(c.mu);
|
|
c.actualExecutionProvider = ep;
|
|
}
|
|
}
|
|
} catch (...) {
|
|
// Keep cached EP on any failure.
|
|
}
|
|
}
|
|
|
|
j["execution_provider"] = ep;
|
|
|
|
auto providers = yams::onnx_util::OrtRuntimeLoader::instance().availableProviders();
|
|
if (!providers.empty()) {
|
|
j["available_providers"] = providers;
|
|
}
|
|
}
|
|
|
|
if (!c.runtimeInfo.version.empty()) {
|
|
j["runtime_version"] = c.runtimeInfo.version;
|
|
}
|
|
if (!c.runtimeInfo.libraryPath.empty()) {
|
|
j["runtime_library"] = c.runtimeInfo.libraryPath;
|
|
}
|
|
if (!c.runtimeInfo.reason.empty()) {
|
|
j["runtime_status"] = c.runtimeInfo.reason;
|
|
}
|
|
if (!c.runtimeInfo.errorMessage.empty()) {
|
|
j["runtime_error"] = c.runtimeInfo.errorMessage;
|
|
}
|
|
if (!c.runtimeInfo.source.empty()) {
|
|
j["runtime_source"] = c.runtimeInfo.source;
|
|
}
|
|
|
|
// Add model states for diagnostics
|
|
{
|
|
std::lock_guard<std::mutex> lk(c.mu);
|
|
if (!c.model_states.empty() || !c.model_failures.empty()) {
|
|
nlohmann::json models = nlohmann::json::object();
|
|
for (const auto& [modelId, state] : c.model_states) {
|
|
nlohmann::json m;
|
|
switch (state) {
|
|
case ProviderCtx::State::Unloaded:
|
|
m["state"] = "unloaded";
|
|
break;
|
|
case ProviderCtx::State::Loading:
|
|
m["state"] = "loading";
|
|
break;
|
|
case ProviderCtx::State::Ready:
|
|
m["state"] = "ready";
|
|
break;
|
|
case ProviderCtx::State::Failed:
|
|
m["state"] = "failed";
|
|
break;
|
|
}
|
|
// Add failure info if present
|
|
auto fit = c.model_failures.find(modelId);
|
|
if (fit != c.model_failures.end() && fit->second.consecutiveFailures > 0) {
|
|
m["failures"] = fit->second.consecutiveFailures;
|
|
m["last_error"] = fit->second.lastError;
|
|
m["in_cooldown"] = c.isInCooldown(modelId);
|
|
}
|
|
models[modelId] = m;
|
|
}
|
|
j["models"] = std::move(models);
|
|
}
|
|
}
|
|
|
|
try {
|
|
json = j.dump();
|
|
} catch (...) {
|
|
json = "{\"status\":\"unknown\"}";
|
|
}
|
|
return json.c_str();
|
|
}
|
|
|
|
#ifdef YAMS_TESTING
|
|
extern "C" void yams_onnx_test_get_pool_config(std::size_t* max_loaded, std::size_t* hot_pool) {
|
|
auto& c = singleton().ctx;
|
|
if (max_loaded)
|
|
*max_loaded = c.configuredMaxLoadedModels;
|
|
if (hot_pool)
|
|
*hot_pool = c.configuredHotPoolSize;
|
|
}
|
|
#endif
|