From 3ce03ab9e6418da3183b3ce8bab5a43d8b538e78 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Sun, 30 Aug 2026 16:37:27 -0700 Subject: [PATCH 1/2] 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 050fd7ebf0459c2e04b84f359546f27a76d95589 Mon Sep 17 00:00:00 2001 From: akshayjadiyanv Date: Sun, 30 Aug 2026 16:37:39 -0700 Subject: [PATCH 2/2] 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 c3c4019c7da2..02323598424c 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -92,6 +92,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 \