Support TensorRT 10 and later in TensorRTEngineHandlerNumPy

CHANGES.md

@@ -85,6 +85,10 @@
 * Portable Java SDK now encodes SchemaCoders in a portable way ([#34672](https://github.com/apache/beam/issues/34672)).
   - Original custom Java coder encoding can still be obtained using [StreamingOptions.setUpdateCompatibilityVersion("2.76")](https://github.com/apache/beam/blob/2cf0930e7ae1aa389c26ce6639b584877a3e31d9/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java#L47) ([#34672](https://github.com/apache/beam/issues/34672)).
   - Fixes ([#36496](https://github.com/apache/beam/issues/36496)), ([#30276](https://github.com/apache/beam/issues/30276)), ([#29245](https://github.com/apache/beam/issues/29245)).
+* (Python) `TensorRTEngineHandlerNumPy` now requires TensorRT 10 or later. TensorRT 8.x is no longer supported, since TensorRT 10 removed the engine binding API the handler was written against ([#36306](https://github.com/apache/beam/issues/36306)).
+  - Engines serialized by TensorRT 8.x must be rebuilt, as an engine can only be deserialized by the major version that built it.
+  - TensorRT 10 and later require a GPU with compute capability 7.5 or higher, which excludes NVIDIA Pascal and Volta GPUs.
+  - If dropping TensorRT 8.x support is a hard blocker for you, please comment on ([#36306](https://github.com/apache/beam/issues/36306)).
 
 ## Deprecations
 
@@ -98,6 +102,7 @@
 * (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)).
 * (Go) Fixed GCS glob matching silently dropping objects when the glob pattern contains multi-byte characters ([#39969](https://github.com/apache/beam/issues/39969)).
+* (Python) Fixed `TensorRTEngineHandlerNumPy` failing with `CUDA_ERROR_INVALID_VALUE` on models with a single-element input or output tensor ([#36306](https://github.com/apache/beam/issues/36306)).
 
 ## Security Fixes
 

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

@@ -83,13 +83,19 @@ pip install torch==1.10.0
 ### TensorRT dependencies
 
 The RunInference API supports TensorRT SDK for high-performance deep learning inference with NVIDIA GPUs.
-To use TensorRT locally, we suggest an environment with TensorRT >= 8.0.1. Install TensorRT as per the
+To use TensorRT locally, we suggest an environment with TensorRT >= 10.0. Install TensorRT as per the
 [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.
 
+TensorRT 10 or later is required. 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. TensorRT 10 and later also require a GPU with
+compute capability 7.5 or higher, for example, T4, L4, A100. The NVIDIA Pascal and Volta GPUs
+such as the Tesla P4, P100 and V100 are no longer supported.
+
 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

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import functools
 import logging
 import threading
 from collections.abc import Callable
@@ -48,9 +49,58 @@
     'runner has tensorrt dependencies installed.'
   LOGGER.warning(msg)
 
+MIN_TRT_MAJOR_VERSION = 10
+
+
+@functools.lru_cache(maxsize=1)
+def _check_trt_version() -> None:
+  """Fails fast if the installed TensorRT is older than we support.
+
+  TensorRT 10 removed the index based "binding" API this module used to be
+  written against. Without this check the failure surfaces as an obscure
+  AttributeError deep inside engine setup.
+
+  Cached rather than checked at import time because the module is importable
+  without TensorRT, so that jobs can be submitted from a machine that does not
+  have it installed.
+  """
+  import tensorrt as trt
+  try:
+    major = int(trt.__version__.split('.')[0])
+  except (AttributeError, IndexError, ValueError):
+    # Fall back to probing for an attribute that only exists from 10 onwards.
+    major = 10 if hasattr(trt.ICudaEngine, 'num_io_tensors') else 8
+  if major < MIN_TRT_MAJOR_VERSION:
+    raise RuntimeError(
+        'RunInference requires TensorRT %d or later, but found %s. Support '
+        'for TensorRT 8.x was removed because TensorRT 10 replaced the '
+        'engine binding API this handler depends on.' %
+        (MIN_TRT_MAJOR_VERSION, getattr(trt, '__version__', 'unknown')))
+
+
+@functools.lru_cache(maxsize=1)
+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.
+
+  Cached because this is called from _assign_or_fail, which runs on every
+  CUDA call.
+  """
+  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
+  # Checked before deserializing, because an engine built by a newer TensorRT
+  # fails to deserialize with an opaque error that hides the real cause.
+  _check_trt_version()
   file = FileSystems.open(engine_path, 'rb')
   runtime = trt.Runtime(TRT_LOGGER)
   engine = runtime.deserialize_cuda_engine(file.read())
@@ -60,9 +110,11 @@ def _load_engine(engine_path):
 
 def _load_onnx(onnx_path):
   import tensorrt as trt
+  _check_trt_version()
   builder = trt.Builder(TRT_LOGGER)
-  network = builder.create_network(
-      flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+  # Explicit batch is the only supported mode from TensorRT 10 onwards, so no
+  # network creation flags are needed.
+  network = builder.create_network()
   parser = trt.OnnxParser(network, TRT_LOGGER)
   with FileSystems.open(onnx_path) as f:
     if not parser.parse(f.read()):
@@ -85,7 +137,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 +163,8 @@ def __init__(self, engine: trt.ICudaEngine):
       engine: trt.ICudaEngine object that contains TensorRT engine
     """
     import tensorrt as trt
-    from cuda import cuda
+    _check_trt_version()
+    cuda = _import_cuda_driver()
     self.engine = engine
     self.context = engine.create_execution_context()
     self.context_lock = threading.RLock()
@@ -120,19 +173,12 @@ 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)
+    # Setup I/O tensors. 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 = {
@@ -144,7 +190,8 @@ def __init__(self, engine: trt.ICudaEngine):
           'size': size
       }
       self.gpu_allocations.append(allocation)
-      if self.engine.binding_is_input(i):
+      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)
@@ -182,7 +229,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,
@@ -195,17 +242,26 @@ def _default_tensorRT_inference_fn(
 
   # Process I/O and execute the network
   with context_lock:
+    # Host buffers are passed as explicit addresses rather than as arrays.
+    # A numpy array holding exactly one element is coerced to a scalar, which
+    # is then read as a null host pointer and fails with CUDA_ERROR_INVALID_
+    # VALUE. Single element outputs are common, for example the num_detections
+    # output of an object detection model.
+    # host_input must stay referenced until the stream is synchronized below,
+    # because the copy is asynchronous.
+    host_input = np.ascontiguousarray(batch)
     _assign_or_fail(
         cuda.cuMemcpyHtoDAsync(
             inputs[0]['allocation'],
-            np.ascontiguousarray(batch),
+            host_input.ctypes.data,
             inputs[0]['size'],
             stream))
-    context.execute_async_v2(gpu_allocations, stream)
+    # Tensor addresses were bound when the engine was created.
+    context.execute_async_v3(stream)
     for output in range(len(cpu_allocations)):
       _assign_or_fail(
           cuda.cuMemcpyDtoHAsync(
-              cpu_allocations[output],
+              cpu_allocations[output].ctypes.data,
               outputs[output]['allocation'],
               outputs[output]['size'],
               stream))

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

@@ -19,6 +19,7 @@
 
 import os
 import unittest
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -37,6 +38,9 @@
   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 _check_trt_version
+  from apache_beam.ml.inference.tensorrt_inference import _import_cuda_driver
 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,
@@ -119,17 +108,18 @@ def _custom_tensorRT_inference_fn(batch, engine, inference_args):
 
   # Process I/O and execute the network
   with context_lock:
+    host_input = np.ascontiguousarray(batch)
     _assign_or_fail(
         cuda.cuMemcpyHtoDAsync(
             inputs[0]['allocation'],
-            np.ascontiguousarray(batch),
+            host_input.ctypes.data,
             inputs[0]['size'],
             stream))
-    context.execute_async_v2(gpu_allocations, stream)
+    context.execute_async_v3(stream)
     for output in range(len(cpu_allocations)):
       _assign_or_fail(
           cuda.cuMemcpyDtoHAsync(
-              cpu_allocations[output],
+              cpu_allocations[output].ctypes.data,
               outputs[output]['allocation'],
               outputs[output]['size'],
               stream))
@@ -189,8 +179,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()
     input_tensor = network.add_input(
         name="input", dtype=trt.float32, shape=(4, 1))
     weight_const = network.add_constant(
@@ -227,8 +216,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()
     input_tensor = network.add_input(
         name="input", dtype=trt.float32, shape=(4, 1))
     weight_const = network.add_constant(
@@ -263,8 +251,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()
     input_tensor = network.add_input(
         name="input", dtype=trt.float32, shape=(4, 2))
     weight_const = network.add_constant(
@@ -349,7 +336,26 @@ 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_supported_tensorrt_exposes_expected_api(self):
+    """The installed TensorRT must expose the API this module is written to.
+
+    TensorRT 10 removed the index based binding API in favour of the name
+    based tensor API. _check_trt_version() rejects anything older, so a passing
+    version check and a missing API would mean the two have drifted apart.
+    """
+    _check_trt_version()
+    self.assertTrue(hasattr(trt.ICudaEngine, 'num_io_tensors'))
+    self.assertTrue(hasattr(trt.IExecutionContext, 'execute_async_v3'))
+
+  def test_version_check_rejects_unsupported_tensorrt(self):
+    """An unsupported TensorRT must fail with a clear message."""
+    with mock.patch.object(trt, '__version__', '8.6.1'):
+      _check_trt_version.cache_clear()
+      with self.assertRaisesRegex(RuntimeError, 'requires TensorRT 10'):
+        _check_trt_version()
+    _check_trt_version.cache_clear()
 
 
 @pytest.mark.uses_tensorrt
@@ -381,7 +387,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,
@@ -394,17 +400,18 @@ def fake_inference_fn(batch, engine, inference_args=None):
 
         # Process I/O and execute the network
         with context_lock:
+          host_input = np.ascontiguousarray(batch)
           _assign_or_fail(
               cuda.cuMemcpyHtoDAsync(
                   inputs[0]['allocation'],
-                  np.ascontiguousarray(batch),
+                  host_input.ctypes.data,
                   inputs[0]['size'],
                   stream))
-          context.execute_async_v2(gpu_allocations, stream)
+          context.execute_async_v3(stream)
           for output in range(len(cpu_allocations)):
             _assign_or_fail(
                 cuda.cuMemcpyDtoHAsync(
-                    cpu_allocations[output],
+                    cpu_allocations[output].ctypes.data,
                     outputs[output]['allocation'],
                     outputs[output]['size'],
                     stream))