mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-21 19:18:55 +00:00
816e8cb2fb
### Description <!-- Describe your changes. --> * Update env to cuda 12.6/ubuntu 22.04 (ubuntu 20.04 uses outdated py38 by default) * Clean old trt8.6 test config ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. -->
284 lines
9 KiB
Python
284 lines
9 KiB
Python
# -------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
# --------------------------------------------------------------------------
|
|
"""
|
|
Builds an Ubuntu-based Docker image with TensorRT.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import pty
|
|
import shlex
|
|
import subprocess
|
|
import sys
|
|
|
|
TRT_DOCKER_FILES = {
|
|
"10.7_cuda11.8_cudnn8": "tools/ci_build/github/linux/docker/Dockerfile.ubuntu_cuda11_tensorrt10",
|
|
"10.7_cuda12.6_cudnn9": "tools/ci_build/github/linux/docker/Dockerfile.ubuntu_cuda12_tensorrt10",
|
|
"BIN": "tools/ci_build/github/linux/docker/Dockerfile.ubuntu_tensorrt_bin",
|
|
}
|
|
|
|
|
|
def run_cmd(cmd: list[str]) -> int | None:
|
|
"""
|
|
Runs a shell command and returns the process's return code.
|
|
|
|
:param cmd: List of command strings.
|
|
|
|
:return: The return code.
|
|
"""
|
|
|
|
escaped_cmd = " ".join(map(shlex.quote, cmd))
|
|
print(f"[CMD] {escaped_cmd}\n")
|
|
|
|
return pty.spawn(cmd)
|
|
|
|
|
|
def get_common_docker_build_args(args: argparse.Namespace) -> list[str]:
|
|
"""
|
|
Returns a list of common 'docker build' command-line arguments/options.
|
|
|
|
:param args: Arguments to the script.
|
|
|
|
:return: A list of common 'docker build' arguments.
|
|
"""
|
|
|
|
command = [
|
|
"--no-cache",
|
|
"-t",
|
|
f"{args.image_name}",
|
|
"--build-arg",
|
|
f"CMAKE_CUDA_ARCHITECTURES={args.cuda_arch}",
|
|
"--build-arg",
|
|
f"ONNXRUNTIME_BRANCH={args.branch}",
|
|
]
|
|
if args.use_tensorrt_oss_parser:
|
|
command.extend(
|
|
[
|
|
"--build-arg",
|
|
"PARSER_CONFIG=--use_tensorrt_oss_parser",
|
|
]
|
|
)
|
|
return command
|
|
|
|
|
|
def is_valid_ver_str(version: str, min_comps: int = 0, max_comps: int = 0) -> bool:
|
|
"""
|
|
Returns a boolean indicating if the string argument is a 'valid' version string.
|
|
|
|
:param version: String representing a version (e.g., 8.2.3).
|
|
:param min_comps: The minimum number of expected version components.
|
|
Set to 0 to ignore.
|
|
:param max_comps: The maximum number of expected version components.
|
|
Set to 0 to ignore.
|
|
|
|
:return: True if the string is valid (i.e., positive integers separated by dots)
|
|
"""
|
|
|
|
if not version:
|
|
return False
|
|
|
|
ver_nums = version.split(".")
|
|
num_comps = len(ver_nums)
|
|
|
|
if num_comps < min_comps:
|
|
return False
|
|
|
|
if num_comps > max_comps > 0:
|
|
return False
|
|
|
|
return all(num.isdecimal() for num in ver_nums)
|
|
|
|
|
|
def docker_build_trt(args: argparse.Namespace):
|
|
"""
|
|
Builds a Docker image that installs TensorRT from a public repository.
|
|
|
|
:param args: The arguments to this script.
|
|
"""
|
|
|
|
if args.trt_version not in TRT_DOCKER_FILES:
|
|
print(f"[ERROR]: TensorRT version '{args.trt_version}' is currently unsupported", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
docker_file = TRT_DOCKER_FILES[args.trt_version]
|
|
docker_file_path = os.path.normpath(os.path.join(args.repo_path, docker_file))
|
|
|
|
if not os.path.isfile(docker_file_path):
|
|
print(f"[ERROR]: Invalid docker file path '{docker_file_path}'", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
common_args = get_common_docker_build_args(args)
|
|
cmd_ret = run_cmd(["docker", "build", *common_args, "-f", f"{docker_file_path}", "."])
|
|
|
|
if cmd_ret != 0:
|
|
print(f"[ERROR]: docker build command failed with return code {cmd_ret}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def docker_build_trt_bin(args: argparse.Namespace):
|
|
"""
|
|
Builds a Docker image that installs TensorRT from a tar.gz package containing binaries.
|
|
See: https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-tar
|
|
|
|
:param args: The arguments to this script.
|
|
"""
|
|
|
|
docker_file = TRT_DOCKER_FILES["BIN"]
|
|
docker_file_path = os.path.normpath(os.path.join(args.repo_path, docker_file))
|
|
|
|
if not is_valid_ver_str(args.trt_version, 4, 4):
|
|
print(
|
|
"[ERROR]: Must specify a valid TensorRT version for binary TensorRT installs (e.g., 8.x.x.x)",
|
|
file=sys.stderr,
|
|
)
|
|
sys.exit(1)
|
|
|
|
if not is_valid_ver_str(args.tar_cuda_version, 2, 2):
|
|
print("[ERROR]: Must specify a valid CUDA version for binary TensorRT installs (e.g., 12.4)", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if not os.path.isfile(docker_file_path):
|
|
print(f"[ERROR]: Invalid docker file path '{docker_file_path}'", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if not args.trt_bins_dir or not os.path.isdir(args.trt_bins_dir):
|
|
print(f"[ERROR]: Invalid TensorRT bin directory '{args.trt_bins_dir}'", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
common_args = get_common_docker_build_args(args)
|
|
cmd_ret = run_cmd(
|
|
[
|
|
"docker",
|
|
"build",
|
|
*common_args,
|
|
"--build-arg",
|
|
f"TAR_TRT_VERSION={args.trt_version}",
|
|
"--build-arg",
|
|
f"TAR_CUDA_VERSION={args.tar_cuda_version}",
|
|
"--build-arg",
|
|
f"TRT_BINS_DIR={args.trt_bins_dir}",
|
|
"-f",
|
|
f"{docker_file_path}",
|
|
".",
|
|
]
|
|
)
|
|
|
|
if cmd_ret != 0:
|
|
print(f"[ERROR]: docker build command failed with return code {cmd_ret}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def overwrite_onnx_tensorrt_commit_id(commit_id):
|
|
"""
|
|
Overwrite onnx-tensorrt commit id in cmake/deps.txt.
|
|
"""
|
|
deps_file_path = "../../../../../../cmake/deps.txt"
|
|
line_index = None
|
|
zip_url = None
|
|
|
|
with open(deps_file_path) as file:
|
|
lines = file.readlines()
|
|
|
|
for i, line in enumerate(lines):
|
|
if line.startswith("onnx_tensorrt"):
|
|
parts = line.split(";")
|
|
zip_url = ";".join([parts[0], f"https://github.com/onnx/onnx-tensorrt/archive/{commit_id}.zip", parts[2]])
|
|
line_index = i
|
|
break
|
|
|
|
if line_index and zip_url:
|
|
wget_command = f"wget {zip_url.split(';')[1]} -O temp.zip"
|
|
subprocess.run(wget_command, shell=True, check=True)
|
|
|
|
sha1sum_command = "sha1sum temp.zip"
|
|
result = subprocess.run(sha1sum_command, shell=True, capture_output=True, text=True, check=True)
|
|
hash_value = result.stdout.split()[0]
|
|
|
|
lines[line_index] = zip_url.split(";")[0] + ";" + zip_url.split(";")[1] + ";" + hash_value + "\n"
|
|
|
|
with open(deps_file_path, "w") as file:
|
|
file.writelines(lines)
|
|
|
|
print(f"Updated deps.txt with new commit id {commit_id} and hash {hash_value}")
|
|
|
|
# Verify updated deps.txt
|
|
try:
|
|
with open(deps_file_path) as file:
|
|
lines = file.readlines()
|
|
for line in lines:
|
|
if line.startswith("onnx_tensorrt"):
|
|
print(line.strip())
|
|
break
|
|
except Exception as e:
|
|
print(f"Failed to read the file: {e}")
|
|
|
|
os.remove("temp.zip")
|
|
else:
|
|
print("onnx_tensorrt commit id overwrite failed, entry not found in deps.txt")
|
|
|
|
|
|
def parse_arguments() -> argparse.Namespace:
|
|
"""
|
|
Parses command-line arguments and returns an object with each argument as a field.
|
|
|
|
:return: An object whose fields represent the parsed command-line arguments.
|
|
"""
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-r", "--repo_path", required=True, help="Path to the onnxruntime repository")
|
|
parser.add_argument("-i", "--image_name", required=True, help="The resulting Docker image name")
|
|
parser.add_argument("-b", "--branch", default="main", help="Name of the onnxruntime git branch to checkout")
|
|
parser.add_argument(
|
|
"-t", "--trt_version", default="8.6.cuda_11_8_cudnn_8", help="TensorRT version (e.g., 8.6.cuda_11_8_cudnn_8)"
|
|
)
|
|
parser.add_argument("-a", "--cuda_arch", default="75", help="CUDA architecture (e.g., 75)")
|
|
parser.add_argument("-o", "--oss_parser_commit_id", default="", help="commit id of onnx-tensorrt")
|
|
|
|
# Command-line options for installing TensorRT from binaries.
|
|
parser.add_argument(
|
|
"--install_bin",
|
|
action="store_true",
|
|
default=False,
|
|
help="Enable to install TensorRT from tar.gz binary package",
|
|
)
|
|
parser.add_argument(
|
|
"--tar_cuda_version",
|
|
default="",
|
|
help="CUDA version (e.g., 12.4) used to find TensorRT EA binary tar.gz package",
|
|
)
|
|
parser.add_argument("--trt_bins_dir", default="", help="Directory containing TensorRT tar.gz package")
|
|
parser.add_argument(
|
|
"--use_tensorrt_oss_parser",
|
|
action="store_true",
|
|
default=False,
|
|
help="Use TensorRT OSS Parser",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
"""
|
|
Script entry point. Builds an Ubuntu-based Docker image with TensorRT.
|
|
|
|
:return: 0 on success, 1 on error.
|
|
"""
|
|
|
|
args = parse_arguments()
|
|
|
|
if args.install_bin:
|
|
docker_build_trt_bin(args)
|
|
else:
|
|
if args.oss_parser_commit_id != "":
|
|
overwrite_onnx_tensorrt_commit_id(args.oss_parser_commit_id)
|
|
docker_build_trt(args)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|