From 9e28170aae85689656ef1962a47a79e682cfe044 Mon Sep 17 00:00:00 2001 From: Akhil Date: Thu, 8 Oct 2026 15:26:23 +0000 Subject: [PATCH] fix: Prevent silent gRPC stream cancellations under backpressure Add --grpc-stream-accept-prefetch (default 16) to keep several stream accept requests outstanding. --- .../inference_protocols.md | 12 ++++ qa/L0_python_api/test_kserve.py | 7 ++- .../ensemble_backpressure_test.py | 48 +++++++++----- qa/L0_simple_ensemble/test.sh | 62 +++++++++++++++++++ src/command_line_parser.cc | 17 +++++ src/grpc/grpc_server.cc | 15 ++++- src/grpc/grpc_server.h | 6 +- src/grpc/stream_infer_handler.cc | 21 +++++++ src/grpc/stream_infer_handler.h | 16 ++++- src/python/tritonfrontend/_api/_kservegrpc.py | 3 +- .../tritonfrontend/_api/_kservegrpc.pyi | 3 +- 11 files changed, 185 insertions(+), 25 deletions(-) diff --git a/docs/customization_guide/inference_protocols.md b/docs/customization_guide/inference_protocols.md index 4aa3eaeae4..bb40b0aa3e 100644 --- a/docs/customization_guide/inference_protocols.md +++ b/docs/customization_guide/inference_protocols.md @@ -150,6 +150,18 @@ In general, using 2 threads per completion queue seems to give the best performa Note: More threads don't always mean better performance. +#### GRPC Streaming Accept Prefetch + +A new stream reaches Triton only after gRPC matches it to an accept request posted by the streaming inference handler. +gRPC cancels a stream that stays unmatched for more than 30 seconds (the gRPC default for `grpc.server_max_unrequested_time_in_server`, which Triton does not currently expose). +Keeping several accept requests outstanding lets new streams be matched right away even while the handler is busy, for example when an ensemble's `max_inflight_requests` limit is holding back requests. + +* `--grpc-stream-accept-prefetch`: 16 by default. + The number of accept requests each streaming inference handler keeps outstanding, so up to this many new streams can be matched immediately. + Increase it if larger bursts of new streams arrive while requests are being held back. + A value of `1` keeps a single outstanding accept request, which was the behavior before this option was added. + Valid range is `1-128`. + ### Limit Endpoint Access (BETA) Triton users may want to restrict access to protocols or APIs that are diff --git a/qa/L0_python_api/test_kserve.py b/qa/L0_python_api/test_kserve.py index 3bd9ce2411..8c7521f01c 100644 --- a/qa/L0_python_api/test_kserve.py +++ b/qa/L0_python_api/test_kserve.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -62,6 +62,7 @@ def test_correct_grpc_parameters(self): reuse_port=True, infer_allocation_pool_size=12, http2_max_pings_without_data=10, + stream_accept_prefetch=32, ) def test_wrong_grpc_parameters(self): @@ -84,6 +85,10 @@ def test_wrong_grpc_parameters(self): KServeGrpc.Options(max_connection_age_ms=-1) with pytest.raises(Exception): KServeGrpc.Options(max_connection_age_grace_ms=-1) + with pytest.raises(Exception): + KServeGrpc.Options(stream_accept_prefetch=0) + with pytest.raises(Exception): + KServeGrpc.Options(stream_accept_prefetch=129) # Wrong data type with pytest.raises(Exception): diff --git a/qa/L0_simple_ensemble/ensemble_backpressure_test.py b/qa/L0_simple_ensemble/ensemble_backpressure_test.py index e57b5b8cc0..d3899e328f 100755 --- a/qa/L0_simple_ensemble/ensemble_backpressure_test.py +++ b/qa/L0_simple_ensemble/ensemble_backpressure_test.py @@ -30,18 +30,18 @@ sys.path.append("../common") -import os -import queue -import threading -import time -import unittest -from contextlib import ExitStack -from functools import partial - -import numpy as np -import test_util as tu -import tritonclient.grpc as grpcclient -from tritonclient.utils import InferenceServerException +import os # noqa: E402 +import queue # noqa: E402 +import threading # noqa: E402 +import time # noqa: E402 +import unittest # noqa: E402 +from contextlib import ExitStack # noqa: E402 +from functools import partial # noqa: E402 + +import numpy as np # noqa: E402 +import test_util as tu # noqa: E402 +import tritonclient.grpc as grpcclient # noqa: E402 +from tritonclient.utils import InferenceServerException # noqa: E402 SERVER_URL = "localhost:8001" DEFAULT_RESPONSE_TIMEOUT = 60 @@ -79,6 +79,14 @@ def prepare_infer_args(input_value, enable_batching=False): return infer_input, outputs +def _describe_error(err): + """Concise, readable description of a tritonclient error for logs and asserts.""" + try: + return f"{err.status()}: {err.message()}" + except Exception: + return str(err) + + def collect_responses(user_data, timeout=DEFAULT_RESPONSE_TIMEOUT): """ Collect responses from user_data until the final response flag is seen. @@ -138,16 +146,22 @@ def _run_inference( model_name=model_name, inputs=inputs, outputs=outputs ) - # Collect and verify responses for all requests + # Collect and verify responses for all requests. Assert on errors + # first so a real failure (e.g. CANCELLED) surfaces instead of a + # misleading "expected N, got 0". for i, ud in enumerate(user_datas): errors, responses = collect_responses(ud) self.assertEqual( - len(responses), - expected_responses_per_request, - f"Request {i}: expected {expected_responses_per_request} responses, got {len(responses)}", + len(errors), + 0, + f"Request {i}: Triton returned error(s): " + f"{[_describe_error(e) for e in errors]}", ) self.assertEqual( - len(errors), 0, f"Request {i}: unexpected errors: {errors}" + len(responses), + expected_responses_per_request, + f"Request {i}: expected {expected_responses_per_request} " + f"response(s), got {len(responses)}", ) # Verify correctness of responses for idx, resp in enumerate(responses): diff --git a/qa/L0_simple_ensemble/test.sh b/qa/L0_simple_ensemble/test.sh index 1e62c91e7b..51622cb75d 100755 --- a/qa/L0_simple_ensemble/test.sh +++ b/qa/L0_simple_ensemble/test.sh @@ -556,6 +556,68 @@ kill $SERVER_PID wait $SERVER_PID +######## Test '--grpc-stream-accept-prefetch' command-line option ######## +# The option sets how many accept requests the gRPC streaming handler keeps +# outstanding. Invalid values must be rejected at startup. For the default and +# the boundary values the handler must post exactly that many accept requests +# at startup (one "New request handler" line each; no streams are opened here). +EMPTY_MODEL_DIR="`pwd`/empty_models" +rm -rf ${EMPTY_MODEL_DIR} && mkdir -p ${EMPTY_MODEL_DIR} +set +e + +for val in 0 129 abc; do + expected_err="invalid argument for --grpc-stream-accept-prefetch. Must be in the range 1 to 128." + if [ "$val" == "abc" ]; then + expected_err="Invalid option value. Got abc" + fi + SERVER_LOG="./grpc_stream_accept_prefetch_${val}_server.log" + SERVER_ARGS="--model-repository=${EMPTY_MODEL_DIR} --grpc-stream-accept-prefetch=${val}" + rm -f $SERVER_LOG + run_server + if [ "$SERVER_PID" != "0" ]; then + echo -e "\n***\n*** FAILED: server started with invalid --grpc-stream-accept-prefetch=${val}\n***" + cat $SERVER_LOG + kill $SERVER_PID + wait $SERVER_PID + RET=1 + elif ! grep -qF "$expected_err" $SERVER_LOG; then + echo -e "\n***\n*** FAILED: missing expected error for --grpc-stream-accept-prefetch=${val}\n***" + cat $SERVER_LOG + RET=1 + fi +done + +# "default" runs without the option and expects 16. +for val in default 1 128; do + SERVER_LOG="./grpc_stream_accept_prefetch_${val}_server.log" + SERVER_ARGS="--model-repository=${EMPTY_MODEL_DIR} --log-verbose=1" + expected=16 + if [ "$val" != "default" ]; then + SERVER_ARGS="${SERVER_ARGS} --grpc-stream-accept-prefetch=${val}" + expected=$val + fi + rm -f $SERVER_LOG + run_server + if [ "$SERVER_PID" == "0" ]; then + echo -e "\n***\n*** FAILED: server did not start with --grpc-stream-accept-prefetch=${val}\n***" + cat $SERVER_LOG + RET=1 + continue + fi + posted=`grep -c "New request handler for ModelStreamInferHandler" $SERVER_LOG` + if ! grep -E -q "ModelStreamInferHandler outstanding accept requests: ${expected}([^0-9]|$)" $SERVER_LOG || \ + [ "$posted" != "$expected" ]; then + echo -e "\n***\n*** FAILED: expected ${expected} outstanding stream accept requests for ${val}, found ${posted} posted\n***" + cat $SERVER_LOG + RET=1 + fi + kill $SERVER_PID + wait $SERVER_PID +done +set -e +rm -rf ${EMPTY_MODEL_DIR} + + ######## Test invalid values for 'max_inflight_requests' config option ######## INVALID_PARAM_MODEL_DIR="`pwd`/invalid_param_test_models" SERVER_ARGS="--model-repository=${INVALID_PARAM_MODEL_DIR}" diff --git a/src/command_line_parser.cc b/src/command_line_parser.cc index f228bcaf77..4705c78ffb 100644 --- a/src/command_line_parser.cc +++ b/src/command_line_parser.cc @@ -307,6 +307,7 @@ enum TritonOptionId { OPTION_GRPC_ADDRESS, OPTION_GRPC_HEADER_FORWARD_PATTERN, OPTION_GRPC_INFER_THREAD_COUNT, + OPTION_GRPC_STREAM_ACCEPT_PREFETCH, OPTION_GRPC_INFER_ALLOCATION_POOL_SIZE, OPTION_GRPC_MAX_RESPONSE_POOL_SIZE, OPTION_GRPC_USE_SSL, @@ -542,6 +543,13 @@ TritonParser::SetupOptions() {OPTION_GRPC_INFER_THREAD_COUNT, "grpc-infer-thread-count", Option::ArgInt, "The number of gRPC inference handler threads. Default is 2."}); + grpc_options_.push_back( + {OPTION_GRPC_STREAM_ACCEPT_PREFETCH, "grpc-stream-accept-prefetch", + Option::ArgInt, + "The number of accept requests each gRPC streaming inference handler " + "keeps outstanding, so up to this many new streams can be matched " + "immediately even while the handler is busy. Must be in the range 1 " + "to 128. Default is 16."}); grpc_options_.push_back( {OPTION_GRPC_INFER_ALLOCATION_POOL_SIZE, "grpc-infer-allocation-pool-size", Option::ArgInt, @@ -1479,6 +1487,15 @@ TritonParser::Parse(int argc, char** argv) "the range 2 to 128."); } break; + case OPTION_GRPC_STREAM_ACCEPT_PREFETCH: + lgrpc_options.stream_accept_prefetch_ = ParseOption(optarg); + if (lgrpc_options.stream_accept_prefetch_ < 1 || + lgrpc_options.stream_accept_prefetch_ > 128) { + throw ParseException( + "invalid argument for --grpc-stream-accept-prefetch. Must be " + "in the range 1 to 128."); + } + break; case OPTION_GRPC_INFER_ALLOCATION_POOL_SIZE: lgrpc_options.infer_allocation_pool_size_ = ParseOption(optarg); break; diff --git a/src/grpc/grpc_server.cc b/src/grpc/grpc_server.cc index daa90870dc..51958529f5 100644 --- a/src/grpc/grpc_server.cc +++ b/src/grpc/grpc_server.cc @@ -2479,7 +2479,7 @@ Server::Server( options.infer_allocation_pool_size_ /* max_state_bucket_count */, options.max_response_pool_size_, options.infer_compression_level_, restricted_kv, options.forward_header_pattern_, &conn_mtx_, &conn_cnt_, - &accepting_new_conn_)); + &accepting_new_conn_, options.stream_accept_prefetch_)); } Server::~Server() @@ -2544,6 +2544,19 @@ Server::GetOptions(Options& options, UnorderedMapType& options_map) RETURN_IF_ERR(GetValue( options_map, "infer_thread_count", &options.infer_thread_count_)); + // Read stream_accept_prefetch only if present, so a caller that builds the + // options map without it keeps the default. + if (options_map.find("stream_accept_prefetch") != options_map.end()) { + RETURN_IF_ERR(GetValue( + options_map, "stream_accept_prefetch", + &options.stream_accept_prefetch_)); + if (options.stream_accept_prefetch_ < 1 || + options.stream_accept_prefetch_ > 128) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "stream_accept_prefetch must be in the range 1 to 128"); + } + } RETURN_IF_ERR(GetValue( options_map, "infer_allocation_pool_size", &options.infer_allocation_pool_size_)); diff --git a/src/grpc/grpc_server.h b/src/grpc/grpc_server.h index 6020bff7ad..d3f8ee8e78 100644 --- a/src/grpc/grpc_server.h +++ b/src/grpc/grpc_server.h @@ -1,4 +1,4 @@ -// Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -88,6 +88,10 @@ struct Options { // The number of gRPC inference handler threads. Useful for // throughput tuning of models that are request handling bounded. int infer_thread_count_{2}; + // The number of accept requests each streaming inference handler keeps + // outstanding, so up to this many new streams are matched without waiting + // for the handler thread to post the next one. + int stream_accept_prefetch_{16}; // The maximum number of inference request/response objects that // remain allocated for reuse. As long as the number of in-flight // requests doesn't exceed this value there will be no diff --git a/src/grpc/stream_infer_handler.cc b/src/grpc/stream_infer_handler.cc index 97c1ab2e98..46fa42d127 100644 --- a/src/grpc/stream_infer_handler.cc +++ b/src/grpc/stream_infer_handler.cc @@ -26,6 +26,7 @@ #include "stream_infer_handler.h" +#include #include namespace triton { namespace server { namespace grpc { @@ -109,6 +110,26 @@ StreamOutputBufferAttributes( void ModelStreamInferHandler::StartNewRequest() +{ + // Keep 'accept_prefetch_' accept requests outstanding rather than one, so up + // to that many new streams are matched at once even while this thread is + // busy. gRPC cancels a call that stays unmatched longer than its + // unrequested-time limit (30s by default). All of them are posted on the + // first call; after that every accepted stream posts exactly one + // replacement, so the number outstanding stays the same. + int count = 1; + if (!accept_prefetch_posted_) { + accept_prefetch_posted_ = true; + count = std::max(accept_prefetch_, 1); + LOG_VERBOSE(1) << Name() << " outstanding accept requests: " << count; + } + for (int i = 0; i < count; ++i) { + PostAcceptRequest(); + } +} + +void +ModelStreamInferHandler::PostAcceptRequest() { auto context = std::make_shared(cq_, NEXT_UNIQUE_ID); context->SetCompressionLevel(compression_level_); diff --git a/src/grpc/stream_infer_handler.h b/src/grpc/stream_infer_handler.h index ac2acdbaa3..daee2702a0 100644 --- a/src/grpc/stream_infer_handler.h +++ b/src/grpc/stream_infer_handler.h @@ -1,4 +1,4 @@ -// Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions @@ -74,13 +74,14 @@ class ModelStreamInferHandler size_t max_response_queue_size, grpc_compression_level compression_level, std::pair restricted_kv, const std::string& header_forward_pattern, std::shared_mutex* conn_mtx, - std::atomic* conn_cnt, bool* accepting_new_conn) + std::atomic* conn_cnt, bool* accepting_new_conn, + int accept_prefetch) : InferHandler( name, tritonserver, service, cq, max_state_bucket_count, max_response_queue_size, restricted_kv, header_forward_pattern, conn_mtx, conn_cnt, accepting_new_conn), trace_manager_(trace_manager), shm_manager_(shm_manager), - compression_level_(compression_level) + compression_level_(compression_level), accept_prefetch_(accept_prefetch) { // Create the allocator that will be used to allocate buffers for // the result tensors. @@ -116,12 +117,21 @@ class ModelStreamInferHandler void* userp); static void StateWriteResponse(InferHandler::State* state); bool Finish(State* state); + // Posts one RequestModelStreamInfer accept request. + void PostAcceptRequest(); TraceManager* trace_manager_; std::shared_ptr shm_manager_; TRITONSERVER_ResponseAllocator* allocator_; grpc_compression_level compression_level_; + + // Number of RequestModelStreamInfer calls this handler keeps outstanding, + // so up to that many new streams can be matched while the thread is busy. + const int accept_prefetch_; + // Whether the extra accept requests have been posted. Only accessed on the + // handler thread. + bool accept_prefetch_posted_{false}; }; }}} // namespace triton::server::grpc diff --git a/src/python/tritonfrontend/_api/_kservegrpc.py b/src/python/tritonfrontend/_api/_kservegrpc.py index ef5fa9e35f..d77922b84c 100644 --- a/src/python/tritonfrontend/_api/_kservegrpc.py +++ b/src/python/tritonfrontend/_api/_kservegrpc.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -80,6 +80,7 @@ class Options: int, Grpc_compression_level ] = Grpc_compression_level.NONE infer_thread_count: int = Field(2, ge=0) + stream_accept_prefetch: int = Field(16, ge=1, le=128) infer_allocation_pool_size: int = Field(8, ge=0) max_response_pool_size: int = Field(2_147_483_647, ge=0) forward_header_pattern: str = "" diff --git a/src/python/tritonfrontend/_api/_kservegrpc.pyi b/src/python/tritonfrontend/_api/_kservegrpc.pyi index cacfe24e56..8739ae0c3d 100644 --- a/src/python/tritonfrontend/_api/_kservegrpc.pyi +++ b/src/python/tritonfrontend/_api/_kservegrpc.pyi @@ -1,4 +1,4 @@ -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -57,6 +57,7 @@ class KServeGrpc: max_connection_age_grace_ms: int infer_compression_level: int | Grpc_compression_level infer_thread_count: int + stream_accept_prefetch: int infer_allocation_pool_size: int max_response_pool_size: int forward_header_pattern: str