From bc4a7cb0c085d7983bb5e1846b2e5dd8efe1b0ea Mon Sep 17 00:00:00 2001 From: Francesco Petrini Date: Tue, 19 Nov 2024 16:07:17 -0800 Subject: [PATCH] build: Support RHEL ORT TensorRT Execution Provider (#285) * Enable ORT TRT extension for RHEL --- tools/gen_ort_dockerfile.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/tools/gen_ort_dockerfile.py b/tools/gen_ort_dockerfile.py index efae8cd..9b511b6 100755 --- a/tools/gen_ort_dockerfile.py +++ b/tools/gen_ort_dockerfile.py @@ -105,7 +105,7 @@ def dockerfile_for_linux(output_file): df += """ # The manylinux container defaults to Python 3.7, but some feature installation # requires a higher version. -ARG PYVER=3.10 +ARG PYVER=3.12 ENV PYTHONPATH=/opt/python/v RUN ln -sf /opt/python/cp${PYVER/./}* ${PYTHONPATH} @@ -572,7 +572,13 @@ def preprocess_gpu_flags(): print("error: linux build requires --cudnn-home and --cuda-home") if FLAGS.tensorrt_home is None: - FLAGS.tensorrt_home = "/usr/src/tensorrt" + if target_platform() == "rhel": + if platform.machine().lower() == "aarch64": + FLAGS.tensorrt_home = "/usr/local/cuda/targets/sbsa-linux/" + else: + FLAGS.tensorrt_home = "/usr/local/cuda/targets/x86_64-linux/" + else: + FLAGS.tensorrt_home = "/usr/src/tensorrt" if __name__ == "__main__": @@ -602,9 +608,8 @@ def preprocess_gpu_flags(): "--target-platform", required=False, default=None, - help='Target for build, can be "linux", "windows" or "igpu". If not specified, build targets the current platform.', + help='Target for build, can be "linux", "windows", "rhel", or "igpu". If not specified, build targets the current platform.', ) - parser.add_argument( "--cuda-version", type=str, required=False, help="Version for CUDA." ) @@ -661,4 +666,4 @@ def preprocess_gpu_flags(): FLAGS.ort_openvino = None dockerfile_for_windows(FLAGS.output) else: - dockerfile_for_linux(FLAGS.output) \ No newline at end of file + dockerfile_for_linux(FLAGS.output)