This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 94482e8da03 Re-enable the TensorRT integration test (#40139)
94482e8da03 is described below
commit 94482e8da035c47d842e50583a21968c64197742
Author: akshayjadiyanv <[email protected]>
AuthorDate: Fri Sep 18 10:44:58 2026 -0700
Re-enable the TensorRT integration test (#40139)
* Re-enable the TensorRT integration test
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 #33078
* Keep the default NVIDIA driver and trigger the Python post-commit
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.
* Request the latest NVIDIA driver for the TensorRT integration test
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.
---
.github/trigger_files/beam_PostCommit_Python.json | 4 ++--
.../apache_beam/ml/inference/tensorrt_inference_test.py | 13 +++++++------
.../test-suites/containers/tensorrt_runinference/README.md | 9 ++++++---
sdks/python/test-suites/dataflow/common.gradle | 9 +++------
4 files changed, 18 insertions(+), 17 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python.json
b/.github/trigger_files/beam_PostCommit_Python.json
index e8079b053ae..9bfa92ff510 100644
--- a/.github/trigger_files/beam_PostCommit_Python.json
+++ b/.github/trigger_files/beam_PostCommit_Python.json
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to
run.",
- "pr": "38701",
- "modification": 57
+ "pr": "40139",
+ "modification": 59
}
diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
index 48931eee399..23b516f502a 100644
--- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
@@ -291,7 +291,7 @@ class TensorRTRunInferenceTest(unittest.TestCase):
min_batch_size=4,
max_batch_size=4,
engine_path=
- 'gs://apache-beam-ml/models/single_tensor_features_engine.trt')
+ 'gs://apache-beam-ml/models/single_tensor_features_engine_trt11.trt')
engine = inference_runner.load_model()
predictions = inference_runner.run_inference(
SINGLE_FEATURE_EXAMPLES, engine)
@@ -314,7 +314,7 @@ class TensorRTRunInferenceTest(unittest.TestCase):
min_batch_size=4,
max_batch_size=4,
engine_path=
- 'gs://apache-beam-ml/models/multiple_tensor_features_engine.trt')
+ 'gs://apache-beam-ml/models/multiple_tensor_features_engine_trt11.trt')
engine = inference_runner.load_model()
predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, engine)
for actual, expected in zip(predictions, TWO_FEATURES_PREDICTIONS):
@@ -367,7 +367,7 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
min_batch_size=4,
max_batch_size=4,
engine_path=
- 'gs://apache-beam-ml/models/single_tensor_features_engine.trt')
+ 'gs://apache-beam-ml/models/single_tensor_features_engine_trt11.trt')
pcoll = pipeline | 'start' >> beam.Create(
SINGLE_FEATURE_EXAMPLES, reshuffle=False)
predictions = pcoll | RunInference(engine_handler)
@@ -428,7 +428,7 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
min_batch_size=4,
max_batch_size=4,
engine_path=
- 'gs://apache-beam-ml/models/single_tensor_features_engine.trt',
+ 'gs://apache-beam-ml/models/single_tensor_features_engine_trt11.trt',
inference_fn=fake_inference_fn,
large_model=True)
pcoll = pipeline | 'start' >> beam.Create(
@@ -447,7 +447,7 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
min_batch_size=4,
max_batch_size=4,
engine_path=
- 'gs://apache-beam-ml/models/single_tensor_features_engine.trt')
+ 'gs://apache-beam-ml/models/single_tensor_features_engine_trt11.trt')
os.environ.pop('FOO', None)
self.assertFalse('FOO' in os.environ)
_ = (
@@ -465,7 +465,8 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
min_batch_size=4,
max_batch_size=4,
engine_path=
- 'gs://apache-beam-ml/models/multiple_tensor_features_engine.trt')
+
'gs://apache-beam-ml/models/multiple_tensor_features_engine_trt11.trt'
+ )
pcoll = pipeline | 'start' >> beam.Create(
TWO_FEATURES_EXAMPLES, reshuffle=False)
predictions = pcoll | RunInference(engine_handler)
diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/README.md
b/sdks/python/test-suites/containers/tensorrt_runinference/README.md
index 35d805c54d9..66f52fda84f 100644
--- a/sdks/python/test-suites/containers/tensorrt_runinference/README.md
+++ b/sdks/python/test-suites/containers/tensorrt_runinference/README.md
@@ -32,9 +32,12 @@ The TensorRT tests load pre-built engines from
`gs://apache-beam-ml/models/`:
| Engine | Used by |
| --- | --- |
-| `single_tensor_features_engine.trt` | `tensorrt_inference_test.py` |
-| `multiple_tensor_features_engine.trt` | `tensorrt_inference_test.py` |
-| `ssd_mobilenet_v2_320x320_coco17_tpu-8.trt` | the `tensorRTtests` Dataflow
integration test |
+| `single_tensor_features_engine_trt11.trt` | `tensorrt_inference_test.py` |
+| `multiple_tensor_features_engine_trt11.trt` | `tensorrt_inference_test.py` |
+| `ssd_mobilenet_v2_320x320_coco17_tpu-8_trt11.trt` | the `tensorRTtests`
Dataflow integration test |
+
+The unsuffixed objects next to them are the original TensorRT 8.x engines,
kept for reference.
+They are not loadable by any currently supported Beam version.
**A serialized TensorRT engine is not portable.** It can only be deserialized
by the
same TensorRT major version and the same GPU architecture that built it. So
these files
diff --git a/sdks/python/test-suites/dataflow/common.gradle
b/sdks/python/test-suites/dataflow/common.gradle
index 4a21391d2fd..7e6919d18ab 100644
--- a/sdks/python/test-suites/dataflow/common.gradle
+++ b/sdks/python/test-suites/dataflow/common.gradle
@@ -453,13 +453,13 @@ def tensorRTTests = tasks.create("tensorRTtests") {
"region": "us-central1",
"input":
"gs://apache-beam-ml/testing/inputs/tensorrt_image_file_names.txt",
"output": "gs://apache-beam-ml/outputs/tensorrt_predictions.txt",
- "engine_path":
"gs://apache-beam-ml/models/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt",
+ "engine_path":
"gs://apache-beam-ml/models/ssd_mobilenet_v2_320x320_coco17_tpu-8_trt11.trt",
"disk_size_gb": 75
]
def cmdArgs = mapToArgString(argMap)
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:latest'
--experiment=no_use_multiple_sdk_containers"
}
}
}
@@ -672,10 +672,7 @@ project.tasks.register("inferencePostCommitIT") {
// suite belongs here rather than in the py310 suite above.
project.tasks.register("inferencePostCommitITPy312") {
dependsOn = [
- // TODO(https://github.com/apache/beam/issues/33078): restore the tensorRT
tests once the staged
- // model is rebuilt. A serialized engine can only be read by the TensorRT
major version that
- // built it, and the staged engine was built with TensorRT 8.x.
- // 'tensorRTtests',
+ 'tensorRTtests',
'vllmTests',
]
}
