Add a test case for Linux aarch64 cu129 build (#7373) (#7383)

#7364 missed this spot and
failed to modify the right variable. Let's add a test case to avoid
this. I can confirm that Linux aarch64 cu129 is showing up now
https://github.com/pytorch/test-infra/actions/runs/18603342105

I also test this locally with:

```
OS=linux-aarch64 python3 tools/scripts/generate_binary_build_matrix.py | jq
{
  "include": [
    {
      "python_version": "3.10",
      "gpu_arch_type": "cpu-aarch64",
      "gpu_arch_version": "",
      "desired_cuda": "cpu",
      "container_image": "pytorch/manylinux2_28_aarch64-builder:cpu-aarch64",
      "package_type": "wheel",
      "build_name": "wheel-py3_10-cpu-aarch64",
      "validation_runner": "linux.arm64.2xlarge",
      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu",
      "channel": "nightly",
      "upload_to_base_bucket": "no",
      "stable_version": "2.9.0"
    },
    {
      "python_version": "3.10",
      "gpu_arch_type": "cuda-aarch64",
      "gpu_arch_version": "12.6-aarch64",
      "desired_cuda": "cu126",
      "container_image": "pytorch/manylinuxaarch64-builder:cuda12.6",
      "package_type": "wheel",
      "build_name": "wheel-py3_10-cuda-aarch6412_6-aarch64",
      "validation_runner": "linux.arm64.m7g.4xlarge",
      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126",
      "channel": "nightly",
      "upload_to_base_bucket": "no",
      "stable_version": "2.9.0"
    },
    {
      "python_version": "3.10",
      "gpu_arch_type": "cuda-aarch64",
      "gpu_arch_version": "12.8-aarch64",
      "desired_cuda": "cu128",
      "container_image": "pytorch/manylinuxaarch64-builder:cuda12.8",
      "package_type": "wheel",
      "build_name": "wheel-py3_10-cuda-aarch6412_8-aarch64",
      "validation_runner": "linux.arm64.m7g.4xlarge",
      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128",
      "channel": "nightly",
      "upload_to_base_bucket": "no",
      "stable_version": "2.9.0"
  1 Update os check
    },
    {
      "python_version": "3.10",
      "gpu_arch_type": "cuda-aarch64",
      "gpu_arch_version": "13.0-aarch64",
      "desired_cuda": "cu130",
      "container_image": "pytorch/manylinuxaarch64-builder:cuda13.0",
      "package_type": "wheel",
      "build_name": "wheel-py3_10-cuda-aarch6413_0-aarch64",
      "validation_runner": "linux.arm64.m7g.4xlarge",
      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130",
      "channel": "nightly",
      "upload_to_base_bucket": "no",
      "stable_version": "2.9.0"
    },
    {
      "python_version": "3.10",
      "gpu_arch_type": "cuda-aarch64",
      "gpu_arch_version": "12.9-aarch64",
      "desired_cuda": "cu129",
      "container_image": "pytorch/manylinuxaarch64-builder:cuda12.9",
      "package_type": "wheel",
      "build_name": "wheel-py3_10-cuda-aarch6412_9-aarch64",
      "validation_runner": "linux.arm64.m7g.4xlarge",
      "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu129",
      "channel": "nightly",
      "upload_to_base_bucket": "no",
      "stable_version": "2.9.0"
    }
  ]
}
```

---------

Signed-off-by: Huy Do <[email protected]>

Author: huydhn

.github/workflows/test_build_wheels_linux_aarch64_with_cuda.yml

@@ -0,0 +1,55 @@
+name: Test Build Linux Aarch64 Wheels with CUDA
+
+on:
+  pull_request:
+    paths:
+      - .github/actions/setup-binary-builds/action.yml
+      - .github/workflows/test_build_wheels_linux.yml
+      - .github/workflows/build_wheels_linux.yml
+      - .github/workflows/generate_binary_build_matrix.yml
+      - .github/workflows/test_build_wheels_linux_aarch64_with_cuda.yml
+      - tools/scripts/generate_binary_build_matrix.py
+  workflow_dispatch:
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    uses: ./.github/workflows/generate_binary_build_matrix.yml
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: ${{ github.repository }}
+      test-infra-ref: ${{ github.ref }}
+      with-cuda: enable
+  test:
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/vision
+            pre-script: packaging/pre_build_script.sh
+            post-script: packaging/post_build_script.sh
+            smoke-test-script: test/smoke_test.py
+            package-name: torchvision
+          - repository: pytorch/audio
+            smoke-test-script: test/smoke_test/smoke_test.py
+            package-name: torchaudio
+    uses: ./.github/workflows/build_wheels_linux.yml
+    name: ${{ matrix.repository }}
+    with:
+      repository: ${{ matrix.repository }}
+      ref: nightly
+      test-infra-repository: ${{ github.repository }}
+      test-infra-ref: ${{ github.ref }}
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      post-script: ${{ matrix.post-script }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      package-name: ${{ matrix.package-name }}
+      trigger-event: "${{ github.event_name }}"
+      architecture: aarch64
+      setup-miniconda: false

tools/scripts/generate_binary_build_matrix.py

@@ -149,15 +149,15 @@ def validation_runner(arch_type: str, os: str) -> str:
 
 
 def initialize_globals(channel: str, os: str, build_python_only: bool) -> None:
-    global CURRENT_VERSION, CUDA_ARCHES, ROCM_ARCHES, PYTHON_ARCHES
+    global CURRENT_VERSION, CUDA_ARCHES, CUDA_AARCH64_ARCHES, ROCM_ARCHES, PYTHON_ARCHES
     global WHEEL_CONTAINER_IMAGES, LIBTORCH_CONTAINER_IMAGES
     if channel == TEST:
         CURRENT_VERSION = CURRENT_CANDIDATE_VERSION
     else:
         CURRENT_VERSION = CURRENT_STABLE_VERSION
 
     CUDA_ARCHES = CUDA_ARCHES_DICT[channel]
-    if channel != "release" and os == LINUX:
+    if channel != "release" and (os == LINUX or os == LINUX_AARCH64):
         # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up
         # in 2.10
         CUDA_ARCHES.append("12.9")