server: tell a streaming client when its slot is parked and restored

tools/server/server-context.cpp

@@ -18,6 +18,7 @@
 #include "mtmd-helper.h"
 
 #include <algorithm>
+#include <set>
 #include <cstddef>
 #include <cstring>
 #include <cinttypes>
@@ -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<server_task_result_preempt_notice>();
+
+        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<server_task_result_cmpl_partial>();
 
@@ -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_res_generator> 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<size_t> 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<server_task_result_preempt_notice*>(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<server_task_result_preempt_notice*>(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<server_task_result_cmpl_partial*>(first_result.get()) != nullptr ||
-            dynamic_cast<server_task_result_cmpl_final*>  (first_result.get()) != nullptr
-        );
+            GGML_ASSERT(
+                dynamic_cast<server_task_result_cmpl_partial*>(first_result.get()) != nullptr ||
+                dynamic_cast<server_task_result_cmpl_final*>  (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_res_generator> 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_res_generator> 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_res_generator> 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<server_task_result_preempt_notice*>(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<server_task_result_cmpl_partial*>(result.get()) != nullptr
                         || dynamic_cast<server_task_result_cmpl_final*>(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) {