From f06eb616334caae202e31db67e56b07eea37fc9e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 3 Sep 2026 21:43:03 +0000 Subject: [PATCH] fix(llama-cpp): stop a generation whose stream is gone grpc::ServerWriter::Write() returns false once the peer is gone, and PredictStream ignored that result at every call site. The handler kept pulling decoded tokens and writing them into a dead stream, so the llama.cpp slot stayed busy until the generation ended on its own terms. A model configured with max_tokens 0 and a large context ends on its own terms only at the context limit. On a 35B model at ~41 t/s a 120k context is about fifty minutes, and a slot held that long is a slot every other request for that model queues behind. Two abandoned requests were enough to make a node with free VRAM and a healthy control plane serve nothing: new requests timed out waiting for a slot, each timeout abandoned another generation, and the node fell further behind the longer it ran. Track the peer instead. The first failed write retires it for good, since a stream never recovers, and the RPC's own cancellation flag folds into the same predicate so the loop has one condition to test. Returning early is what frees the slot: ~server_response_reader() posts SERVER_TASK_TYPE_CANCEL for whatever is still decoding. TTSStream already checked Write(); this brings PredictStream in line. Cancellation stays cooperative and is checked between decoded results, so a batch already in flight may finish before the request stops. Assisted-by: Claude:claude-opus-5 --- backend/cpp/llama-cpp/grpc-server.cpp | 29 +++++++--- backend/cpp/llama-cpp/stream_peer.h | 44 ++++++++++++++ backend/cpp/llama-cpp/stream_peer_test.cpp | 67 ++++++++++++++++++++++ docs/content/features/backends.md | 23 ++++++++ 4 files changed, 156 insertions(+), 7 deletions(-) create mode 100644 backend/cpp/llama-cpp/stream_peer.h create mode 100644 backend/cpp/llama-cpp/stream_peer_test.cpp diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index 95345aaf1203..0205df9796dd 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -56,6 +56,7 @@ #include "thread_params.h" #include "message_content.h" #include "passthrough_options.h" +#include "stream_peer.h" #include "tts_request_options.h" #include #include @@ -2266,6 +2267,11 @@ class BackendServiceImpl final : public backend::Backend::Service { // such concept, so there is nothing to emit — the real tokens arrive in // the loop below. Feeding this null into build_reply_from_json would // throw (uncaught) and surface as a generic RPC error. + // A write that returns false means the peer is gone for good. Track it + // so the loop below stops decoding instead of feeding a dead stream — + // see stream_peer.h for why that matters to everyone else's requests. + llama_grpc::StreamPeer peer; + if (first_res_json.is_null()) { // skip the begin-of-stream marker } else if (first_res_json.is_array()) { @@ -2278,17 +2284,21 @@ class BackendServiceImpl final : public backend::Backend::Service { if (!is_role_init) { attach_chat_deltas(reply, first_result.get()); } - writer->Write(reply); + peer.observe_write(writer->Write(reply)); + if (peer.gone()) { + break; + } } } else { auto reply = build_reply_from_json(first_res_json, first_result.get()); attach_chat_deltas(reply, first_result.get()); - writer->Write(reply); + peer.observe_write(writer->Write(reply)); } // Process subsequent results while (rd.has_next()) { - if (context->IsCancelled()) { + peer.observe_cancelled(context->IsCancelled()); + if (peer.gone()) { break; } @@ -2309,17 +2319,22 @@ class BackendServiceImpl final : public backend::Backend::Service { if (!is_role_init) { attach_chat_deltas(reply, result.get()); } - writer->Write(reply); + peer.observe_write(writer->Write(reply)); + if (peer.gone()) { + break; + } } } else { auto reply = build_reply_from_json(res_json, result.get()); attach_chat_deltas(reply, result.get()); - writer->Write(reply); + peer.observe_write(writer->Write(reply)); } } - // Check if context was cancelled during processing - if (context->IsCancelled()) { + // Returning here is what releases the slot: ~server_response_reader() + // posts SERVER_TASK_TYPE_CANCEL for whatever is still decoding. + peer.observe_cancelled(context->IsCancelled()); + if (peer.gone()) { return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client"); } diff --git a/backend/cpp/llama-cpp/stream_peer.h b/backend/cpp/llama-cpp/stream_peer.h new file mode 100644 index 000000000000..065a31c95688 --- /dev/null +++ b/backend/cpp/llama-cpp/stream_peer.h @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: MIT +#pragma once + +namespace llama_grpc { + +// Tracks whether a server-streaming RPC still has somewhere to send tokens. +// +// grpc::ServerWriter::Write() returns false once the peer is gone, and a +// stream never recovers afterwards. Ignoring that result is not harmless: the +// handler goes on draining decoded tokens into a dead stream, so the llama.cpp +// slot stays busy for the rest of the request's token budget. A model config +// with no max_tokens and a large context turns that into tens of minutes per +// abandoned request, and the slots are exactly what every other request queues +// behind. +// +// Returning as soon as the peer is gone is what frees the slot: the handler's +// server_response_reader then goes out of scope and its destructor posts +// SERVER_TASK_TYPE_CANCEL for whatever is still decoding. +class StreamPeer { +public: + // Records the outcome of a Write(). Once a write has failed the peer stays + // gone -- a later write cannot succeed on a broken stream. + void observe_write(bool ok) noexcept { + if (!ok) { + gone_ = true; + } + } + + // Folds in the RPC's own cancellation flag, so callers have a single + // predicate to test rather than two that can disagree. + void observe_cancelled(bool cancelled) noexcept { + if (cancelled) { + gone_ = true; + } + } + + bool gone() const noexcept { return gone_; } + bool alive() const noexcept { return !gone_; } + +private: + bool gone_ = false; +}; + +} // namespace llama_grpc diff --git a/backend/cpp/llama-cpp/stream_peer_test.cpp b/backend/cpp/llama-cpp/stream_peer_test.cpp new file mode 100644 index 000000000000..2681f182c4b6 --- /dev/null +++ b/backend/cpp/llama-cpp/stream_peer_test.cpp @@ -0,0 +1,67 @@ +#include "stream_peer.h" + +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char *what) { + if (!condition) { + std::fprintf(stderr, "FAIL: %s\n", what); + ++failures; + } +} + +} // namespace + +int main() { + { + llama_grpc::StreamPeer peer; + check(peer.alive(), "a fresh peer is alive"); + check(!peer.gone(), "a fresh peer is not gone"); + } + + { + llama_grpc::StreamPeer peer; + peer.observe_write(true); + peer.observe_write(true); + check(peer.alive(), "successful writes keep the peer alive"); + } + + { + llama_grpc::StreamPeer peer; + peer.observe_write(false); + check(peer.gone(), "a failed write marks the peer gone"); + } + + { + // The whole point of the guard: a stream never comes back, so a later + // success must not resurrect a peer an earlier failure retired. + llama_grpc::StreamPeer peer; + peer.observe_write(false); + peer.observe_write(true); + check(peer.gone(), "a failed write is sticky across later writes"); + } + + { + llama_grpc::StreamPeer peer; + peer.observe_cancelled(false); + check(peer.alive(), "an uncancelled RPC keeps the peer alive"); + peer.observe_cancelled(true); + check(peer.gone(), "cancellation marks the peer gone"); + } + + { + llama_grpc::StreamPeer peer; + peer.observe_cancelled(true); + peer.observe_cancelled(false); + check(peer.gone(), "cancellation is sticky across later checks"); + } + + if (failures != 0) { + std::fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + return 0; +} diff --git a/docs/content/features/backends.md b/docs/content/features/backends.md index 2d04b987c592..8056290b7ef8 100644 --- a/docs/content/features/backends.md +++ b/docs/content/features/backends.md @@ -195,3 +195,26 @@ cannot be retracted; DS4 does not flush incomplete buffered parser state or persist an abandoned request to the disk KV cache. Cancellation is cooperative: DS4 checks it at safe prompt-prefill and decode-loop boundaries, so a GPU kernel already in flight may finish before the request stops. + +### llama.cpp request cancellation + +The llama.cpp backend stops a streaming generation as soon as the response can +no longer be written to the client, not only when the RPC is formally cancelled. +A stream never recovers once a write fails, so the backend treats the first +failed write as final and returns, which releases the slot the generation held. + +This matters most for a model configured without a generation cap. With +`max_tokens: 0` and a large `context_size`, an abandoned request that keeps +decoding occupies its slot until it reaches the context limit — tens of minutes +on a large model — and every other request for that model queues behind it. A +couple of abandoned requests is enough to make a healthy node look wedged. + +Cancellation is cooperative and checked between decoded results, so a batch +already in flight may finish before the request stops. + +{{% notice tip %}} +A generation cap is still worth setting. Cancellation only helps once a client +has actually gone away; a client that waits receives the full context worth of +tokens. Set `max_tokens` on the model config, and keep `repeat_penalty` above +`1` so a repetition loop terminates on its own. +{{% /notice %}}