Skip to content

Commit

Permalink
build: Support RHEL ORT TensorRT Execution Provider (#285)
Browse files Browse the repository at this point in the history
* Enable ORT TRT extension for RHEL
  • Loading branch information
fpetrini15 authored Nov 20, 2024
1 parent 56ce169 commit bc4a7cb
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions tools/gen_ort_dockerfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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__":
Expand Down Expand Up @@ -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."
)
Expand Down Expand Up @@ -661,4 +666,4 @@ def preprocess_gpu_flags():
FLAGS.ort_openvino = None
dockerfile_for_windows(FLAGS.output)
else:
dockerfile_for_linux(FLAGS.output)
dockerfile_for_linux(FLAGS.output)

0 comments on commit bc4a7cb

Please sign in to comment.