Skip to content

[Failing Test]: Fix and re-enable the TensorRT integration test #33078

Closed
#40139
Closed
[Failing Test]: Fix and re-enable the TensorRT integration test#33078
#40139
Assignees
akshayjadiyanv
Labels
P2bugfailing testpermaredpythontests
Milestone
2.78.0 Release

Description

@jrmccluskey
jrmccluskey
opened on Nov 11, 2024

What happened?

The TensorRT integration test went red when the workflow was moved from a base container with Python 3.8 to Python 3.10 as part of the 3.8 support deprecation. The problem is that the model engine staged at gs://apache-beam-ml/models/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt (based on a TF Model Garden config) was built with TensorRT 8.x, and Python 3.10 containers use TensorRT 10.x. Unfortunately the documentation around loading the model from the TF side is somewhat out of date or not necessarily what we need; additionally, we need to convert the model from a TF format to a TensorRT format since we do not use the ONNX route in the test.

The gradle task for the test is :sdks:python:test-suites:dataflow:py310:tensorRTtests and is defined here:

beam/sdks/python/test-suites/dataflow/common.gradle

Line 444 in 2488ca1

def tensorRTTests = tasks.create("tensorRTtests") {
.

When testing the workflow ensure that you're running on dataflow or on a machine with a GPU, as the workflow will fail with CUDA error 35 if there isn't a GPU present.

Issue Failure

Failure: Test is continually failing

Issue Priority

Priority: 2 (backlog / disabled test but we think the product is healthy)

Issue Components

  • Component: Python SDK
  • Component: Java SDK
  • Component: Go SDK
  • Component: Typescript SDK
  • Component: IO connector
  • Component: Beam YAML
  • Component: Beam examples
  • Component: Beam playground
  • Component: Beam katas
  • Component: Website
  • Component: Infrastructure
  • Component: Spark Runner
  • Component: Flink Runner
  • Component: Samza Runner
  • Component: Twister2 Runner
  • Component: Hazelcast Jet Runner
  • Component: Google Cloud Dataflow Runner

Activity

  1. jrmccluskey
    added
    bug
    failing test
    awaiting triage
    on Nov 11, 2024
  2. github-actions
    added
    python
    tests
    P2
    permared
    on Nov 11, 2024

  3. github-actions commented on May 29, 2025

    @github-actions
    github-actionsBot
    on May 29, 2025 – with GitHub Actions
    Contributor

    This issue has been marked as stale due to 150 days of inactivity. It will be closed in 30 days if no further activity occurs. If you think that’s incorrect or this issue still needs to be addressed, please simply write any comment. If closed, you can reopen the issue at any time. Thank you for your contributions.

  4. github-actions
    added
    stale
    on May 29, 2025

  5. github-actions commented on Jun 29, 2025

    @github-actions
    github-actionsBot
    on Jun 29, 2025 – with GitHub Actions
    Contributor

    This issue has been closed due to lack of activity. If you think that is incorrect, you can reopen the issue at any time.

  6. github-actions
    closed this as not plannedon Jun 29, 2025
  7. tvalentyn
    reopened this on Sep 15, 2026

  8. tvalentyn commented on Sep 15, 2026

    @tvalentyn
    tvalentyn
    on Sep 15, 2026
    Contributor

    @akshayjadiyanv would it be possible for you to take a look at this one since you have the most context on this right now? I am happy to run necessary commands to rebuild images you need in the repo for these tests. Feel free to post here what's needed. Thanks!

  9. akshayjadiyanv commented on Sep 15, 2026

    @akshayjadiyanv
    akshayjadiyanv
    on Sep 15, 2026
    Contributor

    @tvalentyn
    here is the full state and what is left. Two of the steps need write access I do not have, and two are a PR I can send.

    A serialized TRT engine can only be deserialized by the same TensorRT major version and GPU architecture that built it. All three .trt files under gs://apache-beam-ml/models/ were built with TensorRT 8.x in 2022:

    1. single_tensor_features_engine.trt
    2. multiple_tensor_features_engine.trt
    3. ssd_mobilenet_v2_320x320_coco17_tpu-8.trt

    Since #39922 the handler requires TensorRT 10+ , none of these load. Rebuilding them is the blocker.

    Already done in #39922 :

    sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py rebuilds all three from the ONNX sources already staged in the same bucket — nothing new needs uploading as a source. It verifies each result by loading it back through TensorRTEngineHandlerNumPy, and only uploads after verification passes. README.md next to it has the commands.

    I have run it on a Tesla T4 with TensorRT 11.0.0.114 in nvcr.io/nvidia/tensorrt:26.06-py3. All three build and verify:

    • single_tensor_features_engine → [2.5, 10.5, -5.5, 20.5], exactly what the unit test asserts
    • multiple_tensor_features_engine → [17.5, 36.5, -27.5, 3.0], likewise
    • ssd_mobilenet_v2_320x320_coco17_tpu-8 → confident detections on both COCO images from tensorrt_image_file_names.txt, top scores 0.77 and 0.72
    1. You can either run the script yourself on a T4:
    sudo apt-get install -y docker.io
sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker

sudo docker run --rm --gpus all -v "$PWD:/w" -w /w nvcr.io/nvidia/tensorrt:26.06-py3 bash -c "\
    pip install -q --break-system-packages 'apache-beam[gcp]' cuda-python pillow && \
    python3 build_test_engines.py --dest gs://SOME_BUCKET_YOU_CAN_WRITE/models"

    …or I can hand over the files I already built and verified, whichever you prefer.

    The script deliberately refuses --dest gs://apache-beam-ml/models, since that bucket has no object versioning and an overwrite could not be undone.

    One thing to decide : The script names its output with a _trt11 suffix, so the 2022 files survive. Can keep the same if we are okay with overwrite.

    1. Rebuild and push the test container image

    us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest, from the updated tensor_rt.dockerfile on master. It now uses nvcr.io/nvidia/tensorrt:26.06-py3 (TensorRT 11.0, CUDA 13.3) and apache/beam_python3.12_sdk. The old pushed image is still TensorRT 8.6 / Python 3.10, so the IT would fail on the image alone.

    What I will send as a PR once those land

    3. Bump the driver request. common.gradle:454 asks for install-nvidia-driver with no version, while the vLLM ITs on 486/490/503 ask for install-nvidia-driver:5xx. TRT 11 with CUDA 13.3 needs a 580+ driver - so without this it fails at CUDA init rather than at engine load.

    4. Re-enable the task. tensorRTtests is commented out in inferencePostCommitITPy312. It already moved to the py312 suite in #39922 ,since the new base image is Python 3.12.

    Plus the path updates from step 1 if we go with the suffixed names.

    Not a blocker, just so it is not a surprise

    1. The IT runs the example pipeline and passes if it completes - it does not diff against a golden output, so nothing needs regenerating.
    2. Also, num_detections is always 100 because the ONNX pins EfficientNMS keepTopK=100; that was equally true of the 2022 engine, so it is not a regression.

  10. tvalentyn commented on Sep 15, 2026

    @tvalentyn
    tvalentyn
    on Sep 15, 2026
    Contributor

    Created:

    gs://apache-beam-ml/models/multiple_tensor_features_engine_trt11.trt
    gs://apache-beam-ml/models/single_tensor_features_engine_trt11.trt
    gs://apache-beam-ml/models/ssd_mobilenet_v2_320x320_coco17_tpu-8_trt11.trt

    we can keep old ones for now

  11. github-actions
    removed
    stale
    on Sep 16, 2026

  12. tvalentyn commented on Sep 16, 2026

    @tvalentyn
    tvalentyn
    on Sep 16, 2026
    Contributor

    Pushed us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest

  13. tvalentyn commented on Sep 16, 2026

    @tvalentyn
    tvalentyn
    on Sep 16, 2026
    Contributor

    @akshayjadiyanv lmk if you need any other assistance and thanks for helping fix these once again.

  14. tvalentyn
    assigned
    akshayjadiyanv
    on Sep 16, 2026
  15. github-actions
    removed
    awaiting triage
    on Sep 16, 2026

  16. akshayjadiyanv commented on Sep 16, 2026

    @akshayjadiyanv
    akshayjadiyanv
    on Sep 16, 2026
    Contributor

    Thanks @tvalentyn , I opened a new PR #40139

  17. tvalentyn
    closed this as completedin #40139on Sep 18, 2026
  18. tvalentyn
    added a commit that references this issue on Sep 18, 2026

    Re-enable the TensorRT integration test (#40139)

    94482e8
  19. github-actions
    added this to the 2.78.0 Release milestone on Sep 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

P2bugfailing testpermaredpythontests

Type

No type

Projects

No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions