Support TensorRT 10 and later in TensorRTEngineHandlerNumPy by akshayjadiyanv · Pull Request #39922 · apache/beam

akshayjadiyanv · 2026-08-30T23:39:07Z

Beam's TensorRT model handler is written against the TensorRT 8.x binding API, which TensorRT 10 removed. Any pipeline using TensorRTEngineHandlerNumPy on a modern TensorRT container fails immediately at engine load:

AttributeError: 'tensorrt.tensorrt.ICudaEngine' object has no attribute 'num_bindings'

Two users have reported this, and both reports were closed as stale without a fix: addresses #36306 and addresses #33946.

Approach: support both APIs, rather than upgrading

This PR selects the API at runtime from the TensorRT major version instead. _trt_major_version() is the single switch:

TensorRT 8.x keeps num_bindings / get_binding_* / execute_async_v2.
TensorRT 10 and later use num_io_tensors / get_tensor_* / set_tensor_address / execute_async_v3.

No currently supported GPU loses support, and Blackwell works for the first time.

Also included

cuda-python 13 compatibility. cuda.cuda was removed in favour of cuda.bindings.driver. _import_cuda_driver() handles both. The new container ships cuda-python 13, so without this the upgrade fails for a second, unrelated reason.
Test container bumped from nvcr.io/nvidia/tensorrt:23.05-py3 (TensorRT 8.6, May 2023) to 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.
Two guard tests so the version check cannot silently drift from the API it selects.
A second, pre-existing bug, in its own commit. Host buffers were passed to cuMemcpyHtoDAsync / cuMemcpyDtoHAsync as numpy arrays. An array holding exactly one element is coerced to a scalar and read as a null host pointer, so the copy fails with CUDA_ERROR_INVALID_VALUE. They are now passed by address. This is not version specific — it reproduces on cuda-python 12.9 and 13.3 alike — and it is why the tensorRTtests integration test could not pass even with a good engine: the SSD MobileNet model's num_detections output has shape (1, 1). The existing unit tests missed it because every tensor in their models holds four elements.
A script to rebuild the staged test engines, sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py, with instructions in the README next to it. A serialized engine is tied to both the TensorRT major version and the GPU architecture that built it, so these files have to be rebuilt whenever either changes. That was previously undocumented manual work, which is part of why the staged engines went stale.
A pre-existing test bug, in its own commit. test_namespace has asserted 'RunInferenceTensorRT' since the original 2022 TensorRT commit, but the handler returns 'BeamML_TensorRT', matching the BeamML_* prefix every other handler uses. It has been wrong for three years and nobody noticed, which is itself a sign of how long this area has gone unexercised.

Testing

Run on a real T4 GPU on GCE in both containers:

Container	TensorRT	Result
`nvcr.io/nvidia/tensorrt:26.06-py3`	11.0.0 (new path)	7 / 7 pass
`nvcr.io/nvidia/tensorrt:23.05-py3`	8.6.1 (old path)	7 / 7 pass

These are the tests that build engines in process. The tests that read pre-built engines from gs://apache-beam-ml/ were not run - see below.

What this PR does not fix

The Dataflow integration test (tensorRTtests) stays disabled in this PR, and #33078 stays open. But the missing piece now exists and only needs a committer to stage it.

A serialized TensorRT engine can only be deserialized by the major version that built it. All three .trt files under gs://apache-beam-ml/models/ were built with TensorRT 8.x in 2022, so none of them load on the new container regardless of this change. Staging replacements needs write access to that bucket, which I do not have.

I have rebuilt and verified all three, and can pass them on to whoever picks up #33078.

Each was rebuilt from the ONNX source already staged in gs://apache-beam-ml/models/ — no new model sources are needed — and verified by the script above, which loads each engine back through the TensorRTEngineHandlerNumPy in this PR:

The two small engines return the exact values the unit tests assert.
The SSD MobileNet engine was run against both COCO images in gs://apache-beam-ml/testing/inputs/tensorrt_image_file_names.txt and returns confident detections, top score 0.77 and 0.72 respectively. Its I/O tensor order is num_detections, detection_boxes, detection_scores, detection_classes, which is exactly what PostProcessor indexes, so no example code change is required. Note this engine only runs through the handler at all because of the single-element buffer fix above.

