Re-enable the TensorRT integration test - #40139
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
|
Is there post commit trigger file we should trivially modify to kick-off this test in the CI? |
|
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). |
| 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" |
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'.
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
There was a problem hiding this comment.
The default driver did not work, I have moved to install-nvidia-driver:latest. There was another failure in that run, leader_board_it_test, is unrelated to this PR
There was a problem hiding this comment.
@tvalentyn install-nvidia-driver:latest did it. Py 3.12 post-commit passed too, Thanks for staging the engines! Other red check seems to be unrelated to this PR.
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.
The default driver is too old for the TensorRT 11 container, which is built on
CUDA 13.3. The post-commit run failed before it ever reached the engine:
[TRT] [E] createInferRuntime: Error Code 6: API Usage Error
(CUDA initialization failure with error: 35)
Python (worker sdk-0-0) exited 4 times: signal: segmentation fault
CUDA error 35 is CUDA_ERROR_INSUFFICIENT_DRIVER. Dataflow accepts only
'default' and 'latest' for this option, so 'latest' is the remaining choice.
|
LGTM, thanks for the help. |
2 participants
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