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 ec27996113..e67859691c 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")