diff --git a/.github/scripts/install-test-deps.sh b/.github/scripts/install-test-deps.sh new file mode 100644 index 0000000000..cf34aab07a --- /dev/null +++ b/.github/scripts/install-test-deps.sh @@ -0,0 +1,19 @@ +#set -exou pipefail +set -x +# Set default values for CHANNEL and CU_VERSION if not already set +CHANNEL=${CHANNEL:-nightly} +CU_VERSION=${CU_VERSION:-cu129} + +PLATFORM=$(python -c "import sys; print(sys.platform)") +TORCH=$(grep "^torch>" ${PWD}/py/requirements.txt) +TORCHVISION=$(grep "^torchvision>" ${PWD}/tests/py/requirements.txt) +INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION} + +# Install all the dependencies required for Torch-TensorRT +pip install --pre -r ${PWD}/tests/py/requirements.txt +# dependencies in the tests/py/requirements.txt might install a different version of torch or torchvision +# eg. timm will install the latest torchvision, however we want to use the torchvision from nightly +# reinstall torch torchvisionto make sure we have the correct version +pip uninstall -y torch torchvision +pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL} +pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL} diff --git a/.github/scripts/install-torch-tensorrt.sh b/.github/scripts/install-torch-tensorrt.sh index 94de5f022a..4e51f2c43f 100755 --- a/.github/scripts/install-torch-tensorrt.sh +++ b/.github/scripts/install-torch-tensorrt.sh @@ -1,9 +1,6 @@ #set -exou pipefail set -x -TORCH=$(grep "^torch>" ${PWD}/py/requirements.txt) -TORCHVISION=$(grep "^torchvision>" ${PWD}/tests/py/requirements.txt) -INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION} PLATFORM=$(python -c "import sys; print(sys.platform)") if [[ $(uname -m) == "aarch64" ]]; then @@ -12,14 +9,8 @@ if [[ $(uname -m) == "aarch64" ]]; then install_cuda_aarch64 fi -# Install all the dependencies required for Torch-TensorRT -pip install --pre -r ${PWD}/tests/py/requirements.txt -# dependencies in the tests/py/requirements.txt might install a different version of torch or torchvision -# eg. timm will install the latest torchvision, however we want to use the torchvision from nightly -# reinstall torch torchvisionto make sure we have the correct version -pip uninstall -y torch torchvision -pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL} -pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL} +# Install all the external dependencies required for testing Torch-TensorRT in CI +source .github/scripts/install-test-deps.sh # Install Torch-TensorRT diff --git a/tests/py/requirements.txt b/tests/py/requirements.txt index c022378f6d..5bfd8a5bcb 100644 --- a/tests/py/requirements.txt +++ b/tests/py/requirements.txt @@ -14,5 +14,5 @@ nvidia-modelopt[all]; python_version >'3.9' and python_version <'3.13' # flashinfer-python is broken on python 3.9 at the moment, so skip it for now flashinfer-python; python_version >'3.9' and python_version <'3.13' --extra-index-url https://download.pytorch.org/whl/nightly/cu129 -torchvision>=0.23.0.dev,<0.24.0 +torchvision>=0.24.0.dev,<0.25.0 timm>=1.0.3 \ No newline at end of file