Support TensorRT 10 and later in TensorRTEngineHandlerNumPy

CHANGES.md

@@ -92,6 +92,7 @@
 * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)).
 * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)).
 * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)).
+* (Python) `TensorRTEngineHandlerNumPy` now works with TensorRT 10 and later, which removed the binding API it was written against. TensorRT 8.x remains supported, so no existing GPU loses support ([#36306](https://github.com/apache/beam/issues/36306)).
 
 ## Security Fixes
 

sdks/python/apache_beam/examples/inference/README.md

@@ -87,9 +87,16 @@ To use TensorRT locally, we suggest an environment with TensorRT >= 8.0.1. Insta
 [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html). You
 will need to make sure the Python bindings for TensorRT are also installed correctly, these are available by installing the python3-libnvinfer and python3-libnvinfer-dev packages on your TensorRT download.
 
+Both the TensorRT 8.x binding API and the TensorRT 10.x and later tensor API are
+supported. Note that a serialized TensorRT engine can only be deserialized by
+the TensorRT major version that built it, so an engine built with TensorRT 8.x
+must be rebuilt before it can be used with TensorRT 10 or later. TensorRT 10 and
+later also require a GPU with compute capability 7.5 or higher, which excludes
+NVIDIA Pascal and Volta GPUs such as the Tesla P4, P100 and V100.
+
 If you would like to use Docker, you can use an NGC image like:
 ```
-docker pull nvcr.io/nvidia/tensorrt:22.04-py3
+docker pull nvcr.io/nvidia/tensorrt:26.06-py3
 ```
 as an existing container base to [build custom Apache Beam container](https://beam.apache.org/documentation/runtime/environments/#modify-existing-base-image).
 

sdks/python/apache_beam/ml/inference/tensorrt_inference.py

@@ -49,6 +49,34 @@
   LOGGER.warning(msg)
 
 
+def _trt_major_version() -> int:
+  """Returns the major version of the installed TensorRT.
+
+  TensorRT 10 replaced the index based "binding" API with a name based
+  "tensor" API, so the major version selects which code path to take.
+  """
+  import tensorrt as trt
+  try:
+    return int(trt.__version__.split('.')[0])
+  except (AttributeError, IndexError, ValueError):
+    # Fall back to probing for an attribute that only exists from 10 onwards.
+    return 10 if hasattr(trt.ICudaEngine, 'num_io_tensors') else 8
+
+
+def _import_cuda_driver():
+  """Imports the CUDA driver bindings.
+
+  ``cuda.bindings.driver`` is the module path used by cuda-python 12.8 and
+  later. It replaced the ``cuda.cuda`` alias, which was removed in
+  cuda-python 13.0, so only fall back to that for older installations.
+  """
+  try:
+    from cuda.bindings import driver as cuda
+  except ImportError:
+    from cuda import cuda
+  return cuda
+
+
 def _load_engine(engine_path):
   import tensorrt as trt
   file = FileSystems.open(engine_path, 'rb')
@@ -58,11 +86,25 @@ def _load_engine(engine_path):
   return engine
 
 
+def _network_creation_flags() -> int:
+  """Returns the ``create_network`` flags for the installed TensorRT.
+
+  Explicit batch is the only supported mode from TensorRT 10 onwards, where
+  the flag is first deprecated and then removed, so it is only passed to
+  TensorRT 8.x.
+  """
+  import tensorrt as trt
+  explicit_batch = getattr(
+      trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH', None)
+  if explicit_batch is None or _trt_major_version() >= 10:
+    return 0
+  return 1 << int(explicit_batch)
+
+
 def _load_onnx(onnx_path):
   import tensorrt as trt
   builder = trt.Builder(TRT_LOGGER)
-  network = builder.create_network(
-      flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+  network = builder.create_network(flags=_network_creation_flags())
   parser = trt.OnnxParser(network, TRT_LOGGER)
   with FileSystems.open(onnx_path) as f:
     if not parser.parse(f.read()):
@@ -85,7 +127,7 @@ def _build_engine(network, builder):
 
 def _assign_or_fail(args):
   """CUDA error checking."""
-  from cuda import cuda
+  cuda = _import_cuda_driver()
   err, ret = args[0], args[1:]
   if isinstance(err, cuda.CUresult):
     if err != cuda.CUresult.CUDA_SUCCESS:
@@ -111,7 +153,7 @@ def __init__(self, engine: trt.ICudaEngine):
       engine: trt.ICudaEngine object that contains TensorRT engine
     """
     import tensorrt as trt
-    from cuda import cuda
+    cuda = _import_cuda_driver()
     self.engine = engine
     self.context = engine.create_execution_context()
     self.context_lock = threading.RLock()
@@ -120,34 +162,59 @@ def __init__(self, engine: trt.ICudaEngine):
     self.gpu_allocations = []
     self.cpu_allocations = []
 
-    # TODO(https://github.com/NVIDIA/TensorRT/issues/2557):
-    # Clean up when fixed upstream.
-    try:
-      _ = np.bool
-    except AttributeError:
-      # numpy >= 1.24.0
-      np.bool = np.bool_  # type: ignore
-
     # Setup I/O bindings.
-    for i in range(self.engine.num_bindings):
-      name = self.engine.get_binding_name(i)
-      dtype = self.engine.get_binding_dtype(i)
-      shape = self.engine.get_binding_shape(i)
-      size = trt.volume(shape) * dtype.itemsize
-      allocation = _assign_or_fail(cuda.cuMemAlloc(size))
-      binding = {
-          'index': i,
-          'name': name,
-          'dtype': np.dtype(trt.nptype(dtype)),
-          'shape': list(shape),
-          'allocation': allocation,
-          'size': size
-      }
-      self.gpu_allocations.append(allocation)
-      if self.engine.binding_is_input(i):
-        self.inputs.append(binding)
-      else:
-        self.outputs.append(binding)
+    if _trt_major_version() >= 10:
+      # TensorRT 10 removed the index based binding API in favour of a name
+      # based tensor API. Device addresses are bound to the context once here
+      # because execute_async_v3 takes no allocation list at execution time.
+      for i in range(self.engine.num_io_tensors):
+        name = self.engine.get_tensor_name(i)
+        dtype = self.engine.get_tensor_dtype(name)
+        shape = self.engine.get_tensor_shape(name)
+        size = trt.volume(shape) * dtype.itemsize
+        allocation = _assign_or_fail(cuda.cuMemAlloc(size))
+        binding = {
+            'index': i,
+            'name': name,
+            'dtype': np.dtype(trt.nptype(dtype)),
+            'shape': list(shape),
+            'allocation': allocation,
+            'size': size
+        }
+        self.gpu_allocations.append(allocation)
+        self.context.set_tensor_address(name, int(allocation))
+        if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+          self.inputs.append(binding)
+        else:
+          self.outputs.append(binding)
+    else:
+      # TODO(https://github.com/NVIDIA/TensorRT/issues/2557):
+      # Clean up when the TensorRT 8.x path is dropped.
+      try:
+        _ = np.bool
+      except AttributeError:
+        # numpy >= 1.24.0
+        np.bool = np.bool_  # type: ignore
+
+      for i in range(self.engine.num_bindings):
+        name = self.engine.get_binding_name(i)
+        dtype = self.engine.get_binding_dtype(i)
+        shape = self.engine.get_binding_shape(i)
+        size = trt.volume(shape) * dtype.itemsize
+        allocation = _assign_or_fail(cuda.cuMemAlloc(size))
+        binding = {
+            'index': i,
+            'name': name,
+            'dtype': np.dtype(trt.nptype(dtype)),
+            'shape': list(shape),
+            'allocation': allocation,
+            'size': size
+        }
+        self.gpu_allocations.append(allocation)
+        if self.engine.binding_is_input(i):
+          self.inputs.append(binding)
+        else:
+          self.outputs.append(binding)
 
     assert self.context
     assert len(self.inputs) > 0
@@ -182,7 +249,7 @@ def _default_tensorRT_inference_fn(
     engine: TensorRTEngine,
     inference_args: Optional[dict[str,
                                   Any]] = None) -> Iterable[PredictionResult]:
-  from cuda import cuda
+  cuda = _import_cuda_driver()
   (
       engine,
       context,
@@ -201,7 +268,11 @@ def _default_tensorRT_inference_fn(
             np.ascontiguousarray(batch),
             inputs[0]['size'],
             stream))
-    context.execute_async_v2(gpu_allocations, stream)
+    if _trt_major_version() >= 10:
+      # Tensor addresses were bound when the engine was created.
+      context.execute_async_v3(stream)
+    else:
+      context.execute_async_v2(gpu_allocations, stream)
     for output in range(len(cpu_allocations)):
       _assign_or_fail(
           cuda.cuMemcpyDtoHAsync(

sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py

@@ -37,6 +37,10 @@
   from apache_beam.ml.inference.base import PredictionResult
   from apache_beam.ml.inference.base import RunInference
   from apache_beam.ml.inference.tensorrt_inference import TensorRTEngineHandlerNumPy
+  from apache_beam.ml.inference.tensorrt_inference import _assign_or_fail
+  from apache_beam.ml.inference.tensorrt_inference import _import_cuda_driver
+  from apache_beam.ml.inference.tensorrt_inference import _network_creation_flags
+  from apache_beam.ml.inference.tensorrt_inference import _trt_major_version
 except ImportError:
   raise unittest.SkipTest('TensorRT dependencies are not installed')
 
@@ -90,23 +94,8 @@ def _compare_prediction_result(a, b):
       for actual, expected in zip(a.inference, b.inference)))
 
 
-def _assign_or_fail(args):
-  """CUDA error checking."""
-  from cuda import cuda
-  err, ret = args[0], args[1:]
-  if isinstance(err, cuda.CUresult):
-    if err != cuda.CUresult.CUDA_SUCCESS:
-      raise RuntimeError("Cuda Error: {}".format(err))
-  else:
-    raise RuntimeError("Unknown error type: {}".format(err))
-  # Special case so that no unpacking is needed at call-site.
-  if len(ret) == 1:
-    return ret[0]
-  return ret
-
-
 def _custom_tensorRT_inference_fn(batch, engine, inference_args):
-  from cuda import cuda
+  cuda = _import_cuda_driver()
   (
       engine,
       context,
@@ -125,7 +114,10 @@ def _custom_tensorRT_inference_fn(batch, engine, inference_args):
             np.ascontiguousarray(batch),
             inputs[0]['size'],
             stream))
-    context.execute_async_v2(gpu_allocations, stream)
+    if _trt_major_version() >= 10:
+      context.execute_async_v3(stream)
+    else:
+      context.execute_async_v2(gpu_allocations, stream)
     for output in range(len(cpu_allocations)):
       _assign_or_fail(
           cuda.cuMemcpyDtoHAsync(
@@ -189,8 +181,7 @@ def test_inference_single_tensor_feature(self):
     inference_runner = TensorRTEngineHandlerNumPy(
         min_batch_size=4, max_batch_size=4)
     builder = trt.Builder(LOGGER)
-    network = builder.create_network(
-        flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    network = builder.create_network(flags=_network_creation_flags())
     input_tensor = network.add_input(
         name="input", dtype=trt.float32, shape=(4, 1))
     weight_const = network.add_constant(
@@ -227,8 +218,7 @@ def test_inference_custom_single_tensor_feature(self):
         max_batch_size=4,
         inference_fn=_custom_tensorRT_inference_fn)
     builder = trt.Builder(LOGGER)
-    network = builder.create_network(
-        flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    network = builder.create_network(flags=_network_creation_flags())
     input_tensor = network.add_input(
         name="input", dtype=trt.float32, shape=(4, 1))
     weight_const = network.add_constant(
@@ -263,8 +253,7 @@ def test_inference_multiple_tensor_features(self):
     inference_runner = TensorRTEngineHandlerNumPy(
         min_batch_size=4, max_batch_size=4)
     builder = trt.Builder(LOGGER)
-    network = builder.create_network(
-        flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    network = builder.create_network(flags=_network_creation_flags())
     input_tensor = network.add_input(
         name="input", dtype=trt.float32, shape=(4, 2))
     weight_const = network.add_constant(
@@ -349,7 +338,30 @@ def test_namespace(self):
     inference_runner = TensorRTEngineHandlerNumPy(
         min_batch_size=4, max_batch_size=4)
     self.assertEqual(
-        'RunInferenceTensorRT', inference_runner.get_metrics_namespace())
+        'BeamML_TensorRT', inference_runner.get_metrics_namespace())
+
+  def test_version_check_matches_installed_api(self):
+    """The branch taken must match the API the installed TensorRT exposes.
+
+    TensorRT 10 removed the index based binding API in favour of the name
+    based tensor API. This guards against the version check drifting away
+    from the API it selects.
+    """
+    if _trt_major_version() >= 10:
+      self.assertTrue(hasattr(trt.ICudaEngine, 'num_io_tensors'))
+      self.assertTrue(hasattr(trt.IExecutionContext, 'execute_async_v3'))
+    else:
+      self.assertTrue(hasattr(trt.ICudaEngine, 'num_bindings'))
+      self.assertTrue(hasattr(trt.IExecutionContext, 'execute_async_v2'))
+
+  def test_network_creation_flags(self):
+    """Explicit batch must only be requested on TensorRT 8.x."""
+    if _trt_major_version() >= 10:
+      self.assertEqual(0, _network_creation_flags())
+    else:
+      self.assertEqual(
+          1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH),
+          _network_creation_flags())
 
 
 @pytest.mark.uses_tensorrt
@@ -381,7 +393,7 @@ def fake_inference_fn(batch, engine, inference_args=None):
           raise Exception(
               f'Loaded engine of type {type(engine)}, was ' +
               'expecting multi_process_shared engine')
-        from cuda import cuda
+        cuda = _import_cuda_driver()
         (
             engine,
             context,
@@ -400,7 +412,10 @@ def fake_inference_fn(batch, engine, inference_args=None):
                   np.ascontiguousarray(batch),
                   inputs[0]['size'],
                   stream))
-          context.execute_async_v2(gpu_allocations, stream)
+          if _trt_major_version() >= 10:
+            context.execute_async_v3(stream)
+          else:
+            context.execute_async_v2(gpu_allocations, stream)
           for output in range(len(cpu_allocations)):
             _assign_or_fail(
                 cuda.cuMemcpyDtoHAsync(

sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile

@@ -14,15 +14,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:23.05-py3
+# 26.06 ships TensorRT 11.0.0, CUDA 13.3 and Ubuntu 24.04 with Python 3.12.
+# TensorRT 8.x cannot target Blackwell GPUs such as the RTX Pro 6000 that
+# Dataflow now offers, so this image tracks the current TensorRT major version.
+# The Python version here must match the Beam SDK image copied in below.
+ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:26.06-py3
+ARG BEAM_SDK_IMAGE=apache/beam_python3.12_sdk:latest
 
-FROM ${BUILD_IMAGE} 
+FROM ${BEAM_SDK_IMAGE} AS beam_sdk
+
+FROM ${BUILD_IMAGE}
 
 ENV PATH="/usr/src/tensorrt/bin:${PATH}"
 
 WORKDIR /workspace
 
-COPY --from=apache/beam_python3.10_sdk:latest /opt/apache/beam /opt/apache/beam
+COPY --from=beam_sdk /opt/apache/beam /opt/apache/beam
 
 RUN pip install --upgrade pip \
     && pip install torch>=1.7.1 \
@@ -32,4 +39,4 @@ RUN pip install --upgrade pip \
     && pip install cuda-python
 
 ENTRYPOINT [ "/opt/apache/beam/boot" ]
-RUN apt-get update && apt-get install -y python3.10-venv 
+RUN apt-get update && apt-get install -y python3.12-venv