fix tensorrt dependencies issue (#3105)

pytorch · Aug 22, 2024 · 3e7be44 · 3e7be44
1 parent 51ccbea
commit 3e7be44
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 14 deletions.
diff --git a/examples/dynamo/requirements.txt b/examples/dynamo/requirements.txt
@@ -1,5 +1,4 @@
 cupy==13.1.0
-tensorrt==10.1.0
 torch>=2.4.0.dev20240503+cu121
 torch-tensorrt>=2.4.0.dev20240503+cu121
 triton==2.3.0
diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh
@@ -17,12 +17,18 @@ INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 # Install all the dependencies required for Torch-TensorRT
 pip uninstall -y torch torchvision
 pip install --force-reinstall --pre ${TORCH_TORCHVISION} --index-url ${INDEX_URL}
-pip install --pre -r tests/py/requirements.txt --use-deprecated legacy-resolver
-
 
 export TORCH_BUILD_NUMBER=$(python -c "import torch, urllib.parse as ul; print(ul.quote_plus(torch.__version__))")
 export TORCH_INSTALL_PATH=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))")
 
+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" \
+         pyproject.toml
+fi
+
 cat toolchains/ci_workspaces/MODULE.bazel.tmpl | envsubst > MODULE.bazel
 
 cat MODULE.bazel

diff --git a/packaging/pre_build_script_windows.sh b/packaging/pre_build_script_windows.sh
@@ -2,12 +2,16 @@ set -exou pipefail
 
 pip install -U numpy packaging pyyaml setuptools wheel
 
-# Install TRT from PyPI
-TRT_VERSION=$(python -c "import yaml; print(yaml.safe_load(open('dev_dep_versions.yml', 'r'))['__tensorrt_version__'])")
-pip install tensorrt==${TRT_VERSION} tensorrt-${CU_VERSION::4}-bindings==${TRT_VERSION} tensorrt-${CU_VERSION::4}-libs==${TRT_VERSION} --extra-index-url https://pypi.nvidia.com
-
 choco install bazelisk -y
 
+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" \
+         pyproject.toml
+fi
+
 #curl -Lo TensorRT.zip https://developer.download.nvidia.com/compute/machine-learning/tensorrt/10.0.1/zip/TensorRT-10.0.1.6.Windows10.win10.cuda-12.4.zip
 #unzip -o TensorRT.zip -d C:/
 TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt)
@@ -16,7 +20,6 @@ INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 # Install all the dependencies required for Torch-TensorRT
 pip uninstall -y torch torchvision
 pip install --force-reinstall --pre ${TORCH_TORCHVISION} --index-url ${INDEX_URL}
-pip install --pre -r tests/py/requirements.txt --use-deprecated legacy-resolver
 
 export CUDA_HOME="$(echo ${CUDA_PATH} | sed -e 's#\\#\/#g')"
 export TORCH_INSTALL_PATH="$(python -c "import torch, os; print(os.path.dirname(torch.__file__))" | sed -e 's#\\#\/#g')"

diff --git a/py/requirements.txt b/py/requirements.txt
@@ -6,4 +6,3 @@ torch>=2.5.0.dev,<2.6.0
 torchvision>=0.20.0.dev,<0.21.0
 --extra-index-url https://pypi.ngc.nvidia.com
 pyyaml
-tensorrt==10.1.0
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,7 +8,7 @@ requires = [
     "cffi>=1.15.1",
     "typing-extensions>=4.7.0",
     "future>=0.18.3",
-    "tensorrt==10.1.0",
+    "tensorrt-cu12==10.1.0",
     "torch >=2.5.0.dev,<2.6.0",
     "pybind11==2.6.2",
     "numpy",
@@ -55,9 +55,9 @@ keywords = [
 ]
 dependencies = [
     "torch >=2.5.0.dev,<2.6.0",
-    "tensorrt==10.1.0",
-    "tensorrt-cu12_bindings==10.1.0",
-    "tensorrt-cu12_libs==10.1.0",
+    "tensorrt-cu12==10.1.0",
+    "tensorrt-cu12-bindings==10.1.0",
+    "tensorrt-cu12-libs==10.1.0",
     "packaging>=23",
     "numpy",
     "typing-extensions>=4.7.0",

diff --git a/tests/py/requirements.txt b/tests/py/requirements.txt
@@ -7,7 +7,6 @@ parameterized>=0.2.0
 pytest>=8.2.1
 pytest-xdist>=3.6.1
 pyyaml
-tensorrt==10.1.0
 timm>=1.0.3
 transformers==4.40.2
 --extra-index-url https://pypi.nvidia.com