Use index-url for all test steps (#4948)

As per instructions: https://pytorch.org/get-started/locally/
pytorch · Feb 26, 2024 · df8eb74 · df8eb74
1 parent 36cc261
commit df8eb74
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 17 deletions.
diff --git a/.github/workflows/test-export-matrix-variables.yml b/.github/workflows/test-export-matrix-variables.yml
@@ -13,7 +13,7 @@ jobs:
       fail-fast: false
       matrix:
         installation:
-          - pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+          - pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
     with:
       binary-matrix: ${{ toJSON(matrix) }}
       script: |
@@ -24,7 +24,7 @@ jobs:
       fail-fast: false
       matrix:
         installation:
-          - pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+          - pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
     with:
       binary-matrix: ${{ toJSON(matrix) }}
       script: |
@@ -35,7 +35,7 @@ jobs:
       fail-fast: false
       matrix:
         installation:
-          - pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+          - pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
     with:
       binary-matrix: ${{ toJSON(matrix) }}
       script: |

diff --git a/.github/workflows/test_linux_job.yml b/.github/workflows/test_linux_job.yml
@@ -48,7 +48,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch
         python3 -c 'import torch'
   test-gpu:
@@ -57,24 +57,24 @@ jobs:
       matrix:
         runner_type: ["linux.4xlarge.nvidia.gpu", "linux.g5.4xlarge.nvidia.gpu"]
     with:
-      job-name: "linux-py3.8-cu116"
+      job-name: "linux-py3.8-cu121"
       runner: ${{ matrix.runner_type }}
       test-infra-repository: ${{ github.repository }}
       test-infra-ref: ${{ github.ref }}
       submodules: ${{ 'true' }}
       gpu-arch-type: cuda
-      gpu-arch-version: "11.6"
+      gpu-arch-version: "12.1"
       timeout: 60
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cu116 --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cu121 --pre torch
         # Can import pytorch, cuda is available
         python3 -c 'import torch;cuda_avail = torch.cuda.is_available();print("CUDA available: " + str(cuda_avail));assert(cuda_avail)'
         python3 -c 'import torch;t = torch.ones([2,2], device="cuda:0");print(t);print("tensor device:" + str(t.device))'
         nvidia-smi
-        nvcc --version | grep "cuda_11.6"
-        [[ "${CUDA_HOME}" == "/usr/local/cuda-11.6" ]] || exit 1
+        nvcc --version | grep "cuda_12.1"
+        [[ "${CUDA_HOME}" == "/usr/local/cuda-12.1" ]] || exit 1
   test-docker-image:
     uses: ./.github/workflows/linux_job.yml
     with:

diff --git a/.github/workflows/test_macos_job.yml b/.github/workflows/test_macos_job.yml
@@ -19,7 +19,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch
         python3 -c 'import torch'
   test-m1:
@@ -34,7 +34,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch, cuda is available
         python3 -c 'import torch'
   test-x86-with-repo:
@@ -49,7 +49,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch
         python3 -c 'import torch'
   test-m1-with-repo:
@@ -66,7 +66,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python3 -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch, cuda is available
         python3 -c 'import torch'
   test-upload-artifact:

diff --git a/.github/workflows/test_windows_job.yml b/.github/workflows/test_windows_job.yml
@@ -19,7 +19,7 @@ jobs:
       script: |
         conda create --yes --quiet -n test python=3.8
         conda activate test
-        python -m pip install --extra-index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
+        python -m pip install --index-url https://download.pytorch.org/whl/nightly/cpu --pre torch
         # Can import pytorch
         python -c 'import torch'
   test-gpu:

diff --git a/tools/pkg-helpers/pytorch_pkg_helpers/wheel.py b/tools/pkg-helpers/pytorch_pkg_helpers/wheel.py
@@ -24,7 +24,7 @@ def get_pytorch_pip_install_command(
     if channel == "nightly":
         pip_install += " --pre"
     extra_index = f"https://download.pytorch.org/whl/{channel}/{gpu_arch_version}"
-    return [f"export PIP_INSTALL_TORCH='{pip_install} --extra-index-url {extra_index}'"]
+    return [f"export PIP_INSTALL_TORCH='{pip_install} --index-url {extra_index}'"]
 
 
 def get_pytorch_s3_bucket_path(

diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py
@@ -273,7 +273,6 @@ def get_wheel_install_command(
     use_only_dl_pytorch_org: bool,
 ) -> str:
 
-    index_url_option = "--index-url" if os != LINUX_AARCH64 else "--extra-index-url"
     if  channel == RELEASE and (not use_only_dl_pytorch_org) and (
         (gpu_arch_version == "12.1" and os == LINUX)
         or (
@@ -289,7 +288,7 @@ def get_wheel_install_command(
             if channel == "nightly"
             else f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}"
         )
-        return f"{whl_install_command} {index_url_option} {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}"
+        return f"{whl_install_command} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}"
 
 
 def generate_conda_matrix(