diff --git a/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/1/model.py b/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/1/model.py index fc8ff2f691..9b2a0a0141 100644 --- a/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/1/model.py +++ b/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/1/model.py @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-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 @@ -25,6 +25,8 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import time + import numpy as np import triton_python_backend_utils as pb_utils @@ -38,13 +40,14 @@ def execute(self, requests): for request in requests: # Get input - number of responses to produce in_tensor = pb_utils.get_input_tensor_by_name(request, "IN") - count = in_tensor.as_numpy()[0] + count = in_tensor.as_numpy().item() response_sender = request.get_response_sender() + out_tensor = pb_utils.Tensor("OUT", np.array([[0.5]], dtype=np.float32)) # Produce 'count' responses, each with 0.5 as the output value for i in range(count): - out_tensor = pb_utils.Tensor("OUT", np.array([0.5], dtype=np.float32)) + time.sleep(0.1) # Simulate some processing delay response = pb_utils.InferenceResponse(output_tensors=[out_tensor]) response_sender.send(response) diff --git a/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/config.pbtxt b/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/config.pbtxt index 4ef3a444ce..0f7e05db0e 100644 --- a/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/config.pbtxt +++ b/qa/L0_simple_ensemble/backpressure_test_models/decoupled_producer/config.pbtxt @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-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,7 +27,7 @@ name: "decoupled_producer" backend: "python" -max_batch_size: 0 +max_batch_size: 1 input [ { @@ -55,4 +55,3 @@ instance_group [ model_transaction_policy { decoupled: true } - diff --git a/qa/L0_simple_ensemble/ensemble_backpressure_test.py b/qa/L0_simple_ensemble/ensemble_backpressure_test.py index efebd03036..5521c4f00e 100755 --- a/qa/L0_simple_ensemble/ensemble_backpressure_test.py +++ b/qa/L0_simple_ensemble/ensemble_backpressure_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-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 @@ -31,6 +31,7 @@ sys.path.append("../common") import queue +import threading import time import unittest from contextlib import ExitStack @@ -61,49 +62,53 @@ def callback(user_data, result, error): user_data._response_queue.put(result) -class EnsembleBackpressureTest(tu.TestResultCollector): +def prepare_infer_args(input_value): """ - Tests for ensemble backpressure feature (max_inflight_requests). + Create InferInput/InferRequestedOutput lists """ + input_data = np.array([input_value], dtype=np.int32) + infer_input = [grpcclient.InferInput("IN", input_data.shape, "INT32")] + infer_input[0].set_data_from_numpy(input_data) + outputs = [grpcclient.InferRequestedOutput("OUT")] + return infer_input, outputs - def _prepare_infer_args(self, input_value): - """ - Create InferInput/InferRequestedOutput lists - """ - input_data = np.array([input_value], dtype=np.int32) - infer_input = [grpcclient.InferInput("IN", input_data.shape, "INT32")] - infer_input[0].set_data_from_numpy(input_data) - outputs = [grpcclient.InferRequestedOutput("OUT")] - return infer_input, outputs - def _collect_responses(self, user_data): - """ - Collect responses from user_data until the final response flag is seen. - """ - responses = [] - while True: - try: - result = user_data._response_queue.get(timeout=DEFAULT_RESPONSE_TIMEOUT) - except queue.Empty: - self.fail( - f"No response received within {DEFAULT_RESPONSE_TIMEOUT} seconds." - ) - - self.assertNotIsInstance( - result, Exception, f"Callback returned an exception: {result}" +def collect_responses(user_data): + """ + Collect responses from user_data until the final response flag is seen. + """ + errors = [] + responses = [] + while True: + try: + result = user_data._response_queue.get(timeout=DEFAULT_RESPONSE_TIMEOUT) + except queue.Empty: + raise Exception( + f"No response received within {DEFAULT_RESPONSE_TIMEOUT} seconds." ) - response = result.get_response() - # Add response to list if it has data (not empty final-only response) - if len(response.outputs) > 0: - responses.append(result) + if type(result) == InferenceServerException: + errors.append(result) + # error responses are final - stream terminates + break + + response = result.get_response() + # Add response to list if it has data (not empty final-only response) + if len(response.outputs) > 0: + responses.append(result) + + # Check if this is the final response + final = response.parameters.get("triton_final_response") + if final and final.bool_param: + break - # Check if this is the final response - final = response.parameters.get("triton_final_response") - if final and final.bool_param: - break + return errors, responses - return responses + +class EnsembleBackpressureTest(tu.TestResultCollector): + """ + Tests for ensemble backpressure feature (max_inflight_requests). + """ def _run_inference(self, model_name, expected_responses_count=32): """ @@ -112,19 +117,24 @@ def _run_inference(self, model_name, expected_responses_count=32): user_data = UserData() with grpcclient.InferenceServerClient(SERVER_URL) as triton_client: try: - inputs, outputs = self._prepare_infer_args(expected_responses_count) + inputs, outputs = prepare_infer_args(expected_responses_count) triton_client.start_stream(callback=partial(callback, user_data)) triton_client.async_stream_infer( model_name=model_name, inputs=inputs, outputs=outputs ) # Collect and verify responses - responses = self._collect_responses(user_data) + errors, responses = collect_responses(user_data) self.assertEqual( len(responses), expected_responses_count, f"Expected {expected_responses_count} responses, got {len(responses)}", ) + self.assertEqual( + len(errors), + 0, + f"Expected no errors during inference, got {len(errors)} errors", + ) # Verify correctness of responses for idx, resp in enumerate(responses): @@ -172,7 +182,7 @@ def test_max_inflight_requests_limit_concurrent_requests(self): for _ in range(num_concurrent) ] - inputs, outputs = self._prepare_infer_args(expected_per_request) + inputs, outputs = prepare_infer_args(expected_per_request) # Start all concurrent requests for i in range(num_concurrent): @@ -183,12 +193,17 @@ def test_max_inflight_requests_limit_concurrent_requests(self): # Collect and verify responses for all requests for i, ud in enumerate(user_datas): - responses = self._collect_responses(ud) + errors, responses = collect_responses(ud) self.assertEqual( len(responses), expected_per_request, f"Request {i}: expected {expected_per_request} responses, got {len(responses)}", ) + self.assertEqual( + len(errors), + 0, + f"Request {i}: Expected no errors during inference, got {len(errors)} errors", + ) # Verify correctness of responses for idx, resp in enumerate(responses): output = resp.as_numpy("OUT") @@ -214,7 +229,7 @@ def test_max_inflight_requests_limit_request_cancellation(self): user_data = UserData() with grpcclient.InferenceServerClient(SERVER_URL) as triton_client: - inputs, outputs = self._prepare_infer_args(input_value) + inputs, outputs = prepare_infer_args(input_value) triton_client.start_stream(callback=partial(callback, user_data)) # Start the request @@ -281,5 +296,201 @@ def test_max_inflight_requests_limit_request_cancellation(self): ) +class EnsembleStepMaxQueueSizeTest(tu.TestResultCollector): + def _run_inference(self, model_name, expected_responses_count): + """ + Helper function for streaming inference. + + For decoupled streaming ensembles with queue limit on internal step: + - Each producer response creates an independent flow through the ensemble + - Flows that complete before error is set send their outputs successfully + - Once error occurs (queue full), stream terminates with error + - Result: 0-N successful responses + 1 error (N depends on timing) + """ + user_data = UserData() + with grpcclient.InferenceServerClient(SERVER_URL) as triton_client: + try: + inputs, outputs = prepare_infer_args(expected_responses_count) + triton_client.start_stream(callback=partial(callback, user_data)) + triton_client.async_stream_infer( + model_name=model_name, inputs=inputs, outputs=outputs + ) + + # Collect and verify responses + errors, responses = collect_responses(user_data) + self.assertGreaterEqual( + len(responses), + 0, + "May have 0 or more successful responses depending on timing", + ) + self.assertLess( + len(responses), + expected_responses_count, + f"Should have fewer than {expected_responses_count} responses (some flows failed)", + ) + self.assertEqual( + len(errors), + 1, + "Expected exactly one error when queue full terminates stream", + ) + + # Verify correctness of successful responses + for idx, resp in enumerate(responses): + output = resp.as_numpy("OUT") + self.assertAlmostEqual( + output[0], + EXPECTED_INFER_OUTPUT, + places=5, + msg=f"Response {idx} has incorrect value - {output[0]}", + ) + + # Verify error is queue-full error + self.assertIn( + "Exceeds maximum queue size", + str(errors[0]), + f"Expected queue size error, got: {str(errors[0])}", + ) + finally: + triton_client.stop_stream() + + def _run_concurrent_inference(self, model_name, expected_responses_count): + """ + Helper function for concurrent independent requests. + Each request either succeeds completely or fails completely. + Returns: (num_successes, num_errors) tuple + """ + user_data = UserData() + with grpcclient.InferenceServerClient(SERVER_URL) as triton_client: + try: + inputs, outputs = prepare_infer_args(expected_responses_count) + triton_client.start_stream(callback=partial(callback, user_data)) + triton_client.async_stream_infer( + model_name=model_name, inputs=inputs, outputs=outputs + ) + + # Collect responses + errors, responses = collect_responses(user_data) + + # For concurrent independent requests with queue limit on internal step: + # - Requests that arrive before queue fills: succeed with all outputs + # - Requests that arrive after queue fills: fail with error + total = len(responses) + len(errors) + self.assertEqual( + total, + expected_responses_count, + f"Expected {expected_responses_count} total responses, got {total}", + ) + + if len(errors) > 0: + # This request failed + self.assertEqual( + len(responses), + 0, + "Failed request should have no successful outputs", + ) + self.assertEqual( + len(errors), 1, "Failed request should have exactly one error" + ) + self.assertIn( + "Exceeds maximum queue size", + str(errors[0]), + f"Expected queue size error, got: {str(errors[0])}", + ) + return (0, 1) # 0 successes, 1 error + else: + # This request succeeded + self.assertEqual( + len(responses), + expected_responses_count, + f"Successful request should have all {expected_responses_count} outputs", + ) + # Verify correctness of successful responses + for idx, resp in enumerate(responses): + output = resp.as_numpy("OUT") + self.assertAlmostEqual( + output[0], + EXPECTED_INFER_OUTPUT, + places=5, + msg=f"Response {idx} has incorrect value - {output[0]}", + ) + return (expected_responses_count, 0) # N successes, 0 errors + finally: + triton_client.stop_stream() + + def test_step1_max_queue_size(self): + """ + Test max_queue_size on step 1 (decoupled_producer). + + Trigger 32 concurrent ensemble requests, each producing 1 response + - Step 1 (producer) has max_queue_size limit + - Some ensemble requests succeed completely (before queue fills) + - Some fail completely (when producer queue is full) + """ + model_name = "ensemble_step1_enabled_max_queue_size" + num_requests = 32 + + # Store results from each thread + results = [] + + def thread_wrapper(model_name, expected_count, results_list): + """Wrapper to capture thread results""" + result = self._run_concurrent_inference(model_name, expected_count) + results_list.append(result) + + # Launch concurrent threads to perform infer requests + threads = [] + for i in range(num_requests): + t = threading.Thread(target=thread_wrapper, args=(model_name, 1, results)) + threads.append(t) + t.start() + + # Wait for all requests to complete + for t in threads: + t.join(timeout=60) + + # Aggregate results from all threads + total_successes = sum(r[0] for r in results) + total_errors = sum(r[1] for r in results) + + # Verify aggregate behavior + self.assertEqual( + total_successes + total_errors, + num_requests, + f"Expected {num_requests} total results (successes + errors), " + f"got {total_successes} successes + {total_errors} errors = {total_successes + total_errors}", + ) + + # Verify at least some errors occurred (queue limit was hit) + self.assertGreater( + total_errors, + 0, + f"Expected some errors due to max_queue_size limit, " + f"but all {num_requests} requests succeeded.", + ) + + # Verify at least some successes occurred (not all rejected) + self.assertGreater( + total_successes, + 0, + f"Expected some successful requests before queue filled, " + f"but all {num_requests} requests failed.", + ) + + def test_step2_max_queue_size(self): + """ + Test max_queue_size on step 2 (slow_consumer). + + Trigger 1 streaming ensemble request producing 32 responses + - Step 1 (producer) generates 32 responses rapidly (every 100ms) + - Step 2 (consumer) has max_queue_size=5 and processes slowly (500ms each) + - Each producer response is an independent request to the second step through + - the ensemble flow. Some requests complete successfully before queue fills + - When queue fills, error is set and stream terminates + - All inflight steps drain, then error response sent to client + """ + model_name = "ensemble_step2_enabled_max_queue_size" + self._run_inference(model_name=model_name, expected_responses_count=32) + + if __name__ == "__main__": unittest.main() diff --git a/qa/L0_simple_ensemble/test.sh b/qa/L0_simple_ensemble/test.sh index 458884533b..be11d0db76 100755 --- a/qa/L0_simple_ensemble/test.sh +++ b/qa/L0_simple_ensemble/test.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2019-2025, NVIDIA CORPORATION. All rights reserved. +# Copyright 2019-2026, NVIDIA CORPORATION. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -147,7 +147,91 @@ set -e kill $SERVER_PID wait $SERVER_PID -######## Test ensemble backpressure feature (max_inflight_requests parameter) +######## Test max_queue_size dynamic batching parameter in ensemble steps ######## +## Ensemble model: step1-decoupled_producer -> step2-slow_consumer +MODEL_DIR="`pwd`/max_queue_size_test_models" +rm -rf ${MODEL_DIR} + +# Enable max_queue_size in the first step (decoupled_producer) +mkdir -p ${MODEL_DIR}/ensemble_step1_enabled_max_queue_size/1 ${MODEL_DIR}/decoupled_producer_enabled_max_queue_size/1 ${MODEL_DIR}/slow_consumer/1 +cp ./backpressure_test_models/ensemble_disabled_max_inflight_requests/config.pbtxt ${MODEL_DIR}/ensemble_step1_enabled_max_queue_size/ +sed -i 's/"decoupled_producer"/"decoupled_producer_enabled_max_queue_size"/g' ${MODEL_DIR}/ensemble_step1_enabled_max_queue_size/config.pbtxt + +cp ../python_models/ground_truth/model.py ${MODEL_DIR}/slow_consumer/1 +cp ../python_models/ground_truth/config.pbtxt ${MODEL_DIR}/slow_consumer/ +sed -i 's/name: "ground_truth"/name: "slow_consumer"/g' ${MODEL_DIR}/slow_consumer/config.pbtxt +sed -i 's/max_batch_size: 64/max_batch_size: 1/g' ${MODEL_DIR}/slow_consumer/config.pbtxt + +cp ./backpressure_test_models/decoupled_producer/1/model.py ${MODEL_DIR}/decoupled_producer_enabled_max_queue_size/1 +cp ./backpressure_test_models/decoupled_producer/config.pbtxt ${MODEL_DIR}/decoupled_producer_enabled_max_queue_size/ +sed -i 's/name: "decoupled_producer"/name: "decoupled_producer_enabled_max_queue_size"/g' ${MODEL_DIR}/decoupled_producer_enabled_max_queue_size/config.pbtxt +# Add dynamic_batching with max_queue_size to decoupled_producer +cat >> ${MODEL_DIR}/decoupled_producer_enabled_max_queue_size/config.pbtxt << 'EOF' + +dynamic_batching { + preferred_batch_size: [ 1 ] + default_queue_policy { + max_queue_size: 4 + } +} +EOF + +# Enable max_queue_size in the second step (slow_consumer) +mkdir -p ${MODEL_DIR}/ensemble_step2_enabled_max_queue_size/1 ${MODEL_DIR}/decoupled_producer/1 ${MODEL_DIR}/slow_consumer_enabled_max_queue_size/1 +cp ./backpressure_test_models/ensemble_disabled_max_inflight_requests/config.pbtxt ${MODEL_DIR}/ensemble_step2_enabled_max_queue_size/ +sed -i 's/"slow_consumer"/"slow_consumer_enabled_max_queue_size"/g' ${MODEL_DIR}/ensemble_step2_enabled_max_queue_size/config.pbtxt + +cp ./backpressure_test_models/decoupled_producer/1/model.py ${MODEL_DIR}/decoupled_producer/1 +cp ./backpressure_test_models/decoupled_producer/config.pbtxt ${MODEL_DIR}/decoupled_producer/ + +cp ../python_models/ground_truth/model.py ${MODEL_DIR}/slow_consumer_enabled_max_queue_size/1 +cp ../python_models/ground_truth/config.pbtxt ${MODEL_DIR}/slow_consumer_enabled_max_queue_size/ +sed -i 's/name: "ground_truth"/name: "slow_consumer_enabled_max_queue_size"/g' ${MODEL_DIR}/slow_consumer_enabled_max_queue_size/config.pbtxt +sed -i 's/max_batch_size: 64/max_batch_size: 1/g' ${MODEL_DIR}/slow_consumer_enabled_max_queue_size/config.pbtxt +# Add dynamic_batching with max_queue_size to slow_consumer +cat >> ${MODEL_DIR}/slow_consumer_enabled_max_queue_size/config.pbtxt << 'EOF' + +dynamic_batching { + preferred_batch_size: [ 1 ] + default_queue_policy { + max_queue_size: 4 + } +} +EOF + +BACKPRESSURE_TEST_PY=./ensemble_backpressure_test.py +TEST_NAME="EnsembleStepMaxQueueSizeTest" +SERVER_LOG="./ensemble_step_max_queue_size_test_server.log" +CLIENT_LOG="./ensemble_step_max_queue_size_test_client.log" +rm -f $SERVER_LOG $CLIENT_LOG + +SERVER_ARGS="--model-repository=${MODEL_DIR}" +run_server +if [ "$SERVER_PID" == "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + cat $SERVER_LOG + exit 1 +fi + +set +e +python $BACKPRESSURE_TEST_PY $TEST_NAME -v >> $CLIENT_LOG 2>&1 +if [ $? -ne 0 ]; then + RET=1 +else + check_test_results $TEST_RESULT_FILE 2 + if [ $? -ne 0 ]; then + cat $CLIENT_LOG + echo -e "\n***\n*** Test Result Verification Failed\n***" + RET=1 + fi +fi +set -e + +kill $SERVER_PID +wait $SERVER_PID + + +######## Test ensemble backpressure feature (max_inflight_requests parameter) ######## MODEL_DIR="`pwd`/backpressure_test_models" mkdir -p ${MODEL_DIR}/ensemble_disabled_max_inflight_requests/1 @@ -173,7 +257,7 @@ sed -i 's/platform: "ensemble"/name: "ensemble_max_inflight_requests_limit_1"\np sed -i 's/ensemble_scheduling {/ensemble_scheduling {\n max_inflight_requests: 1/g' \ ${MODEL_DIR}/ensemble_max_inflight_requests_limit_1/config.pbtxt -BACKPRESSURE_TEST_PY=./ensemble_backpressure_test.py +TEST_NAME="EnsembleBackpressureTest" SERVER_LOG="./ensemble_backpressure_test_server.log" CLIENT_LOG="./ensemble_backpressure_test_client.log" rm -f $SERVER_LOG $CLIENT_LOG @@ -187,7 +271,7 @@ if [ "$SERVER_PID" == "0" ]; then fi set +e -python $BACKPRESSURE_TEST_PY -v >> $CLIENT_LOG 2>&1 +python $BACKPRESSURE_TEST_PY $TEST_NAME -v >> $CLIENT_LOG 2>&1 if [ $? -ne 0 ]; then RET=1 else