akshayjadiyanv opened a new pull request, #40139: URL: https://github.com/apache/beam/pull/40139
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 `.trt` objects are now the TensorRT 11 rebuilds, uploaded alongside the 2022 originals rather than over them: ``` gs://apache-beam-ml/models/single_tensor_features_engine_trt11.trt gs://apache-beam-ml/models/multiple_tensor_features_engine_trt11.trt gs://apache-beam-ml/models/ssd_mobilenet_v2_320x320_coco17_tpu-8_trt11.trt ``` Six references in `tensorrt_inference_test.py` and one in `common.gradle` now point at them. The `.onnx` sources are untouched — ONNX is not version locked, only serialized engines are. **Driver request.** `tensorRTtests` asked for `install-nvidia-driver` with 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 through `TensorRTEngineHandlerNumPy` — 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 ### One thing to watch on the first post-commit run A serialized engine is tied to the **GPU architecture** it was built on, as well as the TensorRT version. This test requests an `nvidia-tesla-t4` (compute capability 7.5), so the staged engines need to have been built on a T4. @tvalentyn — if the post-commit fails at engine deserialization rather than later in the pipeline, that is the likely cause, and a rebuild on a T4 box would fix it. Everything else about the engines checks out. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
