Use cuda docker images

Author: zcbenz

.github/actions/setup-linux/action.yml

@@ -14,18 +14,18 @@ inputs:
 runs:
   using: "composite"
   steps:
+    - name: Install common dependencies
+      shell: bash
+      run: |
+        $(which sudo) apt-get update
+        $(which sudo) apt-get install -y libblas-dev liblapack-dev liblapacke-dev zip git
+
     - name: Use ccache
       uses: hendrikmuhs/ccache-action@v1.2
       with:
         key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}-py${{ inputs.python-version }}
         max-size: 1GB
 
-    - name: Install common dependencies
-      shell: bash
-      run: |
-        sudo apt-get update
-        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev zip
-
     - uses: actions/setup-python@v6
       with:
         python-version: ${{ inputs.python-version }}
@@ -42,43 +42,14 @@ runs:
 
     - name: Install MPI
       shell: bash
-      run: sudo apt-get install -y openmpi-bin openmpi-common libopenmpi-dev
-
-    - name: Install CUDA toolkit
-      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
-      shell: bash
-      env:
-        # Note: the CI machine does not meet CUDA 13's driver requirement.
-        # Compatibility matrix:
-        # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html
-        PACKAGES: |
-          {
-            "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-toolkit-12-6",
-            "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-toolkit-12-9",
-            "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0"
-          }
-      run: |
-        # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is
-        # Jetson specific. SBSA means Arm Server Base System Architecture.
-        ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }}
-        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb
-        sudo dpkg -i cuda-keyring_1.1-1_all.deb
-        sudo apt-get update
-        sudo apt-get install -y \
-            libnccl2 libnccl-dev \
-            ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
-        echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH
+      run: $(which sudo) apt-get install -y openmpi-bin openmpi-common libopenmpi-dev
 
     - name: CUDA packages and driver report
       if: ${{ startsWith(inputs.toolkit, 'cuda') }}
       shell: bash
       run: |
-        sudo apt-get install -y ubuntu-drivers-common dkms
-        echo "NVIDIA Driver Packages Available:"
-        sudo ubuntu-drivers list --gpgpu
         echo "NVIDIA Driver Version:"
         cat /proc/driver/nvidia/version || echo "nvidia driver not found"
-        echo "Installed NVIDIA and CUDA packages:"
         dpkg -l | egrep "cuda|nvidia" -i
         echo "DKMS Status:"
         dkms status || echo "dkms not found"

.github/workflows/build_and_test.yml

@@ -22,6 +22,7 @@ jobs:
     outputs:
       matrix: ${{ steps.generator.outputs.matrix }}
       runner: ${{ steps.generator.outputs.runner }}
+      container: ${{ steps.generator.outputs.container }}
     steps:
       - uses: actions/checkout@v5
       - name: Check lint
@@ -41,19 +42,28 @@ jobs:
                 'aarch64': 'ubuntu-22.04-arm'
               }
             }
+            let container = {
+              'cpu': null
+            }
             if (${{ github.repository == 'ml-explore/mlx' }}) {
-              let cuda_toolkits = ['cuda-12.6', 'cuda-12.9']
+              let cuda_versions = ['12.6', '12.9']
               let cuda_runner = {
-                'x86_64': 'gpu-t4-4-core',
-                'aarch64': 'ubuntu-22.04-arm'
+                 'x86_64': 'gpu-t4-4-core',
+                 'aarch64': 'ubuntu-22.04-arm'
               }
-              for (let toolkit of cuda_toolkits) {
-                matrix.toolkit.push(toolkit)
+              for (let version of cuda_versions) {
+                let toolkit = `cuda-${version}`
+                matrix['toolkit'].push(toolkit)
                 runner[toolkit] = cuda_runner
+                container[toolkit] = `pytorch/pytorch:2.8.0-cuda${version}-cudnn9-devel`
               }
             }
+            console.log(matrix)
+            console.log(runner)
+            console.log(container)
             core.setOutput('matrix', JSON.stringify(matrix))
             core.setOutput('runner', JSON.stringify(runner))
+            core.setOutput('container', JSON.stringify(container))
 
   linux_build_and_test:
     name: Linux (${{ matrix.toolkit }}, ${{ matrix.arch }})
@@ -62,6 +72,7 @@ jobs:
       fail-fast: false
       matrix: ${{ fromJson(needs.prepare.outputs.matrix) }}
     runs-on: ${{ fromJson(needs.prepare.outputs.runner)[matrix.toolkit][matrix.arch] }}
+    container: ${{ fromJson(needs.prepare.outputs.container)[matrix.toolkit] }}
     steps:
       - uses: actions/checkout@v5
       - uses: ./.github/actions/setup-linux