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
This commit is contained in:
Ettore Di Giacinto committed 2026-09-11 23:13:32 +00:00
1 parent 869000ceb7
commit 43c54bc9ba
6 files changed
+166 -7

No files matched your search

+6
View File
@@ -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()
+22 -7
View File
@@ -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 <getopt.h>
#include <grpcpp/ext/proto_server_reflection_plugin.h>
@@ -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");
}
+4
View File
@@ -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/
+44
View File
@@ -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
@@ -0,0 +1,67 @@
#include "stream_peer.h"
#include <cstdio>
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;
}
+23
View File
@@ -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 %}}