diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index b4ef05b92..93e940951 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -70,6 +70,188 @@ struct server_subproc { } }; +struct server_lru_sched { + server_lru_sched(server_models & models) : models(models) {} + + bool has_capacity(std::unique_lock & lk) { + check_lock(lk); + return models.base_params.models_max <= 0 + || count_running() < (size_t) models.base_params.models_max; + } + + // returns "" if no model can be given up + std::string pick_victim(std::unique_lock & lk, const std::string & exclude) { + check_lock(lk); + std::string victim; + int64_t victim_last_used = 0; + for (const auto & m : models.mapping) { + if (m.first == exclude) { + continue; + } + // a busy model is mid-request, one still coming up has no request to finish + if (m.second.req_count != 0 || !m.second.meta.is_ready_or_sleep()) { + continue; + } + if (victim.empty() || m.second.meta.last_used < victim_last_used) { + victim = m.first; + victim_last_used = m.second.meta.last_used; + } + } + return victim; + } + + // requests wanting the same model share one entry, so they all need only one slot + // and all get unblocked by the single load that entry performs + void join(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + if (entry_t * e = find(model_id)) { + e->n_waiters++; + SRV_INF("request for name=%s joined the queue, %d waiting\n", model_id.c_str(), e->n_waiters); + return; + } + queue.push_back({ model_id, 1, false, false }); + SRV_INF("models_max reached, request for name=%s queued at position %zu\n", + model_id.c_str(), queue.size()); + } + + void leave(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + for (auto it = queue.begin(); it != queue.end(); ++it) { + if (it->model_id == model_id) { + if (--it->n_waiters <= 0) { + queue.erase(it); // last one waiting for this model went away + } + return; + } + } + } + + bool queue_empty(std::unique_lock & lk) { + check_lock(lk); + return queue.empty(); + } + + // true if it is this model's turn to load, and nobody is loading it yet + bool try_claim(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + if (queue.empty() || queue.front().model_id != model_id || queue.front().loading) { + return false; + } + if (!has_capacity(lk)) { + return false; + } + queue.front().loading = true; + return true; + } + + // ok means the model is up: drop the entry, the other waiters just watch its status now + void claim_done(std::unique_lock & lk, const std::string & model_id, bool ok) { + check_lock(lk); + for (auto it = queue.begin(); it != queue.end(); ++it) { + if (it->model_id == model_id) { + if (ok) { + queue.erase(it); + } else { + it->loading = false; + } + return; + } + } + } + + // a model is on its way out for this entry, so other requests do not also give up one + void mark_slot_pending(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + if (entry_t * e = find(model_id)) { + e->slot_pending = true; + } + } + + // model_id went idle: give up its slot if a queued request needs one + // thread-safe, caller must NOT hold models.mutex + void on_model_idle(const std::string & model_id) { + if (models.base_params.models_max <= 0) { + return; // no limit, nothing is ever queued + } + { + std::unique_lock lk(models.mutex); + if (queue.empty()) { + return; + } + size_t promised = 0; + bool has_unserved = false; + for (const auto & e : queue) { + if (e.needs_slot()) { + has_unserved = true; + } else { + promised++; + } + } + if (!has_unserved) { + return; + } + if ((int) count_running() - (int) promised < models.base_params.models_max) { + return; // a slot is already on its way + } + // never give up a model that a queued request wants + for (const auto & e : queue) { + if (e.model_id == model_id) { + return; + } + } + auto it = models.mapping.find(model_id); + if (it == models.mapping.end() || it->second.req_count != 0 || !it->second.meta.is_ready_or_sleep()) { + return; + } + for (auto & e : queue) { + if (!e.slot_pending) { + e.slot_pending = true; + break; + } + } + } + SRV_INF("model name=%s went idle, giving up its slot to a queued request\n", model_id.c_str()); + models.unload(model_id); + } + + private: + struct entry_t { + std::string model_id; + int n_waiters; // requests waiting for this model + bool slot_pending; // a model is already being evicted for this entry + bool loading; // one of the waiters is doing the load right now + + // a slot is already coming, or already taken by the load in flight + bool needs_slot() const { return !slot_pending && !loading; } + }; + + entry_t * find(const std::string & model_id) { + for (auto & e : queue) { + if (e.model_id == model_id) { + return &e; + } + } + return nullptr; + } + + void check_lock(std::unique_lock & lk) { + GGML_ASSERT(lk.owns_lock() && lk.mutex() == &models.mutex); + } + + size_t count_running() { + size_t count = 0; + for (const auto & m : models.mapping) { + if (m.second.meta.is_running()) { + count++; + } + } + return count; + } + + server_models & models; + std::deque queue; +}; + // short loopback budget for the resumable stream router to child JSON calls (probe, lookup, // delete). distinct from params.timeout_read/write which only applies to the generation proxy static constexpr int STREAM_LOOKUP_TIMEOUT_MS = 250; @@ -229,7 +411,8 @@ server_models::server_models( : ctx_preset(LLAMA_EXAMPLE_SERVER), base_params(params), base_env(get_environment()), - base_preset(ctx_preset.load_from_args(argc, argv)) { + base_preset(ctx_preset.load_from_args(argc, argv)), + sched(std::make_unique(*this)) { // clean up base preset unset_reserved_args(base_preset, true); // set binary path @@ -241,8 +424,11 @@ server_models::server_models( LOG_WRN("using original argv[0] as fallback: %s\n", argv[0]); } load_models(); + debug_fake_timing = !common_get_env("LLAMA_SERVER_DEBUG_FAKE_TIMING").empty(); } +server_models::~server_models() = default; + void server_models::add_model(server_model_meta && meta) { if (mapping.find(meta.name) != mapping.end()) { throw std::runtime_error(string_format("model '%s' appears multiple times", meta.name.c_str())); @@ -713,24 +899,15 @@ void server_models::unload_lru() { return; // no limit } // remove one of the servers if we passed the models_max (least recently used - LRU) - std::string lru_model_name = ""; - int64_t lru_last_used = ggml_time_ms(); - size_t count_active = 0; + std::string lru_model_name; { std::unique_lock lk(mutex); - for (const auto & m : mapping) { - if (m.second.meta.is_running()) { - count_active++; - // do not evict busy one - bool is_model_idle = m.second.req_count == 0 && m.second.meta.is_ready_or_sleep(); - if (is_model_idle && m.second.meta.last_used < lru_last_used) { - lru_model_name = m.first; - lru_last_used = m.second.meta.last_used; - } - } + if (sched->has_capacity(lk)) { + return; } + lru_model_name = sched->pick_victim(lk, ""); } - if (!lru_model_name.empty() && count_active >= (size_t)base_params.models_max) { + if (!lru_model_name.empty()) { SRV_INF("models_max limit reached, removing LRU name=%s\n", lru_model_name.c_str()); unload(lru_model_name); // wait for unload to complete @@ -741,7 +918,6 @@ void server_models::unload_lru() { }); } } - // TODO @ngxson : if no idle model is found, queue the load request } void server_models::load(const std::string & name) { @@ -749,6 +925,11 @@ void server_models::load(const std::string & name) { } void server_models::load(const std::string & name, const load_options & opts) { + if (debug_fake_timing) { + // do not hold the mutex here, other requests must keep making progress + std::this_thread::sleep_for(std::chrono::seconds(2)); + } + if (!opts.custom_meta.has_value()) { if (!has_model(name)) { throw std::runtime_error("model name=" + name + " is not found"); @@ -1141,7 +1322,7 @@ void server_models::wait(std::unique_lock & lk, const std::string & }); } -bool server_models::ensure_model_ready(const std::string & name) { +bool server_models::ensure_model_ready(const std::string & name, const std::function & should_stop) { auto meta = get_meta(name); if (!meta.has_value()) { throw std::runtime_error("model name=" + name + " is not found"); @@ -1152,25 +1333,112 @@ bool server_models::ensure_model_ready(const std::string & name) { if (meta->status == SERVER_MODEL_STATUS_SLEEPING) { return false; // child is sleeping but still running; new request will wake it up } - if (meta->status == SERVER_MODEL_STATUS_UNLOADED) { - SRV_INF("model name=%s is not loaded, loading...\n", name.c_str()); - load(name); - } - // wait for loading to complete - SRV_INF("waiting until model name=%s is fully loaded...\n", name.c_str()); - wait(name, [&meta](const server_model_meta & new_meta) { - if (new_meta.status != SERVER_MODEL_STATUS_LOADING) { - meta = new_meta; // update meta for final check after wait - return true; + bool queued = false; + bool did_load = false; + std::string victim; + { + std::unique_lock lk(mutex); + auto it = mapping.find(name); + if (it != mapping.end() && it->second.meta.status == SERVER_MODEL_STATUS_UNLOADED) { + bool has_capacity = sched->has_capacity(lk); + if (has_capacity && sched->queue_empty(lk)) { + lk.unlock(); + SRV_INF("model name=%s is not loaded, loading...\n", name.c_str()); + load(name); + did_load = true; + } else { + // also queue when a slot looks free but others wait already, else they starve + sched->join(lk, name); + queued = true; + if (!has_capacity) { + // an idle model may sit here right now, do not wait for a request to end + victim = sched->pick_victim(lk, name); + if (!victim.empty()) { + sched->mark_slot_pending(lk, name); + } + } + } } - return false; - }); - - // check final status - if (!meta.has_value() || meta->is_failed()) { - throw std::runtime_error("model name=" + name + " failed to load"); } + if (!victim.empty()) { + SRV_INF("evicting idle LRU name=%s to make room for name=%s\n", victim.c_str(), name.c_str()); + unload(victim); + } + + // while queued, this is also where the load happens: the head of the queue does it + SRV_INF("waiting until model name=%s is fully loaded...\n", name.c_str()); + std::unique_lock lk(mutex); + auto leave_queue = [this, &queued, &lk, &name]() { + if (queued) { + sched->leave(lk, name); + queued = false; + } + }; + + try { + bool saw_loading = false; + while (true) { + auto it = mapping.find(name); + if (it == mapping.end()) { + break; // removed by another code path, nothing to wait for + } + const server_model_status status = it->second.meta.status; + + if (status == SERVER_MODEL_STATUS_LOADED || status == SERVER_MODEL_STATUS_SLEEPING) { + break; + } + if (status == SERVER_MODEL_STATUS_DOWNLOADING || status == SERVER_MODEL_STATUS_DOWNLOADED) { + break; // do not wait on a download child + } + if (status == SERVER_MODEL_STATUS_LOADING) { + saw_loading = true; + } else if (status == SERVER_MODEL_STATUS_UNLOADED) { + if (did_load || saw_loading) { + // a spawn happened and the instance came back down + if (it->second.meta.is_failed()) { + throw std::runtime_error("model name=" + name + " failed to load"); + } + break; // unloaded by another code path, caller reports "not running" + } + if (!queued) { + break; // not queued, and the load someone else started fell over + } + } + + if (should_stop && should_stop()) { + // if a model was evicted for us, the free slot goes to the next waiter + throw std::runtime_error("request cancelled while waiting for model name=" + name); + } + + // our turn: our model is at the head, and a slot really did free up + if (status == SERVER_MODEL_STATUS_UNLOADED && sched->try_claim(lk, name)) { + lk.unlock(); + bool ok = true; + try { + SRV_INF("slot available, loading queued model name=%s\n", name.c_str()); + load(name); + did_load = true; + } catch (const std::exception & e) { + // lost a race for the slot, stay in line and retry + SRV_WRN("queued load of name=%s did not go through: %s\n", name.c_str(), e.what()); + ok = false; + } + lk.lock(); + sched->claim_done(lk, name, ok); + if (ok) { + queued = false; // entry is gone, the other waiters watch the status now + } + continue; + } + + cv.wait_for(lk, std::chrono::milliseconds(200)); + } + } catch (...) { + leave_queue(); + throw; + } + leave_queue(); return true; } @@ -1190,6 +1458,10 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co } mapping[name].req_count++; } + if (debug_fake_timing) { + // sleep after req_count++, so the model counts as busy while we wait here + std::this_thread::sleep_for(std::chrono::seconds(2)); + } SRV_INF("proxying request to model %s on port %d\n", name.c_str(), meta->port); std::string proxy_path = req.path; if (!req.query_string.empty()) { @@ -1213,10 +1485,17 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co ); proxy->cleanup = [this, name]() { - std::unique_lock lk(mutex); - auto it = mapping.find(name); - if (it != mapping.end() && it->second.req_count > 0) { - it->second.req_count--; + bool went_idle = false; + { + std::unique_lock lk(mutex); + auto it = mapping.find(name); + if (it != mapping.end() && it->second.req_count > 0) { + it->second.req_count--; + went_idle = it->second.req_count == 0; + } + } + if (went_idle) { + sched->on_model_idle(name); } }; @@ -1583,7 +1862,7 @@ void server_models_routes::init_routes() { return error_res; } if (autoload) { - models.ensure_model_ready(name); + models.ensure_model_ready(name, req.should_stop); } return models.proxy_request(req, method, name, false); }; @@ -1603,7 +1882,9 @@ void server_models_routes::init_routes() { // this request instead of leaving an orphan generation std::string conv_id = server_stream_conv_id_from_headers(req.headers); uint64_t ticket = models.conv_models.remember(conv_id, name); - bool waited = autoload && models.ensure_model_ready(name); + // a dead socket must not cancel a session request, only a stop does (checked right below) + auto should_stop = ticket == 0 ? req.should_stop : nullptr; + bool waited = autoload && models.ensure_model_ready(name, should_stop); if (ticket != 0 && !models.conv_models.alive(conv_id, ticket)) { SRV_INF("request for conv_id=%s cancelled while model name=%s was loading\n", conv_id.c_str(), name.c_str()); diff --git a/tools/server/server-models.h b/tools/server/server-models.h index 1c6123cb1..615acb577 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -106,10 +106,12 @@ struct server_model_meta { }; struct server_models_routes; -struct server_subproc; // defined in server-models.cpp +struct server_subproc; // defined in server-models.cpp +struct server_lru_sched; // defined in server-models.cpp struct server_models { friend struct server_models_routes; + friend struct server_lru_sched; private: struct instance_t { @@ -195,6 +197,12 @@ private: std::vector base_env; common_preset base_preset; // base preset from llama-server CLI args + // queue of requests waiting for a models_max slot + std::unique_ptr sched; + + // if true, add some delay to simulate works (useful for testing) + bool debug_fake_timing = false; + void update_meta(const std::string & name, const server_model_meta & meta); // unload least recently used models if the limit is reached @@ -211,6 +219,7 @@ public: conv_model_tracker conv_models; server_models(const common_params & params, int argc, char ** argv); + ~server_models(); server_response sse; // for real-time updates via SSE endpoint @@ -267,7 +276,9 @@ public: // ensure the model is in ready state (thread-safe) // return false if model is ready // otherwise, load the model and blocking wait until it's ready, then return true (meta may need to be refreshed) - bool ensure_model_ready(const std::string & name); + // if models_max is reached, the request waits in a queue until a slot frees up + // throws if the load fails, or if should_stop fires while waiting + bool ensure_model_ready(const std::string & name, const std::function & should_stop = nullptr); // 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, bool detached = false); diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index 94165e520..c503ee342 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -145,6 +145,156 @@ def test_router_models_max_evicts_lru(): assert _get_model_status(first) == "unloaded" +# server_lru_sched tests (relying on LLAMA_SERVER_DEBUG_FAKE_TIMING) + +MODEL_A = "ggml-org/tinygemma3-GGUF:Q8_0" +MODEL_B = "ggml-org/test-model-stories260K:F32" +MODEL_C = "ggml-org/test-model-stories260K-infill:F32" + + +def _tokenize(model_id: str, timeout: float | None = DEFAULT_REQUEST_TIMEOUT) -> ServerResponse: + return server.make_request( + "POST", "/tokenize", data={"model": model_id, "content": "hello world"}, timeout=timeout + ) + + +class _Bg: + """runs one request in a thread, keeps its result, error and finish time""" + + def __init__(self, fn): + self.result = None + self.error: Exception | None = None + self.done_at: float = 0.0 + self._thread = threading.Thread(target=self._run, args=(fn,), daemon=True) + + def _run(self, fn): + try: + self.result = fn() + except Exception as e: + self.error = e + self.done_at = time.time() + + def start(self): + self._thread.start() + return self + + def join(self, timeout: int = 180): + self._thread.join(timeout) + assert not self._thread.is_alive(), "background request did not finish in time" + return self + + def assert_ok(self, what: str): + assert self.error is None, f"{what} raised {self.error!r}" + assert self.result is not None and self.result.status_code == 200, \ + f"{what} failed: {self.result.status_code if self.result else None} {self.result.body if self.result else None}" + + +def test_router_queue_does_not_evict_busy_model(): + """a request that finds no free slot waits, and the model serving a request survives it""" + global server + server.models_max = 1 + server.start() + + _load_model_and_wait(MODEL_A, timeout=120) + + busy = _Bg(lambda: _tokenize(MODEL_A)).start() + time.sleep(0.5) # let the request reach the child and take the only slot + + # no slot free and MODEL_A is busy, so this queues instead of evicting mid-request + queued = _Bg(lambda: _tokenize(MODEL_B)).start() + + busy.join() + queued.join() + + # had MODEL_A been evicted while serving, its own request would have died + busy.assert_ok("request against the busy model") + queued.assert_ok("queued request") + + _wait_for_model_status(MODEL_B, {"loaded"}, timeout=120) + assert _get_model_status(MODEL_A) == "unloaded" + + +def test_router_queue_coalesces_requests_for_same_model(): + """many requests for one missing model share a slot, so only one model is given up""" + global server + server.models_max = 2 + server.start() + + _load_model_and_wait(MODEL_A, timeout=120) + _load_model_and_wait(MODEL_B, timeout=120) + + # keep MODEL_A busy so MODEL_B is the only model that can be given up + busy = _Bg(lambda: _tokenize(MODEL_A)).start() + time.sleep(0.5) + + waiters = [_Bg(lambda: _tokenize(MODEL_C)).start() for _ in range(3)] + + busy.join() + for w in waiters: + w.join() + + busy.assert_ok("request against the busy model") + for i, w in enumerate(waiters): + w.assert_ok(f"queued request {i}") + + _wait_for_model_status(MODEL_C, {"loaded"}, timeout=120) + # one entry for 3 requests means one eviction: MODEL_B goes, MODEL_A is left alone. + # without coalescing the leftover entries still ask for a slot, + # and MODEL_A is taken too as soon as it goes idle + assert _get_model_status(MODEL_A) == "loaded" + assert _get_model_status(MODEL_B) == "unloaded" + + +def test_router_queue_client_disconnect_keeps_model(): + """a client that leaves while queued must not cost a running model its slot""" + global server + server.models_max = 1 + server.start() + + _load_model_and_wait(MODEL_A, timeout=120) + + busy = _Bg(lambda: _tokenize(MODEL_A)).start() + time.sleep(0.5) + + # queues behind MODEL_A, then gives up long before MODEL_A goes idle + with pytest.raises(requests.exceptions.RequestException): + _tokenize(MODEL_B, timeout=1) + + busy.join() + busy.assert_ok("request against the busy model") + + # nobody is waiting anymore, so MODEL_A keeps its slot + time.sleep(3) + assert _get_model_status(MODEL_A) == "loaded" + assert _get_model_status(MODEL_B) == "unloaded" + + +def test_router_queue_is_fifo(): + """the queue is served in arrival order""" + global server + server.models_max = 1 + server.start() + + _load_model_and_wait(MODEL_A, timeout=120) + + busy = _Bg(lambda: _tokenize(MODEL_A)).start() + time.sleep(0.5) + + first = _Bg(lambda: _tokenize(MODEL_B)).start() + time.sleep(1) # keep the arrival order unambiguous + second = _Bg(lambda: _tokenize(MODEL_C)).start() + + busy.join() + first.join() + second.join() + + busy.assert_ok("request against the busy model") + first.assert_ok("first queued request") + second.assert_ok("second queued request") + + assert first.done_at < second.done_at, "queue was not served in arrival order" + + def test_router_no_models_autoload(): global server server.no_models_autoload = True diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index ae56bc70a..8f3a98642 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -132,7 +132,10 @@ class ServerProcess: self.external_server = "DEBUG_EXTERNAL" in os.environ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: - env = {**os.environ} + env = { + **os.environ, + "LLAMA_SERVER_DEBUG_FAKE_TIMING": "1", + } if "LLAMA_CACHE" not in os.environ: env["LLAMA_CACHE"] = "tmp" if self.external_server: