diff --git a/backend/cpp/llama-cpp/CMakeLists.txt b/backend/cpp/llama-cpp/CMakeLists.txt index 7f2570980..50982b45d 100644 --- a/backend/cpp/llama-cpp/CMakeLists.txt +++ b/backend/cpp/llama-cpp/CMakeLists.txt @@ -130,4 +130,10 @@ if(LLAMA_GRPC_BUILD_TESTS) target_include_directories(model_load_error_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) target_compile_features(model_load_error_test PRIVATE cxx_std_17) add_test(NAME model_load_error_test COMMAND model_load_error_test) + + # Dead-stream tracker (standard library only). + add_executable(stream_peer_test stream_peer_test.cpp stream_peer.h) + target_include_directories(stream_peer_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + target_compile_features(stream_peer_test PRIVATE cxx_std_17) + add_test(NAME stream_peer_test COMMAND stream_peer_test) endif() diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index eb44d9a53..9a2240881 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -57,6 +57,7 @@ #include "thread_params.h" #include "message_content.h" #include "passthrough_options.h" +#include "stream_peer.h" #include "tts_request_options.h" #include #include @@ -2268,6 +2269,11 @@ public: // 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()) { @@ -2280,17 +2286,21 @@ public: 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; } @@ -2311,17 +2321,22 @@ public: 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/prepare.sh b/backend/cpp/llama-cpp/prepare.sh index 308379cd2..8bae396cc 100644 --- a/backend/cpp/llama-cpp/prepare.sh +++ b/backend/cpp/llama-cpp/prepare.sh @@ -67,6 +67,10 @@ cp -r thread_params_test.cpp llama.cpp/tools/grpc-server/ # test (run via backend/cpp/run-unit-tests.sh; also buildable under ctest). cp -r parent_watch.h llama.cpp/tools/grpc-server/ cp -r parent_watch_test.cpp llama.cpp/tools/grpc-server/ +# Dead-stream tracker (included by grpc-server.cpp) and its standalone unit +# test (run via backend/cpp/run-unit-tests.sh; also buildable under ctest). +cp -r stream_peer.h llama.cpp/tools/grpc-server/ +cp -r stream_peer_test.cpp llama.cpp/tools/grpc-server/ cp -rfv llama.cpp/vendor/nlohmann/json.hpp llama.cpp/tools/grpc-server/ cp -rfv llama.cpp/vendor/cpp-httplib/httplib.h llama.cpp/tools/grpc-server/ diff --git a/backend/cpp/llama-cpp/stream_peer.h b/backend/cpp/llama-cpp/stream_peer.h new file mode 100644 index 000000000..065a31c95 --- /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 000000000..2681f182c --- /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 c3578fc4f..9ef87ee4d 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 %}}