From 91605fdb5c8c6f4dd9b9e295393af50669ab41c3 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Mon, 20 Oct 2025 10:31:25 -0700 Subject: [PATCH] Add a test case for Linux aarch64 cu129 build (#7373) https://github.com/pytorch/test-infra/pull/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 --- ...t_build_wheels_linux_aarch64_with_cuda.yml | 55 +++++++++++++++++++ tools/scripts/generate_binary_build_matrix.py | 4 +- 2 files changed, 57 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/test_build_wheels_linux_aarch64_with_cuda.yml diff --git a/.github/workflows/test_build_wheels_linux_aarch64_with_cuda.yml b/.github/workflows/test_build_wheels_linux_aarch64_with_cuda.yml new file mode 100644 index 0000000000..7bd4976f84 --- /dev/null +++ b/.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 diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py index 2b91272c15..2336a07348 100755 --- a/tools/scripts/generate_binary_build_matrix.py +++ b/tools/scripts/generate_binary_build_matrix.py @@ -149,7 +149,7 @@ 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 @@ -157,7 +157,7 @@ def initialize_globals(channel: str, os: str, build_python_only: bool) -> None: 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")