upgrade tensorrt dependency to >=10.3.0,<=10.6.0 (#3286)

lanluo-nvidia · web-flow · commit 7ff3ad66ecf1 · 2024-12-12T13:18:25.000-08:00
diff --git a/.github/workflows/docker_builder.yml b/.github/workflows/docker_builder.yml
@@ -7,6 +7,7 @@ on:
       - main
       - nightly
       - release/*
+  workflow_dispatch:
 
 # If pushes to main are made in rapid succession,
 # cancel existing docker builds and use newer commits
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -79,20 +79,20 @@ http_archive(
 http_archive(
     name = "tensorrt",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "adff1cd5abe5d87013806172351e58fd024e5bf0fc61d49ef4b84cd38ed99081",
-    strip_prefix = "TensorRT-10.3.0.26",
+    sha256 = "33d3c2f3f4c84dc7991a4337a6fde9ed33f5c8e5c4f03ac2eb6b994a382b03a0",
+    strip_prefix = "TensorRT-10.6.0.26",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/tars/TensorRT-10.3.0.26.Linux.x86_64-gnu.cuda-12.5.tar.gz",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.6.0/tars/TensorRT-10.6.0.26.Linux.x86_64-gnu.cuda-12.6.tar.gz",
     ],
 )
 
 http_archive(
     name = "tensorrt_win",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "2bb4bcb79e8c33575816d874b0512ea28c302af1c06ee6d224da71aa182f75e0",
-    strip_prefix = "TensorRT-10.3.0.26",
+    sha256 = "6c6d92c108a1b3368423e8f69f08d31269830f1e4c9da43b37ba34a176797254",
+    strip_prefix = "TensorRT-10.6.0.26",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/zip/TensorRT-10.3.0.26.Windows.win10.cuda-12.5.zip",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.6.0/zip/TensorRT-10.6.0.26.Windows.win10.cuda-12.6.zip",
     ],
 )
 
diff --git a/README.md b/README.md
@@ -119,7 +119,7 @@ These are the following dependencies used to verify the testcases. Torch-TensorR
 - Bazel 6.3.2
 - Libtorch 2.5.0.dev (latest nightly) (built with CUDA 12.4)
 - CUDA 12.4
-- TensorRT 10.3.0.26
+- TensorRT 10.6.0.26
 
 ## Deprecation Policy
 
diff --git a/dev_dep_versions.yml b/dev_dep_versions.yml
@@ -1,2 +1,2 @@
 __cuda_version__: "12.4"
-__tensorrt_version__: "10.3.0"
+__tensorrt_version__: ">=10.3.0,<=10.6.0"
diff --git a/docker/README.md b/docker/README.md
@@ -17,14 +17,14 @@ Note: By default the container uses the `pre-cxx11-abi` version of Torch + Torch
 
 ### Instructions
 
-- The example below uses TensorRT 10.3.0.26
+- The example below uses TensorRT 10.6.0.26
 - See <a href="https://github.com/pytorch/TensorRT#dependencies">dependencies</a> for a list of current default dependencies.
 
 > From root of Torch-TensorRT repo
 
 Build:
 ```
-DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=10.3.0 -f docker/Dockerfile -t torch_tensorrt:latest .
+DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=10.6.0 -f docker/Dockerfile -t torch_tensorrt:latest .
 ```
 
 Run:
diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh
@@ -23,18 +23,20 @@ export TORCH_INSTALL_PATH=$(python -c "import torch, os; print(os.path.dirname(t
 
 if [[ ${TENSORRT_VERSION} != "" ]]; then
   # this is the upgraded TensorRT version, replace current tensorrt version to the upgrade tensorRT version in the pyproject.toml
+  # example: __tensorrt_version__: ">=10.3.0,<=10.6.0"
+  # replace: tensorrt-cu12>=10.3.0,<=10.6.0 to tensorrt-cu12==10.7.0
   current_version=$(cat dev_dep_versions.yml | grep __tensorrt_version__ | sed 's/__tensorrt_version__: //g' | sed 's/"//g')
-  sed -i -e "s/tensorrt-cu12==${current_version}/tensorrt-cu12==${TENSORRT_VERSION}/g" \
-         -e "s/tensorrt-cu12-bindings==${current_version}/tensorrt-cu12-bindings==${TENSORRT_VERSION}/g" \
-         -e "s/tensorrt-cu12-libs==${current_version}/tensorrt-cu12-libs==${TENSORRT_VERSION}/g" \
+  sed -i -e "s/tensorrt-cu12${current_version}/tensorrt-cu12==${TENSORRT_VERSION}/g" \
+         -e "s/tensorrt-cu12-bindings${current_version}/tensorrt-cu12-bindings==${TENSORRT_VERSION}/g" \
+         -e "s/tensorrt-cu12-libs${current_version}/tensorrt-cu12-libs==${TENSORRT_VERSION}/g" \
           pyproject.toml
 fi
 
 if [[ "${CU_VERSION::4}" < "cu12" ]]; then
   # replace dependencies from tensorrt-cu12-bindings/libs to tensorrt-cu11-bindings/libs
-  sed -i -e "s/tensorrt-cu12==/tensorrt-${CU_VERSION::4}==/g" \
-         -e "s/tensorrt-cu12-bindings==/tensorrt-${CU_VERSION::4}-bindings==/g" \
-         -e "s/tensorrt-cu12-libs==/tensorrt-${CU_VERSION::4}-libs==/g" \
+  sed -i -e "s/tensorrt-cu12/tensorrt-${CU_VERSION::4}/g" \
+         -e "s/tensorrt-cu12-bindings/tensorrt-${CU_VERSION::4}-bindings/g" \
+         -e "s/tensorrt-cu12-libs/tensorrt-${CU_VERSION::4}-libs/g" \
          pyproject.toml
 fi
 
diff --git a/packaging/pre_build_script_windows.sh b/packaging/pre_build_script_windows.sh
@@ -8,23 +8,23 @@ echo TENSORRT_VERSION=${TENSORRT_VERSION}
 
 if [[ ${TENSORRT_VERSION} != "" ]]; then
   # this is the upgraded TensorRT version, replace current tensorrt version to the upgrade tensorRT version in the pyproject.toml
+  # example: __tensorrt_version__: ">=10.3.0,<=10.6.0"
+  # replace: tensorrt-cu12>=10.3.0,<=10.6.0 to tensorrt-cu12==10.7.0
   current_version=$(cat dev_dep_versions.yml | grep __tensorrt_version__ | sed 's/__tensorrt_version__: //g' | sed 's/"//g')
-  sed -i -e "s/tensorrt-cu12==${current_version}/tensorrt-cu12==${TENSORRT_VERSION}/g" \
-         -e "s/tensorrt-cu12-bindings==${current_version}/tensorrt-cu12-bindings==${TENSORRT_VERSION}/g" \
-         -e "s/tensorrt-cu12-libs==${current_version}/tensorrt-cu12-libs==${TENSORRT_VERSION}/g" \
+  sed -i -e "s/tensorrt-cu12${current_version}/tensorrt-cu12==${TENSORRT_VERSION}/g" \
+         -e "s/tensorrt-cu12-bindings${current_version}/tensorrt-cu12-bindings==${TENSORRT_VERSION}/g" \
+         -e "s/tensorrt-cu12-libs${current_version}/tensorrt-cu12-libs==${TENSORRT_VERSION}/g" \
           pyproject.toml
 fi
 
 if [[ "${CU_VERSION::4}" < "cu12" ]]; then
   # replace dependencies from tensorrt-cu12-bindings/libs to tensorrt-cu11-bindings/libs
-  sed -i -e "s/tensorrt-cu12==/tensorrt-${CU_VERSION::4}==/g" \
-         -e "s/tensorrt-cu12-bindings==/tensorrt-${CU_VERSION::4}-bindings==/g" \
-         -e "s/tensorrt-cu12-libs==/tensorrt-${CU_VERSION::4}-libs==/g" \
+  sed -i -e "s/tensorrt-cu12/tensorrt-${CU_VERSION::4}/g" \
+         -e "s/tensorrt-cu12-bindings/tensorrt-${CU_VERSION::4}-bindings/g" \
+         -e "s/tensorrt-cu12-libs/tensorrt-${CU_VERSION::4}-libs/g" \
          pyproject.toml
 fi
 
-#curl -Lo TensorRT.zip https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/zip/TensorRT-10.3.0.26.Windows.win10.cuda-12.5.zip
-#unzip -o TensorRT.zip -d C:/
 TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt)
 INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,7 +8,9 @@ requires = [
     "cffi>=1.15.1",
     "typing-extensions>=4.7.0",
     "future>=0.18.3",
-    "tensorrt-cu12==10.3.0",
+    "tensorrt-cu12>=10.3.0,<=10.6.0",
+    "tensorrt-cu12-bindings>=10.3.0,<=10.6.0",
+    "tensorrt-cu12-libs>=10.3.0,<=10.6.0",
     "torch>=2.6.0.dev,<2.7.0",
     "pybind11==2.6.2",
     "numpy",
@@ -55,9 +57,9 @@ keywords = [
 ]
 dependencies = [
     "torch>=2.6.0.dev,<2.7.0",
-    "tensorrt-cu12==10.3.0",
-    "tensorrt-cu12-bindings==10.3.0",
-    "tensorrt-cu12-libs==10.3.0",
+    "tensorrt-cu12>=10.3.0,<=10.6.0",
+    "tensorrt-cu12-bindings>=10.3.0,<=10.6.0",
+    "tensorrt-cu12-libs>=10.3.0,<=10.6.0",
     "packaging>=23",
     "numpy",
     "typing-extensions>=4.7.0",
diff --git a/toolchains/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl
@@ -67,20 +67,20 @@ http_archive(
 http_archive(
     name = "tensorrt",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "adff1cd5abe5d87013806172351e58fd024e5bf0fc61d49ef4b84cd38ed99081",
-    strip_prefix = "TensorRT-10.3.0.26",
+    sha256 = "33d3c2f3f4c84dc7991a4337a6fde9ed33f5c8e5c4f03ac2eb6b994a382b03a0",
+    strip_prefix = "TensorRT-10.6.0.26",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/tars/TensorRT-10.3.0.26.Linux.x86_64-gnu.cuda-12.5.tar.gz",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.6.0/tars/TensorRT-10.6.0.26.Linux.x86_64-gnu.cuda-12.6.tar.gz",
     ],
 )
 
 http_archive(
     name = "tensorrt_win",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "2bb4bcb79e8c33575816d874b0512ea28c302af1c06ee6d224da71aa182f75e0",
-    strip_prefix = "TensorRT-10.3.0.26",
+    sha256 = "6c6d92c108a1b3368423e8f69f08d31269830f1e4c9da43b37ba34a176797254",
+    strip_prefix = "TensorRT-10.6.0.26",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/zip/TensorRT-10.3.0.26.Windows.win10.cuda-12.5.zip",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.6.0/zip/TensorRT-10.6.0.26.Windows.win10.cuda-12.6.zip",
     ],
 )
 
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`__cuda_version__: "12.4"`
`2`		`-__tensorrt_version__: "10.3.0"`
	`2`	`+__tensorrt_version__: ">=10.3.0,<=10.6.0"`