Linux aarch64 wheels

.github/workflows/build_ffmpeg.yaml

@@ -48,6 +48,32 @@ jobs:
         mkdir -p "${artifact_dir}"
         mv ffmpeg.tar.gz "${artifact_dir}/${FFMPEG_VERSION}.tar.gz"
 
+  LGPL-Linux-aarch64:
+    strategy:
+      fail-fast: false
+      matrix:
+        ffmpeg-version: ["4.4.4", "5.1.4", "6.1.1", "7.0.1"]
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      job-name: Build
+      runner: linux.arm64.2xlarge
+      upload-artifact: ffmpeg-lgpl-linux_aarch64-${{ matrix.ffmpeg-version }}
+      repository: pytorch/torchcodec
+      script: |
+        export FFMPEG_VERSION="${{ matrix.ffmpeg-version }}"
+        export FFMPEG_ROOT="${PWD}/ffmpeg"
+
+        packaging/build_ffmpeg.sh
+
+        tar -cf ffmpeg.tar.gz ffmpeg/include ffmpeg/lib
+
+        artifact_dir="${RUNNER_ARTIFACT_DIR}/$(date +%Y-%m-%d)/linux_aarch64"
+        mkdir -p "${artifact_dir}"
+        mv ffmpeg.tar.gz "${artifact_dir}/${FFMPEG_VERSION}.tar.gz"
+
   LGPL-macOS:
     strategy:
       fail-fast: false

.github/workflows/linux_aarch64_wheel.yaml

@@ -0,0 +1,127 @@
+name: Build and test Linux aarch64 wheel
+
+on:
+  pull_request:
+  push:
+    branches:
+      - nightly
+      - main
+      - release/*
+    tags:
+        - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: write
+
+defaults:
+  run:
+    shell: bash -l -eo pipefail {0}
+
+jobs:
+
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-xpu: disable
+      with-rocm: disable
+      with-cuda: disable
+      build-python-only: "disable"
+
+  build:
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+    name: Build and Upload Linux wheel
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: pytorch/torchcodec
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: packaging/pre_build_script.sh
+      post-script: packaging/post_build_script.sh
+      smoke-test-script: packaging/fake_smoke_test.py
+      package-name: torchcodec
+      trigger-event: ${{ github.event_name }}
+      build-platform: "python-build-package"
+      build-command: "BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1 python -m build --wheel -vvv --no-isolation"
+      architecture: aarch64
+
+  install-and-test:
+    runs-on: ubuntu-24.04-arm
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.10']
+        ffmpeg-version-for-tests: ['4.4.2', '5.1.2', '6.1.1', '7.0.1']
+    needs: build
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: pytorch_torchcodec__${{ matrix.python-version }}_cpu_aarch64
+          path: pytorch/torchcodec/dist/
+      - name: Setup conda env
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          auto-update-conda: true
+          miniconda-version: "latest"
+          activate-environment: test
+          python-version: ${{ matrix.python-version }}
+      - name: Update pip
+        run: python -m pip install --upgrade pip
+      - name: Install PyTorch
+        run: |
+          python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+      - name: Install torchcodec from the wheel
+        run: |
+          wheel_path=`find pytorch/torchcodec/dist -type f -name "*.whl"`
+          echo Installing $wheel_path
+          python -m pip install $wheel_path -vvv
+
+      - name: Check out repo
+        uses: actions/checkout@v3
+      - name: Install ffmpeg, post build
+        run: |
+          # Ideally we would have checked for that before installing the wheel,
+          # but we need to checkout the repo to access this file, and we don't
+          # want to checkout the repo before installing the wheel to avoid any
+          # side-effect. It's OK.
+          source packaging/helpers.sh
+          assert_ffmpeg_not_installed
+
+          conda install "ffmpeg=${{ matrix.ffmpeg-version-for-tests }}" -c conda-forge
+          ffmpeg -version
+
+      - name: Install test dependencies
+        run: |
+          # Ideally we would find a way to get those dependencies from pyproject.toml
+          python -m pip install numpy pytest pillow
+
+      - name: Delete the src/ folder just for fun
+        run: |
+          # The only reason we checked-out the repo is to get access to the
+          # tests. We don't care about the rest. Out of precaution, we delete
+          # the src/ folder to be extra sure that we're running the code from
+          # the installed wheel rather than from the source.
+          # This is just to be extra cautious and very overkill because a)
+          # there's no way the `torchcodec` package from src/ can be found from
+          # the PythonPath: the main point of `src/` is precisely to protect
+          # against that and b) if we ever were to execute code from
+          # `src/torchcodec`, it would fail loudly because the built .so files
+          # aren't present there.
+          rm -r src/
+          ls
+      - name: Run Python tests
+        run: |
+          pytest --override-ini="addopts=-v" test

test/test_encoders.py

@@ -1,6 +1,7 @@
 import io
 import json
 import os
+import platform
 import re
 import subprocess
 import sys
@@ -305,6 +306,9 @@ def test_against_cli(
             rtol, atol = 0, 1e-3
             if sys.platform == "darwin":
                 assert_close = partial(assert_tensor_close_on_at_least, percentage=99)
+            arch = platform.machine().lower()
+            if sys.platform == "linux" and arch == "aarch64":
+                assert_close = partial(assert_tensor_close_on_at_least, percentage=98)
         elif format == "wav":
             rtol, atol = 0, 1e-4
         elif format == "mp3" and asset is SINE_MONO_S32 and num_channels == 2:

test/utils.py

@@ -2,6 +2,7 @@
 import json
 import os
 import pathlib
+import platform
 import subprocess
 import sys
 
@@ -81,7 +82,12 @@ def assert_frames_equal(*args, **kwargs):
             else:
                 torch.testing.assert_close(*args, **kwargs, atol=atol, rtol=0)
         else:
-            torch.testing.assert_close(*args, **kwargs, atol=0, rtol=0)
+            arch = platform.machine().lower()
+            if arch == "x86_64":
+                torch.testing.assert_close(*args, **kwargs, atol=0, rtol=0)
+            else:
+                # linux aarch64
+                torch.testing.assert_close(*args, **kwargs, atol=3, rtol=0)
     else:
         torch.testing.assert_close(*args, **kwargs, atol=3, rtol=0)