Build environment: TensorRT 11.0.0.114 in nvcr.io/nvidia/tensorrt:26.06-py3 on a Tesla T4 (compute capability 7.5) matching the T4 the integration test requests.

I am happy to hand these over.

github-actions · 2026-08-31T17:14:14Z

Assigning reviewers:

R: @tvalentyn for label python.
R: @kennknowles for label website.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

test_namespace has asserted 'RunInferenceTensorRT' since the original TensorRT commit (a8ca305). The handler was later changed to return 'BeamML_TensorRT' in f477b85, 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.

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 apache#36306 Addresses apache#33946

_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.

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.

tvalentyn · 2026-08-31T22:23:47Z

cc: @jrmccluskey

tvalentyn · 2026-08-31T22:25:31Z

@akshayjadiyanv Thanks for the fix. Are you a Beam TensorRT user?

@jrmccluskey @akshayjadiyanv Should we bump the lower bound on TensorRT and not introduce different codepaths?

tvalentyn · 2026-08-31T22:39:59Z

  LOGGER.warning(msg)


+def _trt_major_version() -> int:


Can we evaluate this once?

we could either do it once per module or add @functools.lru_cache(maxsize=1) here and for _import_cuda_driver

Agreed. I have added @functools.lru_cache(maxsize=1) to both rather than making them module-level constants.

tvalentyn · 2026-08-31T22:51:28Z


 RUN pip install --upgrade pip \
-    && pip install torch==1.13.1 \
+    && pip install torch \


can we pin a newer version?

I have pinned it to torch==2.13.0, the current stable release which supports CUDA 13 and python 3.10+

tvalentyn · 2026-08-31T22:52:13Z

 * (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)).


, which removed the binding API it was written against. TensorRT 8.x remains supported, so no existing GPU loses support

we can remove this

Agreed, trimmed.

tvalentyn · 2026-08-31T22:53:06Z

 * (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)).


, such as the num_detections output of an object detection model

too specific, remove this too?

Agreed, removed.

tvalentyn · 2026-08-31T22:58:11Z

+  """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:


do we need this branching in test code?

No good point. I have dropped the branch and it now calls create_network() directly.

- 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.

jrmccluskey · 2026-09-01T14:39:58Z

@akshayjadiyanv Thanks for the fix. Are you a Beam TensorRT user?

@jrmccluskey @akshayjadiyanv Should we bump the lower bound on TensorRT and not introduce different codepaths?

From a maintainability standpoint I would rather just bump the lower bounds, but I do not know if we would wind up with users hard-stuck on the last supported Beam version for 8.X and unwilling to upgrade.

tvalentyn · 2026-09-01T15:35:42Z

yeah. I am open to bump the bounds as well. the version 8 is pretty old. I am fine say in the release notes (CHANGES.MD) that this is a breaking change and ask users to comment on an issue X if backwards incompatibility with version 8 is a hard blocker for them.

akshayjadiyanv · 2026-09-01T23:12:43Z

@akshayjadiyanv Thanks for the fix. Are you a Beam TensorRT user?

@jrmccluskey @akshayjadiyanv Should we bump the lower bound on TensorRT and not introduce different codepaths?

@damccorm do you have any preference here?

github-actions Bot added python website examples labels Aug 30, 2026

akshayjadiyanv marked this pull request as ready for review August 31, 2026 16:41

github-actions Bot added the Next Action: Reviewers label Aug 31, 2026

akshayjadiyanv added 4 commits August 31, 2026 12:10

akshayjadiyanv force-pushed the tensorrt-version-compat branch from 050fd7e to d71238b Compare August 31, 2026 19:11

tvalentyn reviewed Aug 31, 2026

View reviewed changes

Uh oh!

Conversation

akshayjadiyanv commented Aug 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Approach: support both APIs, rather than upgrading

Also included

Testing

What this PR does not fix

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

tvalentyn commented Aug 31, 2026

Uh oh!

tvalentyn commented Aug 31, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tvalentyn Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jrmccluskey commented Sep 1, 2026

Uh oh!

tvalentyn commented Sep 1, 2026

Uh oh!

akshayjadiyanv commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

akshayjadiyanv commented Aug 30, 2026 •

edited

Loading

tvalentyn Aug 31, 2026 •

edited

Loading