akshayjadiyanv opened a new pull request, #39922: URL: https://github.com/apache/beam/pull/39922
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 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 loading it 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 plausible detections (`person`, `person`, `car` at 0.77 / 0.61 / 0.57 on `000000289594.jpg`). 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. 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. -- 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]
