Use ccache in GitHub Actions (#2773)

* Remove unnecessary steps

* Use ccache

* Log when using ccache

* Set max-size to 1GB

* Pass --no-build-isolation

* Remove more unused things

.github/actions/build-cuda/action.yml

@@ -15,7 +15,7 @@ runs:
       env:
         DEBUG: 1
         CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=${{ inputs.nvcc-location }}
-      run: pip install -e ".[dev]" -v
+      run: pip install --no-build-isolation -e ".[dev]" -v
 
     - name: Build CPP only
       shell: bash

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

@@ -9,8 +9,8 @@ runs:
       env:
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
         DEBUG: 1
-      run: pip install -e ".[dev]" -v
-    
+      run: pip install --no-build-isolation -e ".[dev]" -v
+
     - name: Generate package stubs
       shell: sh
       run: |

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

@@ -18,42 +18,39 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Free disk space
-      shell: sh
-      if: inputs.runner-type == 'linux'
-      run: sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+    - name: Use ccache
+      uses: hendrikmuhs/ccache-action@v1.2
+      with:
+        key: ccache-${{ inputs.runner-type }}-${{ runner.arch }}-py${{ inputs.python-version }}
+        max-size: 1GB
 
     - name: Install common dependencies
-      env:
-        TZ: Etc/UTC
       shell: bash
       run: |
         sudo apt-get update
-        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev tzdata zip
-        sudo apt autoremove -y
+        sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev zip
 
     - uses: actions/setup-python@v6
       with:
         python-version: ${{ inputs.python-version }}
         cache: 'pip'
 
-    - name: setup python venv
+    - name: Setup Python venv
       shell: bash
       run: |
         python -m venv .venv
         source .venv/bin/activate
+        pip install cmake nanobind==2.4.0
         echo PATH=$PATH >> $GITHUB_ENV
-        pip install --upgrade pip cmake
+        # Make cmake search .venv for nanobind
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
 
     - name: Install MPI
       shell: bash
       run: sudo apt-get install -y openmpi-bin openmpi-common libopenmpi-dev
 
     - name: Network CUDA installation from packages
-      id: install-cuda
       if: inputs.runner-type == 'cuda'
-      env:
-        TZ: Etc/UTC
       shell: bash ## Specific to Ubuntu 22.04 & Architecture x86_64
       run: |
         wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb

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

@@ -9,13 +9,18 @@ inputs:
 runs:
   using: "composite"
   steps:
+    - name: Run MPI tests
+      shell: bash
+      run: |
+        echo "::group::MPI tests"
+        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
+        echo "::endgroup::"
+
     - name: Run distributed tests
-      # FIXME: This test fails with CUDA build.
       if: ${{ inputs.cpu-only == 'true' }}
       shell: bash
       run: |
         echo "::group::Distributed tests"
-        mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
         if grep -Fq '[WARN]' stderr.log ; then
           grep -F '[WARN]' stderr.log
@@ -34,7 +39,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run Python tests - GPU
-      if: ${{ !inputs.cpu-only }}
+      if: ${{ inputs.cpu-only == 'false' }}
       shell: bash
       env:
         DEVICE: gpu
@@ -53,7 +58,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run CPP tests - GPU
-      if: ${{ !inputs.cpu-only }}
+      if: ${{ inputs.cpu-only == 'false' }}
       shell: bash
       env:
         DEVICE: gpu

CMakeLists.txt

@@ -74,6 +74,7 @@ endif()
 if(MLX_USE_CCACHE)
   find_program(CCACHE_PROGRAM ccache)
   if(CCACHE_PROGRAM)
+    message(STATUS "Found CCache: ${CCACHE_PROGRAM}")
     set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
     set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
     set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")

cmake/Findnvpl.cmake

@@ -0,0 +1,3 @@
+# This file does nothing but to suppress the cmake warning: "By not providing
+# Findnvpl.cmake in CMAKE_MODULE_PATH...", which is caused by the
+# find_package(nvpl) from cmake's builtin FindLAPACK.cmake module.

mlx/backend/cuda/utils.cpp

@@ -60,7 +60,7 @@ const char* dtype_to_cuda_type(const Dtype& dtype) {
     case float64:
       return "double";
     case complex64:
-      return "complex64_t";
+      return "mlx::core::cu::complex64_t";
     default:
       return "unknown";
   }

pyproject.toml

@@ -2,6 +2,6 @@
 requires = [
   "setuptools>=80",
   "nanobind==2.4.0",
-  "cmake>=3.25,<4.1",
+  "cmake>=3.25",
 ]
 build-backend = "setuptools.build_meta"