From a5355a02269570c4810105168e1bf779d3912a84 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Thu, 16 Apr 2026 13:40:13 +0200 Subject: [PATCH] server: keep router model refcount to avoid unloading models that have running requests this avoids a deadlock when models A and B don't fit together, but both have requests, so the server gets into a loop unloading A, loading B, unloading B, loading A again, and so on --- tools/server/server-http.h | 8 +- tools/server/server-models.cpp | 66 ++++++++++++- tools/server/server-models.h | 7 ++ tools/server/tests/unit/test_router.py | 124 +++++++++++++++++++++++++ 4 files changed, 200 insertions(+), 5 deletions(-) diff --git a/tools/server/server-http.h b/tools/server/server-http.h index 68ae2170cf..42ea8a8e99 100644 --- a/tools/server/server-http.h +++ b/tools/server/server-http.h @@ -28,7 +28,13 @@ struct server_http_res { return next != nullptr; } - virtual ~server_http_res() = default; + std::function on_destroy = nullptr; + + virtual ~server_http_res() { + if (on_destroy) { + on_destroy(); + } + } }; // unique pointer, used by set_chunked_content_provider diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 23ba9c944a..379b01a4f0 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -517,6 +517,19 @@ std::vector server_models::get_all_meta() { return result; } +void server_models::inc_refs(const std::string & name) { + std::lock_guard lk(mutex); + mapping[name].active_refs++; +} + +void server_models::dec_refs(const std::string & name) { + { + std::lock_guard lk(mutex); + mapping[name].active_refs--; + } + cv.notify_all(); +} + int server_models::can_fit(const device_memory_map & dmm_req) const { device_memory_map dmm_total; for (const auto & m : mapping) { @@ -573,7 +586,8 @@ void server_models::unload_lru(const device_memory_map & dmm_req) { for (const auto & m : mapping) { if (m.second.meta.is_running()) { count_active++; - if (m.second.meta.last_used < lru_last_used) { + // Only consider idle models + if (m.second.active_refs == 0 && m.second.meta.last_used < lru_last_used) { lru_model_name = m.first; lru_last_used = m.second.meta.last_used; } @@ -598,6 +612,21 @@ void server_models::unload_lru(const device_memory_map & dmm_req) { return mapping[lru_model_name].meta.status == SERVER_MODEL_STATUS_UNLOADED; }); } + } else if (count_active > 0 && (active_exceeded || memory_exceeded)) { + // No model idle, wait for drain + std::unique_lock lk(mutex); + bool drained = cv.wait_for(lk, std::chrono::seconds(DEFAULT_STOP_TIMEOUT), [this]() { + for (const auto & m : mapping) { + if (m.second.meta.is_running() && m.second.active_refs == 0) { + return true; + } + } + return false; + }); + if (!drained) { + SRV_WRN("%s", "drain timeout, falling back to force eviction\n"); + break; + } } else { break; } @@ -833,6 +862,7 @@ void server_models::_load(const std::string & name, const device_memory_map & dm inst.meta.port = get_free_port(); inst.meta.status = SERVER_MODEL_STATUS_LOADING; inst.meta.last_used = ggml_time_ms(); + inst.active_refs = mapping[name].active_refs; if (inst.meta.port <= 0) { throw std::runtime_error("failed to get a port number"); @@ -1168,10 +1198,18 @@ static bool router_validate_model(std::string & name, server_models & models, bo } // resolve alias to canonical model name name = meta->name; + // To avoid unloading a model before it is loaded, protect with increased ref count before it starts loading + models.inc_refs(name); if (models_autoload) { - models.ensure_model_ready(name); + try { + models.ensure_model_ready(name); + } catch (...) { + models.dec_refs(name); + throw; + } } else { if (!meta->is_running()) { + models.dec_refs(name); res_err(res, format_error_response("model is not loaded", ERROR_TYPE_INVALID_REQUEST)); return false; } @@ -1222,7 +1260,17 @@ void server_models_routes::init_routes() { if (!router_validate_model(name, models, autoload, error_res)) { return error_res; } - return models.proxy_request(req, method, name, false); + server_http_res_ptr proxy; + try { + proxy = models.proxy_request(req, method, name, false); + } catch(...) { + models.dec_refs(name); + throw; + } + proxy->on_destroy = [this, name]() { + this->models.dec_refs(name); + }; + return proxy; }; this->proxy_post = [this](const server_http_req & req) { @@ -1234,7 +1282,17 @@ void server_models_routes::init_routes() { if (!router_validate_model(name, models, autoload, error_res)) { return error_res; } - return models.proxy_request(req, method, name, true); // update last usage for POST request only + server_http_res_ptr proxy; + try { + proxy = models.proxy_request(req, method, name, true); // update last usage for POST request only + } catch(...) { + models.dec_refs(name); + throw; + } + proxy->on_destroy = [this, name]() { + this->models.dec_refs(name); + }; + return proxy; }; this->post_router_models_load = [this](const server_http_req & req) { diff --git a/tools/server/server-models.h b/tools/server/server-models.h index aa6abf7cac..36cd0296f6 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -100,6 +100,7 @@ private: std::thread th; server_model_meta meta; FILE * stdin_file = nullptr; + uint64_t active_refs = 0; }; std::mutex mutex; @@ -174,6 +175,12 @@ public: // proxy an HTTP request to the model instance server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used); + // Increase instance ref counter + void inc_refs(const std::string & name); + + // Decrease instance ref counter + void dec_refs(const std::string & name); + // return true if the current process is a child server instance static bool is_child_server(); diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index 79e60db408..d471ff88b5 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -1,4 +1,5 @@ import pytest +import threading from utils import * server: ServerProcess @@ -205,3 +206,126 @@ def test_router_api_key_required(): ) assert authed.status_code == 200 assert "error" not in authed.body + + +# --- Drain-aware eviction tests --- + + +def _make_completion(model_id: str, max_tokens: int = 16) -> dict: + """Send a non-streaming completion request. Returns {"content": ..., "error": ...}.""" + result = {"content": "", "error": None} + try: + res = server.make_request("POST", "/v1/chat/completions", data={ + "model": model_id, + "max_tokens": max_tokens, + "messages": [{"role": "user", "content": "hi"}], + }) + if res.status_code == 200: + choices = res.body.get("choices", []) + if choices: + result["content"] = choices[0].get("message", {}).get("content", "") + else: + result["error"] = f"status {res.status_code}: {res.body}" + except Exception as e: + result["error"] = str(e) + return result + + +def test_router_concurrent_no_thrashing(): + """Concurrent requests for different models should all succeed, not thrash.""" + global server + server = ServerPreset.router() + server.models_max = 1 + server.start() + + model_a = "ggml-org/tinygemma3-GGUF:Q8_0" + model_b = "ggml-org/test-model-stories260K:F32" + n_per_model = 3 + results = {} + + def send_request(model_id, idx): + results[(model_id, idx)] = _make_completion(model_id) + + threads = [] + for i in range(n_per_model): + threads.append(threading.Thread(target=send_request, args=(model_a, i))) + threads.append(threading.Thread(target=send_request, args=(model_b, i))) + + for t in threads: + t.start() + for t in threads: + t.join(timeout=300) + + failures = [f"{m} #{i}: {r['error']}" for (m, i), r in results.items() if r["error"] is not None] + assert len(failures) == 0, f"{len(failures)} request(s) failed:\n" + "\n".join(failures) + + +def test_router_concurrent_partial_capacity(): + """With models_max=2 and 3 models, concurrent requests should all succeed.""" + global server + server = ServerPreset.router() + server.models_max = 2 + server.start() + + models = [ + "ggml-org/tinygemma3-GGUF:Q8_0", + "ggml-org/test-model-stories260K:F32", + "ggml-org/test-model-stories260K-infill:F32", + ] + results = {} + + def send_request(model_id, idx): + results[(model_id, idx)] = _make_completion(model_id) + + threads = [] + for model in models: + for i in range(2): + threads.append(threading.Thread(target=send_request, args=(model, i))) + + for t in threads: + t.start() + for t in threads: + t.join(timeout=300) + + failures = [f"{m} #{i}: {r['error']}" for (m, i), r in results.items() if r["error"] is not None] + assert len(failures) == 0, f"{len(failures)} request(s) failed:\n" + "\n".join(failures) + + +def test_router_alternating_requests(): + """Repeated alternating requests between two models should all succeed.""" + global server + server = ServerPreset.router() + server.models_max = 1 + server.start() + + model_a = "ggml-org/tinygemma3-GGUF:Q8_0" + model_b = "ggml-org/test-model-stories260K:F32" + + for i in range(3): + result = _make_completion(model_a) + assert result["error"] is None, f"Round {i} model A failed: {result['error']}" + result = _make_completion(model_b) + assert result["error"] is None, f"Round {i} model B failed: {result['error']}" + + +def test_router_concurrent_same_model(): + """Concurrent requests for the same model should all succeed.""" + global server + server = ServerPreset.router() + server.models_max = 1 + server.start() + + model_id = "ggml-org/tinygemma3-GGUF:Q8_0" + results = {} + + def send_request(idx): + results[idx] = _make_completion(model_id) + + threads = [threading.Thread(target=send_request, args=(i,)) for i in range(6)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=300) + + failures = [f"#{i}: {r['error']}" for i, r in results.items() if r["error"] is not None] + assert len(failures) == 0, f"{len(failures)} request(s) failed:\n" + "\n".join(failures)