fix the libtorch version mismatch issue (#3086)

Author: lanluo-nvidia

.github/scripts/install-torch-tensorrt.sh

@@ -1,4 +1,4 @@
-set -eou pipefail
+set -exou pipefail
 
 TORCH_TORCHVISION=$(grep "^torch" ${PWD}/py/requirements.txt)
 INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}

.github/workflows/release-wheel-linux.yml

@@ -199,7 +199,7 @@ jobs:
           mkdir release/tarball
           mkdir release/wheel
           if [[ "${{ inputs.cxx11-tarball-release }}" == "true" ]]; then
-              bazel build //:libtorchtrt --compilation_mode opt --config=default
+              bazel build //:libtorchtrt --compilation_mode opt --config=linux
               cp bazel-bin/libtorchtrt.tar.gz \
               release/tarball/libtorchtrt-${BUILD_VERSION}-tensorrt${TRT_VERSION}-cuda${CU_VERSION:2}-libtorch${PYTORCH_VERSION}-x86_64-linux.tar.gz
           else

.gitignore

@@ -73,4 +73,5 @@ wheelhouse/
 tests/py/dynamo/models/*.ts
 tests/py/dynamo/models/*.ep
 *.deb
-*.tar.xz
+*.tar.xz
+MODULE.bazel.lock

MODULE.bazel.lock

@@ -1,14 +1,29 @@
 #!/bin/bash
 
+set -x
+
 # Install dependencies
 python3 -m pip install pyyaml
+
 yum install -y ninja-build gettext
+
 wget https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64 \
     && mv bazelisk-linux-amd64 /usr/bin/bazel \
     && chmod +x /usr/bin/bazel
 
+TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt)
+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__))")
 
 cat toolchains/ci_workspaces/MODULE.bazel.tmpl | envsubst > MODULE.bazel
+
+cat MODULE.bazel
 export CI_BUILD=1

packaging/pre_build_script.sh

@@ -1,4 +1,4 @@
-set -eou pipefail
+set -exou pipefail
 
 pip install -U numpy packaging pyyaml setuptools wheel
 
@@ -10,9 +10,18 @@ choco install bazelisk -y
 
 #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)
+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')"
 
 cat toolchains/ci_workspaces/MODULE.bazel.tmpl | envsubst > MODULE.bazel
 
+cat MODULE.bazel
 echo "RELEASE=1" >> ${GITHUB_ENV}

packaging/pre_build_script_windows.sh

@@ -534,4 +534,4 @@ def save(
                     exp_program = torch.export.export(
                         module, tuple(arg_inputs), kwargs=kwarg_inputs, strict=False
                     )
-                    torch.export.save(exp_program, file_path)
+                    torch.export.save(exp_program, file_path)

py/torch_tensorrt/_compile.py

@@ -3,6 +3,7 @@
 import logging
 from typing import Any, List, Optional, Sequence
 
+import tensorrt as trt
 import torch
 from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
 from torch_tensorrt._Device import Device
@@ -17,8 +18,6 @@
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
 from torch_tensorrt.dynamo.utils import get_torch_inputs
 
-import tensorrt as trt
-
 logger = logging.getLogger(__name__)
 
 

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -60,20 +60,6 @@ http_archive(
     urls = ["https://download.pytorch.org/libtorch/${CHANNEL}/${CU_VERSION}/libtorch-cxx11-abi-shared-with-deps-latest.zip"],
 )
 
-http_archive(
-    name = "libtorch_pre_cxx11_abi",
-    build_file = "@//third_party/libtorch:BUILD",
-    strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/${CHANNEL}/${CU_VERSION}/libtorch-shared-with-deps-latest.zip"],
-)
-
-http_archive(
-    name = "libtorch_win",
-    build_file = "@//third_party/libtorch:BUILD",
-    strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/${CHANNEL}/${CU_VERSION}/libtorch-win-shared-with-deps-latest.zip"],
-)
-
 # Download these tarballs manually from the NVIDIA website
 # Either place them in the distdir directory in third_party and use the --distdir flag
 # or modify the urls to "file:///<PATH TO TARBALL>/<TARBALL NAME>.tar.gz
@@ -110,17 +96,17 @@ http_archive(
 # x86_64 python distribution. If using NVIDIA's version just point to the root of the package
 # for both versions here and do not use --config=pre-cxx11-abi
 
-#new_local_repository(
-#    name = "libtorch",
-#    path = "${TORCH_INSTALL_PATH}",
-#    build_file = "third_party/libtorch/BUILD"
-#)
+new_local_repository(
+    name = "libtorch_win",
+    path = "${TORCH_INSTALL_PATH}",
+    build_file = "third_party/libtorch/BUILD"
+)
 
-#new_local_repository(
-#    name = "libtorch_pre_cxx11_abi",
-#    path = "${TORCH_INSTALL_PATH}",
-#    build_file = "third_party/libtorch/BUILD"
-#
+new_local_repository(
+    name = "libtorch_pre_cxx11_abi",
+    path = "${TORCH_INSTALL_PATH}",
+    build_file = "third_party/libtorch/BUILD"
+)
 
 #new_local_repository(
 #   name = "tensorrt",