This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new bdf650a75b5 Support TensorRT 10 and later in
TensorRTEngineHandlerNumPy (#39922)
bdf650a75b5 is described below
commit bdf650a75b54ab76c01025622fb31f91420359d2
Author: akshayjadiyanv <[email protected]>
AuthorDate: Tue Sep 15 13:55:51 2026 -0700
Support TensorRT 10 and later in TensorRTEngineHandlerNumPy (#39922)
* 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.
* 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
* 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: <CUresult.CUDA_ERROR_INVALID_VALUE: 1>
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.
* 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.
* 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.
* Require TensorRT 10 or later instead of supporting both APIs
Per review discussion, bump the lower bound rather than carrying two code
paths. TensorRT 8.x is from 2023 and keeping both APIs alive makes the
handler
harder to maintain for a shrinking set of users.
Removes _trt_major_version() and _network_creation_flags() and the branches
they fed, so engine setup always uses the name based tensor API and
execution
always uses execute_async_v3. The numpy monkey patch that only existed for
TensorRT 8.x goes with them.
_check_trt_version() replaces them. It raises a clear error naming the
installed version, rather than letting an unsupported TensorRT surface as an
AttributeError deep inside engine setup. It stays lazy and cached so the
module remains importable without TensorRT, which is what lets jobs be
submitted from a machine that does not have it.
CHANGES.md records this as a breaking change and asks anyone hard blocked by
it to comment on #36306.
* Check the TensorRT version before deserializing an engine
_load_engine() deserialized first and only reached _check_trt_version() when
TensorRTEngine was constructed, so loading a pre-built engine on an
unsupported TensorRT failed with an opaque deserialization error instead of
the message naming the installed version. Loading a pre-built engine is the
most common path, and the one users hitting the new lower bound will take.
* Update sdks/python/apache_beam/examples/inference/README.md
* Update sdks/python/apache_beam/examples/inference/README.md
---------
Co-authored-by: tvalentyn <[email protected]>
---
CHANGES.md | 5 +
.../apache_beam/examples/inference/README.md | 10 +-
.../apache_beam/ml/inference/tensorrt_inference.py | 100 +++++++--
.../ml/inference/tensorrt_inference_test.py | 67 +++---
.../containers/tensorrt_runinference/README.md | 73 +++++-
.../tensorrt_runinference/build_test_engines.py | 250 +++++++++++++++++++++
.../tensorrt_runinference/tensor_rt.dockerfile | 15 +-
sdks/python/test-suites/dataflow/common.gradle | 11 +-
.../en/documentation/ml/tensorrt-runinference.md | 8 +-
9 files changed, 470 insertions(+), 69 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index fc7d62ece33..18351eebd1a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -85,6 +85,10 @@
* Portable Java SDK now encodes SchemaCoders in a portable way
([#34672](https://github.com/apache/beam/issues/34672)).
- Original custom Java coder encoding can still be obtained using
[StreamingOptions.setUpdateCompatibilityVersion("2.76")](https://github.com/apache/beam/blob/2cf0930e7ae1aa389c26ce6639b584877a3e31d9/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java#L47)
([#34672](https://github.com/apache/beam/issues/34672)).
- Fixes ([#36496](https://github.com/apache/beam/issues/36496)),
([#30276](https://github.com/apache/beam/issues/30276)),
([#29245](https://github.com/apache/beam/issues/29245)).
+* (Python) `TensorRTEngineHandlerNumPy` now requires TensorRT 10 or later.
TensorRT 8.x is no longer supported, since TensorRT 10 removed the engine
binding API the handler was written against
([#36306](https://github.com/apache/beam/issues/36306)).
+ - Engines serialized by TensorRT 8.x must be rebuilt, as an engine can only
be deserialized by the major version that built it.
+ - TensorRT 10 and later require a GPU with compute capability 7.5 or higher,
which excludes NVIDIA Pascal and Volta GPUs.
+ - If dropping TensorRT 8.x support is a hard blocker for you, please comment
on ([#36306](https://github.com/apache/beam/issues/36306)).
## Deprecations
@@ -99,6 +103,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)).
* (Go) Fixed GCS glob matching silently dropping objects when the glob pattern
contains multi-byte characters
([#39969](https://github.com/apache/beam/issues/39969)).
+* (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/examples/inference/README.md
b/sdks/python/apache_beam/examples/inference/README.md
index 5eed659d068..c43ad458347 100644
--- a/sdks/python/apache_beam/examples/inference/README.md
+++ b/sdks/python/apache_beam/examples/inference/README.md
@@ -83,13 +83,19 @@ pip install torch==1.10.0
### TensorRT dependencies
The RunInference API supports TensorRT SDK for high-performance deep learning
inference with NVIDIA GPUs.
-To use TensorRT locally, we suggest an environment with TensorRT >= 8.0.1.
Install TensorRT as per the
+To use TensorRT locally, we suggest an environment with TensorRT >= 10.0.
Install TensorRT as per the
[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.
+TensorRT 10 or later is required. 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. TensorRT 10 and later also require a GPU
with
+compute capability 7.5 or higher, for example, T4, L4, A100. The NVIDIA Pascal
and Volta GPUs
+such as the Tesla P4, P100 and V100 are no longer supported.
+
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 333187301b2..b6089788eae 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
@@ -48,9 +49,58 @@ except ModuleNotFoundError:
'runner has tensorrt dependencies installed.'
LOGGER.warning(msg)
+MIN_TRT_MAJOR_VERSION = 10
+
+
[email protected]_cache(maxsize=1)
+def _check_trt_version() -> None:
+ """Fails fast if the installed TensorRT is older than we support.
+
+ TensorRT 10 removed the index based "binding" API this module used to be
+ written against. Without this check the failure surfaces as an obscure
+ AttributeError deep inside engine setup.
+
+ Cached rather than checked 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:
+ major = int(trt.__version__.split('.')[0])
+ except (AttributeError, IndexError, ValueError):
+ # Fall back to probing for an attribute that only exists from 10 onwards.
+ major = 10 if hasattr(trt.ICudaEngine, 'num_io_tensors') else 8
+ if major < MIN_TRT_MAJOR_VERSION:
+ raise RuntimeError(
+ 'RunInference requires TensorRT %d or later, but found %s. Support '
+ 'for TensorRT 8.x was removed because TensorRT 10 replaced the '
+ 'engine binding API this handler depends on.' %
+ (MIN_TRT_MAJOR_VERSION, getattr(trt, '__version__', 'unknown')))
+
+
[email protected]_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
+ except ImportError:
+ from cuda import cuda
+ return cuda
+
def _load_engine(engine_path):
import tensorrt as trt
+ # Checked before deserializing, because an engine built by a newer TensorRT
+ # fails to deserialize with an opaque error that hides the real cause.
+ _check_trt_version()
file = FileSystems.open(engine_path, 'rb')
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(file.read())
@@ -60,9 +110,11 @@ def _load_engine(engine_path):
def _load_onnx(onnx_path):
import tensorrt as trt
+ _check_trt_version()
builder = trt.Builder(TRT_LOGGER)
- network = builder.create_network(
- flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+ # Explicit batch is the only supported mode from TensorRT 10 onwards, so no
+ # network creation flags are needed.
+ network = builder.create_network()
parser = trt.OnnxParser(network, TRT_LOGGER)
with FileSystems.open(onnx_path) as f:
if not parser.parse(f.read()):
@@ -85,7 +137,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 +163,8 @@ class TensorRTEngine:
engine: trt.ICudaEngine object that contains TensorRT engine
"""
import tensorrt as trt
- from cuda import cuda
+ _check_trt_version()
+ cuda = _import_cuda_driver()
self.engine = engine
self.context = engine.create_execution_context()
self.context_lock = threading.RLock()
@@ -120,19 +173,12 @@ class TensorRTEngine:
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)
+ # Setup I/O tensors. 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 = {
@@ -144,7 +190,8 @@ class TensorRTEngine:
'size': size
}
self.gpu_allocations.append(allocation)
- if self.engine.binding_is_input(i):
+ 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)
@@ -182,7 +229,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,
@@ -195,17 +242,26 @@ 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))
- context.execute_async_v2(gpu_allocations, stream)
+ # Tensor addresses were bound when the engine was created.
+ context.execute_async_v3(stream)
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 80a01b8f4d4..48931eee399 100644
--- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
@@ -19,6 +19,7 @@
import os
import unittest
+from unittest import mock
import numpy as np
import pytest
@@ -37,6 +38,9 @@ try:
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 _check_trt_version
+ from apache_beam.ml.inference.tensorrt_inference import _import_cuda_driver
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,
@@ -119,17 +108,18 @@ 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))
- context.execute_async_v2(gpu_allocations, stream)
+ context.execute_async_v3(stream)
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))
@@ -189,8 +179,7 @@ class TensorRTRunInferenceTest(unittest.TestCase):
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()
input_tensor = network.add_input(
name="input", dtype=trt.float32, shape=(4, 1))
weight_const = network.add_constant(
@@ -227,8 +216,7 @@ class TensorRTRunInferenceTest(unittest.TestCase):
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()
input_tensor = network.add_input(
name="input", dtype=trt.float32, shape=(4, 1))
weight_const = network.add_constant(
@@ -263,8 +251,7 @@ class TensorRTRunInferenceTest(unittest.TestCase):
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()
input_tensor = network.add_input(
name="input", dtype=trt.float32, shape=(4, 2))
weight_const = network.add_constant(
@@ -349,7 +336,26 @@ class TensorRTRunInferenceTest(unittest.TestCase):
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())
+
+ def test_supported_tensorrt_exposes_expected_api(self):
+ """The installed TensorRT must expose the API this module is written to.
+
+ TensorRT 10 removed the index based binding API in favour of the name
+ based tensor API. _check_trt_version() rejects anything older, so a passing
+ version check and a missing API would mean the two have drifted apart.
+ """
+ _check_trt_version()
+ self.assertTrue(hasattr(trt.ICudaEngine, 'num_io_tensors'))
+ self.assertTrue(hasattr(trt.IExecutionContext, 'execute_async_v3'))
+
+ def test_version_check_rejects_unsupported_tensorrt(self):
+ """An unsupported TensorRT must fail with a clear message."""
+ with mock.patch.object(trt, '__version__', '8.6.1'):
+ _check_trt_version.cache_clear()
+ with self.assertRaisesRegex(RuntimeError, 'requires TensorRT 10'):
+ _check_trt_version()
+ _check_trt_version.cache_clear()
@pytest.mark.uses_tensorrt
@@ -381,7 +387,7 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
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,
@@ -394,17 +400,18 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
# 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))
- context.execute_async_v2(gpu_allocations, stream)
+ context.execute_async_v3(stream)
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/test-suites/containers/tensorrt_runinference/README.md
b/sdks/python/test-suites/containers/tensorrt_runinference/README.md
index 99fbf83cbd7..35d805c54d9 100644
--- a/sdks/python/test-suites/containers/tensorrt_runinference/README.md
+++ b/sdks/python/test-suites/containers/tensorrt_runinference/README.md
@@ -17,8 +17,77 @@
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`. 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
+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<major>` 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 <name>` 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 00000000000..606a74df398
--- /dev/null
+++
b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py
@@ -0,0 +1,250 @@
+#
+# 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. Requires TensorRT 10 or
+later; the engines built by TensorRT 8 are the ones being replaced.
+"""
+
+# 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 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)
+ # 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()):
+ 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<major>.')
+ 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())
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 c1dc4deb6e6..0e86e79cd2d 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 7b64ace131e..4a21391d2fd 100644
--- a/sdks/python/test-suites/dataflow/common.gradle
+++ b/sdks/python/test-suites/dataflow/common.gradle
@@ -659,22 +659,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 4bae2d3ba7c..49ea6e0e088 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=<path to onnx model> --saveEngine=<path to
save TensorRT engine>
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==2.13.0 \
&& pip install torchvision>=0.8.2 \
&& pip install pillow>=8.0.0 \
&& pip install transformers>=4.18.0 \
