simplify ci env setup (#7411)

.github/scripts/setup-env.sh

@@ -5,7 +5,6 @@ set -euo pipefail
 # Prepare conda
 CONDA_PATH=$(which conda)
 eval "$(${CONDA_PATH} shell.bash hook)"
-conda config --set channel_priority strict
 
 # Setup the OS_TYPE environment variable that should be used for conditions involving the OS below.
 case $(uname) in
@@ -33,43 +32,15 @@ if [[ "${OS_TYPE}" == "macos" && $(uname -m) == x86_64 ]]; then
 fi
 echo '::endgroup::'
 
-echo '::group::Set PyTorch conda channel and wheel index'
-# TODO: Can we maybe have this as environment variable in the job template? For example, `IS_RELEASE`.
-if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
-  CHANNEL_ID=test
-else
-  CHANNEL_ID=nightly
-fi
-PYTORCH_CONDA_CHANNEL=pytorch-"${CHANNEL_ID}"
-echo "PYTORCH_CONDA_CHANNEL=${PYTORCH_CONDA_CHANNEL}"
-
-case $GPU_ARCH_TYPE in
-  cpu)
-    GPU_ARCH_ID="cpu"
-    ;;
-  cuda)
-    VERSION_WITHOUT_DOT=$(echo "${GPU_ARCH_VERSION}" | sed 's/\.//')
-    GPU_ARCH_ID="cu${VERSION_WITHOUT_DOT}"
-    ;;
-  *)
-    echo "Unknown GPU_ARCH_TYPE=${GPU_ARCH_TYPE}"
-    exit 1
-    ;;
-esac
-PYTORCH_WHEEL_INDEX="https://download.pytorch.org/whl/${CHANNEL_ID}/${GPU_ARCH_ID}"
-echo "PYTORCH_WHEEL_INDEX=${PYTORCH_WHEEL_INDEX}"
-echo '::endgroup::'
-
 echo '::group::Create build environment'
 # See https://github.com/pytorch/vision/issues/7296 for ffmpeg
 conda create \
   --name ci \
   --quiet --yes \
   python="${PYTHON_VERSION}" pip \
-  ninja libpng jpeg \
-  'ffmpeg<4.3' \
-  -c "${PYTORCH_CONDA_CHANNEL}" \
-  -c defaults
+  ninja \
+  libpng jpeg \
+  'ffmpeg<4.3'
 conda activate ci
 pip install --progress-bar=off --upgrade setuptools
 
@@ -81,7 +52,18 @@ fi
 echo '::endgroup::'
 
 echo '::group::Install PyTorch'
-pip install --progress-bar=off --pre torch --index-url="${PYTORCH_WHEEL_INDEX}"
+# TODO: Can we maybe have this as environment variable in the job template? For example, `IS_RELEASE`.
+if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
+  CHANNEL=test
+else
+  CHANNEL=nightly
+fi
+
+pip install --progress-bar=off light-the-torch
+ltt install --progress-bar=off \
+  --pytorch-computation-backend="${GPU_ARCH_TYPE}${GPU_ARCH_VERSION}" \
+  --pytorch-channel="${CHANNEL}" \
+  torch
 
 if [[ $GPU_ARCH_TYPE == 'cuda' ]]; then
   python3 -c "import torch; exit(not torch.cuda.is_available())"

.github/workflows/lint.yml

@@ -81,6 +81,8 @@ jobs:
         
         export PYTHON_VERSION=3.8
         export GPU_ARCH_TYPE=cpu
+        export GPU_ARCH_VERSION=''
+
         ./.github/scripts/setup-env.sh
         
         CONDA_PATH=$(which conda)

.github/workflows/test-linux.yml

@@ -51,6 +51,7 @@ jobs:
 
         export PYTHON_VERSION=3.8
         export GPU_ARCH_TYPE=cpu
+        export GPU_ARCH_VERSION=''
 
         ./.github/scripts/setup-env.sh
         
@@ -80,6 +81,7 @@ jobs:
 
         export PYTHON_VERSION=3.8
         export GPU_ARCH_TYPE=cpu
+        export GPU_ARCH_VERSION=''
 
         ./.github/scripts/setup-env.sh
         

.github/workflows/test-macos.yml

@@ -35,5 +35,6 @@ jobs:
 
         export PYTHON_VERSION=${{ matrix.python-version }}
         export GPU_ARCH_TYPE=cpu
+        export GPU_ARCH_VERSION=''
 
         ./.github/scripts/unittest.sh