You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "SaiShashank12 (via GitHub)" <gi...@apache.org> on 2025/09/29 04:15:11 UTC
[PR] Fix TensorRT compatibility with version 10.x [beam]
SaiShashank12 opened a new pull request, #36309:
URL: https://github.com/apache/beam/pull/36309
Perfect — then you should explicitly reference it so the PR auto-closes it. Here’s the finalized PR description with your issue number plugged in:
---
**Description**
This PR introduces a new `TensorRTEngineHandlerNumPy` implementation that supports TensorRT **10.x** with the Tensor API (`num_io_tensors`, `get_tensor_name`, `execute_async_v3`, etc).
It resolves runtime errors such as:
* `AttributeError: 'ICudaEngine' object has no attribute 'num_bindings'`
* `RuntimeError: This handler requires TensorRT 10.x+. Missing attributes on ICudaEngine: [...]`
by replacing the legacy TensorRT < 10 API calls with the new Tensor API.
It also fixes shape validation during dynamic batching and ONNX parsing (`onnx_path provided but build_on_worker=False`) to allow building engines directly on Beam workers if desired.
**Addresses:** fixes #36306
---
**Checklist**
* [x] Mentioned the relevant issue (`fixes #36306`) so it will auto-close once merged.
* [x] Added code that checks for TensorRT 10.x compatibility (`_require_tensorrt_10`) before engine load.
* [x] Replaced `num_bindings` and `execute_async_v2` with `num_io_tensors` and `execute_async_v3`.
* [x] Validated with dummy ResNet-50 ONNX pipeline on DirectRunner.
* [x] Validated with dummy ResNet-50 ONNX pipeline on DataRunner.
* [ ] Update `CHANGES.md` with a new entry under **Python SDK → ML Inference**: “Added TensorRT 10.x handler (`TensorRTEngineHandlerNumPy`).”
* [ ] Write down the test cases.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3442830167
Reminder, please take a look at this pr: @jrmccluskey
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3345054358
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3490894630
Reminder, please take a look at this pr: @shunping
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3461212582
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @shunping for label python.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "SaiShashank12 (via GitHub)" <gi...@apache.org>.
SaiShashank12 commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3393873357
Hi @jrmccluskey,
Thanks for the feedback! I've addressed all the formatting, linting, and whitespace issues in my latest commits. The precommit checks should now be passing:
- Fixed ASF license header formatting
- Corrected lines exceeding 80 characters
- Fixed import ordering and separated imports properly
- Added proper type annotations for mypy
- Cleaned up unused imports and variables
I apologize if the initial submission wasn't quite up to standard—this is my first contribution to an Apache project, and I'm still getting familiar with the code style and precommit requirements. I really appreciate your patience and guidance!
Let me know if there's anything else that needs adjustment.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3511964547
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment `assign set of reviewers`
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "SaiShashank12 (via GitHub)" <gi...@apache.org>.
SaiShashank12 commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3413027813
Great questions @jrmccluskey ! Let me explain the architectural reasoning behind the separate implementation and address your concerns about naming and infrastructure.
**Note:** I actually raised these same architectural concerns on the [dev mailing list](https://lists.apache.org/list.html?dev@beam.apache.org) before implementing this approach - specifically about whether to update the existing handler or create separate files. The feedback I received was to bring this discussion to the PR review, so I'm glad you're raising these points here!
## Why Separate Files? The Core Technical Difference
The separation isn't primarily about the API version change (v2 → v3), but rather a **fundamental workflow improvement** that solves a critical production issue with TensorRT.
### The Environment Alignment Problem
TensorRT engine files must be compiled on the **exact same hardware, CUDA version, and TensorRT version** where inference will run. This is [documented by NVIDIA](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#compatibility-serialized-engines) as a hard requirement. Pre-building engines on a dev machine and deploying to Dataflow workers almost always causes compatibility failures.
### How the Old Handler Works
```python
# tensorrt_inference.py - Manual workflow
handler = TensorRTEngineHandlerNumPy(...)
network, builder = handler.load_onnx() # User must call this
engine = handler.build_engine(network, builder) # Then call this
# Notably: load_onnx() is only used in test files, not production
```
### The New Handler's Key Innovation
```python
# trt_handler_numpy_compact.py - Automatic on-worker building
handler = TensorRTEngineHandlerNumPy(
onnx_path='gs://bucket/model.onnx', # Portable ONNX
build_on_worker=True # Builds IN the Dataflow worker
)
engine = handler.load_model() # Engine built automatically on correct hardware!
```
This eliminates an entire class of deployment failures:
- ✅ ONNX models are fully portable across environments
- ✅ Engine built with the exact GPU/CUDA/TensorRT version used for inference
- ✅ No "works on my machine" issues when deploying to Dataflow
- ✅ Production-ready workflow (vs test-only `load_onnx()`)
### Example Scenario This Solves
```
Developer machine: Dataflow worker:
- GPU: RTX 4090 - GPU: Tesla T4
- CUDA: 12.8 - CUDA: 12.6
- TensorRT: 10.7 - TensorRT: 10.5
Pre-built engine → FAILS with cryptic errors ❌
ONNX + on-worker build → Works perfectly ✅
```
## Addressing Your Specific Concerns
### 1. Naming Inconsistency (`trt_handler_numpy_compact.py`)
You're absolutely right - this doesn't follow the naming convention. I'm happy to rename to maintain consistency:
- Option A: `tensorrt_inference_v10.py` (clear version indicator)
- Option B: `tensorrt_inference_auto.py` (highlights auto-building)
- Preferred: **`tensorrt_inference_v10.py`** to match `tensorrt_inference.py`
### 2. Docker & README Duplication
Agreed - these can absolutely be consolidated into the existing `tensorrt_runinference/` directory. I can restructure as:
```
tensorrt_runinference/
├── README.md (updated with both versions)
├── legacy/
│ └── tensorrt_8x.dockerfile
└── v10/
└── tensorrt_10x.dockerfile
```
### 3. Code Reuse & Future Scalability
There's definitely shared logic (~60-70%) that could be refactored. This was actually one of my key concerns when I reached out to the dev list - whether to prioritize clean separation or code reuse. My concern with merging into the existing handler is:
- Risk of breaking existing users who rely on the manual `load_onnx()`/`build_engine()` pattern
- The API differences are significant enough that conditional logic might make the code harder to follow
## Proposed Refactoring Options
### Option A: Version-Aware Single Handler (Most maintainable)
```python
# tensorrt_inference.py
class TensorRTEngineHandlerNumPy:
def __init__(self, ..., tensorrt_version=None, build_on_worker=False):
if tensorrt_version is None:
tensorrt_version = self._detect_tensorrt_version()
if tensorrt_version >= 10:
self._use_tensor_api = True
else:
self._use_tensor_api = False
self._build_on_worker = build_on_worker
def load_model(self):
if self._build_on_worker and self.onnx_path:
# Automatic ONNX → engine conversion
return self._build_engine_from_onnx()
else:
# Legacy behavior
return self._load_prebuilt_engine()
```
### Option B: Base Class with Version Subclasses
```python
# tensorrt_inference_base.py - Shared logic
# tensorrt_inference_v9.py - Pre-10.x (legacy buffer API)
# tensorrt_inference_v10.py - 10.x+ (Tensor API + auto-building)
```
### Option C: Keep Separate, Fix Naming (Current + improvements)
- Rename to `tensorrt_inference_v10.py`
- Add deprecation notice to old handler encouraging 10.x adoption
- Consolidate Docker infrastructure
- Extract truly shared code (e.g., metrics, batching) to a utility module
## My Recommendation
I'd lean toward **Option A** if the on-worker building feature is valuable enough to backport to the old handler, or **Option C** if we want to keep clean separation for stability.
Since the dev list suggested bringing this architectural discussion here, I'm very open to your expertise on what fits best with Beam's patterns. **Question for you:** Does Beam typically prefer version-aware handlers (like PyTorch's different model loaders) or separate handlers per major version? That would help me align with the project's patterns.
## Summary
The architectural decision for separate files was driven by:
1. **Solving a real production problem** (environment compatibility via on-worker building)
2. **Enabling a new workflow** that wasn't production-ready in the old handler
3. **Avoiding breaking changes** to existing users
However, I completely agree with your concerns about:
- ❌ Naming inconsistency → Will fix
- ❌ Infrastructure duplication → Will consolidate
- ❌ Code duplication → Open to refactoring
I'm happy to implement whichever approach aligns best with Beam's architecture patterns. Let me know your preference and I'll refactor accordingly!
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3934171425
This pull request has been marked as stale due to 60 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@beam.apache.org list. Thank you for your contributions.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "jrmccluskey (via GitHub)" <gi...@apache.org>.
jrmccluskey commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3558850357
If the on-worker builds are that big of an improvement then I don't see why you wouldn't want to back-port them (which has the benefit of keeping things simple for end users, we would just need to document the new path *or* allow the old one if the params to the model handler aren't passed in the same way.)
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "jrmccluskey (via GitHub)" <gi...@apache.org>.
jrmccluskey commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3377083163
Please fix the formatting, linting, and whitespace precommits.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3511955579
R: @jrmccluskey
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "SaiShashank12 (via GitHub)" <gi...@apache.org>.
SaiShashank12 commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3344807843
### Python SDK
#### ML Inference
- Added support for **TensorRT 10.x** with the new `TensorRTEngineHandlerNumPy`.
This replaces legacy TensorRT < 10 APIs (e.g., `num_bindings`, `execute_async_v2`)
with the new Tensor API (`num_io_tensors`, `get_tensor_name`, `execute_async_v3`).
This resolves runtime errors when using TensorRT >= 10 in Beam pipelines.
(See [#36306](https://github.com/apache/beam/issues/36306))
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3368653998
Assigning reviewers:
R: @jrmccluskey for label python.
Note: If you would like to opt out of this review, comment `assign to next reviewer`.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
The PR bot will only process comments in the main thread (not review comments).
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "SaiShashank12 (via GitHub)" <gi...@apache.org>.
SaiShashank12 commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3368653546
assign set of reviewers
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3344794788
## Summary of Changes
Hello @SaiShashank12, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly upgrades Apache Beam's machine learning inference capabilities by introducing a dedicated handler for TensorRT 10.x. It modernizes the integration with TensorRT's latest Tensor API, resolving critical compatibility issues and enhancing the flexibility of deploying ONNX models with dynamic batching within Beam pipelines.
### Highlights
* **TensorRT 10.x Compatibility**: Introduces a new `TensorRTEngineHandlerNumPy` specifically designed to support TensorRT 10.x and its modern Tensor API.
* **API Modernization**: Replaces deprecated TensorRT < 10 API calls (e.g., `num_bindings`, `execute_async_v2`) with the new Tensor API methods (e.g., `num_io_tensors`, `get_tensor_name`, `execute_async_v3`, `set_input_shape`, `set_tensor_address`).
* **Error Resolution**: Fixes common runtime errors encountered with TensorRT 10.x, such as `AttributeError: 'ICudaEngine' object has no attribute 'num_bindings'` and `RuntimeError: This handler requires TensorRT 10.x+`.
* **Enhanced Dynamic Batching and ONNX Parsing**: Improves shape validation during dynamic batching and addresses issues with ONNX parsing, enabling the building of TensorRT engines directly on Beam workers if desired.
<details>
<summary><b>Using Gemini Code Assist</b></summary>
<br>
The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
<b>Invoking Gemini</b>
You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
Feature | Command | Description
--- | --- | ---
Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help | `/gemini help` | Displays a list of available commands.
<b>Customization</b>
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
<b>Limitations & Feedback</b>
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up [here](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
<b>You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the [Gemini Code Assist IDE Extension](https://cloud.google.com/products/gemini/code-assist).</b>
</details>
[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #36309:
URL: https://github.com/apache/beam/pull/36309#discussion_r2423302115
##########
sdks/python/apache_beam/ml/inference/trt_handler_numpy_compact_test.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# pytype: skip-file
+
+import os
+import tempfile
+import unittest
+
+import numpy as np
+import pytest
+
+import apache_beam as beam
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+
+# Protect against environments where TensorRT python library is not available.
+# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
+try:
+ from apache_beam.ml.inference.base import PredictionResult, RunInference
+ from apache_beam.ml.inference.trt_handler_numpy_compact import (
+ TensorRTEngine,
+ TensorRTEngineHandlerNumPy,
+ )
+except ImportError:
+ raise unittest.SkipTest('TensorRT 10.x dependencies are not installed')
+
+try:
+ from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+except ImportError:
+ GCSFileSystem = None # type: ignore
+
+# Test data
+SINGLE_FEATURE_EXAMPLES = np.array([[1.0], [5.0], [-3.0], [10.0]],
+ dtype=np.float32)
+
+SINGLE_FEATURE_PREDICTIONS = [
+ PredictionResult(
+ SINGLE_FEATURE_EXAMPLES[i], [
+ np.array([SINGLE_FEATURE_EXAMPLES[i][0] * 2.0 + 0.5],
+ dtype=np.float32)
+ ]) for i in range(len(SINGLE_FEATURE_EXAMPLES))
+]
+
+TWO_FEATURES_EXAMPLES = np.array([[1, 5], [3, 10], [-14, 0], [0.5, 0.5]],
+ dtype=np.float32)
+
+TWO_FEATURES_PREDICTIONS = [
+ PredictionResult(
+ TWO_FEATURES_EXAMPLES[i],
+ [
+ np.array([
+ TWO_FEATURES_EXAMPLES[i][0] * 2.0 +
+ TWO_FEATURES_EXAMPLES[i][1] * 3 + 0.5
+ ],
+ dtype=np.float32)
+ ]) for i in range(len(TWO_FEATURES_EXAMPLES))
+]
Review Comment:
The test data variables `SINGLE_FEATURE_EXAMPLES`, `SINGLE_FEATURE_PREDICTIONS`, `TWO_FEATURES_EXAMPLES`, and `TWO_FEATURES_PREDICTIONS` are defined but not used in any of the tests. Consider removing them to improve code clarity. If they are intended for future tests, they could be moved to those tests or commented out with a TODO.
##########
sdks/python/apache_beam/ml/inference/trt_handler_numpy_compact.py:
##########
@@ -0,0 +1,540 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+
+import numpy as np
+
+from apache_beam.io.filesystems import FileSystems
+from apache_beam.ml.inference import utils
+from apache_beam.ml.inference.base import ModelHandler
+from apache_beam.ml.inference.base import PredictionResult
+
+LOGGER = logging.getLogger("TensorRTEngineHandlerNumPy_TRT10")
+
+__all__ = [
+ "TensorRTEngine",
+ "TensorRTEngineHandlerNumPy",
+]
+
+
+# ---------------------------------------------------------------------
+# CUDA / TensorRT helpers
+# ---------------------------------------------------------------------
+def _assign_or_fail(args):
+ """CUDA error checking for cuda-python (Driver API)."""
+ from cuda import cuda # lazy import to avoid submit-time dependency
+
+ err, *ret = args
+ if isinstance(err, cuda.CUresult):
+ if err != cuda.CUresult.CUDA_SUCCESS:
+ raise RuntimeError(f"CUDA error: {err}")
+ else:
+ raise RuntimeError(f"Unknown CUDA error type: {err}")
+ return ret[0] if len(ret) == 1 else tuple(ret)
+
+
+def _require_tensorrt_10() -> None:
+ """Assert that TensorRT 10.x Tensor API is available on this worker."""
+ try:
+ import tensorrt as trt # noqa: F401
+ except Exception as e: # pragma: no cover
+ raise RuntimeError("TensorRT is not installed on this worker.") from e
+
+ # TRT 10.x indicators:
+ # - Engine exposes the Tensor API (num_io_tensors / get_tensor_name)
+ # - ExecutionContext exposes execute_async_v3 / set_input_shape /
+ # set_tensor_address
+ engine_reqs = ("num_io_tensors", "get_tensor_name")
+ ctx_reqs = ("execute_async_v3", "set_input_shape", "set_tensor_address")
+
+ missing_engine = [m for m in engine_reqs if not hasattr(trt.ICudaEngine, m)]
+ missing_ctx = [m for m in ctx_reqs if not hasattr(trt.IExecutionContext, m)]
+
+ if missing_engine or missing_ctx:
+ raise RuntimeError(
+ "This handler requires TensorRT 10.x+. "
+ f"Missing on ICudaEngine: {missing_engine or 'OK'}, "
+ f"Missing on IExecutionContext: {missing_ctx or 'OK'}")
+
+
+# ---------------------------------------------------------------------
+# Engine load / build (TRT 10)
+# ---------------------------------------------------------------------
+def _load_engine(engine_path: str):
+ """Deserialize a .engine (plan) from FileSystems into a TRT engine."""
+ _require_tensorrt_10()
+
+ import tensorrt as trt
+
+ with FileSystems.open(engine_path, "rb") as f:
+ blob = f.read()
+
+ logger = trt.Logger(trt.Logger.INFO)
+ trt.init_libnvinfer_plugins(logger, "")
+ rt = trt.Runtime(logger)
+ eng = rt.deserialize_cuda_engine(blob)
+ if eng is None:
+ raise RuntimeError(
+ "Failed to deserialize TensorRT engine. "
+ "The plan may be corrupt or built with an incompatible TRT.")
+ return eng
+
+
+def _load_onnx_build_engine(onnx_path: str):
+ """Parse ONNX and build a TRT engine immediately (Tensor API)."""
+ if onnx_path.lower().endswith(".engine"):
+ raise ValueError(
+ "Provided onnx_path points to .engine; pass it as engine_path instead.")
+
+ _require_tensorrt_10()
+ import tensorrt as trt
+
+ logger = trt.Logger(trt.Logger.INFO)
+ trt.init_libnvinfer_plugins(logger, "")
+
+ builder = trt.Builder(logger)
+ flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+ network = builder.create_network(flags)
+ parser = trt.OnnxParser(network, logger)
+
+ with FileSystems.open(onnx_path, "rb") as f:
+ data = f.read()
+
+ if not parser.parse(data):
+ LOGGER.error("Failed to parse ONNX: %s", onnx_path)
+ for i in range(parser.num_errors):
+ LOGGER.error(parser.get_error(i))
+ raise ValueError(f"Failed to parse ONNX: {onnx_path}")
+
+ config = builder.create_builder_config()
+ # Workbench: ~1GiB workspace (tune for your model/infra)
+ config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
+
+ if getattr(builder, "platform_has_fast_fp16", False):
+ config.set_flag(trt.BuilderFlag.FP16)
+
+ # Generic optimization profile for dynamic inputs.
+ if network.num_inputs > 0:
+ prof = builder.create_optimization_profile()
+ for i in range(network.num_inputs):
+ inp = network.get_input(i)
+ shp = list(inp.shape)
+
+ def _d(v: int, default: int) -> int:
+ return default if v < 0 else v
+
+ if len(shp) == 4:
+ # Assume NCHW; supply defaults where dims are dynamic.
+ min_shape: Tuple[int, ...] = (
+ _d(shp[0], 1), _d(shp[1], 3), _d(shp[2], 224), _d(shp[3], 224))
+ opt_shape: Tuple[int, ...] = (
+ _d(shp[0], 4), _d(shp[1], 3), _d(shp[2], 224), _d(shp[3], 224))
+ max_shape: Tuple[int, ...] = (
+ _d(shp[0], 8), _d(shp[1], 3), _d(shp[2], 224), _d(shp[3], 224))
+ else:
+ # Fallback: make batch dynamic, keep others as-is or 1.
+ min_shape = tuple(_d(x, 1) for x in shp)
+ opt_shape = tuple(_d(x, 4 if j == 0 else 1) for j, x in enumerate(shp))
+ max_shape = tuple(_d(x, 8 if j == 0 else 1) for j, x in enumerate(shp))
+
+ prof.set_shape(inp.name, min=min_shape, opt=opt_shape, max=max_shape)
+
+ config.add_optimization_profile(prof)
+
+ plan = builder.build_serialized_network(network, config)
+ if plan is None:
+ raise RuntimeError(
+ "build_serialized_network() returned None; check ONNX and profiles.")
+
+ rt = trt.Runtime(logger)
+ eng = rt.deserialize_cuda_engine(bytes(plan))
+ if eng is None:
+ raise RuntimeError("Failed to deserialize engine after build.")
+ return eng
+
+
+# ---------------------------------------------------------------------
+# Shape & batch helpers
+# ---------------------------------------------------------------------
+def _resolve_output_shape(shape: Sequence[int] | None,
+ batch_size: int) -> Tuple[int, ...] | None:
+ """Replace a leading -1 (batch) with batch_size; any other -1 is an error."""
+ if shape is None:
+ return None
+ shp = list(shape)
+ if shp and shp[0] < 0:
+ shp[0] = int(batch_size)
+ if any(d < 0 for d in shp[1:]):
+ raise RuntimeError(f"Unresolved non-batch dims in output shape: {shape}")
+ return tuple(shp)
+
+
+def _to_contiguous_batch(x: Sequence[np.ndarray] | np.ndarray) -> np.ndarray:
+ """Accept either an ndarray (already a batch) or a list of ndarrays.
+
+ Concatenates on axis 0. This avoids accidental rank-5 shapes from
+ upstream batching.
+ """
+ if isinstance(x, np.ndarray):
+ return np.ascontiguousarray(x)
+
+ if isinstance(x, (list, tuple)):
+ if len(x) == 1 and isinstance(x[0], np.ndarray):
+ return np.ascontiguousarray(x[0])
+
+ if all(isinstance(a, np.ndarray) for a in x):
+ first = x[0].shape
+ for a in x[1:]:
+ if len(a.shape) != len(first) or any(
+ sa != sb for sa, sb in zip(a.shape[1:], first[1:])):
+ raise ValueError(
+ f"Inconsistent element shapes for concatenation: "
+ f"{first} vs {a.shape}")
+ return np.ascontiguousarray(np.concatenate(x, axis=0))
+
+ raise ValueError(
+ "Batch must be ndarray or sequence of ndarrays of same "
+ "rank/shape (except batch).")
+
+
+# ---------------------------------------------------------------------
+# TRT 10.x engine wrapper (Tensor API only)
+# ---------------------------------------------------------------------
+class TensorRTEngine:
+ """TRT 10.x engine wrapper using the Tensor API."""
+ def __init__(self, engine: Any):
+ import tensorrt as trt
+
+ self.engine = engine
+ self.context = engine.create_execution_context()
+ self.context_lock = threading.RLock()
+
+ # Tensor API enumeration
+ self.input_names: List[str] = []
+ self.output_names: List[str] = []
+ self.dtypes: Dict[str, np.dtype] = {}
+
+ for i in range(engine.num_io_tensors):
+ name = engine.get_tensor_name(i)
+ mode = engine.get_tensor_mode(name)
+ if mode == trt.TensorIOMode.INPUT:
+ self.input_names.append(name)
+ else:
+ self.output_names.append(name)
+ self.dtypes[name] = np.dtype(trt.nptype(engine.get_tensor_dtype(name)))
+
+ # Lazy allocations
+ self._device_ptrs: Dict[str, int] = {} # tensor name -> CUdeviceptr
+ self._host_out: Dict[str, np.ndarray] = {}
+ self._in_nbytes: int = 0
+ self._stream: Optional[int] = None
+ self.profile_index: int = 0
+
+ def _ensure_stream(self) -> None:
+ if self._stream is None:
+ from cuda import cuda
+ self._stream = _assign_or_fail(cuda.cuStreamCreate(0))
+
+ def _free_ptr(self, ptr: Optional[int]) -> None:
+ if not ptr:
+ return
+ from cuda import cuda
+ try:
+ _assign_or_fail(cuda.cuMemFree(ptr))
+ except Exception:
+ pass
Review Comment:
The `except Exception: pass` is too broad and can hide important errors during memory deallocation. It's better to catch a more specific exception, like `RuntimeError` which is raised by `_assign_or_fail`. If the goal is to suppress errors during cleanup, consider logging them as warnings instead of silently ignoring them.
```suggestion
try:
_assign_or_fail(cuda.cuMemFree(ptr))
except RuntimeError as e:
LOGGER.warning('Failed to free CUDA memory pointer: %s', e)
```
##########
sdks/python/apache_beam/ml/inference/trt_handler_numpy_compact_test.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# pytype: skip-file
+
+import os
+import tempfile
+import unittest
+
+import numpy as np
+import pytest
+
+import apache_beam as beam
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+
+# Protect against environments where TensorRT python library is not available.
+# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
+try:
+ from apache_beam.ml.inference.base import PredictionResult, RunInference
+ from apache_beam.ml.inference.trt_handler_numpy_compact import (
+ TensorRTEngine,
+ TensorRTEngineHandlerNumPy,
+ )
+except ImportError:
+ raise unittest.SkipTest('TensorRT 10.x dependencies are not installed')
+
+try:
+ from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+except ImportError:
+ GCSFileSystem = None # type: ignore
+
+# Test data
+SINGLE_FEATURE_EXAMPLES = np.array([[1.0], [5.0], [-3.0], [10.0]],
+ dtype=np.float32)
+
+SINGLE_FEATURE_PREDICTIONS = [
+ PredictionResult(
+ SINGLE_FEATURE_EXAMPLES[i], [
+ np.array([SINGLE_FEATURE_EXAMPLES[i][0] * 2.0 + 0.5],
+ dtype=np.float32)
+ ]) for i in range(len(SINGLE_FEATURE_EXAMPLES))
+]
+
+TWO_FEATURES_EXAMPLES = np.array([[1, 5], [3, 10], [-14, 0], [0.5, 0.5]],
+ dtype=np.float32)
+
+TWO_FEATURES_PREDICTIONS = [
+ PredictionResult(
+ TWO_FEATURES_EXAMPLES[i],
+ [
+ np.array([
+ TWO_FEATURES_EXAMPLES[i][0] * 2.0 +
+ TWO_FEATURES_EXAMPLES[i][1] * 3 + 0.5
+ ],
+ dtype=np.float32)
+ ]) for i in range(len(TWO_FEATURES_EXAMPLES))
+]
+
+
+def _compare_prediction_result(a, b):
+ """Compare two PredictionResult objects."""
+ return ((a.example == b.example).all() and all(
+ np.array_equal(actual, expected)
+ for actual, expected in zip(a.inference, b.inference)))
+
+
+def _build_simple_onnx_model(input_size=1, output_path=None):
+ """Build a simple ONNX model for testing: y = 2x + 0.5"""
+ try:
+ from onnx import helper, TensorProto
+ except ImportError:
+ raise unittest.SkipTest('ONNX dependencies are not installed')
+
+ # Create a simple linear model: y = 2*x + 0.5
+ input_tensor = helper.make_tensor_value_info(
+ 'input', TensorProto.FLOAT, [None, input_size])
+ output_tensor = helper.make_tensor_value_info(
+ 'output', TensorProto.FLOAT, [None, input_size])
+
+ # Weight tensor: 2.0
+ weight_init = helper.make_tensor(
+ 'weight',
+ TensorProto.FLOAT, [input_size, input_size],
+ [2.0] * (input_size * input_size))
+ # Bias tensor: 0.5
+ bias_init = helper.make_tensor(
+ 'bias', TensorProto.FLOAT, [input_size], [0.5] * input_size)
+
+ # MatMul node
+ matmul_node = helper.make_node('MatMul', ['input', 'weight'], ['matmul_out'])
+ # Add node
+ add_node = helper.make_node('Add', ['matmul_out', 'bias'], ['output'])
+
+ # Create graph
+ graph = helper.make_graph([matmul_node, add_node],
+ 'simple_linear', [input_tensor], [output_tensor],
+ [weight_init, bias_init])
+
+ # Create model
+ model = helper.make_model(graph, producer_name='trt_test')
+ model.opset_import[0].version = 13
+
+ if output_path:
+ with open(output_path, 'wb') as f:
+ f.write(model.SerializeToString())
+
+ return model
+
+
+@pytest.mark.uses_tensorrt
+class TensorRTEngineHandlerNumPyTest(unittest.TestCase):
+ """Tests for TensorRTEngineHandlerNumPy with TensorRT 10.x Tensor API."""
+ def test_handler_initialization(self):
+ """Test that handler initializes correctly with required parameters."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=1, max_batch_size=4, engine_path='/tmp/test.engine')
+ self.assertEqual(handler.min_batch_size, 1)
+ self.assertEqual(handler.max_batch_size, 4)
+ self.assertEqual(handler.engine_path, '/tmp/test.engine')
+
+ def test_handler_initialization_both_paths_error(self):
+ """Test that providing both engine_path and onnx_path raises an error."""
+ with self.assertRaises(ValueError):
+ TensorRTEngineHandlerNumPy(
+ min_batch_size=1,
+ max_batch_size=4,
+ engine_path='/tmp/test.engine',
+ onnx_path='/tmp/test.onnx')
+
+ def test_handler_initialization_no_path_error(self):
+ """Test that providing neither engine_path nor onnx_path raises an error."""
+ with self.assertRaises(ValueError):
+ TensorRTEngineHandlerNumPy(min_batch_size=1, max_batch_size=4)
+
+ def test_handler_initialization_invalid_engine_path(self):
+ """Test that providing a non-.engine path raises an error."""
+ with self.assertRaises(ValueError):
+ TensorRTEngineHandlerNumPy(
+ min_batch_size=1, max_batch_size=4, engine_path='/tmp/test.onnx')
+
+ def test_handler_initialization_invalid_onnx_path(self):
+ """Test that providing a .engine path as onnx_path raises an error."""
+ with self.assertRaises(ValueError):
+ TensorRTEngineHandlerNumPy(
+ min_batch_size=1, max_batch_size=4, onnx_path='/tmp/test.engine')
+
+ def test_batch_elements_kwargs(self):
+ """Test that batch_elements_kwargs returns correct values."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=2,
+ max_batch_size=8,
+ max_batch_duration_secs=10,
+ engine_path='/tmp/test.engine')
+ kwargs = handler.batch_elements_kwargs()
+ self.assertEqual(kwargs['min_batch_size'], 2)
+ self.assertEqual(kwargs['max_batch_size'], 8)
+ self.assertEqual(kwargs['max_batch_duration_secs'], 10)
+
+ def test_get_num_bytes_ndarray(self):
+ """Test get_num_bytes with numpy ndarray."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=1, max_batch_size=4, engine_path='/tmp/test.engine')
+ batch = np.array([[1, 2], [3, 4]], dtype=np.float32)
+ num_bytes = handler.get_num_bytes(batch)
+ self.assertEqual(num_bytes, batch.nbytes)
+
+ def test_get_num_bytes_list(self):
+ """Test get_num_bytes with list of ndarrays."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=1, max_batch_size=4, engine_path='/tmp/test.engine')
+ batch = [
+ np.array([1, 2], dtype=np.float32), np.array([3, 4], dtype=np.float32)
+ ]
+ num_bytes = handler.get_num_bytes(batch)
+ expected = sum(a.nbytes for a in batch)
+ self.assertEqual(num_bytes, expected)
+
+ def test_get_metrics_namespace(self):
+ """Test that metrics namespace is correct."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=1, max_batch_size=4, engine_path='/tmp/test.engine')
+ self.assertEqual(handler.get_metrics_namespace(), 'BeamML_TensorRT10')
+
+ def test_share_model_across_processes(self):
+ """Test share_model_across_processes flag."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=1,
+ max_batch_size=4,
+ engine_path='/tmp/test.engine',
+ large_model=True)
+ self.assertTrue(handler.share_model_across_processes())
+
+ def test_model_copies(self):
+ """Test model_copies parameter."""
+ handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=1,
+ max_batch_size=4,
+ engine_path='/tmp/test.engine',
+ model_copies=3)
+ self.assertEqual(handler.model_copies(), 3)
+
+ @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed')
Review Comment:
The tests `test_inference_with_onnx_build_on_worker`, `test_engine_initialization` (line 264), and `test_pipeline_with_onnx_model` (line 292) are skipped if `GCSFileSystem` is not available. However, these tests use local temporary files and do not seem to depend on GCS. The `FileSystems.open` call in the implementation will use `LocalFileSystem` for local paths. The `skipIf` condition seems unnecessary and could be removed from these tests.
##########
sdks/python/apache_beam/ml/inference/trt_handler_numpy_compact.py:
##########
@@ -0,0 +1,540 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+
+import numpy as np
+
+from apache_beam.io.filesystems import FileSystems
+from apache_beam.ml.inference import utils
+from apache_beam.ml.inference.base import ModelHandler
+from apache_beam.ml.inference.base import PredictionResult
+
+LOGGER = logging.getLogger("TensorRTEngineHandlerNumPy_TRT10")
+
+__all__ = [
+ "TensorRTEngine",
+ "TensorRTEngineHandlerNumPy",
+]
+
+
+# ---------------------------------------------------------------------
+# CUDA / TensorRT helpers
+# ---------------------------------------------------------------------
+def _assign_or_fail(args):
+ """CUDA error checking for cuda-python (Driver API)."""
+ from cuda import cuda # lazy import to avoid submit-time dependency
+
+ err, *ret = args
+ if isinstance(err, cuda.CUresult):
+ if err != cuda.CUresult.CUDA_SUCCESS:
+ raise RuntimeError(f"CUDA error: {err}")
+ else:
+ raise RuntimeError(f"Unknown CUDA error type: {err}")
+ return ret[0] if len(ret) == 1 else tuple(ret)
+
+
+def _require_tensorrt_10() -> None:
+ """Assert that TensorRT 10.x Tensor API is available on this worker."""
+ try:
+ import tensorrt as trt # noqa: F401
+ except Exception as e: # pragma: no cover
+ raise RuntimeError("TensorRT is not installed on this worker.") from e
+
+ # TRT 10.x indicators:
+ # - Engine exposes the Tensor API (num_io_tensors / get_tensor_name)
+ # - ExecutionContext exposes execute_async_v3 / set_input_shape /
+ # set_tensor_address
+ engine_reqs = ("num_io_tensors", "get_tensor_name")
+ ctx_reqs = ("execute_async_v3", "set_input_shape", "set_tensor_address")
+
+ missing_engine = [m for m in engine_reqs if not hasattr(trt.ICudaEngine, m)]
+ missing_ctx = [m for m in ctx_reqs if not hasattr(trt.IExecutionContext, m)]
+
+ if missing_engine or missing_ctx:
+ raise RuntimeError(
+ "This handler requires TensorRT 10.x+. "
+ f"Missing on ICudaEngine: {missing_engine or 'OK'}, "
+ f"Missing on IExecutionContext: {missing_ctx or 'OK'}")
+
+
+# ---------------------------------------------------------------------
+# Engine load / build (TRT 10)
+# ---------------------------------------------------------------------
+def _load_engine(engine_path: str):
+ """Deserialize a .engine (plan) from FileSystems into a TRT engine."""
+ _require_tensorrt_10()
+
+ import tensorrt as trt
+
+ with FileSystems.open(engine_path, "rb") as f:
+ blob = f.read()
+
+ logger = trt.Logger(trt.Logger.INFO)
+ trt.init_libnvinfer_plugins(logger, "")
+ rt = trt.Runtime(logger)
+ eng = rt.deserialize_cuda_engine(blob)
+ if eng is None:
+ raise RuntimeError(
+ "Failed to deserialize TensorRT engine. "
+ "The plan may be corrupt or built with an incompatible TRT.")
+ return eng
+
+
+def _load_onnx_build_engine(onnx_path: str):
+ """Parse ONNX and build a TRT engine immediately (Tensor API)."""
+ if onnx_path.lower().endswith(".engine"):
+ raise ValueError(
+ "Provided onnx_path points to .engine; pass it as engine_path instead.")
+
+ _require_tensorrt_10()
+ import tensorrt as trt
+
+ logger = trt.Logger(trt.Logger.INFO)
+ trt.init_libnvinfer_plugins(logger, "")
+
+ builder = trt.Builder(logger)
+ flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+ network = builder.create_network(flags)
+ parser = trt.OnnxParser(network, logger)
+
+ with FileSystems.open(onnx_path, "rb") as f:
+ data = f.read()
+
+ if not parser.parse(data):
+ LOGGER.error("Failed to parse ONNX: %s", onnx_path)
+ for i in range(parser.num_errors):
+ LOGGER.error(parser.get_error(i))
+ raise ValueError(f"Failed to parse ONNX: {onnx_path}")
+
+ config = builder.create_builder_config()
+ # Workbench: ~1GiB workspace (tune for your model/infra)
+ config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
+
+ if getattr(builder, "platform_has_fast_fp16", False):
+ config.set_flag(trt.BuilderFlag.FP16)
+
+ # Generic optimization profile for dynamic inputs.
+ if network.num_inputs > 0:
+ prof = builder.create_optimization_profile()
+ for i in range(network.num_inputs):
+ inp = network.get_input(i)
+ shp = list(inp.shape)
+
+ def _d(v: int, default: int) -> int:
+ return default if v < 0 else v
+
+ if len(shp) == 4:
+ # Assume NCHW; supply defaults where dims are dynamic.
+ min_shape: Tuple[int, ...] = (
+ _d(shp[0], 1), _d(shp[1], 3), _d(shp[2], 224), _d(shp[3], 224))
+ opt_shape: Tuple[int, ...] = (
+ _d(shp[0], 4), _d(shp[1], 3), _d(shp[2], 224), _d(shp[3], 224))
+ max_shape: Tuple[int, ...] = (
+ _d(shp[0], 8), _d(shp[1], 3), _d(shp[2], 224), _d(shp[3], 224))
+ else:
+ # Fallback: make batch dynamic, keep others as-is or 1.
+ min_shape = tuple(_d(x, 1) for x in shp)
+ opt_shape = tuple(_d(x, 4 if j == 0 else 1) for j, x in enumerate(shp))
+ max_shape = tuple(_d(x, 8 if j == 0 else 1) for j, x in enumerate(shp))
+
+ prof.set_shape(inp.name, min=min_shape, opt=opt_shape, max=max_shape)
+
+ config.add_optimization_profile(prof)
+
+ plan = builder.build_serialized_network(network, config)
+ if plan is None:
+ raise RuntimeError(
+ "build_serialized_network() returned None; check ONNX and profiles.")
+
+ rt = trt.Runtime(logger)
+ eng = rt.deserialize_cuda_engine(bytes(plan))
+ if eng is None:
+ raise RuntimeError("Failed to deserialize engine after build.")
+ return eng
+
+
+# ---------------------------------------------------------------------
+# Shape & batch helpers
+# ---------------------------------------------------------------------
+def _resolve_output_shape(shape: Sequence[int] | None,
+ batch_size: int) -> Tuple[int, ...] | None:
+ """Replace a leading -1 (batch) with batch_size; any other -1 is an error."""
+ if shape is None:
+ return None
+ shp = list(shape)
+ if shp and shp[0] < 0:
+ shp[0] = int(batch_size)
+ if any(d < 0 for d in shp[1:]):
+ raise RuntimeError(f"Unresolved non-batch dims in output shape: {shape}")
+ return tuple(shp)
+
+
+def _to_contiguous_batch(x: Sequence[np.ndarray] | np.ndarray) -> np.ndarray:
+ """Accept either an ndarray (already a batch) or a list of ndarrays.
+
+ Concatenates on axis 0. This avoids accidental rank-5 shapes from
+ upstream batching.
+ """
+ if isinstance(x, np.ndarray):
+ return np.ascontiguousarray(x)
+
+ if isinstance(x, (list, tuple)):
+ if len(x) == 1 and isinstance(x[0], np.ndarray):
+ return np.ascontiguousarray(x[0])
+
+ if all(isinstance(a, np.ndarray) for a in x):
+ first = x[0].shape
+ for a in x[1:]:
+ if len(a.shape) != len(first) or any(
+ sa != sb for sa, sb in zip(a.shape[1:], first[1:])):
+ raise ValueError(
+ f"Inconsistent element shapes for concatenation: "
+ f"{first} vs {a.shape}")
+ return np.ascontiguousarray(np.concatenate(x, axis=0))
+
+ raise ValueError(
+ "Batch must be ndarray or sequence of ndarrays of same "
+ "rank/shape (except batch).")
+
+
+# ---------------------------------------------------------------------
+# TRT 10.x engine wrapper (Tensor API only)
+# ---------------------------------------------------------------------
+class TensorRTEngine:
+ """TRT 10.x engine wrapper using the Tensor API."""
+ def __init__(self, engine: Any):
+ import tensorrt as trt
+
+ self.engine = engine
+ self.context = engine.create_execution_context()
+ self.context_lock = threading.RLock()
+
+ # Tensor API enumeration
+ self.input_names: List[str] = []
+ self.output_names: List[str] = []
+ self.dtypes: Dict[str, np.dtype] = {}
+
+ for i in range(engine.num_io_tensors):
+ name = engine.get_tensor_name(i)
+ mode = engine.get_tensor_mode(name)
+ if mode == trt.TensorIOMode.INPUT:
+ self.input_names.append(name)
+ else:
+ self.output_names.append(name)
+ self.dtypes[name] = np.dtype(trt.nptype(engine.get_tensor_dtype(name)))
+
+ # Lazy allocations
+ self._device_ptrs: Dict[str, int] = {} # tensor name -> CUdeviceptr
+ self._host_out: Dict[str, np.ndarray] = {}
+ self._in_nbytes: int = 0
+ self._stream: Optional[int] = None
+ self.profile_index: int = 0
+
+ def _ensure_stream(self) -> None:
+ if self._stream is None:
+ from cuda import cuda
+ self._stream = _assign_or_fail(cuda.cuStreamCreate(0))
+
+ def _free_ptr(self, ptr: Optional[int]) -> None:
+ if not ptr:
+ return
+ from cuda import cuda
+ try:
+ _assign_or_fail(cuda.cuMemFree(ptr))
+ except Exception:
+ pass
+
+ def _select_profile(self) -> None:
+ # Pick optimization profile (sync or async depending on TRT)
+ if hasattr(self.context, "set_optimization_profile_async"):
+ self._ensure_stream()
+ self.context.set_optimization_profile_async(
+ self.profile_index, self._stream)
+ elif hasattr(self.context, "set_optimization_profile"):
+ self.context.set_optimization_profile(self.profile_index)
+
+ def _check_shape_in_profile(self, name: str, shape: Sequence[int]) -> None:
+ mi, _, ma = self.engine.get_tensor_profile_shape(name, self.profile_index)
+
+ def ok(dim: int, lo: int, hi: int) -> bool:
+ if lo < 0:
+ lo = dim
+ if hi < 0:
+ hi = dim
+ return lo <= dim <= hi
+
+ if len(shape) != len(mi):
+ raise RuntimeError(
+ f"Input '{name}' rank mismatch: given {tuple(shape)}, "
+ f"profile[{self.profile_index}] min={tuple(mi)} max={tuple(ma)}")
+ for i, dim in enumerate(shape):
+ if not ok(dim, mi[i], ma[i]):
+ raise RuntimeError(
+ f"Input '{name}' dim {i}={dim} outside "
+ f"profile[{self.profile_index}] bounds "
+ f"[min={mi[i]}, max={ma[i]}]. Given={tuple(shape)}, "
+ f"min={tuple(mi)}, max={tuple(ma)}")
+
+ def ensure_buffers(
+ self,
+ batch: np.ndarray,
+ input_shapes: Optional[Dict[str, Sequence[int]]] = None,
+ ) -> None:
+ """
+ Validate shapes, set input shapes, (re)allocate device + host buffers,
+ and set tensor addresses for Tensor API execution.
+ """
+ from cuda import cuda
+
+ self._select_profile()
+
+ # Derive shapes for inputs
+ shapes: Dict[str, List[int]] = {}
+ if len(self.input_names) == 1:
+ shapes[self.input_names[0]] = list(batch.shape)
+ else:
+ if not input_shapes:
+ raise RuntimeError(
+ f"Engine expects multiple inputs {self.input_names}; "
+ "provide shapes via "
+ "inference_args={'input_shapes': {name: shape, ...}}")
+ for name in self.input_names:
+ if name not in input_shapes:
+ raise RuntimeError(f"Missing shape for input tensor '{name}'")
+ shapes[name] = list(map(int, input_shapes[name]))
+
+ # Validate and set shapes
+ for name, shp in shapes.items():
+ self._check_shape_in_profile(name, shp)
+ self.context.set_input_shape(name, shp)
+
+ # Allocate first input device buffer (we copy only this from 'batch')
+ in_name = self.input_names[0]
+ in_dtype = self.dtypes[in_name]
+ in_nbytes = int(np.prod(shapes[in_name])) * in_dtype.itemsize
+ if self._device_ptrs.get(in_name) is None or self._in_nbytes != in_nbytes:
+ if self._device_ptrs.get(in_name) is not None:
+ self._free_ptr(self._device_ptrs[in_name])
+ self._device_ptrs[in_name] = _assign_or_fail(cuda.cuMemAlloc(in_nbytes))
+ self._in_nbytes = in_nbytes
Review Comment:
The current implementation for multi-input models appears to be incomplete.
1. While it derives shapes for multiple inputs from `inference_args['input_shapes']`, it only uses the `batch` data for the first input tensor (`self.input_names[0]`). There is no mechanism to provide data for the other input tensors. The `_trt10_inference_fn` also only copies data for the first input.
2. The device buffer allocation in `ensure_buffers` only happens for the first input tensor. It should iterate over all `self.input_names` and allocate buffers for each.
A common pattern for multi-input models in Beam is to have `PCollection` elements as dictionaries mapping input names to numpy arrays. The `ModelHandler` and `run_inference` function would need to be adapted to handle this. For example, `ExampleT` could be `Dict[str, np.ndarray]`, and `run_inference` would receive a batch of such dictionaries.
This is a significant limitation for models with multiple inputs.
##########
sdks/python/apache_beam/ml/inference/trt_handler_numpy_compact.py:
##########
@@ -0,0 +1,540 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import Tuple
+
+import numpy as np
+
+from apache_beam.io.filesystems import FileSystems
+from apache_beam.ml.inference import utils
+from apache_beam.ml.inference.base import ModelHandler
+from apache_beam.ml.inference.base import PredictionResult
+
+LOGGER = logging.getLogger("TensorRTEngineHandlerNumPy_TRT10")
+
+__all__ = [
+ "TensorRTEngine",
+ "TensorRTEngineHandlerNumPy",
+]
+
+
+# ---------------------------------------------------------------------
+# CUDA / TensorRT helpers
+# ---------------------------------------------------------------------
+def _assign_or_fail(args):
+ """CUDA error checking for cuda-python (Driver API)."""
+ from cuda import cuda # lazy import to avoid submit-time dependency
+
+ err, *ret = args
+ if isinstance(err, cuda.CUresult):
+ if err != cuda.CUresult.CUDA_SUCCESS:
+ raise RuntimeError(f"CUDA error: {err}")
+ else:
+ raise RuntimeError(f"Unknown CUDA error type: {err}")
+ return ret[0] if len(ret) == 1 else tuple(ret)
+
+
+def _require_tensorrt_10() -> None:
+ """Assert that TensorRT 10.x Tensor API is available on this worker."""
+ try:
+ import tensorrt as trt # noqa: F401
+ except Exception as e: # pragma: no cover
+ raise RuntimeError("TensorRT is not installed on this worker.") from e
+
+ # TRT 10.x indicators:
+ # - Engine exposes the Tensor API (num_io_tensors / get_tensor_name)
+ # - ExecutionContext exposes execute_async_v3 / set_input_shape /
+ # set_tensor_address
+ engine_reqs = ("num_io_tensors", "get_tensor_name")
+ ctx_reqs = ("execute_async_v3", "set_input_shape", "set_tensor_address")
+
+ missing_engine = [m for m in engine_reqs if not hasattr(trt.ICudaEngine, m)]
+ missing_ctx = [m for m in ctx_reqs if not hasattr(trt.IExecutionContext, m)]
+
+ if missing_engine or missing_ctx:
+ raise RuntimeError(
+ "This handler requires TensorRT 10.x+. "
+ f"Missing on ICudaEngine: {missing_engine or 'OK'}, "
+ f"Missing on IExecutionContext: {missing_ctx or 'OK'}")
+
+
+# ---------------------------------------------------------------------
+# Engine load / build (TRT 10)
+# ---------------------------------------------------------------------
+def _load_engine(engine_path: str):
+ """Deserialize a .engine (plan) from FileSystems into a TRT engine."""
+ _require_tensorrt_10()
+
+ import tensorrt as trt
+
+ with FileSystems.open(engine_path, "rb") as f:
+ blob = f.read()
+
+ logger = trt.Logger(trt.Logger.INFO)
+ trt.init_libnvinfer_plugins(logger, "")
+ rt = trt.Runtime(logger)
+ eng = rt.deserialize_cuda_engine(blob)
+ if eng is None:
+ raise RuntimeError(
+ "Failed to deserialize TensorRT engine. "
+ "The plan may be corrupt or built with an incompatible TRT.")
+ return eng
+
+
+def _load_onnx_build_engine(onnx_path: str):
+ """Parse ONNX and build a TRT engine immediately (Tensor API)."""
+ if onnx_path.lower().endswith(".engine"):
+ raise ValueError(
+ "Provided onnx_path points to .engine; pass it as engine_path instead.")
+
+ _require_tensorrt_10()
+ import tensorrt as trt
+
+ logger = trt.Logger(trt.Logger.INFO)
+ trt.init_libnvinfer_plugins(logger, "")
+
+ builder = trt.Builder(logger)
+ flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+ network = builder.create_network(flags)
+ parser = trt.OnnxParser(network, logger)
+
+ with FileSystems.open(onnx_path, "rb") as f:
+ data = f.read()
+
+ if not parser.parse(data):
+ LOGGER.error("Failed to parse ONNX: %s", onnx_path)
+ for i in range(parser.num_errors):
+ LOGGER.error(parser.get_error(i))
+ raise ValueError(f"Failed to parse ONNX: {onnx_path}")
+
+ config = builder.create_builder_config()
+ # Workbench: ~1GiB workspace (tune for your model/infra)
+ config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
Review Comment:
The workspace size for the ONNX builder is hardcoded to 1GB. While the comment suggests tuning it, it would be more flexible for users if this could be configured through the `TensorRTEngineHandlerNumPy` constructor. Consider adding a parameter like `onnx_builder_config_args` to `__init__` that can be used to pass arguments to the TensorRT builder config, for example to set the memory pool limit.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "SaiShashank12 (via GitHub)" <gi...@apache.org>.
SaiShashank12 commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3393865756
/gemini review
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix TensorRT compatibility with version 10.x [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #36309:
URL: https://github.com/apache/beam/pull/36309#issuecomment-3511267671
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @damccorm for label python.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org