From 6009998a74ba95a6a7edcd8f6c61ed77104e6c8a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 00:12:09 +0000 Subject: [PATCH 1/4] server: preempt a slot instead of ending every conversation when the KV pool fills With --parallel N --kv-unified there is one pool of cells and every slot believes it owns all of them. When the pool fills, llama_decode returns 1, the retry ladder in decode() halves n_batch down to 1, and the server calls send_error on EVERY processing slot: "Context size has been exceeded". Four chats sharing a 8192-cell pool on Qwen3.5-4B-MTP die together after six seconds, none of them anywhere near its own 8192 limit. The code already says what should happen instead: "TODO: try to terminate only the largest active slot/sequence and continue with the rest". Terminate nothing. Once per update_slots(), before the batch is built, compare what the pool holds against what the next decode will ask for. If it does not fit, take the cells back from one slot: copy its sequence out with llama_state_seq_get_data_ext, release the cells, and park the slot in a new SLOT_STATE_PREEMPTED. When the pool has room the copy goes back with llama_state_seq_set_data_ext and the slot carries on. The task, the sampler, the generated text and the position the stream has reached never left the slot, so the continuation is the one the slot would have produced without the pause, and a streaming client sees a gap and nothing else. The check sits before the batch is built on purpose: at that point every slot is at a token boundary, prompt.tokens is exactly what the cache holds for it, and no draft is in flight, so a slot can be removed without unpicking a half-decoded batch. The speculative draft is dropped with the cells, which costs the step its speedup and nothing else. Victim policy: keep the slot that is furthest along, since it is the closest to finishing and to giving its cells back, and among the rest prefer one that has not been preempted three times already, then the smallest. A prompt cached on an idle slot is cheaper than a conversation waiting to continue, so try_clear_idle_slots() is asked first, both before preempting anyone and before deciding a resume does not fit. Measured on Qwen3.5-4B-UD-Q4_K_XL with an embedded MTP head, --parallel 4 --kv-unified -c 8192, four streaming chats with 1000-token prompts at temperature 0: base 4 of 4 chats killed by "Context size has been exceeded" after 6.7 s with this 4 of 4 chats completed, 0 errors, 7 preemptions, 7 resumes, 22944 tokens in 66.9 s (343 tok/s aggregate) and the retry ladder never fires at all. At -c 16384 the same load still kills all four on the base and still completes all four here. LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens regardless of pressure. With one request on an idle server the batch has the same shape at every step, so it isolates the resume from batch nondeterminism: over 91 forced preemptions across four prompts, every continuation is byte-identical to the same prompt run without any. --- tools/server/server-context.cpp | 401 ++++++++++++++++++++++++++++++++ 1 file changed, 401 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..85e9494563d1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -59,8 +59,27 @@ enum slot_state { SLOT_STATE_PROCESSING_PROMPT, SLOT_STATE_DONE_PROMPT, SLOT_STATE_GENERATING, + SLOT_STATE_PREEMPTED, // [TAG_PREEMPT] cells released, everything needed to resume is in host RAM }; +// [TAG_PREEMPT] server-side request preemption +// +// With --kv-unified the cells are one pool shared by every slot, and each slot believes it +// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves +// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size +// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the +// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and +// continue with the rest". +// +// Nothing is terminated here. The cells of one slot are taken back and given to it again +// later: its sequence is copied to host RAM, its cells are released, and when the pool has +// room the copy goes back and the slot carries on with the same sampler, the same generated +// text and the same open stream. A streaming client sees a pause, not an error. +constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation +constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected +constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on +constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked + struct server_slot; // forward declaration struct server_batch { @@ -293,6 +312,115 @@ struct server_slot { prompt.clear(); } + // [TAG_PREEMPT] state of a slot whose cells were taken back + // + // Only the KV cells leave. The task, the sampler, the generated text and the position + // the stream has reached stay on the slot, so a resume is a memcpy and not a new + // request: no retokenisation, no replayed prompt, no seam in the output. + slot_state state_before_preempt = SLOT_STATE_IDLE; + std::vector preempt_state_tgt; + std::vector preempt_state_dft; + int32_t n_preempt = 0; // times the CURRENT task has been preempted + int32_t n_preempt_fail = 0; // consecutive failed restores + int64_t t_preempt_us = 0; // when it was parked + + size_t preempt_state_size() const { + return preempt_state_tgt.size() + preempt_state_dft.size(); + } + + void preempt_state_free() { + preempt_state_tgt.clear(); + preempt_state_tgt.shrink_to_fit(); + preempt_state_dft.clear(); + preempt_state_dft.shrink_to_fit(); + } + + // copy the sequence out of the cache and release its cells + bool preempt_save() { + const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + + try { + preempt_state_tgt.resize(size_tgt); + preempt_state_dft.resize(size_dft); + } catch (const std::bad_alloc & e) { + SLT_ERR(*this, "failed to allocate %.3f MiB for the preemption state: %s\n", + (size_tgt + size_dft) / (1024.0 * 1024.0), e.what()); + preempt_state_free(); + return false; + } + + if (llama_state_seq_get_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to copy the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_get_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to copy the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + // The draft is a prediction, not a result, so it goes with the cells. Preemption + // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already + // holds exactly the tokens the state above covers -- including the rollback done by + // the checkpoint path when a draft was only partially accepted. + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + + // note: prompt.tokens is deliberately kept. It is the mirror of the state just + // copied out, and the resume needs it to know how many cells to ask for. + mem.seq_rm(id, -1, -1); + + state_before_preempt = state; + state = SLOT_STATE_PREEMPTED; + t_preempt_us = ggml_time_us(); + + n_preempt++; + + return true; + } + + // put the sequence back; the slot then continues from the token it was about to decode + bool preempt_restore() { + const size_t size_tgt = preempt_state_tgt.size(); + const size_t size_dft = preempt_state_dft.size(); + + if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + // no room after all: drop the half-written sequence and stay parked + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + if (size_dft > 0 && + llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + preempt_state_free(); + + n_preempt_fail = 0; + + state = state_before_preempt; + + // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks + // that the draft context is where it should be, which the restore above ensures + if (can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + + return true; + } + std::vector lora; int32_t alora_invocation_start = -1; @@ -351,6 +479,13 @@ struct server_slot { n_predict_max = -1; + // [TAG_PREEMPT] + preempt_state_free(); + state_before_preempt = SLOT_STATE_IDLE; + n_preempt = 0; + n_preempt_fail = 0; + t_preempt_us = 0; + llama_set_sampler(ctx_tgt, id, nullptr); // clear alora start @@ -503,6 +638,14 @@ struct server_slot { t_last_used = ggml_time_us(); + // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be + // released while parked), so the mirror of them must not outlive them: the next + // task on this slot would otherwise take a prefix match against an empty cache + if (state == SLOT_STATE_PREEMPTED) { + preempt_state_free(); + prompt_clear(); + } + state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough @@ -1249,6 +1392,16 @@ struct server_context_impl { } } + { + const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); + preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; + + if (preempt_test_every > 0) { + SRV_WRN("LLAMA_SERVER_PREEMPT_EVERY = %d (test knob: preempting every slot every %d tokens)\n", + preempt_test_every, preempt_test_every); + } + } + { const char * LLAMA_SERVER_SLOTS_N_DIFF = getenv("LLAMA_SERVER_SLOTS_N_DIFF"); slots_n_diff = LLAMA_SERVER_SLOTS_N_DIFF ? atoi(LLAMA_SERVER_SLOTS_N_DIFF) : 0; @@ -2674,6 +2827,251 @@ struct server_context_impl { }; #endif + // + // [TAG_PREEMPT] server-side request preemption + // + + int64_t n_preempt_total = 0; + int64_t n_resume_total = 0; + + // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, + // whether or not the pool is under pressure. It exists to answer the only question that + // matters about a resume: with one request on an idle server the batch has the same + // shape at every step, so a preempted continuation that is not byte-identical to an + // uninterrupted one is the preemption's fault and nothing else's. + int32_t preempt_test_every = 0; + + int32_t preempt_n_spec_max() const { + return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; + } + + // Cells the pool is holding right now. A released slot keeps its prompt in the cache + // for the next request to reuse as a prefix, so idle slots count too: the first version + // of this counted only the running ones, decided a pool holding 8185 cached cells was + // empty, and every resume failed against a cache that was actually full. + int32_t preempt_kv_used() const { + int32_t res = 0; + + for (const auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTED) { + continue; // parked: its cells are in host RAM, not in the pool + } + + res += slot.prompt.n_tokens(); + } + + return res; + } + + // cells those slots are about to ask for on the next decode + int32_t preempt_kv_reserve() const { + const int32_t n_spec = preempt_n_spec_max(); + const int32_t n_batch = llama_n_batch(ctx_tgt); + + int32_t res = 0; + int32_t res_pmt = 0; + + for (const auto & slot : slots) { + switch (slot.state) { + case SLOT_STATE_GENERATING: + case SLOT_STATE_DONE_PROMPT: + { + res += 1 + n_spec; + } break; + case SLOT_STATE_STARTED: + case SLOT_STATE_PROCESSING_PROMPT: + { + const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + + res_pmt += std::max(1, std::min(n_batch, n_left)); + } break; + default: + break; + } + } + + // one batch is all the prompt slots get between them, however many are waiting + return res + std::min(res_pmt, n_batch); + } + + // Keep the slot that is furthest along -- it is the closest to finishing and to giving + // its cells back -- and among the rest prefer one that has not been preempted + // PREEMPT_N_STARVED times already, then the smallest. + server_slot * preempt_pick_victim() { + server_slot * leader = nullptr; + int32_t n_running = 0; + + for (auto & slot : slots) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + n_running++; + + if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) { + leader = &slot; + } + } + } + + if (n_running < 2) { + // a single conversation that does not fit the pool on its own is a real context + // overflow and not a scheduling problem - leave it to the existing error path + return nullptr; + } + + server_slot * victim = nullptr; + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_GENERATING) { + continue; // a generating slot is the one with a clean point to stop at + } + + if (&slot == leader || slot.prompt.n_tokens() == 0) { + continue; + } + + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + continue; // n_cmpl > 1 slots share one sequence, out of scope here + } + + const bool starved = slot.n_preempt >= PREEMPT_N_STARVED; + const bool starved_cur = victim && victim->n_preempt >= PREEMPT_N_STARVED; + + if (!victim || + (starved_cur && !starved) || + (starved_cur == starved && slot.prompt.n_tokens() < victim->prompt.n_tokens())) { + victim = &slot; + } + } + + return victim; + } + + // called once per update_slots(), before the batch is built: at that point every slot is + // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no + // draft is in flight, so a slot can be removed from the picture without unpicking a + // half-decoded batch + void update_preemption() { + if (!params_base.kv_unified || slots.size() < 2) { + return; // with a cache per slot, no slot can take another one's cells + } + + const int32_t n_cells = n_ctx; + + // put back what fits, the most-preempted slot first + for (;;) { + server_slot * best = nullptr; + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_PREEMPTED) { + continue; + } + + if (!best || + slot.n_preempt > best->n_preempt || + (slot.n_preempt == best->n_preempt && slot.t_preempt_us < best->t_preempt_us)) { + best = &slot; + } + } + + if (!best) { + break; + } + + const int32_t n_need = best->prompt.n_tokens() + 1 + preempt_n_spec_max(); + + // room for the sequence AND for the next step of everything already running, + // so that a resume cannot immediately trigger the preemption of someone else. + // A cached prompt on an idle slot is worth less than a conversation waiting to + // continue, so give those cells up first - same call the KV-full path makes. + while (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells && + try_clear_idle_slots()) { + } + + if (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells) { + break; + } + + const int64_t t_start = ggml_time_us(); + + if (!best->preempt_restore()) { + // update_slots() runs in a tight loop while tasks are pending, so a counter + // alone burns its whole budget in a couple of milliseconds. Give up only on + // a slot that has been failing for a while, and keep the log quiet. + if (best->n_preempt_fail % 64 == 1) { + SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", + best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); + } + + if (best->n_preempt_fail >= PREEMPT_N_FAIL_MAX && + ggml_time_us() - best->t_preempt_us > PREEMPT_FAIL_US) { + send_error(*best, "failed to restore the preempted sequence"); + best->release(); + } + + break; + } + + n_resume_total++; + + SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", + (ggml_time_us() - best->t_preempt_us) / 1e6, + best->prompt.n_tokens(), + (ggml_time_us() - t_start) / 1e3, + preempt_kv_used(), n_cells, + best->n_preempt); + } + + // forced preemption, for the determinism test only + if (preempt_test_every > 0) { + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_GENERATING && + (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && + slot.preempt_save()) { + n_preempt_total++; + + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", + (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + } + } + } + + // and take cells back until the next decode fits + for (;;) { + const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); + + if (n_used + PREEMPT_N_MARGIN <= n_cells) { + break; + } + + // a prompt cached on an idle slot is the cheapest thing in the pool to give up + if (try_clear_idle_slots()) { + continue; + } + + server_slot * victim = preempt_pick_victim(); + + if (!victim) { + SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted\n", n_used, n_cells); + break; + } + + const int32_t n_tokens = victim->prompt.n_tokens(); + const int64_t t_start = ggml_time_us(); + + if (!victim->preempt_save()) { + break; // could not park it; the existing retry ladder is still behind us + } + + n_preempt_total++; + + SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + } + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -2715,6 +3113,9 @@ struct server_context_impl { } } + // [TAG_PREEMPT] make the pool fit the step that is about to be built + update_preemption(); + try { scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); From 41cbff49d02c21998c66215ad473b1686c96c2e7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 00:41:11 +0000 Subject: [PATCH 2/4] server: test that preemption keeps the output and finishes every slot Two tests on the two-slot unified pool. The first runs one request with LLAMA_SERVER_PREEMPT_EVERY=8 and asserts the tokens match the same request without the knob. The second runs two requests that each fit alone but not together and asserts both finish with no context error. Both fail on master: the knob is unknown there, and the second request dies with Context size has been exceeded. --- tools/server/tests/unit/test_preempt.py | 110 ++++++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 tools/server/tests/unit/test_preempt.py diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py new file mode 100644 index 000000000000..479b62b91ef5 --- /dev/null +++ b/tools/server/tests/unit/test_preempt.py @@ -0,0 +1,110 @@ +import os +import tempfile +import pytest +from utils import * + +# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked +# (its sequence copied to host RAM, its cells released) instead of every slot being +# terminated. Both tests need more than one slot and --kv-unified, which is the only +# configuration where one slot can take another one's cells. + +server = ServerPreset.tinyllama2() + + +class LogReader: + def __init__(self, path): + self.path = path + self.pos = 0 + + def drain(self): + with open(self.path) as f: + f.seek(self.pos) + content = f.read() + self.pos = f.tell() + return content + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + server.server_slots = True + server.temperature = 0.0 + server.seed = 42 + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + + +def _complete(n_predict: int, prompt: str = "Hi how are you"): + res = server.make_request("POST", "/completion", data={ + "n_predict": n_predict, + "prompt": prompt, + "ignore_eos": True, + "return_tokens": True, + "temperature": 0.0, + "seed": 42, + }) + return res + + +def test_forced_preemption_does_not_change_the_output(): + # Park and restore the only running slot every 8 tokens. With one request the batch + # has the same shape at every step whether or not the slot was parked in between, so + # any difference in the output is the preemption's fault and nothing else's. + global server + server.n_ctx = 512 + server.start() + reference = _complete(64) + assert reference.status_code == 200 + assert reference.body["timings"]["predicted_n"] == 64 + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_EVERY = 8" in log.drain() + + preempted = _complete(64) + assert preempted.status_code == 200 + assert preempted.body["timings"]["predicted_n"] == 64 + + text = log.drain() + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 + + assert preempted.body["content"] == reference.body["content"] + assert preempted.body["tokens"] == reference.body["tokens"] + + +def test_two_slots_that_overflow_the_pool_together_both_finish(): + # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well + # under 256. Together they do not, 336 against 256. Without preemption the retry + # ladder ends with "Context size has been exceeded" on every processing slot; with it + # the smaller slot is parked until the leader finishes and its cells are purged, and + # then it resumes from the token it was parked on. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + From 32c0a77e16ff62630a335f4caeb64263f9b57484 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 01:03:22 +0000 Subject: [PATCH 3/4] server: park prompt-processing slots too, restore whatever fits first, and bound the parked state with --preempt-ram A slot still processing its prompt is between two chunks of it, which is as clean a boundary as between two sampled tokens, so it is a victim too: two prompts that do not fit together no longer fail together, and a large prompt arriving beside a running chat waits for it instead of ending it. A slot that has not started yet holds at most a cached prefix and is parked the same way, which is how it waits. Restoring takes the most-preempted parked slot first, but one that does not fit yet no longer holds up a smaller one that does: the smaller one is the first to be parked again if the pool fills, so the head of the line loses nothing. --preempt-ram N (LLAMA_ARG_PREEMPT_RAM) bounds the host RAM parked sequences may hold, default 8192 MiB like --cache-ram. A slot that would not fit under the budget is not parked, and when nothing can be parked the KV-full path runs as before. --preempt-ram 0 disables preemption. The prompt batching pass skips parked slots explicitly. Speculation is only restarted on restore for a slot that was generating; one parked mid-prompt starts it when its prompt is done, as it always did. Tests: two prompts that overflow the pool together, a generating slot beside a large prompt, and --preempt-ram 0 restoring the old behaviour. --- common/arg.cpp | 8 ++ common/common.h | 1 + tools/server/README.md | 1 + tools/server/server-context.cpp | 124 +++++++++++++++++++----- tools/server/tests/unit/test_preempt.py | 119 +++++++++++++++++++++++ 5 files changed, 230 insertions(+), 23 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d0..5bfa4adcdf0d 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1709,6 +1709,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.cache_ram_mib = value; } ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--preempt-ram"}, "N", + string_format("with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; " + "N is the maximum host RAM for parked sequences in MiB (default: %d, -1 - no limit, 0 - disable)", params.preempt_ram_mib), + [](common_params & params, int value) { + params.preempt_ram_mib = value; + } + ).set_env("LLAMA_ARG_PREEMPT_RAM").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index de49dac9f63a..c99269f9a967 100644 --- a/common/common.h +++ b/common/common.h @@ -614,6 +614,7 @@ struct common_params { int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = disable preemption std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa9..52e5c1b8c7db 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -164,6 +164,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | +| `--preempt-ram N` | with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; N is the maximum host RAM for parked sequences in MiB (default: 8192, -1 - no limit, 0 - disable)
(env: LLAMA_ARG_PREEMPT_RAM) | | `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)
(env: LLAMA_ARG_KV_UNIFIED) | | `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)
(env: LLAMA_ARG_CACHE_IDLE_SLOTS) | | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 85e9494563d1..a41e5feda5d6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -335,6 +335,12 @@ struct server_slot { preempt_state_dft.shrink_to_fit(); } + // bytes preempt_save() would need for this slot right now + size_t preempt_state_required() const { + return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) + + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0); + } + // copy the sequence out of the cache and release its cells bool preempt_save() { const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); @@ -413,8 +419,10 @@ struct server_slot { state = state_before_preempt; // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks - // that the draft context is where it should be, which the restore above ensures - if (can_speculate()) { + // that the draft context is where it should be, which the restore above ensures. + // A slot parked while still processing its prompt makes that transition itself + // once the prompt is done. + if (state == SLOT_STATE_GENERATING && can_speculate()) { common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } @@ -2845,6 +2853,43 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } + // host RAM the parked sequences hold right now + size_t preempt_ram_used() const { + size_t res = 0; + + for (const auto & slot : slots) { + res += slot.preempt_state_size(); + } + + return res; + } + + // whether parking this slot stays under --preempt-ram + bool preempt_fits_budget(const server_slot & slot) const { + if (params_base.preempt_ram_mib < 0) { + return true; + } + + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + + return preempt_ram_used() + slot.preempt_state_required() <= budget; + } + + // cells the slot will ask for on its next step once it is back in the pool + int32_t preempt_n_need(const server_slot & slot) const { + int32_t res = slot.prompt.n_tokens(); + + if (slot.state_before_preempt == SLOT_STATE_GENERATING) { + res += 1 + preempt_n_spec_max(); + } else { + const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + + res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); + } + + return res; + } + // Cells the pool is holding right now. A released slot keeps its prompt in the cache // for the next request to reuse as a prefix, so idle slots count too: the first version // of this counted only the running ones, decided a pool holding 8185 cached cells was @@ -2920,11 +2965,18 @@ struct server_context_impl { server_slot * victim = nullptr; for (auto & slot : slots) { - if (slot.state != SLOT_STATE_GENERATING) { - continue; // a generating slot is the one with a clean point to stop at + // Before the batch is built every one of these is at a token boundary: a + // generating slot between two sampled tokens, a prompt-processing slot between + // two chunks of its prompt, a started slot with only a cached prefix (or + // nothing) in the pool. A slot holding no cells is still worth parking - it + // is about to ask for a whole batch of them. + if (slot.state != SLOT_STATE_GENERATING && + slot.state != SLOT_STATE_PROCESSING_PROMPT && + slot.state != SLOT_STATE_STARTED) { + continue; } - if (&slot == leader || slot.prompt.n_tokens() == 0) { + if (&slot == leader) { continue; } @@ -2932,6 +2984,10 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + if (!preempt_fits_budget(slot)) { + continue; + } + const bool starved = slot.n_preempt >= PREEMPT_N_STARVED; const bool starved_cur = victim && victim->n_preempt >= PREEMPT_N_STARVED; @@ -2954,39 +3010,57 @@ struct server_context_impl { return; // with a cache per slot, no slot can take another one's cells } + if (params_base.preempt_ram_mib == 0) { + return; // --preempt-ram 0: the KV-full retry ladder, as before + } + const int32_t n_cells = n_ctx; - // put back what fits, the most-preempted slot first + // Put back what fits: the most-preempted slot first, then the one parked longest. + // A slot that does not fit yet must not hold up a smaller one that does: it keeps + // its place at the head of the line, and the smaller one is the first to be parked + // again if the pool fills, so letting it through costs the head nothing. for (;;) { - server_slot * best = nullptr; + std::vector parked; for (auto & slot : slots) { - if (slot.state != SLOT_STATE_PREEMPTED) { - continue; - } - - if (!best || - slot.n_preempt > best->n_preempt || - (slot.n_preempt == best->n_preempt && slot.t_preempt_us < best->t_preempt_us)) { - best = &slot; + if (slot.state == SLOT_STATE_PREEMPTED) { + parked.push_back(&slot); } } - if (!best) { + if (parked.empty()) { break; } - const int32_t n_need = best->prompt.n_tokens() + 1 + preempt_n_spec_max(); + std::sort(parked.begin(), parked.end(), [](const server_slot * a, const server_slot * b) { + if (a->n_preempt != b->n_preempt) { + return a->n_preempt > b->n_preempt; + } - // room for the sequence AND for the next step of everything already running, + return a->t_preempt_us < b->t_preempt_us; + }); + + server_slot * best = nullptr; + + // Room for the sequence AND for the next step of everything already running, // so that a resume cannot immediately trigger the preemption of someone else. // A cached prompt on an idle slot is worth less than a conversation waiting to // continue, so give those cells up first - same call the KV-full path makes. - while (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells && - try_clear_idle_slots()) { + for (;;) { + for (auto * slot : parked) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + best = slot; + break; + } + } + + if (best || !try_clear_idle_slots()) { + break; + } } - if (preempt_kv_used() + preempt_kv_reserve() + n_need + PREEMPT_N_MARGIN > n_cells) { + if (!best) { break; } @@ -3025,6 +3099,7 @@ struct server_context_impl { for (auto & slot : slots) { if (slot.state == SLOT_STATE_GENERATING && (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && + preempt_fits_budget(slot) && slot.preempt_save()) { n_preempt_total++; @@ -3050,7 +3125,8 @@ struct server_context_impl { server_slot * victim = preempt_pick_victim(); if (!victim) { - SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted\n", n_used, n_cells); + SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted (parked %.1f MiB of the %d MiB --preempt-ram budget)\n", + n_used, n_cells, preempt_ram_used() / (1024.0 * 1024.0), params_base.preempt_ram_mib); break; } @@ -3401,7 +3477,9 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - if (!slot.is_processing()) { + // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to + // batch; it takes no part in this pass until it is restored + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { return; } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 479b62b91ef5..2e57ebbdfd97 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -1,4 +1,5 @@ import os +import time import tempfile import pytest from utils import * @@ -37,6 +38,7 @@ def create_server(): os.close(fd) yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) def _complete(n_predict: int, prompt: str = "Hi how are you"): @@ -108,3 +110,120 @@ def test_two_slots_that_overflow_the_pool_together_both_finish(): assert res.body["truncated"] is False assert len(res.body["tokens"]) == n_predict + + +_WORDS = ( + "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor " + "incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud " + "exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure " + "dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. " + "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt " + "mollit anim id est laborum. " +) * 4 + + +def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: + """A prompt whose token count is in [n_tokens - 12, n_tokens], measured on the server.""" + words = (salt + " " + _WORDS).split() + while words: + text = " ".join(words) + res = server.make_request("POST", "/tokenize", data={"content": text}) + assert res.status_code == 200 + n = len(res.body["tokens"]) + if n <= n_tokens: + assert n >= n_tokens - 12, f"could not land near {n_tokens} tokens, got {n}" + return text, n + # about four tokens per word on this model's vocabulary + words = words[: len(words) - max(1, (n - n_tokens) // 8)] + raise AssertionError("empty prompt") + + +def test_two_prompts_that_overflow_the_pool_together_both_finish(): + # Neither slot ever generates before the pool is full: both are still processing their + # prompts. A prompt-processing slot is between two chunks of its prompt, which is as + # clean a boundary as the one between two sampled tokens, so it is parked the same way. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + prompt_a, n_a = _prompt_of_about(150, "Alpha") + prompt_b, n_b = _prompt_of_about(150, "Bravo") + n_predict = 16 + assert n_a + n_predict <= 256 and n_b + n_predict <= 256 + assert n_a + n_b + 2 * n_predict > 256 + + results = parallel_function_calls([ + (_complete, (n_predict, prompt_a)), + (_complete, (n_predict, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert len(res.body["tokens"]) == n_predict + + +def test_a_generating_slot_and_a_large_prompt_both_finish(): + # One slot is generating a long answer to a short prompt when a large prompt arrives + # beside it. Together they need far more than the pool has. The prompt is admitted + # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish. + # This model produces a thousand tokens a second, so the second request is sent right + # behind the first rather than after a delay: its prompt takes several batches to + # process, which is enough for the two to overlap however fast the first one runs. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + prompt_b, n_b = _prompt_of_about(150, "Charlie") + # b lives long enough for the two to collide: the first run of this used 16 tokens + # and b was finished and purged before a had grown into it + n_predict_a = 230 + n_predict_b = 90 + assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 + assert 8 + n_predict_a + n_b + n_predict_b > 256 + + def _late(n_predict, prompt): + time.sleep(0.02) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_a, "Hi how are you")), + (_late, (n_predict_b, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + + assert results[0].status_code == 200 + assert results[0].body["timings"]["predicted_n"] == n_predict_a + assert results[1].status_code == 200 + assert results[1].body["timings"]["predicted_n"] == n_predict_b + + +def test_preempt_ram_zero_disables_preemption(): + # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the + # KV-full path ends the requests the way it always did. + global server + server.n_ctx = 256 + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "preempted:" not in text + assert "Context size has been exceeded" in text + assert any(res.status_code != 200 for res in results) From 63cdac4d7d9c7679d08e2dd3ee37120c75774b5c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 01:13:28 +0000 Subject: [PATCH 4/4] server: report preemption through /metrics and /slots Counters n_preempt_total and n_resume_total, gauges requests_preempted and preempt_ram_bytes, and is_preempted plus n_preempt on each /slots entry, so a client can tell a parked request from a slow one and an operator can see the parked host RAM. A parked slot no longer counts as busy in n_busy_slots_per_decode, since it took no part in the decode. --- tools/server/README.md | 4 +++ tools/server/server-common.h | 4 +++ tools/server/server-context.cpp | 19 ++++++++---- tools/server/server-task.cpp | 16 ++++++++++ tools/server/server-task.h | 2 ++ tools/server/tests/unit/test_preempt.py | 40 +++++++++++++++++++++++++ 6 files changed, 79 insertions(+), 6 deletions(-) diff --git a/tools/server/README.md b/tools/server/README.md index 52e5c1b8c7db..7b4a0330340f 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -1139,6 +1139,10 @@ In *router mode* the query param `?model={model_id}` has to be set. This endpoin | `llamacpp:spec_decode_num_accepted_tokens_total` | Counter | Total draft tokens accepted by the target model (0 when spec-decode is off). | | `llamacpp:spec_decode_num_drafts_total` | Counter | Total speculative decoding verification steps (0 when spec-decode is off). | | `llamacpp:spec_decode_num_accepted_tokens_per_pos_total` | Counter | Accepted tokens per draft position (labeled `position="N"`; absent when spec-decode is off or before the first completed speculative request). | +| `llamacpp:n_preempt_total` | Counter | Slots parked to make room in the unified KV cache (0 unless `--kv-unified` with more than one slot). | +| `llamacpp:n_resume_total` | Counter | Parked slots put back. | +| `llamacpp:requests_preempted` | Gauge | Requests currently parked, waiting for room in the unified KV cache. | +| `llamacpp:preempt_ram_bytes` | Gauge | Host RAM held by parked sequences. | ### POST `/slots/{id_slot}?action=save`: Save the prompt cache of the specified slot to a file. diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f8ea82ef4cf5..f0cf76b8c501 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,6 +467,10 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; + // [TAG_PREEMPT] slots parked to make room in the unified KV pool, and put back + uint64_t n_preempt = 0; + uint64_t n_resume = 0; + uint64_t n_draft_tokens = 0; // Total draft tokens generated uint64_t n_draft_accepted = 0; // Draft tokens actually accepted uint64_t n_draft_verif_steps = 0; // Total draft token verification steps by the target model diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a41e5feda5d6..6723c51397ed 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -796,6 +796,8 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, + {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"n_preempt", n_preempt}, }; const auto & ptask = task ? task : task_prev; @@ -2546,11 +2548,15 @@ struct server_context_impl { case SERVER_TASK_TYPE_METRICS: { int n_processing_slots = 0; + int n_preempted_slots = 0; for (server_slot & slot : slots) { if (slot.is_processing()) { n_processing_slots++; } + if (slot.state == SLOT_STATE_PREEMPTED) { + n_preempted_slots++; + } } SRV_DBG("n_processing_slots = %d\n", n_processing_slots); @@ -2558,6 +2564,8 @@ struct server_context_impl { res->id = task.id; res->n_processing_slots = n_processing_slots; res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size(); + res->n_preempted_slots = n_preempted_slots; + res->preempt_ram_bytes = preempt_ram_used(); res->metrics = metrics; if (task.metrics_reset_bucket) { @@ -2839,8 +2847,6 @@ struct server_context_impl { // [TAG_PREEMPT] server-side request preemption // - int64_t n_preempt_total = 0; - int64_t n_resume_total = 0; // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, // whether or not the pool is under pressure. It exists to answer the only question that @@ -3084,7 +3090,7 @@ struct server_context_impl { break; } - n_resume_total++; + metrics.n_resume++; SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", (ggml_time_us() - best->t_preempt_us) / 1e6, @@ -3101,7 +3107,7 @@ struct server_context_impl { (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && preempt_fits_budget(slot) && slot.preempt_save()) { - n_preempt_total++; + metrics.n_preempt++; SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); @@ -3137,7 +3143,7 @@ struct server_context_impl { break; // could not park it; the existing retry ladder is still behind us } - n_preempt_total++; + metrics.n_preempt++; SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, @@ -4422,7 +4428,8 @@ struct server_context_impl { void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { metrics.n_decode++; for (const auto & slot : slots) { - if (slot.is_processing()) { + // [TAG_PREEMPT] a parked slot is processing but took no part in this decode + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { metrics.n_busy_slots++; } metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens()); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..9afe3c7f06a8 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1562,6 +1562,14 @@ std::string server_task_result_metrics::to_metrics() { "spec_decode_num_drafts_total", "Speculative: Total speculative decoding verification steps", (double) metrics.n_draft_verif_steps + }, { + "n_preempt_total", + "Preemption: Total slots parked to make room in the unified KV cache", + (double) metrics.n_preempt + }, { + "n_resume_total", + "Preemption: Total parked slots put back", + (double) metrics.n_resume }, }; @@ -1586,6 +1594,14 @@ std::string server_task_result_metrics::to_metrics() { "n_busy_slots_per_decode", "Average number of busy slots per llama_decode() call", (double) metrics.n_busy_slots / std::max((double) metrics.n_decode, 1.0) + }, { + "requests_preempted", + "Preemption: Number of requests currently parked, waiting for room in the unified KV cache", + (double) n_preempted_slots + }, { + "preempt_ram_bytes", + "Preemption: Host RAM held by parked sequences", + (double) preempt_ram_bytes }, }; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..00734924bc63 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -494,6 +494,8 @@ struct server_task_result_metrics : server_task_result { // these are immediate stats, not accumulated (server_metrics is cumulative) int n_processing_slots = 0; int n_tasks_deferred = 0; + int n_preempted_slots = 0; // [TAG_PREEMPT] processing slots currently parked + size_t preempt_ram_bytes = 0; // [TAG_PREEMPT] host RAM their parked sequences hold server_metrics metrics; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 2e57ebbdfd97..0da885bcafd5 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -227,3 +227,43 @@ def test_preempt_ram_zero_disables_preemption(): assert "preempted:" not in text assert "Context size has been exceeded" in text assert any(res.status_code != 200 for res in results) + + +def test_metrics_and_slots_report_the_parked_state(): + # A client that wants to tell a parked chat from a slow one reads /slots, and an + # operator reads /metrics. Both must show the preemption happening, and the counters + # must survive the requests finishing. + global server + server.n_ctx = 256 + server.server_metrics = True + server.start() + + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + for slot in res.body: + assert slot["is_preempted"] is False + assert slot["n_preempt"] == 0 + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + for res in results: + assert res.status_code == 200 + + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + metrics = {} + for line in res.body.splitlines(): + if line.startswith("llamacpp:"): + name, value = line.split(" ", 1) + metrics[name[len("llamacpp:"):]] = float(value) + assert metrics["n_preempt_total"] >= 1 + assert metrics["n_resume_total"] == metrics["n_preempt_total"] + assert metrics["requests_preempted"] == 0 + assert metrics["preempt_ram_bytes"] == 0 + + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot"