Re-enable the TensorRT integration test - #40139
Open
akshayjadiyanv wants to merge 2 commits into
Open
akshayjadiyanv wants to merge 2 commits into
akshayjadiyanv wants to merge 2 commits into
Conversation
The staged engines have been rebuilt with TensorRT 11 and uploaded alongside the 2022 originals, so the tests can be pointed at them and the Dataflow integration test can come back. - Point the engine paths at the _trt11 objects. The ONNX sources are unchanged, since ONNX is not version locked. - Request install-nvidia-driver:5xx rather than an unversioned driver. The test container is now TensorRT 11 on CUDA 13.3, which needs a 580 or newer driver, and would otherwise fail at CUDA init rather than at engine load. This matches what the vLLM integration tests already request. - Re-enable tensorRTtests in the py312 post-commit suite. Fixes apache#33078
Contributor
|
Is there post commit trigger file we should trivially modify to kick-off this test in the CI? |
Contributor
|
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). |
Contributor
tvalentyn
reviewed
Sep 16, 2026
| exec { | ||
| executable 'sh' | ||
| args '-c', ". ${envdir}/bin/activate && pip install pillow && python -m apache_beam.examples.inference.tensorrt_object_detection $cmdArgs --experiment='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver' --experiment=no_use_multiple_sdk_containers" | ||
| args '-c', ". ${envdir}/bin/activate && pip install pillow && python -m apache_beam.examples.inference.tensorrt_object_detection $cmdArgs --experiment='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver:5xx' --experiment=no_use_multiple_sdk_containers" |
Contributor
There was a problem hiding this comment.
We don't use 5xx anymore.
https://docs.cloud.google.com/dataflow/docs/gpu/use-gpus#drivers
I would try the default driver without any mods, and if not working, then use 'latest'.
Contributor
Author
There was a problem hiding this comment.
Good catch - thanks! And yes, beam_PostCommit_Python.json looks to be the right trigger file. Made both changes
Dataflow only accepts 'default' and 'latest' for install-nvidia-driver, and an unrecognized value makes the driver install fail, so the 5xx request this PR introduced would have broken the test rather than fixed it. Reverted to the unversioned default per review; 'latest' is the fallback if the default driver turns out to be too old for TensorRT 11 on CUDA 13.3. Also bump .github/trigger_files/beam_PostCommit_Python.json so the Python post-commit runs against this PR. That suite reaches tensorRTtests through python312PostCommit -> inferencePostCommitITPy312, which is the only way to exercise the rebuilt engines before merging.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-enables the TensorRT Dataflow integration test, which has been disabled since #33078.
Follow-up to #39922, which raised the handler's minimum to TensorRT 10 and added the script that rebuilds the staged engines. @tvalentyn has now staged the rebuilt engines and pushed the updated test container image, so the last two pieces are in place.
Fixes #33078
What changed
Engine paths. The staged
.trtobjects are now the TensorRT 11 rebuilds, uploaded alongside the 2022 originals rather than over them:Six references in
tensorrt_inference_test.pyand one incommon.gradlenow point at them. The.onnxsources are untouched — ONNX is not version locked, only serialized engines are.Driver request.
tensorRTtestsasked forinstall-nvidia-driverwith no version. The test container is now TensorRT 11 on CUDA 13.3, which needs a 580 or newer driver, so without this the job would fail at CUDA init rather than at engine load. This matches what the vLLM integration tests already request on the neighbouring lines.Re-enabled the task in the py312 post-commit suite. It moved to py312 in #39922 because the new base image is Python 3.12.
How the engines were produced
With
sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py, added in #39922. It rebuilds each engine from the ONNX source already staged beside it, and verifies the result by loading it back throughTensorRTEngineHandlerNumPy— the two small engines against the exact values the unit tests assert, and the object detection engine against the same COCO images this integration test uses. Nothing is uploaded until verification passes.On my own run, on a Tesla T4 with TensorRT 11.0.0.114 in
nvcr.io/nvidia/tensorrt:26.06-py3:single_tensor_features_engine→[2.5, 10.5, -5.5, 20.5]multiple_tensor_features_engine→[17.5, 36.5, -27.5, 3.0]ssd_mobilenet_v2_320x320_coco17_tpu-8→ top scores 0.77 and 0.72 on the two COCO images