diff --git a/CHANGES.md b/CHANGES.md index 0c519076f0db..e680716ee5b7 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -94,6 +94,8 @@ * (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)). +* (Python) Fixed `TensorRTEngineHandlerNumPy` failing with `CUDA_ERROR_INVALID_VALUE` on models with a single-element input or output tensor, such as the `num_detections` output of an object detection model ([#36306](https://github.com/apache/beam/issues/36306)). ## Security Fixes diff --git a/sdks/python/apache_beam/examples/inference/README.md b/sdks/python/apache_beam/examples/inference/README.md index 5eed659d068c..e9fc8b55a742 100644 --- a/sdks/python/apache_beam/examples/inference/README.md +++ b/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). diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py index 333187301b29..b76f4312ab25 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py +++ b/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, @@ -195,17 +262,29 @@ 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) + 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( - cpu_allocations[output], + cpu_allocations[output].ctypes.data, outputs[output]['allocation'], outputs[output]['size'], stream)) diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py index 80a01b8f4d4c..65e493bb602e 100644 --- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py +++ b/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, @@ -119,17 +108,21 @@ 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) + 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( - cpu_allocations[output], + cpu_allocations[output].ctypes.data, outputs[output]['allocation'], outputs[output]['size'], stream)) @@ -189,8 +182,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 +219,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 +254,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 +339,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 +394,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 +407,21 @@ 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) + 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( - cpu_allocations[output], + cpu_allocations[output].ctypes.data, outputs[output]['allocation'], outputs[output]['size'], stream)) diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/README.md b/sdks/python/test-suites/containers/tensorrt_runinference/README.md index 99fbf83cbd74..fcacf736b81f 100644 --- a/sdks/python/test-suites/containers/tensorrt_runinference/README.md +++ b/sdks/python/test-suites/containers/tensorrt_runinference/README.md @@ -17,8 +17,76 @@ under the License. --> -# TensorRT Dockerfile for Beam +# TensorRT test resources for Beam -This directory contains the Dockerfiles required to run Beam pipelines that use TensorRT. +This directory contains the Dockerfile required to run Beam pipelines that use TensorRT, +and the script that rebuilds the TensorRT engines those tests load from GCS. + +## Container image To build the image, run `docker build -f tensor_rt.dockerfile -t us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest .` + +## Rebuilding the test engines + +The TensorRT tests load pre-built engines from `gs://apache-beam-ml/models/`: + +| Engine | Used by | +| --- | --- | +| `single_tensor_features_engine.trt` | `tensorrt_inference_test.py` | +| `multiple_tensor_features_engine.trt` | `tensorrt_inference_test.py` | +| `ssd_mobilenet_v2_320x320_coco17_tpu-8.trt` | the `tensorRTtests` Dataflow integration test | + +**A serialized TensorRT engine is not portable.** It can only be deserialized by the +same TensorRT major version and the same GPU architecture that built it. So these files +must be rebuilt whenever either of the following changes: + +* the TensorRT version in `tensor_rt.dockerfile`, or +* the GPU that the `tensorRTtests` task requests in + `sdks/python/test-suites/dataflow/common.gradle`. + +`build_test_engines.py` does that. It rebuilds each engine from the ONNX source already +staged next to it in the same bucket, so no new model sources are needed, and it verifies +each result by loading it back through `TensorRTEngineHandlerNumPy` — the small engines +against the exact values the unit tests assert, and the object detection engine against +the same COCO images the integration test uses. + +It needs a GPU, so it cannot run as part of the test suite. Run it in the same container +and on the same GPU type the tests use. As of writing that is +`nvcr.io/nvidia/tensorrt:26.06-py3` on an `nvidia-tesla-t4`. + +The host needs a driver new enough for that container (580 or later), plus Docker and +the NVIDIA container toolkit. On a GCE deep learning VM image the toolkit is already +present but Docker may not be: + +``` +sudo apt-get install -y docker.io +sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker +``` + +Then, from a directory containing `build_test_engines.py`: + +``` +sudo docker run --rm --gpus all -v "$PWD:/w" -w /w nvcr.io/nvidia/tensorrt:26.06-py3 bash -c "\ + pip install -q --break-system-packages 'apache-beam[gcp]' cuda-python pillow && \ + python3 build_test_engines.py --dest gs://YOUR_BUCKET/models" +``` + +`--break-system-packages` is required because the container's Python environment is +marked externally managed. Credentials are picked up from the VM's service account, so +no extra authentication step is needed. + +The verification step imports the model handler from the installed `apache-beam`, so +that version has to support the TensorRT major version you are building for. To verify +against an unreleased change, copy your working tree's +`apache_beam/ml/inference/tensorrt_inference.py` over the installed one inside the +container before running the script. + +Engines are written with a `_trt` suffix, for example +`single_tensor_features_engine_trt11.trt`, so the engines built by earlier TensorRT +versions stay in place and anyone on an older branch is unaffected. + +Point `--dest` at a bucket you can write to. Staging the results under +`gs://apache-beam-ml/models/` is a separate, deliberate step for someone with write +access to that bucket; note it has no object versioning, so an overwrite cannot be undone. + +Pass `--only ` to rebuild a single engine, and `--help` for the full options. diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py new file mode 100644 index 000000000000..deee092848dd --- /dev/null +++ b/sdks/python/test-suites/containers/tensorrt_runinference/build_test_engines.py @@ -0,0 +1,255 @@ +# +# 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. +# + +"""Rebuilds the TensorRT engines that the TensorRT tests load from GCS. + +A serialized TensorRT engine can only be deserialized by the TensorRT major +version and GPU architecture that built it. The engines the tests load are +therefore not portable, and have to be rebuilt whenever the TensorRT version +in the test container changes or the tests move to a different GPU. + +This script rebuilds them from the ONNX sources already staged alongside them, +and verifies each result by loading it back through the model handler the +tests use. It cannot run as part of the test suite because it needs a GPU. + +Run it in the same container and on the same GPU type the tests use. See +README.md in this directory for the exact commands. +""" + +# pytype: skip-file + +import argparse +import io +import logging +import os +import sys +import tempfile + +import numpy as np +import tensorrt as trt + +TRT_LOGGER = trt.Logger(trt.Logger.INFO) +TRT_MAJOR = int(trt.__version__.split('.')[0]) + +SOURCE = 'gs://apache-beam-ml/models' +COCO_IMAGES = [ + 'gs://apache-beam-ml/datasets/coco/raw-data/val2017/000000289594.jpg', + 'gs://apache-beam-ml/datasets/coco/raw-data/val2017/000000000139.jpg', +] + + +def _copy(src, dst): + """Copies between any two paths Beam's FileSystems understands. + + The TensorRT container has no gcloud CLI, but apache-beam[gcp] is installed + for the verification step anyway, so reuse it rather than shelling out. + """ + from apache_beam.io.filesystems import FileSystems + logging.info('copy %s -> %s', src, dst) + with FileSystems.open(src) as fin, FileSystems.create(dst) as fout: + while True: + chunk = fin.read(8 << 20) + if not chunk: + break + fout.write(chunk) + + +def _network_creation_flags(): + """Explicit batch is only a flag on TensorRT 8.x; it is the default after.""" + explicit_batch = getattr( + trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH', None) + if explicit_batch is None or TRT_MAJOR >= 10: + return 0 + return 1 << int(explicit_batch) + + +def build_engine(onnx_path, engine_path): + """Parses an ONNX file and serializes an engine for this GPU.""" + # The SSD MobileNet ONNX contains an EfficientNMS_TRT node, so the bundled + # plugins have to be registered before the parser will accept it. + trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") + + builder = trt.Builder(TRT_LOGGER) + network = builder.create_network(flags=_network_creation_flags()) + parser = trt.OnnxParser(network, TRT_LOGGER) + with open(onnx_path, 'rb') as f: + if not parser.parse(f.read()): + for i in range(parser.num_errors): + logging.error(parser.get_error(i)) + raise ValueError(f'Failed to parse {onnx_path}') + + config = builder.create_builder_config() + plan = builder.build_serialized_network(network, config) + if plan is None: + raise RuntimeError(f'Engine build produced no plan for {onnx_path}') + with open(engine_path, 'wb') as f: + f.write(plan) + logging.info('built %s (%d bytes)', engine_path, os.path.getsize(engine_path)) + + +def _handler(engine_path, batch_size): + from apache_beam.ml.inference.tensorrt_inference import ( + TensorRTEngineHandlerNumPy) + return TensorRTEngineHandlerNumPy( + min_batch_size=batch_size, + max_batch_size=batch_size, + engine_path=engine_path) + + +def verify_linear(engine_path, examples, expected): + """Checks a small linear engine against the values the unit tests assert.""" + handler = _handler(engine_path, len(examples)) + results = handler.run_inference(list(examples), handler.load_model()) + actual = np.array([r.inference[0] for r in results]).reshape(-1) + if not np.allclose(actual, np.asarray(expected).reshape(-1), atol=1e-4): + raise AssertionError(f'{engine_path}: expected {expected}, got {actual}') + logging.info('verified %s -> %s', os.path.basename(engine_path), actual) + + +def verify_ssd(engine_path): + """Runs the object detection engine on the images the Dataflow IT uses. + + The outputs are checked for the shape and ordering the example's + PostProcessor indexes by, and for at least one confident detection. + """ + from apache_beam.io.filesystems import FileSystems + from PIL import Image + + handler = _handler(engine_path, 1) + engine = handler.load_model() + + for image_path in COCO_IMAGES: + with FileSystems.open(image_path) as f: + image = Image.open(io.BytesIO(f.read())).convert('RGB') + # Mirrors preprocess_image() in the tensorrt_object_detection example. + image = image.resize((300, 300), resample=Image.Resampling.BILINEAR) + batch = [np.expand_dims(np.asarray(image, dtype=np.float32), axis=0)] + + inference = list(handler.run_inference(batch, engine))[0].inference + if len(inference) != 4: + raise AssertionError( + f'{engine_path}: expected 4 outputs, got {len(inference)}') + _, boxes, scores, classes = inference + if boxes.shape[-1] != 4 or scores.shape != classes.shape: + raise AssertionError( + f'{engine_path}: unexpected output shapes; the engine tensor order ' + f'must be num_detections, boxes, scores, classes. Got ' + f'{[np.asarray(o).shape for o in inference]}') + if float(np.max(scores)) < 0.3: + raise AssertionError( + f'{engine_path}: no confident detection for {image_path}; top score ' + f'was {float(np.max(scores)):.3f}') + logging.info( + 'verified %s on %s -> top score %.2f', + os.path.basename(engine_path), + os.path.basename(image_path), + float(np.max(scores))) + + +# The inputs and outputs below mirror the constants in +# apache_beam/ml/inference/tensorrt_inference_test.py, so a rebuilt engine is +# checked against exactly what the tests will assert. +SINGLE_EXAMPLES = [np.float32(v) for v in (1, 5, -3, 10)] +SINGLE_EXPECTED = [2.5, 10.5, -5.5, 20.5] # y = 2x + 0.5 + +MULTI_EXAMPLES = np.array([[1, 5], [3, 10], [-14, 0], [0.5, 0.5]], + dtype=np.float32) +MULTI_EXPECTED = [17.5, 36.5, -27.5, 3.0] # y = 2*x0 + 3*x1 + 0.5 + + +def verify_single(engine_path): + verify_linear(engine_path, SINGLE_EXAMPLES, SINGLE_EXPECTED) + + +def verify_multiple(engine_path): + verify_linear(engine_path, MULTI_EXAMPLES, MULTI_EXPECTED) + + +# Each entry rebuilds one staged .trt file from its staged .onnx source. +ENGINES = { + 'single_tensor_features_engine': { + 'onnx': 'single_tensor_features_model.onnx', + 'verify': verify_single, + }, + 'multiple_tensor_features_engine': { + 'onnx': 'multiple_tensor_features_model.onnx', + 'verify': verify_multiple, + }, + 'ssd_mobilenet_v2_320x320_coco17_tpu-8': { + 'onnx': 'ssd_mobilenet_v2_320x320_coco17_tpu-8.onnx', + 'verify': verify_ssd, + }, +} + + +def main(argv=None): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + '--dest', + required=True, + help='Where to write the rebuilt engines, e.g. gs://my-bucket/models or ' + 'a local directory. Staging them under the shared bucket is a separate, ' + 'deliberate step for someone with write access.') + parser.add_argument( + '--suffix', + default=None, + help='Name suffix for the rebuilt engines, so the engines built by ' + 'earlier TensorRT versions can stay in place. Defaults to _trt.') + parser.add_argument( + '--only', + action='append', + choices=sorted(ENGINES), + help='Rebuild only the named engine. May be repeated. Defaults to all.') + args = parser.parse_args(argv) + + if args.dest.rstrip('/') == SOURCE: + parser.error( + f'Refusing to write to {SOURCE}. That bucket has no object ' + 'versioning, so overwriting a staged engine could not be undone.') + + suffix = args.suffix if args.suffix is not None else f'_trt{TRT_MAJOR}' + names = args.only or sorted(ENGINES) + logging.info( + 'TensorRT %s, suffix %r, building: %s', + trt.__version__, + suffix, + ', '.join(names)) + + written = [] + with tempfile.TemporaryDirectory() as tmp: + for name in names: + spec = ENGINES[name] + onnx_local = os.path.join(tmp, spec['onnx']) + _copy(f'{SOURCE}/{spec["onnx"]}', onnx_local) + + engine_local = os.path.join(tmp, f'{name}{suffix}.trt') + build_engine(onnx_local, engine_local) + spec['verify'](engine_local) + + dest = f'{args.dest.rstrip("/")}/{os.path.basename(engine_local)}' + _copy(engine_local, dest) + written.append(dest) + + print('\nRebuilt and verified with TensorRT %s:' % trt.__version__) + for dest in written: + print(f' {dest}') + + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO, format='%(levelname)s %(message)s') + sys.exit(main()) diff --git a/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile b/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile index c1dc4deb6e69..0e86e79cd2d5 100644 --- a/sdks/python/test-suites/containers/tensorrt_runinference/tensor_rt.dockerfile +++ b/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 diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index abe7867181ec..905f03f8f703 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -651,22 +651,23 @@ task mockAPITests { } // add all RunInference E2E tests that run on DataflowRunner -// As of now, this test suite is enable in py310 suite as the base NVIDIA image used for Tensor RT -// contains Python 3.10. // TODO: https://github.com/apache/beam/issues/22651 project.tasks.register("inferencePostCommitIT") { dependsOn = [ - // TODO(https://github.com/apache/beam/issues/33078): restore the tensorRT tests once the staged - // model is fixed. - // 'tensorRTtests', 'vertexAIInferenceTest', 'geminiInferenceTest', 'mockAPITests', ] } +// The base NVIDIA image used for TensorRT contains Python 3.12, so the TensorRT +// suite belongs here rather than in the py310 suite above. project.tasks.register("inferencePostCommitITPy312") { dependsOn = [ + // TODO(https://github.com/apache/beam/issues/33078): restore the tensorRT tests once the staged + // model is rebuilt. A serialized engine can only be read by the TensorRT major version that + // built it, and the staged engine was built with TensorRT 8.x. + // 'tensorRTtests', 'vllmTests', ] } diff --git a/website/www/site/content/en/documentation/ml/tensorrt-runinference.md b/website/www/site/content/en/documentation/ml/tensorrt-runinference.md index 4bae2d3ba7ce..78bdd95b30fa 100644 --- a/website/www/site/content/en/documentation/ml/tensorrt-runinference.md +++ b/website/www/site/content/en/documentation/ml/tensorrt-runinference.md @@ -66,7 +66,7 @@ trtexec --onnx= --saveEngine= To use `trtexec`, follow the steps in the blog post [Simplifying and Accelerating Machine Learning Predictions in Apache Beam with NVIDIA TensorRT](https://developer.nvidia.com/blog/simplifying-and-accelerating-machine-learning-predictions-in-apache-beam-with-nvidia-tensorrt/). The post explains how to build a docker image from a DockerFile that can be used for conversion. We use the following Docker file, which is similar to the file used in the blog post: ``` -ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:22.05-py3 +ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:26.06-py3 FROM ${BUILD_IMAGE} @@ -75,11 +75,11 @@ ENV PATH="/usr/src/tensorrt/bin:${PATH}" WORKDIR /workspace RUN apt-get update -y && apt-get install -y python3-venv -RUN pip install --no-cache-dir apache-beam[gcp]==2.44.0 -COPY --from=apache/beam_python3.8_sdk:2.44.0 /opt/apache/beam /opt/apache/beam +RUN pip install --no-cache-dir apache-beam[gcp]==2.76.0 +COPY --from=apache/beam_python3.12_sdk:2.76.0 /opt/apache/beam /opt/apache/beam RUN pip install --upgrade pip \ - && pip install torch==1.13.1 \ + && pip install torch \ && pip install torchvision>=0.8.2 \ && pip install pillow>=8.0.0 \ && pip install transformers>=4.18.0 \