Test final 2.6.0 wheel (#589)

GoogleCloudPlatform · Feb 3, 2025 · bd3982a · bd3982a
1 parent 9ee7608
commit bd3982a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 12 deletions.
diff --git a/dags/legacy_test/tests/pytorch/r2.6/common.libsonnet b/dags/legacy_test/tests/pytorch/r2.6/common.libsonnet
@@ -18,15 +18,13 @@ local mixins = import 'templates/mixins.libsonnet';
 local utils = import 'templates/utils.libsonnet';
 local volumes = import 'templates/volumes.libsonnet';
 
-local rcVersion = 'rc10';
-
 {
   local r2_6 = {
     frameworkPrefix: 'pt-2-6',
     tpuSettings+: {
       softwareVersion: 'tpu-ubuntu2204-base',
     },
-    imageTag: 'r2.6.0-%(rc)s_3.10' % {rc: rcVersion},
+    imageTag: 'r2.6.0_3.10',
   },
   PyTorchTest:: common.PyTorchTest + r2_6 {
     local config = self,
@@ -109,13 +107,13 @@ local rcVersion = 'rc10';
         pip install torch==2.6 --index-url https://download.pytorch.org/whl/test/cpu
         # torchvision commit reference: https://github.com/pytorch/pytorch/blob/release/2.6/.github/ci_commit_pins/vision.txt
         pip install --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@d23a6e1664d20707c11781299611436e1f0c104f"
-        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0%(rc)s-cp310-cp310-manylinux_2_28_x86_64.whl
+        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0-cp310-cp310-manylinux_2_28_x86_64.whl
         pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
         pip install pillow
         git clone --depth=1 https://github.com/pytorch/pytorch.git
         cd pytorch
-        git clone -b v2.6.0-%(rc)s https://github.com/pytorch/xla.git
-      ||| % {rc: rcVersion},
+        git clone -b v2.6.0 https://github.com/pytorch/xla.git
+      |||,
     },
     podTemplate+:: {
       spec+: {
@@ -152,16 +150,16 @@ local rcVersion = 'rc10';
         pip uninstall -y torch torchvision
         pip install torch==2.6 --index-url https://download.pytorch.org/whl/test/cpu
         pip install --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@d23a6e1664d20707c11781299611436e1f0c104f"
-        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0%(rc)s-cp310-cp310-manylinux_2_28_x86_64.whl
+        pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0-cp310-cp310-manylinux_2_28_x86_64.whl
 
         mkdir -p pytorch/xla
-        git clone -b v2.6.0-%(rc)s https://github.com/pytorch/xla.git pytorch/xla
+        git clone -b v2.6.0 https://github.com/pytorch/xla.git pytorch/xla
 
         %(cmd)s
 
         # Run whatever is in `command` here
         "${@:0}"
-      ||| % {cmd: config.tpuSettings.tpuVmExports, rc: rcVersion},
+      ||| % {cmd: config.tpuSettings.tpuVmExports},
     ],
     command: [
       'torchrun',

diff --git a/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py b/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py
@@ -111,12 +111,12 @@ class R2_5_1(enum.Enum):
     TORCH_XLA_REPO_BRANCH = "-b v2.5.1"
 
   class R2_6(enum.Enum):
-    TORCH_XLA_TPU_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0rc10+cxx11-cp311-cp311-manylinux_2_28_x86_64.whl"
-    TORCH_XLA_CUDA_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.6.0rc10+cxx11-cp310-cp310-linux_x86_64.whl"
+    TORCH_XLA_TPU_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0+cxx11-cp311-cp311-manylinux_2_28_x86_64.whl"
+    TORCH_XLA_CUDA_WHEEL = "https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla-2.6.0+cxx11-cp310-cp310-linux_x86_64.whl"
     TORCH = "torch==2.6.0"
     TORCHVISION = "torchvision"
     TORCHAUDIO = "torchaudio"
-    TORCH_XLA_GPU_DOCKER = "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0rc10_3.10_cuda_12.1"
+    TORCH_XLA_GPU_DOCKER = "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_cuda_12.1"
     TORCH_INDEX_CPU_URL = "https://download.pytorch.org/whl/test/cpu"
     TORCH_INDEX_CUDA_URL = "https://download.pytorch.org/whl/test/cu121"
     TORCH_REPO_BRANCH = "-b release/2.6"