Add GPU runner for linux-aarch64 (NVIDIA#289)

* add linux-aarch64 GPU runner

* fix test skip condition

* check system

* try to run on bare image

* fix

* change to ubuntu container

* Update gh-build-and-test.yml

* Update action.yml

* fix apt install

Author: leofang

.github/actions/test/action.yml

@@ -14,6 +14,13 @@ runs:
       shell: bash --noprofile --norc -xeuo pipefail {0}
       run: nvidia-smi
 
+    # The cache action needs this
+    - name: Install zstd
+      shell: bash --noprofile --norc -xeuo pipefail {0}
+      run: |
+        apt update
+        apt install zstd
+
     - name: Download bindings build artifacts
       uses: actions/download-artifact@v4
       with:

.github/workflows/gh-build-and-test.yml

@@ -76,17 +76,19 @@ jobs:
   test:
     # TODO: improve the name once a separate test matrix is defined
     name: Test (CUDA ${{ inputs.cuda-version }})
-    # TODO: enable testing once linux-aarch64 & win-64 GPU runners are up
+    # TODO: enable testing once win-64 GPU runners are up
     if: ${{ (github.repository_owner == 'nvidia') &&
-             startsWith(inputs.host-platform, 'linux-x64') }}
+             startsWith(inputs.host-platform, 'linux') }}
     permissions:
       id-token: write # This is required for configure-aws-credentials
       contents: read  # This is required for actions/checkout
-    runs-on: ${{ (inputs.host-platform == 'linux-x64' && 'linux-amd64-gpu-v100-latest-1') }}
-    # TODO: use a different (nvidia?) container, or just run on bare image
+    runs-on: ${{ (inputs.host-platform == 'linux-x64' && 'linux-amd64-gpu-v100-latest-1') ||
+                 (inputs.host-platform == 'linux-aarch64' && 'linux-arm64-gpu-a100-latest-1') }}
+    # Our self-hosted runners require a container
+    # TODO: use a different (nvidia?) container
     container:
       options: -u root --security-opt seccomp=unconfined --privileged --shm-size 16g
-      image: condaforge/miniforge3:latest
+      image: ubuntu:22.04
       env:
         NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
     needs: