diff --git a/qa/L0_request_cancellation/cancellation_test_utils.py b/qa/L0_request_cancellation/cancellation_test_utils.py new file mode 100644 index 0000000000..593ac85831 --- /dev/null +++ b/qa/L0_request_cancellation/cancellation_test_utils.py @@ -0,0 +1,187 @@ +# Copyright 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 +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import os +import re +import threading +import time + +import requests +from tritonclient.utils import InferenceServerException + + +class CancellationTest: + """Shared synchronization and server-state helpers for cancellation tests.""" + + # Build an event-backed callback and mutable response state for async infer. + def _generate_callback_and_response_pair(self): + response = { + "completed": threading.Event(), + "responded": False, + "result": None, + "error": None, + } + + def callback(result, error): + response["result"] = result + response["error"] = error + response["responded"] = True + response["completed"].set() + + return callback, response + + # Wait for an async inference callback with a bounded timeout. + def _wait_for_response(self, response, timeout=30): + self.assertTrue( + response["completed"].wait(timeout), + f"inference callback was not invoked within {timeout}s", + ) + + # Verify that an async request completed with the gRPC CANCELLED status. + def _assert_response_is_cancelled(self, response, timeout=30): + self._wait_for_response(response, timeout) + self.assertTrue(response["responded"]) + self.assertIsNone(response["result"]) + self.assertIsInstance(response["error"], InferenceServerException) + self.assertEqual(response["error"].status(), "StatusCode.CANCELLED") + + # Read one labeled Prometheus metric value independent of label ordering. + def _metric_value(self, metric_name, expected_labels): + response = requests.get("http://localhost:8002/metrics", timeout=5) + response.raise_for_status() + for line in response.text.splitlines(): + if not line.startswith(f"{metric_name}{{"): + continue + + fields = line.split() + if len(fields) < 2: + continue + + metric_and_labels = fields[0] + label_text = metric_and_labels[ + len(metric_name) + 1 : metric_and_labels.rfind("}") + ] + labels = dict(re.findall(r'(\w+)="([^"]*)"', label_text)) + if all(labels.get(key) == value for key, value in expected_labels.items()): + return int(float(fields[1])) + + return None + + # Return the backend execution count reported for a model. + def _execution_count(self, model_name): + stats = self._triton.get_inference_statistics( + model_name=model_name, as_json=True + ) + model_stats = stats.get("model_stats", []) + if not model_stats: + return 0 + return int(model_stats[0].get("execution_count", 0)) + + # Return the failure count for a model and failure reason. + def _failure_count(self, model_name, reason): + value = self._metric_value( + "nv_inference_request_failure", + {"model": model_name, "reason": reason, "version": "1"}, + ) + return 0 if value is None else value + + # Wait for a failure metric to increase by the expected amount. + def _assert_metrics( + self, model_name, reason, expected_count_increase, initial_count + ): + expected_count = initial_count + expected_count_increase + self._wait_until( + lambda: self._failure_count(model_name, reason) == expected_count, + f"{model_name} {reason} failure count to reach {expected_count}", + ) + + # Poll a condition until it succeeds or reaches a bounded timeout. + def _wait_until(self, predicate, description, timeout=30, interval=0.1): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(interval) + self.fail(f"timed out after {timeout}s waiting for {description}") + + # Cancel a request and wait for cancellation log message. + def _cancel_and_wait(self, request, request_id): + def cancellation_count(): + with open(os.environ["SERVER_LOG"], encoding="utf-8") as server_log: + return server_log.read().count( + f"[request id: {request_id}] Cancellation issued" + ) + + cancellations_before = cancellation_count() + request.cancel() + self._wait_until( + lambda: cancellation_count() > cancellations_before, + "the server to issue cancellation", + ) + + # Wait until a model's pending-request count is stable at the expected value. + def _wait_until_pending(self, model_name, expected, timeout=30, stable_for=0.5): + deadline = time.monotonic() + timeout + stable_since = None + last_count = None + + while time.monotonic() < deadline: + now = time.monotonic() + last_count = self._metric_value( + "nv_inference_pending_request_count", + {"model": model_name, "version": "1"}, + ) + last_count = 0 if last_count is None else last_count + if last_count == expected: + if stable_since is None: + stable_since = now + elif now - stable_since >= stable_for: + return + else: + stable_since = None + time.sleep(0.1) + + self.fail( + f"'{model_name}' pending count did not remain at {expected} for " + f"{stable_for}s within {timeout}s; last count was {last_count}" + ) + + # Start two holders and wait until one is queued behind the other. + def _start_holders(self, pool, model_name, input_factory): + holders = [ + pool.submit(self._triton.infer, model_name, input_factory(index)) + for index in range(2) + ] + self._wait_until_pending(model_name, 1) + return holders + + # Wait for an exact execution count so extra backend execution is detected. + def _wait_for_execution_count(self, model_name, expected): + self._wait_until( + lambda: self._execution_count(model_name) >= expected, + f"{model_name} execution count to reach {expected}", + ) + self.assertEqual(self._execution_count(model_name), expected) diff --git a/qa/L0_request_cancellation/scheduler_test.py b/qa/L0_request_cancellation/scheduler_test.py index 900073ea7d..38624bea54 100755 --- a/qa/L0_request_cancellation/scheduler_test.py +++ b/qa/L0_request_cancellation/scheduler_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2020-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 @@ -27,45 +27,29 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import concurrent.futures -import re import time import unittest import numpy as np -import requests import tritonclient.grpc as grpcclient +from cancellation_test_utils import CancellationTest from tritonclient.utils import InferenceServerException -class TestScheduler(unittest.TestCase): +class TestScheduler(CancellationTest, unittest.TestCase): def setUp(self): # Initialize client self._triton = grpcclient.InferenceServerClient("localhost:8001") - def _get_inputs(self, batch_size): + # Build identity-model input, optionally with a value unique to a cache key. + def _get_inputs(self, batch_size, value=1.0): self.assertIsInstance(batch_size, int) self.assertGreater(batch_size, 0) shape = [batch_size, 8] inputs = [grpcclient.InferInput("INPUT0", shape, "FP32")] - inputs[0].set_data_from_numpy(np.ones(shape, dtype=np.float32)) + inputs[0].set_data_from_numpy(np.full(shape, value, dtype=np.float32)) return inputs - def _generate_callback_and_response_pair(self): - response = {"responded": False, "result": None, "error": None} - - def callback(result, error): - response["responded"] = True - response["result"] = result - response["error"] = error - - return callback, response - - def _assert_response_is_cancelled(self, response): - self.assertTrue(response["responded"]) - self.assertEqual(response["result"], None) - self.assertIsInstance(response["error"], InferenceServerException) - self.assertEqual(response["error"].status(), "StatusCode.CANCELLED") - def _generate_streaming_callback_and_response_pair(self): response = [] # [{"result": result, "error": error}, ...] @@ -86,29 +70,6 @@ def _assert_streaming_response_is_cancelled(self, response): cancelled_count += 1 self.assertEqual(cancelled_count, 1) - def _get_metrics(self): - metrics_url = "http://localhost:8002/metrics" - r = requests.get(metrics_url) - r.raise_for_status() - return r.text - - def _metrics_before_test(self, model, reason): - pattern = rf'nv_inference_request_failure\{{model="{model}",reason="{reason}",version="1"\}} (\d+)' - metrics = self._get_metrics() - match = re.search(pattern, metrics) - if match: - return int(match.group(1)) - else: - raise Exception(f"Failure metrics for model='{model}' not found") - - def _assert_metrics( - self, model_name, reason, expected_count_increase, initial_count - ): - metrics = self._get_metrics() - # Add initial count + expected count for the the test - expected_metric = f'nv_inference_request_failure{{model="{model_name}",reason="{reason}",version="1"}} {expected_count_increase + initial_count}' - self.assertIn(expected_metric, metrics) - # Test queued requests on dynamic batch scheduler can be cancelled def test_dynamic_batch_scheduler_request_cancellation(self): model_name = "dynamic_batch" @@ -130,16 +91,122 @@ def test_dynamic_batch_scheduler_request_cancellation(self): self.assertFalse(response["responded"]) # Cancel the queued request queue_future.cancel() - time.sleep(2) # ensure the cancellation is delivered self._assert_response_is_cancelled(response) # Join saturating thread - saturate_thread_1.result() - saturate_thread_2.result() + saturate_thread_1.result(timeout=60) + saturate_thread_2.result(timeout=60) + + # Return the response-cache hit count reported for a model. + def _cache_hit_count(self, model_name): + stats = self._triton.get_inference_statistics( + model_name=model_name, as_json=True + ) + model_stats = stats.get("model_stats", []) + if not model_stats: + return 0 + infer_stats = model_stats[0]["inference_stats"] + return int(infer_stats.get("cache_hit", {}).get("count", 0)) + + # Test the no-dynamic-batching path where a cancelled request is already + # queued in the rate limiter waiting for a model instance. + def test_no_batcher_queued_request_cancellation(self): + model_name = "no_batching" + request_id = "no-batcher-queued-cancel" + failures_before = self._failure_count(model_name, "CANCELED") + executions_before = self._execution_count(model_name) + + with concurrent.futures.ThreadPoolExecutor() as pool: + holders = self._start_holders( + pool, model_name, lambda _: self._get_inputs(batch_size=1) + ) + + callback, response = self._generate_callback_and_response_pair() + request = self._triton.async_infer( + model_name, + self._get_inputs(batch_size=1), + callback, + request_id=request_id, + ) + self._wait_until_pending(model_name, 2) + self.assertFalse(response["responded"]) + + self._cancel_and_wait(request, request_id) + for holder in holders: + holder.result(timeout=60) + self._assert_response_is_cancelled(response) + + self._assert_metrics(model_name, "CANCELED", 1, failures_before) + self._wait_for_execution_count(model_name, executions_before + 2) + + # Test a queued request on a model with no batcher is cancelled before + # backend execution and its error response is not cached. + def test_no_batcher_cancelled_request_is_not_executed_or_cached(self): + model_name = "no_batching_cache" + request_id = "no-batcher-cache-cancel" + cancelled_value = 2.0 + failures_before = self._failure_count(model_name, "CANCELED") + executions_before = self._execution_count(model_name) + with concurrent.futures.ThreadPoolExecutor() as pool: + # Unique holder inputs avoid cache hits while the instance is held. + holders = self._start_holders( + pool, + model_name, + lambda index: self._get_inputs(batch_size=1, value=10.0 + index), + ) + + callback, response = self._generate_callback_and_response_pair() + queue_future = self._triton.async_infer( + model_name, + self._get_inputs(batch_size=1, value=cancelled_value), + callback, + request_id=request_id, + ) + self._wait_until_pending(model_name, 2) + self.assertFalse(response["responded"]) + + self._cancel_and_wait(queue_future, request_id) + for holder in holders: + holder.result(timeout=60) + self._assert_response_is_cancelled(response) + + self._assert_metrics(model_name, "CANCELED", 1, failures_before) + self._wait_for_execution_count(model_name, executions_before + 2) + hits_after_cancellation = self._cache_hit_count(model_name) + + # The first identical request must execute. If the cancelled response + # was cached, this request would incorrectly be served as a cache hit. + first_result = self._triton.infer( + model_name, self._get_inputs(batch_size=1, value=cancelled_value) + ) + first_output = first_result.as_numpy("OUTPUT0") + self.assertIsNotNone(first_output) + np.testing.assert_allclose( + first_output, np.full([1, 8], cancelled_value, dtype=np.float32) + ) + self._wait_for_execution_count(model_name, executions_before + 3) + hits_after_first_request = self._cache_hit_count(model_name) + self.assertEqual(hits_after_first_request, hits_after_cancellation) + + # The successful response must now be cached. A second identical request + # must be a cache hit and must not execute the backend again. + second_result = self._triton.infer( + model_name, self._get_inputs(batch_size=1, value=cancelled_value) + ) + second_output = second_result.as_numpy("OUTPUT0") + self.assertIsNotNone(second_output) + np.testing.assert_allclose( + second_output, np.full([1, 8], cancelled_value, dtype=np.float32) + ) + self._wait_until( + lambda: self._cache_hit_count(model_name) == hits_after_cancellation + 1, + f"{model_name} cache hit count to increase by one", + ) + self.assertEqual(self._execution_count(model_name), executions_before + 3) # Test backlogged requests on sequence batch scheduler can be cancelled def test_sequence_batch_scheduler_backlog_request_cancellation(self): model_name = "sequence_direct" - initial_metrics_value = self._metrics_before_test(model_name, "CANCELED") + initial_metrics_value = self._failure_count(model_name, "CANCELED") with concurrent.futures.ThreadPoolExecutor() as pool: # Saturate the single sequence slot saturate_thread = pool.submit( @@ -169,12 +236,10 @@ def test_sequence_batch_scheduler_backlog_request_cancellation(self): self.assertFalse(backlog_requests[1]["response"]["responded"]) # Cancelling any backlogged request cancels the entire sequence backlog_requests[0]["future"].cancel() - time.sleep(2) # ensure the cancellation is delivered - time.sleep(2) # ensure reaper thread has responded self._assert_response_is_cancelled(backlog_requests[0]["response"]) self._assert_response_is_cancelled(backlog_requests[1]["response"]) # Join saturating thread - saturate_thread.result() + saturate_thread.result(timeout=60) expected_count_increase = 2 self._assert_metrics( model_name, @@ -186,7 +251,7 @@ def test_sequence_batch_scheduler_backlog_request_cancellation(self): # Test queued requests on direct sequence batch scheduler can be cancelled def test_direct_sequence_batch_scheduler_request_cancellation(self): model_name = "sequence_direct" - initial_metrics_value = self._metrics_before_test(model_name, "CANCELED") + initial_metrics_value = self._failure_count(model_name, "CANCELED") self._test_sequence_batch_scheduler_queued_request_cancellation(model_name) expected_count_increase = 2 self._assert_metrics( @@ -201,6 +266,57 @@ def test_oldest_sequence_batch_scheduler_request_cancellation(self): model_name = "sequence_oldest" self._test_sequence_batch_scheduler_queued_request_cancellation(model_name) + # Test an oldest-first request cancelled while waiting for an instance. + def test_oldest_sequence_rate_limiter_request_cancellation(self): + model_name = "sequence_oldest" + request_id = "oldest-rate-limiter-cancel" + failures_before = self._failure_count(model_name, "CANCELED") + executions_before = self._execution_count(model_name) + + with concurrent.futures.ThreadPoolExecutor() as pool: + holders = [ + pool.submit( + self._triton.infer, + model_name, + self._get_inputs(batch_size=1), + sequence_id=100 + index, + sequence_start=True, + sequence_end=True, + ) + for index in range(2) + ] + self._wait_until_pending(model_name, 1) + + callback, response = self._generate_callback_and_response_pair() + request = self._triton.async_infer( + model_name, + self._get_inputs(batch_size=1), + callback, + request_id=request_id, + sequence_id=102, + sequence_start=True, + sequence_end=True, + ) + self._wait_until_pending(model_name, 2) + self.assertFalse(response["responded"]) + + self._cancel_and_wait(request, request_id) + live_request = pool.submit( + self._triton.infer, + model_name, + self._get_inputs(batch_size=1), + sequence_id=103, + sequence_start=True, + sequence_end=True, + ) + for holder in holders: + holder.result(timeout=60) + live_request.result(timeout=60) + self._assert_response_is_cancelled(response) + + self._assert_metrics(model_name, "CANCELED", 1, failures_before) + self._wait_for_execution_count(model_name, executions_before + 3) + # Helper function def _test_sequence_batch_scheduler_queued_request_cancellation(self, model_name): with concurrent.futures.ThreadPoolExecutor() as pool: @@ -226,12 +342,10 @@ def _test_sequence_batch_scheduler_queued_request_cancellation(self, model_name) self.assertFalse(queue_requests[1]["response"]["responded"]) # Cancelling any queued request cancels the entire sequence queue_requests[0]["future"].cancel() - time.sleep(2) # ensure the cancellation is delivered - time.sleep(2) # ensure reaper thread has responded self._assert_response_is_cancelled(queue_requests[0]["response"]) self._assert_response_is_cancelled(queue_requests[1]["response"]) # Join start thread - start_thread.result() + start_thread.result(timeout=60) # Test ensemble scheduler will propagate cancellation request to child def test_ensemble_scheduler_request_cancellation(self): @@ -243,7 +357,6 @@ def test_ensemble_scheduler_request_cancellation(self): time.sleep(2) # ensure the inference has started self.assertFalse(response["responded"]) infer_future.cancel() - time.sleep(2) # ensure the cancellation is delivered self._assert_response_is_cancelled(response) # Test cancellation on multiple gRPC streaming sequences diff --git a/qa/L0_request_cancellation/test.sh b/qa/L0_request_cancellation/test.sh index 2b92c12027..961c558c29 100755 --- a/qa/L0_request_cancellation/test.sh +++ b/qa/L0_request_cancellation/test.sh @@ -40,6 +40,8 @@ fi export CUDA_VISIBLE_DEVICES=0 +DATADIR=/data/inferenceserver/${REPO_VERSION} + SERVER=/opt/tritonserver/bin/tritonserver source ../common/util.sh CANCEL_LOG_LINE="Cancellation notification received for " @@ -275,7 +277,24 @@ mkdir -p models/sequence_oldest/1 && (cd models/sequence_oldest && \ echo -e 'input [{ name: "INPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ echo -e 'output [{ name: "OUTPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ echo -e 'instance_group [{ count: 1 \n kind: KIND_CPU }]' >> config.pbtxt && \ - echo -e 'sequence_batching { oldest { max_candidate_sequences: 1 } \n max_sequence_idle_microseconds: 6000000 }' >> config.pbtxt && \ + echo -e 'sequence_batching { oldest { max_candidate_sequences: 3 } \n max_sequence_idle_microseconds: 6000000 }' >> config.pbtxt && \ + echo -e 'parameters [{ key: "execute_delay_ms" \n value: { string_value: "6000" } }]' >> config.pbtxt) +mkdir -p models/no_batching/1 && (cd models/no_batching && \ + echo 'name: "no_batching"' >> config.pbtxt && \ + echo 'backend: "identity"' >> config.pbtxt && \ + echo 'max_batch_size: 1' >> config.pbtxt && \ + echo -e 'input [{ name: "INPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ + echo -e 'output [{ name: "OUTPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ + echo -e 'instance_group [{ count: 1 \n kind: KIND_CPU }]' >> config.pbtxt && \ + echo -e 'parameters [{ key: "execute_delay_ms" \n value: { string_value: "6000" } }]' >> config.pbtxt) +mkdir -p models/no_batching_cache/1 && (cd models/no_batching_cache && \ + echo 'name: "no_batching_cache"' >> config.pbtxt && \ + echo 'backend: "identity"' >> config.pbtxt && \ + echo 'max_batch_size: 1' >> config.pbtxt && \ + echo -e 'input [{ name: "INPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ + echo -e 'output [{ name: "OUTPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ + echo -e 'instance_group [{ count: 1 \n kind: KIND_CPU }]' >> config.pbtxt && \ + echo -e 'response_cache { enable: true }' >> config.pbtxt && \ echo -e 'parameters [{ key: "execute_delay_ms" \n value: { string_value: "6000" } }]' >> config.pbtxt) mkdir -p models/ensemble_model/1 && (cd models/ensemble_model && \ echo 'name: "ensemble_model"' >> config.pbtxt && \ @@ -291,7 +310,8 @@ mkdir -p models/ensemble_model/1 && (cd models/ensemble_model && \ TEST_LOG="scheduler_test.log" SERVER_LOG="./scheduler_test.server.log" -SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=2" +# Cache required by the 'no_batching_cache' model +SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=2 --cache-config=local,size=1048576" run_server if [ "$SERVER_PID" == "0" ]; then echo -e "\n***\n*** Failed to start $SERVER\n***" @@ -300,7 +320,7 @@ if [ "$SERVER_PID" == "0" ]; then fi set +e -python scheduler_test.py > $TEST_LOG 2>&1 +SERVER_LOG=$SERVER_LOG python scheduler_test.py > $TEST_LOG 2>&1 if [ $? -ne 0 ]; then echo -e "\n***\n*** Scheduler Tests Failed\n***" cat $TEST_LOG @@ -311,6 +331,86 @@ set -e kill $SERVER_PID wait $SERVER_PID +# +# TensorRT cancellation while waiting in the rate limiter +# +# 'resource_holder' is slow and holds the global resource both models need, so +# the TensorRT request stays queued on the rate limiter long enough to cancel. +# +TRT_MODEL_SRC=$DATADIR/qa_model_repository/plan_float32_float32_float32 + +rm -rf models && mkdir models +cp -r $TRT_MODEL_SRC models/plan_no_batching +rm -rf models/plan_no_batching/2 models/plan_no_batching/3 + +set +e +python3 - <<'PYEOF' +from google.protobuf import text_format +import tritonclient.grpc.model_config_pb2 as model_config_pb2 + +path = "models/plan_no_batching/config.pbtxt" +with open(path) as config_file: + config = text_format.Parse(config_file.read(), model_config_pb2.ModelConfig()) + +config.name = "plan_no_batching" +config.ClearField("version_policy") +config.ClearField("instance_group") +config.ClearField("dynamic_batching") +config.ClearField("sequence_batching") +config.ClearField("ensemble_scheduling") + +instance_group = config.instance_group.add() +instance_group.count = 1 +instance_group.kind = model_config_pb2.ModelInstanceGroup.KIND_GPU +resource = instance_group.rate_limiter.resources.add() +resource.name = "SHARED" +setattr(resource, "global", True) +resource.count = 1 + +with open(path, "w") as config_file: + config_file.write(text_format.MessageToString(config)) +PYEOF +TRT_CONFIG_RC=$? +set -e + +if [ $TRT_CONFIG_RC -ne 0 ]; then + echo -e "\n***\n*** Failed to prepare the TensorRT model config\n***" + RET=1 +else + + mkdir -p models/resource_holder/1 && (cd models/resource_holder && \ + echo 'name: "resource_holder"' >> config.pbtxt && \ + echo 'backend: "identity"' >> config.pbtxt && \ + echo 'max_batch_size: 1' >> config.pbtxt && \ + echo -e 'input [{ name: "INPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ + echo -e 'output [{ name: "OUTPUT0" \n data_type: TYPE_FP32 \n dims: [ -1 ] }]' >> config.pbtxt && \ + echo -e 'instance_group [{ count: 1 \n kind: KIND_CPU \n rate_limiter { resources [{ name: "SHARED" \n global: true \n count: 1 }] } }]' >> config.pbtxt && \ + echo -e 'parameters [{ key: "execute_delay_ms" \n value: { string_value: "6000" } }]' >> config.pbtxt) + + TEST_LOG="trt_cancellation_test.log" + SERVER_LOG="./trt_cancellation_test.server.log" + + SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=2 --rate-limit=execution_count" + run_server + if [ "$SERVER_PID" == "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + cat $SERVER_LOG + exit 1 + fi + + set +e + SERVER_LOG=$SERVER_LOG python trt_cancellation_test.py > $TEST_LOG 2>&1 + if [ $? -ne 0 ]; then + echo -e "\n***\n*** TensorRT Cancellation Tests Failed\n***" + cat $TEST_LOG + RET=1 + fi + set -e + + kill $SERVER_PID + wait $SERVER_PID +fi + # # Implicit state tests # diff --git a/qa/L0_request_cancellation/trt_cancellation_test.py b/qa/L0_request_cancellation/trt_cancellation_test.py new file mode 100755 index 0000000000..1be0c21fcd --- /dev/null +++ b/qa/L0_request_cancellation/trt_cancellation_test.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 + +# Copyright 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 +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import concurrent.futures +import unittest + +import numpy as np +import tritonclient.grpc as grpcclient +from cancellation_test_utils import CancellationTest + +# A TensorRT request cancelled while waiting in the rate limiter must not run +# inference. Both models need the one global "SHARED" resource, so +# 'resource_holder' keeps the TensorRT requests queued. +TRT_MODEL = "plan_no_batching" +HOLDER_MODEL = "resource_holder" + +# Plan model from qa_model_repository: OUTPUT0 = INPUT0 + INPUT1. +TRT_SHAPE = [1, 16] +HOLDER_SHAPE = [1, 8] + + +class TestTrtRequestCancellation(CancellationTest, unittest.TestCase): + def setUp(self): + self._triton = grpcclient.InferenceServerClient("localhost:8001") + + # Build the two TensorRT inputs with a predictable summed output. + def _trt_inputs(self, value): + inputs = [ + grpcclient.InferInput("INPUT0", TRT_SHAPE, "FP32"), + grpcclient.InferInput("INPUT1", TRT_SHAPE, "FP32"), + ] + for model_input in inputs: + model_input.set_data_from_numpy(np.full(TRT_SHAPE, value, dtype=np.float32)) + return inputs + + # Build identity-model input used only to occupy the shared resource. + def _holder_inputs(self): + inputs = [grpcclient.InferInput("INPUT0", HOLDER_SHAPE, "FP32")] + inputs[0].set_data_from_numpy(np.ones(HOLDER_SHAPE, dtype=np.float32)) + return inputs + + # A cancelled TRT request must not run, while an adjacent live request does. + def test_trt_rate_limited_cancellation_skips_only_cancelled_request(self): + request_id = "trt-rate-limited-cancel" + executions_before = self._execution_count(TRT_MODEL) + failures_before = self._failure_count(TRT_MODEL, "CANCELED") + + with concurrent.futures.ThreadPoolExecutor() as pool: + holders = self._start_holders( + pool, HOLDER_MODEL, lambda _: self._holder_inputs() + ) + + ( + cancelled_callback, + cancelled_response, + ) = self._generate_callback_and_response_pair() + cancelled_request = self._triton.async_infer( + TRT_MODEL, + self._trt_inputs(value=1.0), + cancelled_callback, + request_id=request_id, + ) + + live_callback, live_response = self._generate_callback_and_response_pair() + live_request = self._triton.async_infer( + TRT_MODEL, self._trt_inputs(value=2.0), live_callback + ) + + self._wait_until_pending(TRT_MODEL, 2) + self.assertFalse( + cancelled_response["responded"], + "the cancelled request was not held by the rate limiter", + ) + self.assertFalse( + live_response["responded"], + "the live request was not held by the rate limiter", + ) + self.assertIsNotNone(live_request) + + self._cancel_and_wait(cancelled_request, request_id) + + # Let the holders finish so the queued TRT requests can be scheduled. + for holder in holders: + holder.result(timeout=60) + + self._assert_response_is_cancelled(cancelled_response) + self._wait_for_response(live_response) + self.assertIsNone(live_response["error"]) + self.assertIsNotNone(live_response["result"]) + np.testing.assert_allclose( + live_response["result"].as_numpy("OUTPUT0"), + np.full(TRT_SHAPE, 4.0, dtype=np.float32), + ) + + self._assert_metrics(TRT_MODEL, "CANCELED", 1, failures_before) + self._wait_for_execution_count(TRT_MODEL, executions_before + 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/grpc/infer_handler.h b/src/grpc/infer_handler.h index 7d8c988af0..99f7d101bd 100644 --- a/src/grpc/infer_handler.h +++ b/src/grpc/infer_handler.h @@ -878,13 +878,20 @@ class InferHandlerState { // Note that request may or may not be valid at this point. // Assuming if RequestComplete callback is run asynchronously // before this point. + const std::string request_id = state->request_.id().empty() + ? "" + : state->request_.id(); TRITONSERVER_Error* err = nullptr; err = TRITONSERVER_InferenceRequestCancel( state->inference_request_.get()); - // TODO: Add request id to the message if (err != nullptr) { - LOG_INFO << "Failed to cancel the request: " + LOG_INFO << "[request id: " << request_id + << "] Failed to cancel the request: " << TRITONSERVER_ErrorMessage(err); + TRITONSERVER_ErrorDelete(err); + } else { + LOG_VERBOSE(1) + << "[request id: " << request_id << "] Cancellation issued"; } state->step_ = Steps::CANCELLATION_ISSUED; } else if (state->step_ == Steps::COMPLETE) {