Support TensorRT 10 and later in TensorRTEngineHandlerNumPy - #39922
Conversation
|
Assigning reviewers: R: @tvalentyn for label python. Note: If you would like to opt out of this review, comment Available commands:
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.
050fd7e to
d71238b
Compare
|
cc: @jrmccluskey |
|
@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? |
| LOGGER.warning(msg) | ||
|
|
||
|
|
||
| def _trt_major_version() -> int: |
There was a problem hiding this comment.
Can we evaluate this once?
There was a problem hiding this comment.
we could either do it once per module or add @functools.lru_cache(maxsize=1) here and for _import_cuda_driver
There was a problem hiding this comment.
Agreed. I have added @functools.lru_cache(maxsize=1) to both rather than making them module-level constants.
|
|
||
| RUN pip install --upgrade pip \ | ||
| && pip install torch==1.13.1 \ | ||
| && pip install torch \ |
There was a problem hiding this comment.
can we pin a newer version?
There was a problem hiding this comment.
I have pinned it to torch==2.13.0, the current stable release which supports CUDA 13 and python 3.10+
| * (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)). |
There was a problem hiding this comment.
, which removed the binding API it was written against. TensorRT 8.x remains supported, so no existing GPU loses support
we can remove this
There was a problem hiding this comment.
Agreed, trimmed.
| * (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)). |
There was a problem hiding this comment.
, such as the
num_detectionsoutput of an object detection model
too specific, remove this too?
There was a problem hiding this comment.
Agreed, removed.
| """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: |
There was a problem hiding this comment.
do we need this branching in test code?
There was a problem hiding this comment.
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.
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. |
|
yeah. I am open to bump the bounds as well. the version 8 is pretty old. I am fine to 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. |
@damccorm do you have any preference here? |
|
pinging @damccorm here. |
|
R: @tvalentyn |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
|
Sorry, I missed this - I am +1 to just bumping the minimum version. If users get stuck on the old version of Beam for this, I think that is ok |
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 apache#36306.
Resolves a CHANGES.md conflict and picks up the updated pin for gradle/actions/setup-gradle. The branch still referenced the older SHA, which the Apache org action allowlist now rejects, so every Gradle based precommit failed at "Setup environment" before running any of this PR's code.
_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.
|
Thanks @tvalentyn @jrmccluskey @damccorm. Switched to bumping the minimum to TensorRT 10 per the discussion above |
6 participants
Beam's TensorRT model handler is written against the TensorRT 8.x binding API, which TensorRT 10 removed. Any pipeline using
TensorRTEngineHandlerNumPyon a modern TensorRT container fails immediately at engine load:Two users have reported this, and both reports were closed as stale without a fix: addresses #36306 and addresses #33946.
Approach: raise the minimum to TensorRT 10, per review discussion
This PR originally kept both APIs alive behind a runtime version check, so that no existing user lost support. Reviewers asked whether we should just bump the lower bound instead, and agreed that we should:
So the dual code path is gone.
TensorRTEngineHandlerNumPynow requires TensorRT 10 or later:num_io_tensors/get_tensor_*/set_tensor_address.execute_async_v3._trt_major_version()and_network_creation_flags()and every branch they fed are removed, along with the numpy monkey patch that only existed for TensorRT 8.x._check_trt_version()replaces them. It raises a clear error naming the installed version, rather than letting an unsupported TensorRT surface as anAttributeErrordeep inside engine setup. It is lazy and cached, so the module stays importable without TensorRT and jobs can still be submitted from a machine that does not have it.Per @tvalentyn's suggestion,
CHANGES.mdrecords this under Breaking Changes and asks anyone hard blocked by the drop of 8.x to comment on #36306.Also included
cuda.cudawas removed in favour ofcuda.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.cuMemcpyHtoDAsync/cuMemcpyDtoHAsyncas numpy arrays. An array holding exactly one element is coerced to a scalar and read as a null host pointer, so the copy fails withCUDA_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 thetensorRTtestsintegration test could not pass even with a good engine: the SSD MobileNet model'snum_detectionsoutput has shape(1, 1). The existing unit tests missed it because every tensor in their models holds four elements.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. Added in response to review feedback asking for this to be reproducible rather than hand-built.nvcr.io/nvidia/tensorrt:23.05-py3(TensorRT 8.6, May 2023) to26.06-py3(TensorRT 11.0, CUDA 13.3, Python 3.12). Because that image is Python 3.12, the disabledtensorRTteststask moves from the py310 suite to the py312 suite.test_namespacehas asserted'RunInferenceTensorRT'since the original 2022 TensorRT commit, but the handler returns'BeamML_TensorRT', matching theBeamML_*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.What users need to know
Testing
Verified on a real T4 GPU on GCE, in
nvcr.io/nvidia/tensorrt:26.06-py3(TensorRT 11.0.0.114, CUDA 13.3, Python 3.12, cuda-python 13.3.1):9 passed, 6 deselected. The 6 deselected are the ones that read pre-built engines from GCS, covered below.The tests that read the pre-built engines currently staged in
gs://apache-beam-ml/were not run, because those engines are TensorRT 8.x artifacts — 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.All three
.trtfiles undergs://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 by the script above from the ONNX source already staged in
gs://apache-beam-ml/models/— no new model sources are needed — and verified by loading it back through theTensorRTEngineHandlerNumPyin this PR:gs://apache-beam-ml/testing/inputs/tensorrt_image_file_names.txtand returns confident detections, top score 0.77 and 0.72 respectively. Its I/O tensor order isnum_detections, detection_boxes, detection_scores, detection_classes, which is exactly whatPostProcessorindexes, so no example code change is required.Build environment: TensorRT 11.0.0.114 in
nvcr.io/nvidia/tensorrt:26.06-py3on a Tesla T4 (compute capability 7.5), matching the T4 the integration test requests.One note for whoever re-enables that test: it requests
install-nvidia-driverwith no version, while the vLLM integration tests requestinstall-nvidia-driver:5xx. TensorRT 11 with CUDA 13.3 needs a 580+ driver, so without that it would fail at CUDA init rather than at engine load.