diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b18fa4e2f23a..f57baff24dfc 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -18,6 +18,7 @@ #include "mtmd-helper.h" #include +#include #include #include #include @@ -77,6 +78,7 @@ enum slot_state { // 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 int64_t PREEMPT_KEEPALIVE_MS = 2000; // SSE keepalive period while a streaming slot is parked constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected // [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody @@ -87,6 +89,18 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is // LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest // parked, and a smaller slot may pass a head that does not fit. // LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. +// [TAG_PREEMPT] the SSE comment for a park or a resume. A request with several prompts +// streams them through one reader, so the comment names the prompt it is about, except for +// prompt 0, whose comment stays the bare form a single-prompt client matches on. +static std::string preempt_notice_comment(const server_task_result_preempt_notice & notice) { + std::string res = notice.parked ? ": preempted" : ": resumed"; + + if (notice.index > 0) { + res += " " + std::to_string(notice.index); + } + + return res + "\n\n"; +} 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 constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long @@ -2189,6 +2203,25 @@ struct server_context_impl { queue_results.send(std::move(res)); } + // [TAG_PREEMPT] tell a streaming client that its slot was parked or restored. The + // HTTP layer turns this into an SSE comment, so a client that does not know about + // preemption sees nothing, and one that does can show a pause instead of a stall. + void send_preempt_notice(server_slot & slot, bool parked) { + if (!slot.task || !slot.task->params.stream) { + return; + } + + auto res = std::make_unique(); + + res->id = slot.task->id; + res->index = slot.task->index; + res->id_slot = slot.id; + res->parked = parked; + res->n_preempt = slot.n_preempt; + + queue_results.send(std::move(res)); + } + void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) { auto res = std::make_unique(); @@ -3240,6 +3273,18 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + // a started slot whose request the STARTED block is about to reject gets its + // error on its own pass, and nothing before it: a park notice would open the + // stream and turn that error into 200 plus an in-stream one + if (slot.state == SLOT_STATE_STARTED) { + std::string msg; + error_type type = ERROR_TYPE_SERVER; + + if (slot_prompt_rejected(slot, msg, type)) { + continue; + } + } + if (!preempt_fits_budget(slot)) { continue; } @@ -3449,6 +3494,8 @@ struct server_context_impl { metrics.n_preempt++; + send_preempt_notice(slot, true); + SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn%s, preemptions %d\n", slot.n_ctx_shift, slot.prompt.n_tokens(), slot.preempt_state_size() / (1024.0 * 1024.0), @@ -3489,6 +3536,8 @@ struct server_context_impl { metrics.n_resume++; + send_preempt_notice(*best, false); + 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(), @@ -3506,6 +3555,8 @@ struct server_context_impl { slot.preempt_save()) { metrics.n_preempt++; + send_preempt_notice(slot, true); + 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)); } @@ -3546,6 +3597,8 @@ struct server_context_impl { metrics.n_preempt++; + send_preempt_notice(*victim, true); + 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, @@ -3555,6 +3608,60 @@ struct server_context_impl { } } + // the checks a slot's request has to pass before its prompt is processed, run from the + // SLOT_STATE_STARTED block below. true when the request is rejected, with the message and + // the type of the error it gets. The empty prompt is not here: it is a final response and + // not an error. + // [TAG_PREEMPT] the planner asks the same question before it parks a started slot, so a + // request that is about to be errored is never given a park notice ahead of its error: a + // notice opens the stream, and the client would get 200 plus an in-stream error where the + // non-stream 4xx belongs. + bool slot_prompt_rejected(const server_slot & slot, std::string & msg, error_type & type) const { + if (!slot.task) { + return false; + } + + // TODO: support memory-less logits computation + if (slot.task->need_logits() && !llama_get_memory(ctx_tgt)) { + msg = "the current context does not logits computation. skipping"; + type = ERROR_TYPE_SERVER; + return true; + } + + if (!slot.can_split()) { + const int32_t n_ubatch = llama_n_ubatch(ctx_tgt); + + if (slot.task->n_tokens() > n_ubatch) { + msg = string_format( + "input (%d tokens) is too large to process. increase the physical batch " + "size (current batch size: %d)", + slot.task->n_tokens(), n_ubatch); + type = ERROR_TYPE_SERVER; + return true; + } + + if (slot.task->n_tokens() > slot.n_ctx) { + msg = string_format( + "input (%d tokens) is larger than the max context size (%d tokens). skipping", + slot.task->n_tokens(), slot.n_ctx); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; + } + + return false; + } + + if (slot.task->n_tokens() >= slot.n_ctx) { + msg = string_format( + "request (%d tokens) exceeds the available context size (%d tokens), try increasing it", + slot.task->n_tokens(), slot.n_ctx); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; + } + + return false; + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -3961,46 +4068,18 @@ struct server_context_impl { return; } - // TODO: support memory-less logits computation - if (slot.task->need_logits() && !llama_get_memory(ctx_tgt)) { - send_error(slot, "the current context does not logits computation. skipping", ERROR_TYPE_SERVER); - slot.release(); - return; - } - - if (!slot.can_split()) { - if (slot.task->n_tokens() > n_ubatch) { - send_error(slot, - string_format( - "input (%d tokens) is too large to process. increase the physical batch " - "size (current batch size: %d)", - slot.task->n_tokens(), n_ubatch), - ERROR_TYPE_SERVER); - slot.release(); - return; - } + { + std::string msg; + error_type type = ERROR_TYPE_SERVER; - if (slot.task->n_tokens() > slot.n_ctx) { - send_error( - slot, - string_format( - "input (%d tokens) is larger than the max context size (%d tokens). skipping", - slot.task->n_tokens(), slot.n_ctx), - ERROR_TYPE_EXCEED_CONTEXT_SIZE); - slot.release(); - return; - } - } else { - if (slot.task->n_tokens() >= slot.n_ctx) { - send_error(slot, - string_format("request (%d tokens) exceeds the available context size (%d " - "tokens), try increasing it", - slot.task->n_tokens(), slot.n_ctx), - ERROR_TYPE_EXCEED_CONTEXT_SIZE); + if (slot_prompt_rejected(slot, msg, type)) { + send_error(slot, msg, type); slot.release(); return; } + } + if (slot.can_split()) { if (slot.task->params.cache_prompt) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); @@ -4496,6 +4575,8 @@ struct server_context_impl { metrics.n_preempt++; n_parked++; + send_preempt_notice(*victim, true); + SLT_WRN(*victim, "preempted as a last resort: %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, @@ -5275,37 +5356,59 @@ std::unique_ptr server_routes::handle_completions_impl( // in streaming mode, the first error must be treated as non-stream response // this is to match the OAI API behavior // ref: https://github.com/ggml-org/llama.cpp/pull/16486#discussion_r2419657309 + // [TAG_PREEMPT] a slot can be parked while still processing its prompt, before any + // token exists. Those notices arrive ahead of the first real result; keep them and + // send them in front of it, so the client learns about the wait it just had. + std::string preempt_prefix; + std::set parked_idx; // prompts of this request that are parked right now auto first_result = rd.next(req.should_stop); - if (first_result == nullptr) { - GGML_ASSERT(req.should_stop()); - return res; // connection is closed - } + if (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { + // [TAG_PREEMPT] parked before any token exists. The stream starts now, with the + // notice, so the parked keepalive runs through the wait instead of the client + // seeing nothing until the slot resumes; the first ordinary result follows in + // the stream, an error included, since the response has already begun. + const auto * notice = static_cast(first_result.get()); + preempt_prefix = preempt_notice_comment(*notice); + if (notice->parked) { + parked_idx.insert(notice->index); + } else { + parked_idx.erase(notice->index); + } + first_result.reset(); + } else { + if (first_result == nullptr) { + GGML_ASSERT(req.should_stop()); + return res; // connection is closed + } - if (first_result->is_error()) { - res->error(first_result->to_json()); - return res; - } + if (first_result->is_error()) { + res->error(first_result->to_json()); + return res; + } - GGML_ASSERT( - dynamic_cast(first_result.get()) != nullptr || - dynamic_cast (first_result.get()) != nullptr - ); + GGML_ASSERT( + dynamic_cast(first_result.get()) != nullptr || + dynamic_cast (first_result.get()) != nullptr + ); + } - // next responses are streamed - // to be sent immediately - json first_result_json = first_result->to_json(); + json first_result_json = first_result ? first_result->to_json() : json(nullptr); if (first_result_json == nullptr) { - res->data = ""; // simply send HTTP headers and status code + res->data = preempt_prefix; // simply send HTTP headers and status code } else if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { - res->data = format_anthropic_sse(first_result_json); + res->data = preempt_prefix + format_anthropic_sse(first_result_json); } else if (res_type == TASK_RESPONSE_TYPE_OAI_RESP) { - res->data = format_oai_resp_sse(first_result_json); + res->data = preempt_prefix + format_oai_resp_sse(first_result_json); } else { - res->data = format_oai_sse(first_result_json); + res->data = preempt_prefix + format_oai_sse(first_result_json); } res->status = 200; res->content_type = "text/event-stream"; - res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool { + res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked_idx](std::string & output) mutable -> bool { + // [TAG_PREEMPT] the keepalive runs while ANY prompt of the request is parked: with + // several prompts in one stream, one resuming does not mean the others did + const bool parked = !parked_idx.empty(); + static auto format_error = [](task_response_type res_type, const json & res_json) { if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { return format_anthropic_sse({ @@ -5356,10 +5459,17 @@ std::unique_ptr server_routes::handle_completions_impl( // receive subsequent results bool timeout = false; int64_t start_time = ggml_time_ms(); - auto result = rd.next([&timeout, &start_time, sse_ping_interval, &effective_should_stop]() { + // [TAG_PREEMPT] a parked slot produces nothing for as long as the pool is + // full, so while parked the ping runs at least every 2 s whether or not + // --sse-ping asked for one, and is named, so a client can tell "waiting for + // cells" from "slow". A shorter interval the request asked for is kept: a + // client that wants a ping every second wants it most while nothing else comes. + const int64_t ping_cfg = sse_ping_interval > 0 ? (int64_t) sse_ping_interval * 1000 : -1; + const int64_t ping_ms = parked ? (ping_cfg > 0 ? std::min(ping_cfg, PREEMPT_KEEPALIVE_MS) : PREEMPT_KEEPALIVE_MS) : ping_cfg; + auto result = rd.next([&timeout, &start_time, ping_ms, &effective_should_stop]() { if (effective_should_stop()) { return true; // should_stop condition met - } else if (sse_ping_interval > 0 && ggml_time_ms() - start_time > (int64_t)sse_ping_interval * 1000) { + } else if (ping_ms > 0 && ggml_time_ms() - start_time > ping_ms) { timeout = true; return true; // timeout } @@ -5369,7 +5479,7 @@ std::unique_ptr server_routes::handle_completions_impl( if (timeout) { // some clients may time out (e.g. undici) will time out if no data is received for a while, so we need to send a ping to keep the connection alive SRV_DBG("%s", "sending SSE ping\n"); - output = ":\n\n"; + output = parked ? ": preempt-keepalive\n\n" : ":\n\n"; return true; } @@ -5385,12 +5495,28 @@ std::unique_ptr server_routes::handle_completions_impl( output = format_error(res_type, res_json); SRV_DBG("%s", "error received during streaming, terminating stream\n"); return false; // terminate on error + } else if (const auto * notice = dynamic_cast(result.get())) { + // [TAG_PREEMPT] an SSE comment: invisible to clients that do not know + // about preemption, a pause indicator for the ones that do + if (notice->parked) { + parked_idx.insert(notice->index); + } else { + parked_idx.erase(notice->index); + } + output = preempt_notice_comment(*notice); } else { GGML_ASSERT( dynamic_cast(result.get()) != nullptr || dynamic_cast(result.get()) != nullptr ); json res_json = result->to_json(); + if (res_json.is_null()) { + // [TAG_PREEMPT] the signal a prompt sends before its first token, so + // that the headers go out, carries no data. Normally it is the first + // result and only opens the stream; after a notice opened the stream + // it has nothing to add, and the sender skips an empty chunk. + return true; + } if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { output = format_anthropic_sse(res_json); } else if (res_type == TASK_RESPONSE_TYPE_OAI_RESP) { diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d86..bd74db89408e 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -448,6 +448,18 @@ server_task_result_ptr server_response::recv(const std::unordered_set & id_ } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { + // [TAG_PREEMPT] The timeout is a deadline, not a per-wait duration. + // + // send() notify_all()s on one condition variable for every result of every task, so a + // reader waiting on a task that is producing nothing is woken by every token every other + // task produces. With wait_for() each of those wakeups restarted the wait, and on a busy + // server the timeout was never reached at all: whoever was waiting for a quiet task + // waited forever, however small the timeout they asked for. That is exactly the + // situation of a parked slot, which by definition exists because the others are busy, so + // neither its 2 s keepalive nor the ordinary --sse-ping could ever fire for it. Waiting + // until a fixed point instead makes the timeout mean what every caller reads it as. + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout); + while (true) { std::unique_lock lock(mutex_results); @@ -459,7 +471,7 @@ server_task_result_ptr server_response::recv_with_timeout(const std::unordered_s } } - std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout)); + std::cv_status cr_res = condition_results.wait_until(lock, deadline); if (!running) { RES_DBG("%s : queue result stop\n", __func__); std::terminate(); // we cannot return here since the caller is HTTP code diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 9afe3c7f06a8..48fc77c960ff 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1023,6 +1023,13 @@ void server_task_result_cmpl_partial::update(task_result_state & state) { } } +json server_task_result_preempt_notice::to_json() { + return json { + {"preempted", parked}, + {"n_preempt", n_preempt}, + }; +} + json server_task_result_cmpl_partial::to_json() { GGML_ASSERT(is_updated && "update() must be called before to_json()"); if (is_begin) { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 00734924bc63..f20f891655f5 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -392,6 +392,20 @@ struct server_task_result_cmpl_final : server_task_result { json to_json_anthropic_stream(); }; +// [TAG_PREEMPT] out-of-band notice for a streaming task whose slot was parked or restored. +// Serialised as an SSE comment (": preempted", ": resumed"), which every existing client +// ignores, so the body of the response is unchanged by preemption. Never sent to a +// non-streaming task. +struct server_task_result_preempt_notice : server_task_result { + bool parked = false; // true when the slot was just parked, false when restored + int32_t n_preempt = 0; // how many times this task has been parked so far + + virtual bool is_stop() override { + return false; + } + virtual json to_json() override; +}; + struct server_task_result_cmpl_partial : server_task_result { std::string content; llama_tokens tokens; diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py new file mode 100644 index 000000000000..6fa71c55473d --- /dev/null +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -0,0 +1,286 @@ +import os +import tempfile +import threading +import time +import pytest +import requests +from utils import * + +# [TAG_PREEMPT] A streaming client is told when its slot is parked and when it is +# restored, as SSE comments, and the body is byte for byte what it is without any park. +# Comments are legal SSE that every existing client ignores; a client that knows about +# preemption can show "paused" instead of a dead stream, and a keepalive every 2 s while +# parked keeps proxies and read timeouts from giving up on a wait that is by design long. + +server = ServerPreset.tinyllama2() + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + server.temperature = 0.0 + server.seed = 42 + # A build without libcurl cannot fetch the model itself; point it at a local copy. + local = os.environ.get("LLAMA_SERVER_TEST_MODEL") + if local: + server.model_hf_repo = None + server.model_hf_file = None + server.model_file = local + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + + +def _stream_raw(path: str, data: dict) -> tuple[list[str], list[str]]: + """The SSE lines of one streaming request: (comment lines, data lines).""" + url = f"http://{server.server_host}:{server.server_port}{path}" + res = requests.post(url, json=data, stream=True) + assert res.status_code == 200 + comments, datas = [], [] + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line.startswith(":"): + comments.append(line) + elif line.startswith("data: "): + datas.append(line[6:]) + return comments, datas + + +def _content(datas: list[str]) -> str: + out = "" + for d in datas: + if d == "[DONE]": + break + j = json.loads(d) + if "content" in j: + out += j["content"] + for ch in j.get("choices", []) or []: + delta = ch.get("delta") or {} + out += delta.get("content") or "" + return out + + +def _completion_payload(n_predict: int) -> dict: + return { + "n_predict": n_predict, + "prompt": "Hi how are you", + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + "stream": True, + } + + +def _chat_payload(n_predict: int) -> dict: + return { + "max_tokens": n_predict, + "messages": [{"role": "user", "content": "Hi how are you"}], + "temperature": 0.0, + "seed": 42, + "stream": True, + } + + +def test_a_stream_announces_its_parks_and_the_body_is_unchanged(): + global server + server.n_ctx = 512 + server.start() + ref_comments, ref_datas = _stream_raw("/completion", _completion_payload(64)) + assert not any(c.startswith(": preempted") or c.startswith(": resumed") for c in ref_comments) + assert _content(ref_datas) + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + comments, datas = _stream_raw("/completion", _completion_payload(64)) + + parked = [c for c in comments if c == ": preempted"] + resumed = [c for c in comments if c == ": resumed"] + assert len(parked) >= 6, comments + assert len(resumed) == len(parked), comments + # Every park is followed by its resume before the next park. + seq = [c for c in comments if c in (": preempted", ": resumed")] + assert seq == [": preempted", ": resumed"] * len(parked), seq + # The generated text is byte for byte the unparked text, token by token. Only the + # final chunk's wall-clock timings differ between the two runs. + def _pieces(ds): + return [json.loads(d).get("content") for d in ds if d != "[DONE]"] + + assert _pieces(datas) == _pieces(ref_datas) + assert _content(datas) == _content(ref_datas) + final, ref_final = json.loads(datas[-1]), json.loads(ref_datas[-1]) + assert final["tokens_predicted"] == ref_final["tokens_predicted"] == 64 + + +def test_the_oai_chat_stream_carries_the_same_comments(): + global server + server.n_ctx = 512 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + comments, datas = _stream_raw("/v1/chat/completions", _chat_payload(48)) + assert ": preempted" in comments and ": resumed" in comments + assert datas[-1] == "[DONE]" + assert _content(datas) + + +def test_non_streaming_requests_see_nothing(): + global server + server.n_ctx = 512 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + res = server.make_request("POST", "/completion", data={ + "n_predict": 32, + "prompt": "Hi how are you", + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + }) + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == 32 + assert "preempted" not in res.body + + +def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): + # The pair from test_preempt: each alone fits, together they do not, so one is + # parked until the other finishes. The parked stream must carry the comments and + # finish with its full output. + global server + server.n_ctx = 256 + server.start() + + n_predict = 160 + p1 = _completion_payload(n_predict) | {"prompt": "Once upon a time there was a brave knight who"} + p2 = _completion_payload(n_predict) | {"prompt": "The quick brown fox jumps over the lazy dog and"} + results = parallel_function_calls([ + (_stream_raw, ("/completion", p1)), + (_stream_raw, ("/completion", p2)), + ]) + announced = 0 + for comments, datas in results: + final = json.loads([d for d in datas if d != "[DONE]"][-1]) + assert final["timings"]["predicted_n"] == n_predict + assert final["truncated"] is False + if ": preempted" in comments: + announced += 1 + assert ": resumed" in comments + assert announced >= 1, [r[0] for r in results] + + +def test_a_stream_parked_before_its_first_token_starts_with_the_notice(): + # A request parked while still processing its prompt has no token to send yet. The + # response must not wait for one: it starts with the notice, so the client sees + # "paused" and gets the keepalive at once, instead of a silent connection that only + # opens when the slot resumes. + global server + # The resident keeps growing towards the whole pool; the newcomer's prompt is larger + # than what is free beside it, so the planner parks the newcomer before it has a token. + global server + server.n_ctx = 512 + server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once + server.start() + url = f"http://{server.server_host}:{server.server_port}/completion" + first = _completion_payload(390) | {"prompt": " ".join(["Once upon a time there was a brave knight who"] * 6)} + second = _completion_payload(32) | {"prompt": " ".join(["The quick brown fox jumps over the lazy dog and"] * 14)} + + timeline = [] + lock = threading.Lock() + + def _run(name, payload, started=None): + res = requests.post(url, json=payload, stream=True) + assert res.status_code == 200 + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if not line: + continue + with lock: + timeline.append((time.monotonic(), name, line)) + if started is not None and line.startswith("data: "): + started.set() + + started = threading.Event() + t = threading.Thread(target=_run, args=("first", first, started)) + t.start() + assert started.wait(30) + _run("second", second) + t.join(60) + + second_lines = [(ts, line) for ts, name, line in timeline if name == "second"] + first_end = max(ts for ts, name, _ in timeline if name == "first") + # The notice is the very first thing on the wire, and it arrives while the other + # stream is still running, not when it has finished and the parked slot resumes. + assert second_lines[0][1] == ": preempted", second_lines[:3] + assert second_lines[0][0] < first_end + events = [line for _, line in second_lines if line in (": preempted", ": resumed") or line.startswith("data: ")] + assert events[0] == ": preempted" and events[1] == ": resumed" and events[2].startswith("data: "), events[:3] + datas = [line[6:] for _, line in second_lines if line.startswith("data: ")] + assert _content(datas) + final = json.loads([d for d in datas if d != "[DONE]"][-1]) + assert final["tokens_predicted"] == 32 + + +def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): + # The rotation from test_preempt: a resident cycling through context shifts holds the + # pool, and after the head has waited its turn the resident is parked in its place. + # That park is a park like any other, so its stream must say so, and every notice + # must be paired: no stream ends with a park it was never told about. + global server + server.n_ctx = 256 + server.enable_ctx_shift = True + server.start() + n_predict = 12000 + p1 = _completion_payload(n_predict) | {"prompt": "Once upon a time there was a brave knight who"} + p2 = _completion_payload(n_predict) | {"prompt": "The quick brown fox jumps over the lazy dog and"} + results = parallel_function_calls([ + (_stream_raw, ("/completion", p1)), + (_stream_raw, ("/completion", p2)), + ]) + n_parked = 0 + for comments, datas in results: + final = json.loads([d for d in datas if d != "[DONE]"][-1]) + assert final["tokens_predicted"] == n_predict + seq = [c for c in comments if c in (": preempted", ": resumed")] + assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq + n_parked += len(seq) // 2 + # Both streams took turns: at least one park each, so at least two in all. + assert n_parked >= 2, [r[0] for r in results] + + +def test_an_oversized_prompt_is_errored_instead_of_parked(): + # A slot that has just been given a task has not passed the prompt checks yet: they + # run on its first pass through update_slots. Parked before that, it would be told + # ": preempted" first, and the notice opens the stream, so a prompt larger than the + # context would come back as 200 plus an in-stream error instead of the plain error + # response it gets with nothing running. The planner leaves such a slot alone. + global server + server.n_ctx = 512 + server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + url = f"http://{server.server_host}:{server.server_port}/completion" + resident = _completion_payload(390) | {"prompt": " ".join(["Once upon a time there was a brave knight who"] * 6)} + oversized = _completion_payload(16) | {"prompt": " ".join(["The quick brown fox jumps over the lazy dog and"] * 80)} + + started = threading.Event() + + def _run_resident(): + res = requests.post(url, json=resident, stream=True) + assert res.status_code == 200 + for raw in res.iter_lines(): + if raw.decode("utf-8").startswith("data: "): + started.set() + + t = threading.Thread(target=_run_resident) + t.start() + try: + assert started.wait(60) + res = requests.post(url, json=oversized, stream=True) + body = res.text + assert res.status_code != 200, body + assert not body.lstrip().startswith(":"), body + assert "error" in json.loads(body), body + finally: + t.join(120)