From 35d0333219fe3da569851b784771fe0b3e0af226 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Sun, 30 Aug 2026 16:37:27 -0700 Subject: [PATCH 1/5] Fix stale metrics namespace assertion in TensorRT test test_namespace has asserted 'RunInferenceTensorRT' since the original TensorRT commit (a8ca3057c0b). The handler was later changed to return 'BeamML_TensorRT' in f477b85f230, matching the BeamML_* prefix that every other model handler uses, but the test was never updated. The mismatch went unnoticed because the TensorRT suite does not run in any active CI job. --- sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py index 80a01b8f4d4c..bf190da4e5d4 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py @@ -349,7 +349,7 @@ def test_namespace(self): inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4) self.assertEqual( - 'RunInferenceTensorRT', inference_runner.get_metrics_namespace()) + 'BeamML_TensorRT', inference_runner.get_metrics_namespace()) @pytest.mark.uses_tensorrt From bb37768546da010fa9e334197fd6b0bd074ba388 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Sun, 30 Aug 2026 16:37:39 -0700 Subject: [PATCH 2/5] Support TensorRT 10 and later in TensorRTEngineHandlerNumPy TensorRT 10 removed the index based binding API that the handler was written against, so RunInference fails at engine load time with: AttributeError: 'ICudaEngine' object has no attribute 'num_bindings' Select the API at runtime from the TensorRT major version rather than picking one of them. TensorRT 8.x keeps the binding API and execute_async_v2, while TensorRT 10 and later use the name based tensor API and execute_async_v3. No currently supported GPU loses support. Supporting both versions is necessary rather than merely convenient. Dataflow now offers Blackwell GPUs (RTX Pro 6000, compute capability 12.0) that no TensorRT 8.x release can target, while TensorRT 10 and later require compute capability 7.5 or higher and so cannot target the Pascal and Volta GPUs that Dataflow still offers. No single TensorRT version covers the whole range. Also handle cuda-python 13, which removed the cuda.cuda alias in favour of cuda.bindings.driver, and move the test container to nvcr.io/nvidia/tensorrt:26.06-py3 (TensorRT 11.0, CUDA 13.3, Python 3.12). Because that image is Python 3.12, the disabled tensorRTtests task moves from the py310 suite to the py312 suite. The Dataflow integration test stays disabled. Every .trt engine staged under gs://apache-beam-ml/models/ was built with TensorRT 8.x, and a serialized engine can only be read by the major version that built it. Rebuilt and verified replacements are available, but staging them needs write access to that bucket; see the pull request description. Verified on a T4 GPU on GCE: 7/7 tests pass under TensorRT 11.0 (nvcr.io/nvidia/tensorrt:26.06-py3) and 7/7 under TensorRT 8.6.1 (23.05-py3). Addresses #36306 Addresses #33946 --- CHANGES.md | 1 + .../apache_beam/examples/inference/README.md | 9 +- .../ml/inference/tensorrt_inference.py | 137 +++++++++++++----- .../ml/inference/tensorrt_inference_test.py | 65 +++++---- .../tensor_rt.dockerfile | 15 +- .../python/test-suites/dataflow/common.gradle | 11 +- .../documentation/ml/tensorrt-runinference.md | 8 +- 7 files changed, 174 insertions(+), 72 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 0c519076f0db..e41b5b59d9e2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -94,6 +94,7 @@ * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). +* (Python) `TensorRTEngineHandlerNumPy` now works with TensorRT 10 and later, which removed the binding API it was written against. TensorRT 8.x remains supported, so no existing GPU loses support ([#36306](https://github.com/apache/beam/issues/36306)). ## Security Fixes diff --git a/sdks/python/apache_beam/examples/inference/README.md b/sdks/python/apache_beam/examples/inference/README.md index 5eed659d068c..e9fc8b55a742 100644 --- a/sdks/python/apache_beam/examples/inference/README.md +++ b/sdks/python/apache_beam/examples/inference/README.md @@ -87,9 +87,16 @@ To use TensorRT locally, we suggest an environment with TensorRT >= 8.0.1. Insta [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html). You will need to make sure the Python bindings for TensorRT are also installed correctly, these are available by installing the python3-libnvinfer and python3-libnvinfer-dev packages on your TensorRT download. +Both the TensorRT 8.x binding API and the TensorRT 10.x and later tensor API are +supported. Note that a serialized TensorRT engine can only be deserialized by +the TensorRT major version that built it, so an engine built with TensorRT 8.x +must be rebuilt before it can be used with TensorRT 10 or later. TensorRT 10 and +later also require a GPU with compute capability 7.5 or higher, which excludes +NVIDIA Pascal and Volta GPUs such as the Tesla P4, P100 and V100. + If you would like to use Docker, you can use an NGC image like: ``` -docker pull nvcr.io/nvidia/tensorrt:22.04-py3 +docker pull nvcr.io/nvidia/tensorrt:26.06-py3 ``` as an existing container base to [build custom Apache Beam container](https://beam.apache.org/documentation/runtime/environments/#modify-existing-base-image). diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py index 333187301b29..214847070944 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py +++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py @@ -49,6 +49,34 @@ LOGGER.warning(msg) +def _trt_major_version() -> int: + """Returns the major version of the installed TensorRT. + + TensorRT 10 replaced the index based "binding" API with a name based + "tensor" API, so the major version selects which code path to take. + """ + import tensorrt as trt + try: + return int(trt.__version__.split('.')[0]) + except (AttributeError, IndexError, ValueError): + # Fall back to probing for an attribute that only exists from 10 onwards. + return 10 if hasattr(trt.ICudaEngine, 'num_io_tensors') else 8 + + +def _import_cuda_driver(): + """Imports the CUDA driver bindings. + + ``cuda.bindings.driver`` is the module path used by cuda-python 12.8 and + later. It replaced the ``cuda.cuda`` alias, which was removed in + cuda-python 13.0, so only fall back to that for older installations. + """ + try: + from cuda.bindings import driver as cuda + except ImportError: + from cuda import cuda + return cuda + + def _load_engine(engine_path): import tensorrt as trt file = FileSystems.open(engine_path, 'rb') @@ -58,11 +86,25 @@ def _load_engine(engine_path): return engine +def _network_creation_flags() -> int: + """Returns the ``create_network`` flags for the installed TensorRT. + + Explicit batch is the only supported mode from TensorRT 10 onwards, where + the flag is first deprecated and then removed, so it is only passed to + TensorRT 8.x. + """ + import tensorrt as trt + explicit_batch = getattr( + trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH', None) + if explicit_batch is None or _trt_major_version() >= 10: + return 0 + return 1 << int(explicit_batch) + + def _load_onnx(onnx_path): import tensorrt as trt builder = trt.Builder(TRT_LOGGER) - network = builder.create_network( - flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + network = builder.create_network(flags=_network_creation_flags()) parser = trt.OnnxParser(network, TRT_LOGGER) with FileSystems.open(onnx_path) as f: if not parser.parse(f.read()): @@ -85,7 +127,7 @@ def _build_engine(network, builder): def _assign_or_fail(args): """CUDA error checking.""" - from cuda import cuda + cuda = _import_cuda_driver() err, ret = args[0], args[1:] if isinstance(err, cuda.CUresult): if err != cuda.CUresult.CUDA_SUCCESS: @@ -111,7 +153,7 @@ def __init__(self, engine: trt.ICudaEngine): engine: trt.ICudaEngine object that contains TensorRT engine """ import tensorrt as trt - from cuda import cuda + cuda = _import_cuda_driver() self.engine = engine self.context = engine.create_execution_context() self.context_lock = threading.RLock() @@ -120,34 +162,59 @@ def __init__(self, engine: trt.ICudaEngine): self.gpu_allocations = [] self.cpu_allocations = [] - # TODO(https://github.com/NVIDIA/TensorRT/issues/2557): - # Clean up when fixed upstream. - try: - _ = np.bool - except AttributeError: - # numpy >= 1.24.0 - np.bool = np.bool_ # type: ignore - # Setup I/O bindings. - for i in range(self.engine.num_bindings): - name = self.engine.get_binding_name(i) - dtype = self.engine.get_binding_dtype(i) - shape = self.engine.get_binding_shape(i) - size = trt.volume(shape) * dtype.itemsize - allocation = _assign_or_fail(cuda.cuMemAlloc(size)) - binding = { - 'index': i, - 'name': name, - 'dtype': np.dtype(trt.nptype(dtype)), - 'shape': list(shape), - 'allocation': allocation, - 'size': size - } - self.gpu_allocations.append(allocation) - if self.engine.binding_is_input(i): - self.inputs.append(binding) - else: - self.outputs.append(binding) + if _trt_major_version() >= 10: + # TensorRT 10 removed the index based binding API in favour of a name + # based tensor API. Device addresses are bound to the context once here + # because execute_async_v3 takes no allocation list at execution time. + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + dtype = self.engine.get_tensor_dtype(name) + shape = self.engine.get_tensor_shape(name) + size = trt.volume(shape) * dtype.itemsize + allocation = _assign_or_fail(cuda.cuMemAlloc(size)) + binding = { + 'index': i, + 'name': name, + 'dtype': np.dtype(trt.nptype(dtype)), + 'shape': list(shape), + 'allocation': allocation, + 'size': size + } + self.gpu_allocations.append(allocation) + self.context.set_tensor_address(name, int(allocation)) + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + self.inputs.append(binding) + else: + self.outputs.append(binding) + else: + # TODO(https://github.com/NVIDIA/TensorRT/issues/2557): + # Clean up when the TensorRT 8.x path is dropped. + try: + _ = np.bool + except AttributeError: + # numpy >= 1.24.0 + np.bool = np.bool_ # type: ignore + + for i in range(self.engine.num_bindings): + name = self.engine.get_binding_name(i) + dtype = self.engine.get_binding_dtype(i) + shape = self.engine.get_binding_shape(i) + size = trt.volume(shape) * dtype.itemsize + allocation = _assign_or_fail(cuda.cuMemAlloc(size)) + binding = { + 'index': i, + 'name': name, + 'dtype': np.dtype(trt.nptype(dtype)), + 'shape': list(shape), + 'allocation': allocation, + 'size': size + } + self.gpu_allocations.append(allocation) + if self.engine.binding_is_input(i): + self.inputs.append(binding) + else: + self.outputs.append(binding) assert self.context assert len(self.inputs) > 0 @@ -182,7 +249,7 @@ def _default_tensorRT_inference_fn( engine: TensorRTEngine, inference_args: Optional[dict[str, Any]] = None) -> Iterable[PredictionResult]: - from cuda import cuda + cuda = _import_cuda_driver() ( engine, context, @@ -201,7 +268,11 @@ def _default_tensorRT_inference_fn( np.ascontiguousarray(batch), inputs[0]['size'], stream)) - context.execute_async_v2(gpu_allocations, stream) + if _trt_major_version() >= 10: + # Tensor addresses were bound when the engine was created. + context.execute_async_v3(stream) + else: + context.execute_async_v2(gpu_allocations, stream) for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py index bf190da4e5d4..2cca02e704a1 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py @@ -37,6 +37,10 @@ from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorrt_inference import TensorRTEngineHandlerNumPy + from apache_beam.ml.inference.tensorrt_inference import _assign_or_fail + from apache_beam.ml.inference.tensorrt_inference import _import_cuda_driver + from apache_beam.ml.inference.tensorrt_inference import _network_creation_flags + from apache_beam.ml.inference.tensorrt_inference import _trt_major_version except ImportError: raise unittest.SkipTest('TensorRT dependencies are not installed') @@ -90,23 +94,8 @@ def _compare_prediction_result(a, b): for actual, expected in zip(a.inference, b.inference))) -def _assign_or_fail(args): - """CUDA error checking.""" - from cuda import cuda - err, ret = args[0], args[1:] - if isinstance(err, cuda.CUresult): - if err != cuda.CUresult.CUDA_SUCCESS: - raise RuntimeError("Cuda Error: {}".format(err)) - else: - raise RuntimeError("Unknown error type: {}".format(err)) - # Special case so that no unpacking is needed at call-site. - if len(ret) == 1: - return ret[0] - return ret - - def _custom_tensorRT_inference_fn(batch, engine, inference_args): - from cuda import cuda + cuda = _import_cuda_driver() ( engine, context, @@ -125,7 +114,10 @@ def _custom_tensorRT_inference_fn(batch, engine, inference_args): np.ascontiguousarray(batch), inputs[0]['size'], stream)) - context.execute_async_v2(gpu_allocations, stream) + if _trt_major_version() >= 10: + context.execute_async_v3(stream) + else: + context.execute_async_v2(gpu_allocations, stream) for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( @@ -189,8 +181,7 @@ def test_inference_single_tensor_feature(self): inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4) builder = trt.Builder(LOGGER) - network = builder.create_network( - flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + network = builder.create_network(flags=_network_creation_flags()) input_tensor = network.add_input( name="input", dtype=trt.float32, shape=(4, 1)) weight_const = network.add_constant( @@ -227,8 +218,7 @@ def test_inference_custom_single_tensor_feature(self): max_batch_size=4, inference_fn=_custom_tensorRT_inference_fn) builder = trt.Builder(LOGGER) - network = builder.create_network( - flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + network = builder.create_network(flags=_network_creation_flags()) input_tensor = network.add_input( name="input", dtype=trt.float32, shape=(4, 1)) weight_const = network.add_constant( @@ -263,8 +253,7 @@ def test_inference_multiple_tensor_features(self): inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4) builder = trt.Builder(LOGGER) - network = builder.create_network( - flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + network = builder.create_network(flags=_network_creation_flags()) input_tensor = network.add_input( name="input", dtype=trt.float32, shape=(4, 2)) weight_const = network.add_constant( @@ -351,6 +340,29 @@ def test_namespace(self): self.assertEqual( 'BeamML_TensorRT', inference_runner.get_metrics_namespace()) + def test_version_check_matches_installed_api(self): + """The branch taken must match the API the installed TensorRT exposes. + + TensorRT 10 removed the index based binding API in favour of the name + based tensor API. This guards against the version check drifting away + from the API it selects. + """ + if _trt_major_version() >= 10: + self.assertTrue(hasattr(trt.ICudaEngine, 'num_io_tensors')) + self.assertTrue(hasattr(trt.IExecutionContext, 'execute_async_v3')) + else: + self.assertTrue(hasattr(trt.ICudaEngine, 'num_bindings')) + self.assertTrue(hasattr(trt.IExecutionContext, 'execute_async_v2')) + + def test_network_creation_flags(self): + """Explicit batch must only be requested on TensorRT 8.x.""" + if _trt_major_version() >= 10: + self.assertEqual(0, _network_creation_flags()) + else: + self.assertEqual( + 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH), + _network_creation_flags()) + @pytest.mark.uses_tensorrt class TensorRTRunInferencePipelineTest(unittest.TestCase): @@ -381,7 +393,7 @@ def fake_inference_fn(batch, engine, inference_args=None): raise Exception( f'Loaded engine of type {type(engine)}, was ' + 'expecting multi_process_shared engine') - from cuda import cuda + cuda = _import_cuda_driver() ( engine, context, @@ -400,7 +412,10 @@ def fake_inference_fn(batch, engine, inference_args=None): np.ascontiguousarray(batch), inputs[0]['size'], stream)) - context.execute_async_v2(gpu_allocations, stream) + if _trt_major_version() >= 10: + context.execute_async_v3(stream) + else: + context.execute_async_v2(gpu_allocations, stream) for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile b/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile index c1dc4deb6e69..0e86e79cd2d5 100644 --- a/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile +++ b/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile @@ -14,15 +14,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:23.05-py3 +# 26.06 ships TensorRT 11.0.0, CUDA 13.3 and Ubuntu 24.04 with Python 3.12. +# TensorRT 8.x cannot target Blackwell GPUs such as the RTX Pro 6000 that +# Dataflow now offers, so this image tracks the current TensorRT major version. +# The Python version here must match the Beam SDK image copied in below. +ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:26.06-py3 +ARG BEAM_SDK_IMAGE=apache/beam_python3.12_sdk:latest -FROM ${BUILD_IMAGE} +FROM ${BEAM_SDK_IMAGE} AS beam_sdk + +FROM ${BUILD_IMAGE} ENV PATH="/usr/src/tensorrt/bin:${PATH}" WORKDIR /workspace -COPY --from=apache/beam_python3.10_sdk:latest /opt/apache/beam /opt/apache/beam +COPY --from=beam_sdk /opt/apache/beam /opt/apache/beam RUN pip install --upgrade pip \ && pip install torch>=1.7.1 \ @@ -32,4 +39,4 @@ RUN pip install --upgrade pip \ && pip install cuda-python ENTRYPOINT [ "/opt/apache/beam/boot" ] -RUN apt-get update && apt-get install -y python3.10-venv +RUN apt-get update && apt-get install -y python3.12-venv diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index abe7867181ec..905f03f8f703 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -651,22 +651,23 @@ task mockAPITests { } // add all RunInference E2E tests that run on DataflowRunner -// As of now, this test suite is enable in py310 suite as the base NVIDIA image used for Tensor RT -// contains Python 3.10. // TODO: https://github.com/apache/beam/issues/22651 project.tasks.register("inferencePostCommitIT") { dependsOn = [ - // TODO(https://github.com/apache/beam/issues/33078): restore the tensorRT tests once the staged - // model is fixed. - // 'tensorRTtests', 'vertexAIInferenceTest', 'geminiInferenceTest', 'mockAPITests', ] } +// The base NVIDIA image used for TensorRT contains Python 3.12, so the TensorRT +// suite belongs here rather than in the py310 suite above. project.tasks.register("inferencePostCommitITPy312") { dependsOn = [ + // TODO(https://github.com/apache/beam/issues/33078): restore the tensorRT tests once the staged + // model is rebuilt. A serialized engine can only be read by the TensorRT major version that + // built it, and the staged engine was built with TensorRT 8.x. + // 'tensorRTtests', 'vllmTests', ] } diff --git a/website/www/site/content/en/documentation/ml/tensorrt-runinference.md b/website/www/site/content/en/documentation/ml/tensorrt-runinference.md index 4bae2d3ba7ce..78bdd95b30fa 100644 --- a/website/www/site/content/en/documentation/ml/tensorrt-runinference.md +++ b/website/www/site/content/en/documentation/ml/tensorrt-runinference.md @@ -66,7 +66,7 @@ trtexec --onnx= --saveEngine= To use `trtexec`, follow the steps in the blog post [Simplifying and Accelerating Machine Learning Predictions in Apache Beam with NVIDIA TensorRT](https://developer.nvidia.com/blog/simplifying-and-accelerating-machine-learning-predictions-in-apache-beam-with-nvidia-tensorrt/). The post explains how to build a docker image from a DockerFile that can be used for conversion. We use the following Docker file, which is similar to the file used in the blog post: ``` -ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:22.05-py3 +ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:26.06-py3 FROM ${BUILD_IMAGE} @@ -75,11 +75,11 @@ ENV PATH="/usr/src/tensorrt/bin:${PATH}" WORKDIR /workspace RUN apt-get update -y && apt-get install -y python3-venv -RUN pip install --no-cache-dir apache-beam[gcp]==2.44.0 -COPY --from=apache/beam_python3.8_sdk:2.44.0 /opt/apache/beam /opt/apache/beam +RUN pip install --no-cache-dir apache-beam[gcp]==2.76.0 +COPY --from=apache/beam_python3.12_sdk:2.76.0 /opt/apache/beam /opt/apache/beam RUN pip install --upgrade pip \ - && pip install torch==1.13.1 \ + && pip install torch \ && pip install torchvision>=0.8.2 \ && pip install pillow>=8.0.0 \ && pip install transformers>=4.18.0 \ From 6072108b03b4d65e63d616720775383117c671e8 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Mon, 31 Aug 2026 12:10:25 -0700 Subject: [PATCH 3/5] Pass host buffers to CUDA by address in TensorRT inference _default_tensorRT_inference_fn passed numpy arrays directly to cuMemcpyHtoDAsync and cuMemcpyDtoHAsync. An array holding exactly one element is coerced to a scalar rather than being handled through the buffer protocol, so the value is read as a null host pointer and the copy fails: RuntimeError: Cuda Error: Pass the buffer address explicitly instead. This is not specific to a TensorRT or cuda-python version; it reproduces on cuda-python 12.9 and 13.3 alike, and depends only on an input or output tensor having a single element. Single element tensors are common. The ssd_mobilenet_v2_320x320_coco17_tpu-8 model behind the tensorRTtests integration test has a num_detections output of shape (1, 1), so that test cannot pass without this fix. The existing unit tests did not catch it because every tensor in their models holds four elements. The two copies of this loop in the test file are updated to match, so they do not keep demonstrating the broken pattern. --- CHANGES.md | 1 + .../apache_beam/ml/inference/tensorrt_inference.py | 12 ++++++++++-- .../ml/inference/tensorrt_inference_test.py | 10 ++++++---- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index e41b5b59d9e2..e680716ee5b7 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -95,6 +95,7 @@ * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). * (Python) `TensorRTEngineHandlerNumPy` now works with TensorRT 10 and later, which removed the binding API it was written against. TensorRT 8.x remains supported, so no existing GPU loses support ([#36306](https://github.com/apache/beam/issues/36306)). +* (Python) Fixed `TensorRTEngineHandlerNumPy` failing with `CUDA_ERROR_INVALID_VALUE` on models with a single-element input or output tensor, such as the `num_detections` output of an object detection model ([#36306](https://github.com/apache/beam/issues/36306)). ## Security Fixes diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py index 214847070944..b76f4312ab25 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py +++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py @@ -262,10 +262,18 @@ def _default_tensorRT_inference_fn( # Process I/O and execute the network with context_lock: + # Host buffers are passed as explicit addresses rather than as arrays. + # A numpy array holding exactly one element is coerced to a scalar, which + # is then read as a null host pointer and fails with CUDA_ERROR_INVALID_ + # VALUE. Single element outputs are common, for example the num_detections + # output of an object detection model. + # host_input must stay referenced until the stream is synchronized below, + # because the copy is asynchronous. + host_input = np.ascontiguousarray(batch) _assign_or_fail( cuda.cuMemcpyHtoDAsync( inputs[0]['allocation'], - np.ascontiguousarray(batch), + host_input.ctypes.data, inputs[0]['size'], stream)) if _trt_major_version() >= 10: @@ -276,7 +284,7 @@ def _default_tensorRT_inference_fn( for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( - cpu_allocations[output], + cpu_allocations[output].ctypes.data, outputs[output]['allocation'], outputs[output]['size'], stream)) diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py index 2cca02e704a1..65e493bb602e 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py @@ -108,10 +108,11 @@ def _custom_tensorRT_inference_fn(batch, engine, inference_args): # Process I/O and execute the network with context_lock: + host_input = np.ascontiguousarray(batch) _assign_or_fail( cuda.cuMemcpyHtoDAsync( inputs[0]['allocation'], - np.ascontiguousarray(batch), + host_input.ctypes.data, inputs[0]['size'], stream)) if _trt_major_version() >= 10: @@ -121,7 +122,7 @@ def _custom_tensorRT_inference_fn(batch, engine, inference_args): for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( - cpu_allocations[output], + cpu_allocations[output].ctypes.data, outputs[output]['allocation'], outputs[output]['size'], stream)) @@ -406,10 +407,11 @@ def fake_inference_fn(batch, engine, inference_args=None): # Process I/O and execute the network with context_lock: + host_input = np.ascontiguousarray(batch) _assign_or_fail( cuda.cuMemcpyHtoDAsync( inputs[0]['allocation'], - np.ascontiguousarray(batch), + host_input.ctypes.data, inputs[0]['size'], stream)) if _trt_major_version() >= 10: @@ -419,7 +421,7 @@ def fake_inference_fn(batch, engine, inference_args=None): for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( - cpu_allocations[output], + cpu_allocations[output].ctypes.data, outputs[output]['allocation'], outputs[output]['size'], stream)) From d71238be550e32e58786a4cd09ba27aa151e1dd3 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Mon, 31 Aug 2026 12:10:32 -0700 Subject: [PATCH 4/5] Add a script to rebuild the staged TensorRT test engines A serialized TensorRT engine can only be deserialized by the same TensorRT major version and GPU architecture that built it, so the engines the tests load from gs://apache-beam-ml/models/ have to be rebuilt whenever the TensorRT version in tensor_rt.dockerfile changes, or the tensorRTtests task moves to a different GPU. Until now that was undocumented manual work, which is part of why the staged engines went stale. build_test_engines.py rebuilds each of the three from the ONNX source already staged beside it, so no new model sources are needed, and verifies the result by loading it back through TensorRTEngineHandlerNumPy: the two small engines against the exact values the unit tests assert, and the object detection engine against the same COCO images the integration test uses. Engines are only uploaded once verification passes. The script needs a GPU, so it cannot run as part of the test suite. README.md covers how to run it and when it needs running. --- .../tensorrt_runinference/README.md | 72 ++++- .../build_test_engines.py | 255 ++++++++++++++++++ 2 files changed, 325 insertions(+), 2 deletions(-) create mode 100644 sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/README.md b/sdks/python/test-suites/containers/tensorrt_runinference/README.md index 99fbf83cbd74..fcacf736b81f 100644 --- a/sdks/python/test-suites/containers/tensorrt_runinference/README.md +++ b/sdks/python/test-suites/containers/tensorrt_runinference/README.md @@ -17,8 +17,76 @@ under the License. --> -# TensorRT Dockerfile for Beam +# TensorRT test resources for Beam -This directory contains the Dockerfiles required to run Beam pipelines that use TensorRT. +This directory contains the Dockerfile required to run Beam pipelines that use TensorRT, +and the script that rebuilds the TensorRT engines those tests load from GCS. + +## Container image To build the image, run `docker build -f tensor_rt.dockerfile -t us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest .` + +## Rebuilding the test engines + +The TensorRT tests load pre-built engines from `gs://apache-beam-ml/models/`: + +| Engine | Used by | +| --- | --- | +| `single_tensor_features_engine.trt` | `tensorrt_inference_test.py` | +| `multiple_tensor_features_engine.trt` | `tensorrt_inference_test.py` | +| `ssd_mobilenet_v2_320x320_coco17_tpu-8.trt` | the `tensorRTtests` Dataflow integration test | + +**A serialized TensorRT engine is not portable.** It can only be deserialized by the +same TensorRT major version and the same GPU architecture that built it. So these files +must be rebuilt whenever either of the following changes: + +* the TensorRT version in `tensor_rt.dockerfile`, or +* the GPU that the `tensorRTtests` task requests in + `sdks/python/test-suites/dataflow/common.gradle`. + +`build_test_engines.py` does that. It rebuilds each engine from the ONNX source already +staged next to it in the same bucket, so no new model sources are needed, and it verifies +each result by loading it back through `TensorRTEngineHandlerNumPy` — the small engines +against the exact values the unit tests assert, and the object detection engine against +the same COCO images the integration test uses. + +It needs a GPU, so it cannot run as part of the test suite. Run it in the same container +and on the same GPU type the tests use. As of writing that is +`nvcr.io/nvidia/tensorrt:26.06-py3` on an `nvidia-tesla-t4`. + +The host needs a driver new enough for that container (580 or later), plus Docker and +the NVIDIA container toolkit. On a GCE deep learning VM image the toolkit is already +present but Docker may not be: + +``` +sudo apt-get install -y docker.io +sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker +``` + +Then, from a directory containing `build_test_engines.py`: + +``` +sudo docker run --rm --gpus all -v "$PWD:/w" -w /w nvcr.io/nvidia/tensorrt:26.06-py3 bash -c "\ + pip install -q --break-system-packages 'apache-beam[gcp]' cuda-python pillow && \ + python3 build_test_engines.py --dest gs://YOUR_BUCKET/models" +``` + +`--break-system-packages` is required because the container's Python environment is +marked externally managed. Credentials are picked up from the VM's service account, so +no extra authentication step is needed. + +The verification step imports the model handler from the installed `apache-beam`, so +that version has to support the TensorRT major version you are building for. To verify +against an unreleased change, copy your working tree's +`apache_beam/ml/inference/tensorrt_inference.py` over the installed one inside the +container before running the script. + +Engines are written with a `_trt` suffix, for example +`single_tensor_features_engine_trt11.trt`, so the engines built by earlier TensorRT +versions stay in place and anyone on an older branch is unaffected. + +Point `--dest` at a bucket you can write to. Staging the results under +`gs://apache-beam-ml/models/` is a separate, deliberate step for someone with write +access to that bucket; note it has no object versioning, so an overwrite cannot be undone. + +Pass `--only ` to rebuild a single engine, and `--help` for the full options. diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py new file mode 100644 index 000000000000..deee092848dd --- /dev/null +++ b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py @@ -0,0 +1,255 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Rebuilds the TensorRT engines that the TensorRT tests load from GCS. + +A serialized TensorRT engine can only be deserialized by the TensorRT major +version and GPU architecture that built it. The engines the tests load are +therefore not portable, and have to be rebuilt whenever the TensorRT version +in the test container changes or the tests move to a different GPU. + +This script rebuilds them from the ONNX sources already staged alongside them, +and verifies each result by loading it back through the model handler the +tests use. It cannot run as part of the test suite because it needs a GPU. + +Run it in the same container and on the same GPU type the tests use. See +README.md in this directory for the exact commands. +""" + +# pytype: skip-file + +import argparse +import io +import logging +import os +import sys +import tempfile + +import numpy as np +import tensorrt as trt + +TRT_LOGGER = trt.Logger(trt.Logger.INFO) +TRT_MAJOR = int(trt.__version__.split('.')[0]) + +SOURCE = 'gs://apache-beam-ml/models' +COCO_IMAGES = [ + 'gs://apache-beam-ml/datasets/coco/raw-data/val2017/000000289594.jpg', + 'gs://apache-beam-ml/datasets/coco/raw-data/val2017/000000000139.jpg', +] + + +def _copy(src, dst): + """Copies between any two paths Beam's FileSystems understands. + + The TensorRT container has no gcloud CLI, but apache-beam[gcp] is installed + for the verification step anyway, so reuse it rather than shelling out. + """ + from apache_beam.io.filesystems import FileSystems + logging.info('copy %s -> %s', src, dst) + with FileSystems.open(src) as fin, FileSystems.create(dst) as fout: + while True: + chunk = fin.read(8 << 20) + if not chunk: + break + fout.write(chunk) + + +def _network_creation_flags(): + """Explicit batch is only a flag on TensorRT 8.x; it is the default after.""" + explicit_batch = getattr( + trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH', None) + if explicit_batch is None or TRT_MAJOR >= 10: + return 0 + return 1 << int(explicit_batch) + + +def build_engine(onnx_path, engine_path): + """Parses an ONNX file and serializes an engine for this GPU.""" + # The SSD MobileNet ONNX contains an EfficientNMS_TRT node, so the bundled + # plugins have to be registered before the parser will accept it. + trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") + + builder = trt.Builder(TRT_LOGGER) + network = builder.create_network(flags=_network_creation_flags()) + parser = trt.OnnxParser(network, TRT_LOGGER) + with open(onnx_path, 'rb') as f: + if not parser.parse(f.read()): + for i in range(parser.num_errors): + logging.error(parser.get_error(i)) + raise ValueError(f'Failed to parse {onnx_path}') + + config = builder.create_builder_config() + plan = builder.build_serialized_network(network, config) + if plan is None: + raise RuntimeError(f'Engine build produced no plan for {onnx_path}') + with open(engine_path, 'wb') as f: + f.write(plan) + logging.info('built %s (%d bytes)', engine_path, os.path.getsize(engine_path)) + + +def _handler(engine_path, batch_size): + from apache_beam.ml.inference.tensorrt_inference import ( + TensorRTEngineHandlerNumPy) + return TensorRTEngineHandlerNumPy( + min_batch_size=batch_size, + max_batch_size=batch_size, + engine_path=engine_path) + + +def verify_linear(engine_path, examples, expected): + """Checks a small linear engine against the values the unit tests assert.""" + handler = _handler(engine_path, len(examples)) + results = handler.run_inference(list(examples), handler.load_model()) + actual = np.array([r.inference[0] for r in results]).reshape(-1) + if not np.allclose(actual, np.asarray(expected).reshape(-1), atol=1e-4): + raise AssertionError(f'{engine_path}: expected {expected}, got {actual}') + logging.info('verified %s -> %s', os.path.basename(engine_path), actual) + + +def verify_ssd(engine_path): + """Runs the object detection engine on the images the Dataflow IT uses. + + The outputs are checked for the shape and ordering the example's + PostProcessor indexes by, and for at least one confident detection. + """ + from apache_beam.io.filesystems import FileSystems + from PIL import Image + + handler = _handler(engine_path, 1) + engine = handler.load_model() + + for image_path in COCO_IMAGES: + with FileSystems.open(image_path) as f: + image = Image.open(io.BytesIO(f.read())).convert('RGB') + # Mirrors preprocess_image() in the tensorrt_object_detection example. + image = image.resize((300, 300), resample=Image.Resampling.BILINEAR) + batch = [np.expand_dims(np.asarray(image, dtype=np.float32), axis=0)] + + inference = list(handler.run_inference(batch, engine))[0].inference + if len(inference) != 4: + raise AssertionError( + f'{engine_path}: expected 4 outputs, got {len(inference)}') + _, boxes, scores, classes = inference + if boxes.shape[-1] != 4 or scores.shape != classes.shape: + raise AssertionError( + f'{engine_path}: unexpected output shapes; the engine tensor order ' + f'must be num_detections, boxes, scores, classes. Got ' + f'{[np.asarray(o).shape for o in inference]}') + if float(np.max(scores)) < 0.3: + raise AssertionError( + f'{engine_path}: no confident detection for {image_path}; top score ' + f'was {float(np.max(scores)):.3f}') + logging.info( + 'verified %s on %s -> top score %.2f', + os.path.basename(engine_path), + os.path.basename(image_path), + float(np.max(scores))) + + +# The inputs and outputs below mirror the constants in +# apache_beam/ml/inference/tensorrt_inference_test.py, so a rebuilt engine is +# checked against exactly what the tests will assert. +SINGLE_EXAMPLES = [np.float32(v) for v in (1, 5, -3, 10)] +SINGLE_EXPECTED = [2.5, 10.5, -5.5, 20.5] # y = 2x + 0.5 + +MULTI_EXAMPLES = np.array([[1, 5], [3, 10], [-14, 0], [0.5, 0.5]], + dtype=np.float32) +MULTI_EXPECTED = [17.5, 36.5, -27.5, 3.0] # y = 2*x0 + 3*x1 + 0.5 + + +def verify_single(engine_path): + verify_linear(engine_path, SINGLE_EXAMPLES, SINGLE_EXPECTED) + + +def verify_multiple(engine_path): + verify_linear(engine_path, MULTI_EXAMPLES, MULTI_EXPECTED) + + +# Each entry rebuilds one staged .trt file from its staged .onnx source. +ENGINES = { + 'single_tensor_features_engine': { + 'onnx': 'single_tensor_features_model.onnx', + 'verify': verify_single, + }, + 'multiple_tensor_features_engine': { + 'onnx': 'multiple_tensor_features_model.onnx', + 'verify': verify_multiple, + }, + 'ssd_mobilenet_v2_320x320_coco17_tpu-8': { + 'onnx': 'ssd_mobilenet_v2_320x320_coco17_tpu-8.onnx', + 'verify': verify_ssd, + }, +} + + +def main(argv=None): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + '--dest', + required=True, + help='Where to write the rebuilt engines, e.g. gs://my-bucket/models or ' + 'a local directory. Staging them under the shared bucket is a separate, ' + 'deliberate step for someone with write access.') + parser.add_argument( + '--suffix', + default=None, + help='Name suffix for the rebuilt engines, so the engines built by ' + 'earlier TensorRT versions can stay in place. Defaults to _trt.') + parser.add_argument( + '--only', + action='append', + choices=sorted(ENGINES), + help='Rebuild only the named engine. May be repeated. Defaults to all.') + args = parser.parse_args(argv) + + if args.dest.rstrip('/') == SOURCE: + parser.error( + f'Refusing to write to {SOURCE}. That bucket has no object ' + 'versioning, so overwriting a staged engine could not be undone.') + + suffix = args.suffix if args.suffix is not None else f'_trt{TRT_MAJOR}' + names = args.only or sorted(ENGINES) + logging.info( + 'TensorRT %s, suffix %r, building: %s', + trt.__version__, + suffix, + ', '.join(names)) + + written = [] + with tempfile.TemporaryDirectory() as tmp: + for name in names: + spec = ENGINES[name] + onnx_local = os.path.join(tmp, spec['onnx']) + _copy(f'{SOURCE}/{spec["onnx"]}', onnx_local) + + engine_local = os.path.join(tmp, f'{name}{suffix}.trt') + build_engine(onnx_local, engine_local) + spec['verify'](engine_local) + + dest = f'{args.dest.rstrip("/")}/{os.path.basename(engine_local)}' + _copy(engine_local, dest) + written.append(dest) + + print('\nRebuilt and verified with TensorRT %s:' % trt.__version__) + for dest in written: + print(f' {dest}') + + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO, format='%(levelname)s %(message)s') + sys.exit(main()) From cfb60daaa26604a8470007c3c2ff5ac0c1c83149 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Mon, 31 Aug 2026 17:22:58 -0700 Subject: [PATCH 5/5] Address review feedback on TensorRT version support - Cache _trt_major_version() and _import_cuda_driver() with functools.lru_cache. _import_cuda_driver() is called from _assign_or_fail(), so it ran on every CUDA call. Caching rather than resolving at import time keeps the module importable without TensorRT, so jobs can still be submitted from a machine that does not have it. - Pin torch in the TensorRT documentation rather than leaving it unpinned. - Trim the two CHANGES.md entries to the change itself. - Drop the TensorRT 8 branch from build_test_engines.py. The script only rebuilds engines for the container the tests currently use, so it can require TensorRT 10 or later. The equivalent branch in the handler stays, since the handler does support both. --- CHANGES.md | 4 ++-- .../ml/inference/tensorrt_inference.py | 10 ++++++++++ .../containers/tensorrt_runinference/README.md | 3 ++- .../tensorrt_runinference/build_test_engines.py | 17 ++++++----------- .../documentation/ml/tensorrt-runinference.md | 2 +- 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index e680716ee5b7..679721ce0032 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -94,8 +94,8 @@ * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). -* (Python) `TensorRTEngineHandlerNumPy` now works with TensorRT 10 and later, which removed the binding API it was written against. TensorRT 8.x remains supported, so no existing GPU loses support ([#36306](https://github.com/apache/beam/issues/36306)). -* (Python) Fixed `TensorRTEngineHandlerNumPy` failing with `CUDA_ERROR_INVALID_VALUE` on models with a single-element input or output tensor, such as the `num_detections` output of an object detection model ([#36306](https://github.com/apache/beam/issues/36306)). +* (Python) `TensorRTEngineHandlerNumPy` now works with TensorRT 10 and later ([#36306](https://github.com/apache/beam/issues/36306)). +* (Python) Fixed `TensorRTEngineHandlerNumPy` failing with `CUDA_ERROR_INVALID_VALUE` on models with a single-element input or output tensor ([#36306](https://github.com/apache/beam/issues/36306)). ## Security Fixes diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py index b76f4312ab25..6e41ef586869 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py +++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py @@ -19,6 +19,7 @@ from __future__ import annotations +import functools import logging import threading from collections.abc import Callable @@ -49,11 +50,16 @@ LOGGER.warning(msg) +@functools.lru_cache(maxsize=1) def _trt_major_version() -> int: """Returns the major version of the installed TensorRT. TensorRT 10 replaced the index based "binding" API with a name based "tensor" API, so the major version selects which code path to take. + + Cached rather than resolved at import time because the module is importable + without TensorRT, so that jobs can be submitted from a machine that does not + have it installed. """ import tensorrt as trt try: @@ -63,12 +69,16 @@ def _trt_major_version() -> int: return 10 if hasattr(trt.ICudaEngine, 'num_io_tensors') else 8 +@functools.lru_cache(maxsize=1) def _import_cuda_driver(): """Imports the CUDA driver bindings. ``cuda.bindings.driver`` is the module path used by cuda-python 12.8 and later. It replaced the ``cuda.cuda`` alias, which was removed in cuda-python 13.0, so only fall back to that for older installations. + + Cached because this is called from _assign_or_fail, which runs on every + CUDA call. """ try: from cuda.bindings import driver as cuda diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/README.md b/sdks/python/test-suites/containers/tensorrt_runinference/README.md index fcacf736b81f..35d805c54d9c 100644 --- a/sdks/python/test-suites/containers/tensorrt_runinference/README.md +++ b/sdks/python/test-suites/containers/tensorrt_runinference/README.md @@ -52,7 +52,8 @@ the same COCO images the integration test uses. It needs a GPU, so it cannot run as part of the test suite. Run it in the same container and on the same GPU type the tests use. As of writing that is -`nvcr.io/nvidia/tensorrt:26.06-py3` on an `nvidia-tesla-t4`. +`nvcr.io/nvidia/tensorrt:26.06-py3` on an `nvidia-tesla-t4`. It requires TensorRT 10 or +later, since the TensorRT 8 engines are the ones being replaced. The host needs a driver new enough for that container (580 or later), plus Docker and the NVIDIA container toolkit. On a GCE deep learning VM image the toolkit is already diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py index deee092848dd..606a74df3988 100644 --- a/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py +++ b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py @@ -27,7 +27,8 @@ tests use. It cannot run as part of the test suite because it needs a GPU. Run it in the same container and on the same GPU type the tests use. See -README.md in this directory for the exact commands. +README.md in this directory for the exact commands. Requires TensorRT 10 or +later; the engines built by TensorRT 8 are the ones being replaced. """ # pytype: skip-file @@ -68,15 +69,6 @@ def _copy(src, dst): fout.write(chunk) -def _network_creation_flags(): - """Explicit batch is only a flag on TensorRT 8.x; it is the default after.""" - explicit_batch = getattr( - trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH', None) - if explicit_batch is None or TRT_MAJOR >= 10: - return 0 - return 1 << int(explicit_batch) - - def build_engine(onnx_path, engine_path): """Parses an ONNX file and serializes an engine for this GPU.""" # The SSD MobileNet ONNX contains an EfficientNMS_TRT node, so the bundled @@ -84,7 +76,10 @@ def build_engine(onnx_path, engine_path): trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") builder = trt.Builder(TRT_LOGGER) - network = builder.create_network(flags=_network_creation_flags()) + # Explicit batch is the default from TensorRT 10 onwards, so no creation + # flags are needed. Engines are only ever rebuilt for the container the + # tests currently use, so there is no reason to support TensorRT 8 here. + network = builder.create_network() parser = trt.OnnxParser(network, TRT_LOGGER) with open(onnx_path, 'rb') as f: if not parser.parse(f.read()): diff --git a/website/www/site/content/en/documentation/ml/tensorrt-runinference.md b/website/www/site/content/en/documentation/ml/tensorrt-runinference.md index 78bdd95b30fa..49ea6e0e0881 100644 --- a/website/www/site/content/en/documentation/ml/tensorrt-runinference.md +++ b/website/www/site/content/en/documentation/ml/tensorrt-runinference.md @@ -79,7 +79,7 @@ RUN pip install --no-cache-dir apache-beam[gcp]==2.76.0 COPY --from=apache/beam_python3.12_sdk:2.76.0 /opt/apache/beam /opt/apache/beam RUN pip install --upgrade pip \ - && pip install torch \ + && pip install torch==2.13.0 \ && pip install torchvision>=0.8.2 \ && pip install pillow>=8.0.0 \ && pip install transformers>=4.18.0 \