From d3ebc59185d528faad08ffce8a1a7fd7ab0f7450 Mon Sep 17 00:00:00 2001 From: mattwittwer Date: Thu, 11 Jun 2026 09:03:14 -0700 Subject: [PATCH] feat: Enable PyTorch2 Batching Tests (#8814) --- qa/L0_torch_aoti/test.sh | 57 +++- qa/L0_torch_aoti/torch_aoti_infer_test.py | 351 ++++++++++++++++++--- qa/common/gen_qa_implicit_models.py | 366 +++++++++++++++++++++- qa/common/gen_qa_model_repository | 2 + qa/common/gen_qa_models.py | 165 +++++++++- 5 files changed, 875 insertions(+), 66 deletions(-) diff --git a/qa/L0_torch_aoti/test.sh b/qa/L0_torch_aoti/test.sh index f37751c55e..9a09921e10 100755 --- a/qa/L0_torch_aoti/test.sh +++ b/qa/L0_torch_aoti/test.sh @@ -62,6 +62,7 @@ DATADIR=${DATADIR:="/data/inferenceserver/${REPO_VERSION}"} TRITON_DIR=${TRITON_DIR:="/opt/tritonserver"} SERVER=${TRITON_DIR}/bin/tritonserver BACKEND_DIR=${TRITON_DIR}/backends +SERVER_TIMEOUT=${SERVER_TIMEOUT:=120} # PyTorch on SBSA requires libgomp to be loaded first. See the following # GitHub issue for more information: @@ -79,7 +80,9 @@ export BACKENDS # Copy the models into the model repository echo -e "${COLOR_DARK}Setting up model repository in ${MODELDIR}${COLOR_RESET}" -rm -rf ${MODELDIR} && mkdir -p ${MODELDIR} +BAD_MODELDIR=`pwd`/bad_models +rm -rf ${MODELDIR} ${BAD_MODELDIR} +mkdir -p ${MODELDIR} models=( "torch_aoti_complex_index" "torch_aoti_complex_named" @@ -89,6 +92,8 @@ models=( "torch_aoti_int64_int64" "torch_aoti_float16_float16" "torch_aoti_float32_float32" + "torch_aoti_variable_float32" + "torch_aoti_multi_instance_float32" "torchvision_aoti" ) for model in "${models[@]}"; do @@ -96,6 +101,19 @@ for model in "${models[@]}"; do echo -e "${COLOR_DARK}ls ${MODELDIR}/${model}${COLOR_RESET}" ls -lha ${MODELDIR}/${model} done + +# Sequence-batching AOTI models live in the implicit-state sequence repository. +sequence_models=( + "torch_aoti_sequence_float32" + "torch_aoti_sequence_initstate_float32" + "torch_aoti_sequence_forward_float32" +) +for model in "${sequence_models[@]}"; do + cp -r ${DATADIR}/qa_sequence_implicit_model_repository/${model} ${MODELDIR}/${model} + echo -e "${COLOR_DARK}ls ${MODELDIR}/${model}${COLOR_RESET}" + ls -lha ${MODELDIR}/${model} +done + echo -e "${COLOR_DARK}ls ${MODELDIR}${COLOR_RESET}" ls -lha ${MODELDIR} @@ -133,11 +151,42 @@ fi echo -e "${COLOR_DARK}Killing server (pid: ${SERVER_PID})${COLOR_RESET}" kill -s SIGINT ${SERVER_PID} wait ${SERVER_PID} || true -echo -e "${COLOR_DARK}Removing model repository${COLOR_RESET}" -for model in "${models[@]}"; do - rm -rf ${MODELDIR}/${model} + +# Negative tests: these models declare unsupported types (TYPE_STRING CORRID / +# state) and must fail to load. Start a separate server (exit-on-error=false) so +# it stays up despite the load failures, then assert the models are not ready. +echo -e "${COLOR_DARK}Negative (load-failure) tests${COLOR_RESET}" +mkdir -p ${BAD_MODELDIR} +bad_models=( + "torch_aoti_sequence_bad_corrid" + "torch_aoti_sequence_bad_state" +) +for model in "${bad_models[@]}"; do + cp -r ${DATADIR}/qa_sequence_implicit_model_repository/${model} ${BAD_MODELDIR}/${model} done +SERVER_ARGS="--model-repository=${BAD_MODELDIR} --exit-on-error=false --log-verbose=1" +SERVER_LOG="./torch_aoti_negative-server.log" +run_server_tolive +if [[ "${SERVER_PID}" -eq 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Failed to start ${SERVER} (negative phase)\n***${COLOR_RESET}" 1>&2 + cat ${SERVER_LOG} 1>&2 + RET=1 +else + wait_for_model_stable ${SERVER_TIMEOUT} + for model in "${bad_models[@]}"; do + code=$(curl -s -o /dev/null -w "%{http_code}" localhost:8000/v2/models/${model}/ready) + if [[ "${code}" == "200" ]]; then + echo -e "${COLOR_ERROR}*** Negative model '${model}' unexpectedly loaded (ready)${COLOR_RESET}" 1>&2 + RET=1 + else + echo -e "${COLOR_INFO}*** Negative model '${model}' correctly failed to load${COLOR_RESET}" + fi + done + kill -s SIGINT ${SERVER_PID} + wait ${SERVER_PID} || true +fi + # Report results and exit. if [[ ${RET} -ne 0 ]]; then echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" &1>2 diff --git a/qa/L0_torch_aoti/torch_aoti_infer_test.py b/qa/L0_torch_aoti/torch_aoti_infer_test.py index 2b93f31a48..60cae41360 100755 --- a/qa/L0_torch_aoti/torch_aoti_infer_test.py +++ b/qa/L0_torch_aoti/torch_aoti_infer_test.py @@ -30,10 +30,13 @@ sys.path.append("../common") import unittest +from concurrent.futures import ThreadPoolExecutor +import numpy as np import test_util as tu import torch import tritonclient.http as http +from tritonclient.utils import InferenceServerException class TorchAotiTest(tu.TestResultCollector): @@ -203,81 +206,325 @@ def test_simple_model(self): torch.float16, torch.float32, ] + # The simple AOTI add/sub model is compiled with a dynamic batch + # dimension and configured with max_batch_size: 8. Exercise a range of + # batch sizes (including 1) so we validate that batched inputs are + # assembled and batched outputs are scattered back per-row correctly. + batch_sizes = [1, 4, 8] for io_type in io_types: MODEL_NAME = self._get_simple_model_name(io_type) - INPUT_SHAPE = (16,) - OUTPUT_SHAPE = (16,) TRITON_IO_TYPE = self._dtype_to_triton_dtype(io_type) - input_data = ( - self._get_simple_input_data(INPUT_SHAPE, io_type), - self._get_simple_input_data(INPUT_SHAPE, io_type), - ) + for batch_size in batch_sizes: + INPUT_SHAPE = (batch_size, 16) + OUTPUT_SHAPE = (batch_size, 16) - with http.InferenceServerClient("localhost:8000") as client: - inputs = [ - http.InferInput("ARGS[0]", input_data[0].shape, TRITON_IO_TYPE), - http.InferInput("ARGS[1]", input_data[1].shape, TRITON_IO_TYPE), - ] + input_data = ( + self._get_simple_input_data(INPUT_SHAPE, io_type), + self._get_simple_input_data(INPUT_SHAPE, io_type), + ) - inputs[0].set_data_from_numpy(input_data[0], binary_data=True) - inputs[1].set_data_from_numpy(input_data[1], binary_data=True) + with http.InferenceServerClient("localhost:8000") as client: + inputs = [ + http.InferInput("ARGS[0]", input_data[0].shape, TRITON_IO_TYPE), + http.InferInput("ARGS[1]", input_data[1].shape, TRITON_IO_TYPE), + ] - output_names = [ - "RESULT", - ] + inputs[0].set_data_from_numpy(input_data[0], binary_data=True) + inputs[1].set_data_from_numpy(input_data[1], binary_data=True) - outputs = [] - for output_name in output_names: - outputs.append( - http.InferRequestedOutput(output_name, binary_data=True) - ) + output_names = [ + "RESULT", + ] - output_data = [] - results = client.infer(MODEL_NAME, inputs, outputs=outputs) + outputs = [] + for output_name in output_names: + outputs.append( + http.InferRequestedOutput(output_name, binary_data=True) + ) + + output_data = [] + results = client.infer(MODEL_NAME, inputs, outputs=outputs) - for output_name in output_names: - output_data.append(results.as_numpy(output_name)) + for output_name in output_names: + output_data.append(results.as_numpy(output_name)) - self.assertEqual(len(outputs), len(output_data)) - for data in output_data: - self.assertEqual(data.shape, OUTPUT_SHAPE) - self.assertTrue((data == input_data[0] + input_data[1]).all()) + self.assertEqual(len(outputs), len(output_data)) + for data in output_data: + self.assertEqual(data.shape, OUTPUT_SHAPE) + self.assertTrue((data == input_data[0] + input_data[1]).all()) def test_torchvision(self): + # torchvision_aoti is exported with a dynamic batch dim (max_batch_size + # 8), so exercise batching of a real, higher-rank [N,3,224,224] model. MODEL_NAME = "torchvision_aoti" - INPUT_SHAPE = (1, 3, 224, 224) - OUTPUT_SHAPE = (1, 1000) - - input_data = self._get_torchvision_input_data(INPUT_SHAPE) - input_data[0][0] = 1.0 + with http.InferenceServerClient("localhost:8000") as client: + for batch_size in (1, 4, 8): + input_data = self._get_torchvision_input_data((batch_size, 3, 224, 224)) + inputs = [http.InferInput("ARGS[0]", input_data.shape, "FP32")] + inputs[0].set_data_from_numpy(input_data, binary_data=True) + outputs = [http.InferRequestedOutput("RESULT", binary_data=True)] + results = client.infer(MODEL_NAME, inputs, outputs=outputs) + data = results.as_numpy("RESULT") + self.assertEqual(data.shape, (batch_size, 1000)) + output_tensor = torch.from_numpy(data) + self.assertTrue(torch.isfinite(output_tensor).all().item()) + def test_batch_size_limit(self): + # A request whose batch exceeds max_batch_size (8) must be rejected; + # exactly max_batch_size must succeed. + MODEL_NAME = "torch_aoti_float32_float32" with http.InferenceServerClient("localhost:8000") as client: + ok = self._get_simple_input_data((8, 16), torch.float32) inputs = [ - http.InferInput("ARGS[0]", input_data.shape, "FP32"), + http.InferInput("ARGS[0]", ok.shape, "FP32"), + http.InferInput("ARGS[1]", ok.shape, "FP32"), ] - - inputs[0].set_data_from_numpy(input_data, binary_data=True) - - output_names = [ - "RESULT", + inputs[0].set_data_from_numpy(ok, binary_data=True) + inputs[1].set_data_from_numpy(ok, binary_data=True) + outputs = [http.InferRequestedOutput("RESULT", binary_data=True)] + client.infer(MODEL_NAME, inputs, outputs=outputs) # batch == max OK + + too_big = self._get_simple_input_data((16, 16), torch.float32) + big_inputs = [ + http.InferInput("ARGS[0]", too_big.shape, "FP32"), + http.InferInput("ARGS[1]", too_big.shape, "FP32"), ] + big_inputs[0].set_data_from_numpy(too_big, binary_data=True) + big_inputs[1].set_data_from_numpy(too_big, binary_data=True) + with self.assertRaises(InferenceServerException): + client.infer(MODEL_NAME, big_inputs, outputs=outputs) - outputs = [] - for output_name in output_names: - outputs.append(http.InferRequestedOutput(output_name, binary_data=True)) + def _infer_add(self, model_name, a, b, triton_type="FP32"): + # Run the two-input add model and return its RESULT output. + with http.InferenceServerClient("localhost:8000") as client: + inputs = [ + http.InferInput("ARGS[0]", a.shape, triton_type), + http.InferInput("ARGS[1]", b.shape, triton_type), + ] + inputs[0].set_data_from_numpy(a, binary_data=True) + inputs[1].set_data_from_numpy(b, binary_data=True) + outputs = [http.InferRequestedOutput("RESULT", binary_data=True)] + return client.infer(model_name, inputs, outputs=outputs).as_numpy("RESULT") - output_data = [] - results = client.infer(MODEL_NAME, inputs, outputs=outputs) + def _execution_count(self, model_name): + with http.InferenceServerClient("localhost:8000") as client: + stats = client.get_inference_statistics(model_name=model_name) + return int(stats["model_stats"][0]["execution_count"]) + + def _infer_one_row(self, model_name): + a = self._get_simple_input_data((1, 16), torch.float32) + b = self._get_simple_input_data((1, 16), torch.float32) + out = self._infer_add(model_name, a, b) + self.assertTrue((out == a + b).all()) + + def test_dynamic_batching_coalescing(self): + # Fire many concurrent single-row requests and confirm the dynamic + # batcher coalesced them into far fewer backend executions than requests. + MODEL_NAME = "torch_aoti_float32_float32" + num_requests = 200 + before = self._execution_count(MODEL_NAME) + with ThreadPoolExecutor(max_workers=32) as pool: + futures = [ + pool.submit(self._infer_one_row, MODEL_NAME) + for _ in range(num_requests) + ] + for future in futures: + future.result() + executions = self._execution_count(MODEL_NAME) - before + self.assertGreater(executions, 0) + self.assertLess(executions, num_requests) + + def test_multi_instance(self): + # Concurrent requests against a 2-instance model must all be correct. + MODEL_NAME = "torch_aoti_multi_instance_float32" + with ThreadPoolExecutor(max_workers=16) as pool: + futures = [pool.submit(self._infer_one_row, MODEL_NAME) for _ in range(64)] + for future in futures: + future.result() + + def test_variable_shape_batching(self): + # The variable model has dims [-1]; exercise different feature lengths + # across batch sizes. + MODEL_NAME = "torch_aoti_variable_float32" + for batch_size in (1, 4, 8): + for feature in (1, 16, 64): + a = np.random.randn(batch_size, feature).astype(np.float32) + b = np.random.randn(batch_size, feature).astype(np.float32) + out = self._infer_add(MODEL_NAME, a, b) + self.assertEqual(out.shape, (batch_size, feature)) + self.assertTrue((out == a + b).all()) + + +class TorchAotiSequenceTest(tu.TestResultCollector): + # The AOTI sequence model (see gen_qa_implicit_models.py) is a running + # accumulator that resets on the sequence start and adds the correlation id + # to each emitted output: + # new_state = INPUT0 + INPUT_STATE * (1 - START) + # OUTPUT0 = (new_state + CORRID) * READY + # Triton's sequence scheduler synthesizes the START / READY / CORRID control + # tensors and manages the implicit state, so the client only sends INPUT__0 + # (the correlation id is supplied via sequence_id). + MODEL_NAME = "torch_aoti_sequence_float32" + + def _infer_step( + self, + client, + seq_id, + value, + start, + end, + model=None, + in_name="INPUT__0", + out_name="OUTPUT__0", + ): + data = np.full((1, 1), value, dtype=np.float32) + inputs = [http.InferInput(in_name, data.shape, "FP32")] + inputs[0].set_data_from_numpy(data, binary_data=True) + outputs = [http.InferRequestedOutput(out_name, binary_data=True)] + result = client.infer( + model or self.MODEL_NAME, + inputs, + outputs=outputs, + sequence_id=seq_id, + sequence_start=start, + sequence_end=end, + ) + return result.as_numpy(out_name) + + def test_single_sequence(self): + seq_id = 100 + steps = [2.0, 3.0, 4.0, 5.0] + # Output is the running sum plus the correlation id. + expected = np.cumsum(steps) + seq_id + with http.InferenceServerClient("localhost:8000") as client: + for i, value in enumerate(steps): + out = self._infer_step( + client, + seq_id=seq_id, + value=value, + start=(i == 0), + end=(i == len(steps) - 1), + ) + self.assertEqual(out.shape, (1, 1)) + self.assertAlmostEqual(float(out[0, 0]), float(expected[i]), places=3) + + def test_interleaved_sequences(self): + # Two concurrent sequences must keep independent state. Each output is + # the per-sequence running sum plus that sequence's correlation id. + seqs = { + 201: {"steps": [1.0, 1.0, 1.0, 1.0], "sum": 0.0}, + 202: {"steps": [10.0, 20.0, 30.0, 40.0], "sum": 0.0}, + } + with http.InferenceServerClient("localhost:8000") as client: + num_steps = len(next(iter(seqs.values()))["steps"]) + for i in range(num_steps): + for seq_id, st in seqs.items(): + value = st["steps"][i] + st["sum"] += value + out = self._infer_step( + client, + seq_id=seq_id, + value=value, + start=(i == 0), + end=(i == num_steps - 1), + ) + self.assertAlmostEqual( + float(out[0, 0]), st["sum"] + seq_id, places=3 + ) - for output_name in output_names: - output_data.append(results.as_numpy(output_name)) + def test_many_concurrent_sequences(self): + # Fill every batch slot with a distinct live sequence (max_batch_size is + # 8); each must keep independent state (+ its own correlation id) as they + # are stepped in lockstep. + seq_ids = list(range(300, 308)) # 8 concurrent sequences == slots + sums = {s: 0.0 for s in seq_ids} + num_steps = 5 + with http.InferenceServerClient("localhost:8000") as client: + for i in range(num_steps): + for s in seq_ids: + value = float(s % 7 + 1) + sums[s] += value + out = self._infer_step( + client, + seq_id=s, + value=value, + start=(i == 0), + end=(i == num_steps - 1), + ) + self.assertAlmostEqual(float(out[0, 0]), sums[s] + s, places=3) + + def test_staggered_sequences(self): + # Sequences of different lengths begin and end at different ticks on a + # shared timeline, so they overlap: some finish (freeing their batch + # slot) while others stay live and new ones start into the freed slots. + # Each sequence keeps independent state that resets on its own START. + # plan: seq_id -> (first_tick, values) + plans = { + 400: (0, [1.0, 2.0]), # ticks 0-1 + 401: (0, [3.0, 4.0, 5.0, 6.0]), # ticks 0-3 + 402: (2, [10.0, 10.0, 10.0]), # ticks 2-4 (starts after 400 ends) + 403: (3, [7.0, 8.0]), # ticks 3-4 + } + last_tick = max(start + len(values) - 1 for start, values in plans.values()) + running = {seq_id: 0.0 for seq_id in plans} + with http.InferenceServerClient("localhost:8000") as client: + for tick in range(last_tick + 1): + for seq_id, (first_tick, values) in plans.items(): + idx = tick - first_tick + if idx < 0 or idx >= len(values): + continue # sequence not live at this tick + running[seq_id] += values[idx] + out = self._infer_step( + client, + seq_id=seq_id, + value=values[idx], + start=(idx == 0), + end=(idx == len(values) - 1), + ) + self.assertAlmostEqual( + float(out[0, 0]), running[seq_id] + seq_id, places=3 + ) - self.assertEqual(len(outputs), len(output_data)) - for data in output_data: - self.assertEqual(data.shape, OUTPUT_SHAPE) - output_tensor = torch.from_numpy(data) - self.assertTrue(torch.isfinite(output_tensor).all().item()) + def test_initial_state_sequence(self): + # Model relies on a declared zero initial_state (no START reset). Output + # is the running sum (no correlation id for this variant). + model = "torch_aoti_sequence_initstate_float32" + steps = [2.0, 3.0, 4.0, 5.0] + expected = np.cumsum(steps) + with http.InferenceServerClient("localhost:8000") as client: + for i, value in enumerate(steps): + out = self._infer_step( + client, + seq_id=500, + value=value, + start=(i == 0), + end=(i == len(steps) - 1), + model=model, + ) + self.assertAlmostEqual(float(out[0, 0]), float(expected[i]), places=3) + + def test_forward_interface_sequence(self): + # Same sequence artifact, but control/state addressed via the forward + # interface (ARGS[...] / RESULT[...]). Behaviour must match the ordinal + # model: running sum + correlation id. + model = "torch_aoti_sequence_forward_float32" + seq_id = 600 + steps = [2.0, 3.0, 4.0] + expected = np.cumsum(steps) + seq_id + with http.InferenceServerClient("localhost:8000") as client: + for i, value in enumerate(steps): + out = self._infer_step( + client, + seq_id=seq_id, + value=value, + start=(i == 0), + end=(i == len(steps) - 1), + model=model, + in_name="ARGS[0]", + out_name="RESULT[0]", + ) + self.assertAlmostEqual(float(out[0, 0]), float(expected[i]), places=3) if __name__ == "__main__": diff --git a/qa/common/gen_qa_implicit_models.py b/qa/common/gen_qa_implicit_models.py index 4b91ca24ea..d3b28da96a 100755 --- a/qa/common/gen_qa_implicit_models.py +++ b/qa/common/gen_qa_implicit_models.py @@ -1246,6 +1246,352 @@ def create_plan_modelconfig(models_dir, max_batch, dtype, shape): cfile.write(config) +def create_torch_aoti_modelfile(models_dir, model_version, max_batch, dtype, shape): + # AOT Inductor (PT2) sequence model. The forward arguments map positionally + # to the model's ordinal inputs/outputs, which the config addresses directly: + # INPUT__0 = INPUT0 (data), INPUT__1 = INPUT_STATE (implicit state), + # INPUT__2 = START (control), INPUT__3 = READY (control), + # INPUT__4 = CORRID (control) + # OUTPUT__0 = out (data), OUTPUT__1 = new_state (implicit state) + if dtype not in (np.float32, np.int32): + return + + torch_dtype = np_to_torch_dtype(dtype) + model_name = tu.get_sequence_model_name("torch_aoti", dtype) + shape = [abs(ips) for ips in shape] + + class SequenceNet(nn.Module): + def __init__(self): + super(SequenceNet, self).__init__() + + def forward(self, INPUT0, INPUT_STATE, START, READY, CORRID): + # On sequence START, reset the running state to INPUT0; otherwise + # accumulate onto the carried state. The emitted output adds the + # correlation id so tests can confirm CORRID delivery, and READY + # gates it (active batch slots have READY == 1). + keep = (1 - START).to(INPUT_STATE.dtype) + new_state = INPUT0 + INPUT_STATE * keep + out = (new_state + CORRID.to(new_state.dtype)) * READY.to(new_state.dtype) + return out, new_state + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = SequenceNet().to(device).eval() + + # Export with a dynamic first (batch) dimension so the AOTI artifact accepts + # any batch size in [1, max_batch]. The correlation id is delivered as an + # INT32 tensor (see config), independent of the model's data type. + export_batch = 2 if max_batch > 0 else 1 + data_shape = [export_batch] + list(shape) + ctrl_shape = [export_batch, 1] + sample_inputs = ( + torch.zeros(data_shape, dtype=torch_dtype, device=device), + torch.zeros(data_shape, dtype=torch_dtype, device=device), + torch.zeros(ctrl_shape, dtype=torch_dtype, device=device), + torch.zeros(ctrl_shape, dtype=torch_dtype, device=device), + torch.zeros(ctrl_shape, dtype=torch.int32, device=device), + ) + + dynamic_shapes = None + if max_batch > 0: + batch = torch.export.Dim("batch", min=1, max=max_batch) + dynamic_shapes = ( + {0: batch}, + {0: batch}, + {0: batch}, + {0: batch}, + {0: batch}, + ) + + model_version_dir = models_dir + "/" + model_name + "/" + str(model_version) + try: + os.makedirs(model_version_dir) + except OSError: + pass # ignore existing dir + + exported_model = torch.export.export( + model, sample_inputs, dynamic_shapes=dynamic_shapes + ) + torch._inductor.aoti_compile_and_package( + exported_model, package_path=model_version_dir + "/model.pt2" + ) + + +def create_torch_aoti_modelconfig(models_dir, max_batch, dtype, shape): + if dtype not in (np.float32, np.int32): + return + + model_name = tu.get_sequence_model_name("torch_aoti", dtype) + config_dir = models_dir + "/" + model_name + control_type = "int32" if dtype == np.int32 else "fp32" + + config = f""" +name: "{model_name}" +backend: "pytorch" +platform: "torch_aoti" +max_batch_size: {max_batch} +sequence_batching {{ + max_sequence_idle_microseconds: 5000000 + control_input [ + {{ + name: "INPUT__2" + control [ + {{ + kind: CONTROL_SEQUENCE_START + {control_type}_false_true: [ 0, 1 ] + }} + ] + }}, + {{ + name: "INPUT__3" + control [ + {{ + kind: CONTROL_SEQUENCE_READY + {control_type}_false_true: [ 0, 1 ] + }} + ] + }}, + {{ + name: "INPUT__4" + control [ + {{ + kind: CONTROL_SEQUENCE_CORRID + data_type: TYPE_INT32 + }} + ] + }} + ] + state [ + {{ + input_name: "INPUT__1" + output_name: "OUTPUT__1" + data_type: {np_to_model_dtype(dtype)} + dims: [ {tu.shape_to_dims_str(shape)} ] + }} + ] +}} +input [ + {{ + name: "INPUT__0" + data_type: {np_to_model_dtype(dtype)} + dims: [ {tu.shape_to_dims_str(shape)} ] + }} +] +output [ + {{ + name: "OUTPUT__0" + data_type: {np_to_model_dtype(dtype)} + dims: [ {tu.shape_to_dims_str(shape)} ] + }} +] +instance_group [ + {{ + kind: {"KIND_GPU" if torch.cuda.is_available() else "KIND_CPU"} + }} +] +""" + + try: + os.makedirs(config_dir) + except OSError: + pass # ignore existing dir + + with open(config_dir + "/config.pbtxt", "w") as cfile: + cfile.write(config) + + +def create_torch_aoti_forward_modelconfig(models_dir, model_version): + # Config-only variant of the float32 sequence model that addresses the + # control/state tensors via the forward-argument interface (ARGS[...] / + # RESULT[...]) instead of the ordinal INPUT__N / OUTPUT__N names. It reuses + # the float32 sequence artifact (5 positional inputs -> ARGS[0..4], two + # outputs -> RESULT[0], RESULT[1]). + import shutil + + src = models_dir + "/" + tu.get_sequence_model_name("torch_aoti", np.float32) + src_pt2 = src + "/" + str(model_version) + "/model.pt2" + if not os.path.exists(src_pt2): + print(f"warning: {src_pt2} not found; skipping forward-interface model") + return + + model_name = "torch_aoti_sequence_forward_float32" + dst_dir = models_dir + "/" + model_name + "/" + str(model_version) + try: + os.makedirs(dst_dir) + except OSError: + pass # ignore existing dir + shutil.copy(src_pt2, dst_dir + "/model.pt2") + + config = f""" +name: "{model_name}" +backend: "pytorch" +platform: "torch_aoti" +max_batch_size: 8 +sequence_batching {{ + max_sequence_idle_microseconds: 5000000 + control_input [ + {{ name: "ARGS[2]" control [{{ kind: CONTROL_SEQUENCE_START fp32_false_true: [ 0, 1 ] }}] }}, + {{ name: "ARGS[3]" control [{{ kind: CONTROL_SEQUENCE_READY fp32_false_true: [ 0, 1 ] }}] }}, + {{ name: "ARGS[4]" control [{{ kind: CONTROL_SEQUENCE_CORRID data_type: TYPE_INT32 }}] }} + ] + state [ + {{ + input_name: "ARGS[1]" + output_name: "RESULT[1]" + data_type: {np_to_model_dtype(np.float32)} + dims: [ 1 ] + }} + ] +}} +input [ + {{ name: "ARGS[0]" data_type: {np_to_model_dtype(np.float32)} dims: [ 1 ] }} +] +output [ + {{ name: "RESULT[0]" data_type: {np_to_model_dtype(np.float32)} dims: [ 1 ] }} +] +instance_group [{{ kind: {"KIND_GPU" if torch.cuda.is_available() else "KIND_CPU"} }}] +""" + with open(models_dir + "/" + model_name + "/config.pbtxt", "w") as f: + f.write(config) + print(f"Created forward-interface sequence model {model_name}") + + +def create_torch_aoti_initstate_model(models_dir, model_version, max_batch=8): + # Sequence model that relies on a declared zero initial_state rather than a + # START-driven reset: new_state = INPUT0 + INPUT_STATE (accumulate). On the + # first step the state input is the zero initial_state, so it behaves as a + # running sum that resets when the sequence (and its state) is recycled. + model_name = "torch_aoti_sequence_initstate_float32" + dst_dir = models_dir + "/" + model_name + "/" + str(model_version) + try: + os.makedirs(dst_dir) + except OSError: + pass # ignore existing dir + + class SequenceNet(nn.Module): + def forward(self, INPUT0, INPUT_STATE, START, READY): + new_state = INPUT0 + INPUT_STATE + out = new_state * READY.to(new_state.dtype) + return out, new_state + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = SequenceNet().to(device).eval() + batch = torch.export.Dim("batch", min=1, max=max_batch) + sample = ( + torch.zeros(2, 1, dtype=torch.float32, device=device), + torch.zeros(2, 1, dtype=torch.float32, device=device), + torch.zeros(2, 1, dtype=torch.float32, device=device), + torch.zeros(2, 1, dtype=torch.float32, device=device), + ) + ds = ({0: batch}, {0: batch}, {0: batch}, {0: batch}) + ep = torch.export.export(model, sample, dynamic_shapes=ds) + torch._inductor.aoti_compile_and_package(ep, package_path=dst_dir + "/model.pt2") + + config = f""" +name: "{model_name}" +backend: "pytorch" +platform: "torch_aoti" +max_batch_size: {max_batch} +sequence_batching {{ + max_sequence_idle_microseconds: 5000000 + control_input [ + {{ name: "INPUT__2" control [{{ kind: CONTROL_SEQUENCE_START fp32_false_true: [ 0, 1 ] }}] }}, + {{ name: "INPUT__3" control [{{ kind: CONTROL_SEQUENCE_READY fp32_false_true: [ 0, 1 ] }}] }} + ] + state [ + {{ + input_name: "INPUT__1" + output_name: "OUTPUT__1" + data_type: {np_to_model_dtype(np.float32)} + dims: [ 1 ] + initial_state: {{ + name: "zero state" + data_type: {np_to_model_dtype(np.float32)} + dims: [ 1 ] + zero_data: true + }} + }} + ] +}} +input [ + {{ name: "INPUT__0" data_type: {np_to_model_dtype(np.float32)} dims: [ 1 ] }} +] +output [ + {{ name: "OUTPUT__0" data_type: {np_to_model_dtype(np.float32)} dims: [ 1 ] }} +] +instance_group [{{ kind: {"KIND_GPU" if torch.cuda.is_available() else "KIND_CPU"} }}] +""" + with open(models_dir + "/" + model_name + "/config.pbtxt", "w") as f: + f.write(config) + print(f"Created initial-state sequence model {model_name}") + + +def create_torch_aoti_negative_configs(models_dir, model_version): + # Config-only negative models that reuse the float32 sequence artifact but + # declare an unsupported TYPE_STRING correlation id / state. These must fail + # to load; the L0 test starts a dedicated server and asserts the failure. + import shutil + + src = models_dir + "/" + tu.get_sequence_model_name("torch_aoti", np.float32) + src_pt2 = src + "/" + str(model_version) + "/model.pt2" + if not os.path.exists(src_pt2): + print(f"warning: {src_pt2} not found; skipping negative models") + return + + fp32 = np_to_model_dtype(np.float32) + gpu = "KIND_GPU" if torch.cuda.is_available() else "KIND_CPU" + variants = { + "torch_aoti_sequence_bad_corrid": f""" +name: "torch_aoti_sequence_bad_corrid" +backend: "pytorch" +platform: "torch_aoti" +max_batch_size: 8 +sequence_batching {{ + max_sequence_idle_microseconds: 5000000 + control_input [ + {{ name: "INPUT__2" control [{{ kind: CONTROL_SEQUENCE_START fp32_false_true: [ 0, 1 ] }}] }}, + {{ name: "INPUT__3" control [{{ kind: CONTROL_SEQUENCE_READY fp32_false_true: [ 0, 1 ] }}] }}, + {{ name: "INPUT__4" control [{{ kind: CONTROL_SEQUENCE_CORRID data_type: TYPE_STRING }}] }} + ] + state [ + {{ input_name: "INPUT__1" output_name: "OUTPUT__1" data_type: {fp32} dims: [ 1 ] }} + ] +}} +input [ {{ name: "INPUT__0" data_type: {fp32} dims: [ 1 ] }} ] +output [ {{ name: "OUTPUT__0" data_type: {fp32} dims: [ 1 ] }} ] +instance_group [{{ kind: {gpu} }}] +""", + "torch_aoti_sequence_bad_state": f""" +name: "torch_aoti_sequence_bad_state" +backend: "pytorch" +platform: "torch_aoti" +max_batch_size: 8 +sequence_batching {{ + max_sequence_idle_microseconds: 5000000 + control_input [ + {{ name: "INPUT__2" control [{{ kind: CONTROL_SEQUENCE_START fp32_false_true: [ 0, 1 ] }}] }}, + {{ name: "INPUT__3" control [{{ kind: CONTROL_SEQUENCE_READY fp32_false_true: [ 0, 1 ] }}] }} + ] + state [ + {{ input_name: "INPUT__1" output_name: "OUTPUT__1" data_type: TYPE_STRING dims: [ 1 ] }} + ] +}} +input [ {{ name: "INPUT__0" data_type: {fp32} dims: [ 1 ] }} ] +output [ {{ name: "OUTPUT__0" data_type: {fp32} dims: [ 1 ] }} ] +instance_group [{{ kind: {gpu} }}] +""", + } + for model_name, config in variants.items(): + dst_dir = models_dir + "/" + model_name + "/" + str(model_version) + try: + os.makedirs(dst_dir) + except OSError: + pass # ignore existing dir + shutil.copy(src_pt2, dst_dir + "/model.pt2") + with open(models_dir + "/" + model_name + "/config.pbtxt", "w") as f: + f.write(config) + print(f"Created negative sequence model {model_name}") + + def create_models(models_dir, dtype, shape, initial_state, no_batch=True): model_version = 1 @@ -1288,6 +1634,18 @@ def create_models(models_dir, dtype, shape, initial_state, no_batch=True): models_dir, model_version, 0, dtype, shape + suffix, initial_state ) + if FLAGS.torch_aoti: + # AOTI sequence models are generated with first-dim batching enabled. + create_torch_aoti_modelconfig(models_dir, 8, dtype, shape) + create_torch_aoti_modelfile(models_dir, model_version, 8, dtype, shape) + # Generate the float32-only variants once (they reuse / extend the + # float32 sequence artifact): forward-interface naming, declared zero + # initial_state, and the negative (unsupported-type) load-failure models. + if dtype == np.float32 and no_batch: + create_torch_aoti_forward_modelconfig(models_dir, model_version) + create_torch_aoti_initstate_model(models_dir, model_version, 8) + create_torch_aoti_negative_configs(models_dir, model_version) + if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -1328,6 +1686,12 @@ def create_models(models_dir, dtype, shape, initial_state, no_batch=True): action="store_true", help="Generate Pytorch LibTorch models", ) + parser.add_argument( + "--torch-aoti", + required=False, + action="store_true", + help="Generate PyTorch AOT Inductor (PT2) sequence models", + ) parser.add_argument( "--openvino", required=False, @@ -1356,7 +1720,7 @@ def create_models(models_dir, dtype, shape, initial_state, no_batch=True): if FLAGS.tensorrt: import tensorrt as trt - if FLAGS.libtorch: + if FLAGS.libtorch or FLAGS.torch_aoti: import torch from torch import nn diff --git a/qa/common/gen_qa_model_repository b/qa/common/gen_qa_model_repository index ca873fe8f1..7952ef5cb8 100755 --- a/qa/common/gen_qa_model_repository +++ b/qa/common/gen_qa_model_repository @@ -277,6 +277,8 @@ python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --libtorch --variable --m chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --libtorch --models_dir=$TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL +python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --torch-aoti --models_dir=$TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL +chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --libtorch --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL diff --git a/qa/common/gen_qa_models.py b/qa/common/gen_qa_models.py index 2008cd824a..bf34cb4eb9 100755 --- a/qa/common/gen_qa_models.py +++ b/qa/common/gen_qa_models.py @@ -1383,6 +1383,7 @@ def np_to_dtype(np_dtype): def create_torch_aoti_model_file( models_dir, + max_batch, model_version, input_shape, input_dtype, @@ -1451,11 +1452,32 @@ def forward(self, INPUT0: torch.Tensor, INPUT1: torch.Tensor) -> torch.Tensor: model.to(device) model = model.eval() - sample_inputs = generate_torch_aoti_sample_inputs(input_shape, input_dtype, device) + # When batching is enabled, the AOTI artifact must be compiled with a + # dynamic first (batch) dimension. Otherwise the compiled model specializes + # to a static shape and rejects any batch size other than the one used at + # export time. We export with a representative batch > 1 and declare dim 0 + # of each input as dynamic over [1, max_batch]. + if max_batch > 0: + export_input_shape = [2] + list(input_shape) + else: + export_input_shape = list(input_shape) + + sample_inputs = generate_torch_aoti_sample_inputs( + export_input_shape, input_dtype, device + ) + + dynamic_shapes = None + if max_batch > 0: + batch_dim = torch.export.Dim("batch", min=1, max=max_batch) + # One spec per positional argument of AddSubNet.forward(INPUT0, INPUT1). + dynamic_shapes = ({0: batch_dim}, {0: batch_dim}) + package_path = os.path.join(model_version_dir, "model.pt2") try: - exported_model = torch.export.export(model, sample_inputs) + exported_model = torch.export.export( + model, sample_inputs, dynamic_shapes=dynamic_shapes + ) torch._inductor.aoti_compile_and_package( exported_model, package_path=package_path, @@ -1640,15 +1662,22 @@ def create_torchvision_aoti_model_file( model = model.to(device) model = model.eval() - SHAPE = (max_batch, 3, 224, 224) + # When batching is enabled, export with a dynamic first (batch) dimension so + # the AOTI artifact accepts any batch size in [1, max_batch]. A batch>=2 + # sample is used to avoid 0/1 specialization. + if max_batch > 0: + SHAPE = (2, 3, 224, 224) + dynamic_shapes = ({0: torch.export.Dim("batch", min=1, max=max_batch)},) + else: + SHAPE = (1, 3, 224, 224) + dynamic_shapes = None - # Example input tensor with batch size 1 and 3 color channels (RGB), height and width of 224 sample_inputs = (torch.zeros(SHAPE, dtype=torch.float32, device=device),) package_path = os.path.join(model_version_dir, "model.pt2") try: - ep = torch.export.export(model, sample_inputs) + ep = torch.export.export(model, sample_inputs, dynamic_shapes=dynamic_shapes) torch._inductor.aoti_compile_and_package(ep, package_path=package_path) except Exception as e: print( @@ -1759,6 +1788,7 @@ def create_libtorch_modelconfig( def create_torch_aoti_model_config( models_dir, + max_batch, input_shape, output_shape, input_dtype, @@ -1777,7 +1807,6 @@ def create_torch_aoti_model_config( else: version_policy_str = "{ all { }}" - # Use a different model name for the non-batching variant model_name = tu.get_model_name( "torch_aoti", input_dtype, @@ -1793,7 +1822,12 @@ def create_torch_aoti_model_config( backend: "pytorch" name: "{model_name}" platform: "torch_aoti" +max_batch_size: {max_batch} version_policy: {version_policy_str} +dynamic_batching {{ + preferred_batch_size: [ 4, 8 ] + max_queue_delay_microseconds: 1000 +}} input [ {{ name: "ARGS[0]" @@ -1836,6 +1870,112 @@ def create_torch_aoti_model_config( print(f"Created {label_path}") +def create_torch_aoti_variable_model(models_dir, max_batch=8): + # AOTI add model exported with two dynamic dimensions (batch + feature) so + # batching can be exercised on a model whose non-batch shape is variable + # (dims: [-1]). Float32 only. + model_name = "torch_aoti_variable_float32" + model_version_dir = os.path.join(models_dir, model_name, "1") + try: + os.makedirs(model_version_dir) + except OSError: + pass # ignore existing dir + + print(f"{_color_green}Creating model {model_name}{_color_reset}") + + class AddNet(nn.Module): + def forward(self, INPUT0, INPUT1): + return INPUT0 + INPUT1 + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = AddNet().to(device).eval() + + # Sample uses batch>=2 and feature>=2 to avoid 0/1 specialization. + sample = ( + torch.zeros(2, 4, dtype=torch.float32, device=device), + torch.zeros(2, 4, dtype=torch.float32, device=device), + ) + batch = torch.export.Dim("batch", min=1, max=max_batch) + feature = torch.export.Dim("feature", min=1, max=512) + dynamic_shapes = ( + {0: batch, 1: feature}, + {0: batch, 1: feature}, + ) + try: + ep = torch.export.export(model, sample, dynamic_shapes=dynamic_shapes) + torch._inductor.aoti_compile_and_package( + ep, package_path=os.path.join(model_version_dir, "model.pt2") + ) + except Exception as e: + print( + f"{_color_red}error: Failed to create model {model_name}: {e}{_color_reset}", + file=sys.stderr, + ) + return + + config = f""" +backend: "pytorch" +name: "{model_name}" +platform: "torch_aoti" +max_batch_size: {max_batch} +dynamic_batching {{ max_queue_delay_microseconds: 1000 }} +input [ + {{ name: "ARGS[0]" data_type: TYPE_FP32 dims: [ -1 ] }}, + {{ name: "ARGS[1]" data_type: TYPE_FP32 dims: [ -1 ] }} +] +output [ + {{ name: "RESULT" data_type: TYPE_FP32 dims: [ -1 ] }} +] +instance_group [{{ kind: {"KIND_GPU" if torch.cuda.is_available() else "KIND_CPU"} }}] +""" + with open(os.path.join(models_dir, model_name, "config.pbtxt"), "w") as f: + f.write(config) + print(f"Created config for {model_name}") + + +def create_torch_aoti_multi_instance_model(models_dir, max_batch=8): + # Config-only variant that reuses the float32 add model artifact but runs + # with two instances, to exercise batching across multiple model instances. + import shutil + + src_name = tu.get_model_name("torch_aoti", np.float32, np.float32, None) + src_pt2 = os.path.join(models_dir, src_name, "1", "model.pt2") + if not os.path.exists(src_pt2): + print( + f"{_color_yellow}warning: {src_pt2} not found; skipping multi-instance " + f"model{_color_reset}" + ) + return + + model_name = "torch_aoti_multi_instance_float32" + model_version_dir = os.path.join(models_dir, model_name, "1") + try: + os.makedirs(model_version_dir) + except OSError: + pass # ignore existing dir + shutil.copy(src_pt2, os.path.join(model_version_dir, "model.pt2")) + + print(f"{_color_green}Creating model {model_name}{_color_reset}") + config = f""" +backend: "pytorch" +name: "{model_name}" +platform: "torch_aoti" +max_batch_size: {max_batch} +dynamic_batching {{ max_queue_delay_microseconds: 1000 }} +input [ + {{ name: "ARGS[0]" data_type: TYPE_FP32 dims: [ 16 ] }}, + {{ name: "ARGS[1]" data_type: TYPE_FP32 dims: [ 16 ] }} +] +output [ + {{ name: "RESULT" data_type: TYPE_FP32 dims: [ 16 ] }} +] +instance_group [{{ count: 2 kind: {"KIND_GPU" if torch.cuda.is_available() else "KIND_CPU"} }}] +""" + with open(os.path.join(models_dir, model_name, "config.pbtxt"), "w") as f: + f.write(config) + print(f"Created config for {model_name}") + + def create_torch_aoti_complex_model_config( models_dir, ): @@ -2405,6 +2545,7 @@ def create_models( # max-batch 8 if create_torch_aoti_model_file( models_dir, + 8, model_version, input_shape, input_dtype, @@ -2412,6 +2553,7 @@ def create_models( ): create_torch_aoti_model_config( models_dir, + 8, input_shape, output0_shape, input_dtype, @@ -3061,9 +3203,14 @@ def create_fixed_models( ) if create_torch_aoti_complex_model_file(FLAGS.models_dir): create_torch_aoti_complex_model_config(FLAGS.models_dir) + # Batching coverage models: variable non-batch dim, and a multi-instance + # variant that reuses the float32 add model artifact. + create_torch_aoti_variable_model(FLAGS.models_dir) + create_torch_aoti_multi_instance_model(FLAGS.models_dir) if FLAGS.torchvision_aoti: - # TODO: Add support for variable batch size and version policy for torchvision AOTI models. print(f"{_color_blue}TorchVision AOTI model generation requested{_color_reset}") - if create_torchvision_aoti_model_file(FLAGS.models_dir, 1): - create_torchvision_aoti_model_config(FLAGS.models_dir, 1) + # Export with a dynamic batch dimension (max_batch_size 8) to exercise + # batching of a real, higher-rank ([N,3,224,224]) model. + if create_torchvision_aoti_model_file(FLAGS.models_dir, 8): + create_torchvision_aoti_model_config(FLAGS.models_dir, 8)