Use uv in macOS CI

.github/actions/build-macos-release/action.yml

@@ -8,29 +8,56 @@ inputs:
     default: '15.0'
   build-backend:
     description: 'Build the backend mlx-metal package'
-    type: boolean
     required: false
-    default: false
+    default: 'false'
+  python-version:
+    description: 'Python version to use'
+    required: false
+    default: '3.10'
 
 runs:
   using: "composite"
   steps:
+    - name: Clear build environment
+      shell: bash
+      run: |
+        echo "::group::Clear build environment"
+        uv venv --clear
+        uv pip install build setuptools
+        python setup.py clean --all
+        source .venv/bin/activate
+        echo "::endgroup::"
+
     - name: Build Python package
-      shell: bash -l {0}
+      shell: bash
       env:
         DEVELOPER_DIR: /Applications/Xcode-latest.app
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
-        pip install build
-        python setup.py clean --all
+        echo "::group::Build Python package"
         MLX_BUILD_STAGE=1 python -m build -w
+        python setup.py clean --all
+        echo "::endgroup::"
 
     - name: Build backend package
-      if: ${{ inputs.build-backend }}
-      shell: bash -l {0}
+      if: ${{ inputs.build-backend == 'true' }}
+      shell: bash
       env:
         DEVELOPER_DIR: /Applications/Xcode-latest.app
         MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
-        python setup.py clean --all
+        echo "::group::Build backend package"
         MLX_BUILD_STAGE=2 python -m build -w
+        python setup.py clean --all
+        echo "::endgroup::"
+
+    - name: Test local packages
+      if: ${{ inputs.build-backend == 'true' && inputs.macos-target == '26.0' }}
+      shell: bash
+      run: |
+        echo "::group::Test local packages"
+        uv pip install numpy torch psutil
+        uv pip install dist/mlx_metal-*-macosx_26_0_arm64.whl
+        uv pip install dist/mlx-*-macosx_26_0_arm64.whl
+        python -m unittest discover -v python/tests
+        echo "::endgroup::"

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

@@ -4,61 +4,72 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
+    - name: Install Python package
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
-        pip install -e ".[dev]" -v
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"
 
     - name: Install tests dependencies
-      shell: bash -l {0}
+      shell: bash
       run: |
-        pip install tensorflow
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"
 
     - name: Run Python tests
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
       run: |
+        echo "::group::Run Python tests"
         DEVICE=cpu python -m unittest discover -v python/tests
         DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests
         mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
         mlx.launch --verbose -n 8 python python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
         if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
-
+        echo "::endgroup::"
+
     - name: Build example extension
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build example extension"
         cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-
+        uv pip install -r requirements.txt
+        uv run --no-project setup.py build_ext --inplace
+        uv run --no-project test.py
+        echo "::endgroup::"
+
     - name: Build CPP only
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build CPP only"
         mkdir -p build
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
+        echo "::endgroup::"
 
     - name: Run CPP tests
-      shell: bash -l {0}
+      shell: bash
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run CPP tests"
         ./build/tests/tests
         ./build/tests/test_teardown
+        echo "::endgroup::"
 
     - name: Build small binary with JIT
-      shell: bash -l {0}
+      shell: bash
       run: |
+        echo "::group::Build small binary with JIT"
         mkdir -p build
         cd build
         cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -68,15 +79,18 @@ runs:
           -DMLX_BUILD_GGUF=OFF \
           -DMLX_METAL_JIT=ON
         make -j $(sysctl -n hw.ncpu)
-
+        echo "::endgroup::"
+
     - name: Run Python tests with JIT
-      shell: bash -l {0}
+      shell: bash
       env:
         LOW_MEMORY: 1
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
       run: |
+        echo "::group::Run Python tests with JIT"
         CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
+          uv pip install -e . -v
         python -m unittest discover -v python/tests
+        echo "::endgroup::"

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

@@ -13,12 +13,20 @@ runs:
     - name: Install Homebrew packages
       shell: sh
       run: /opt/homebrew/bin/brew install openmpi
-    
+
     - name: Verify MetalToolchain installed
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - uses: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        uv venv --python ${{ inputs.python-version }} --managed-python
+        source .venv/bin/activate
+        echo PATH=$PATH >> $GITHUB_ENV
+        # Search python packages in .venv
+        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
+        echo "::endgroup::"

.github/workflows/release.yml

@@ -93,27 +93,20 @@ jobs:
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools typing_extensions
-          pip install -e . -v
       - name: Build macOS 14 package
         uses: ./.github/actions/build-macos-release
         with:
-          macos-target: 14.0
+          macos-target: '14.0'
           build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Build macOS 15 package
         uses: ./.github/actions/build-macos-release
         with:
-          macos-target: 15.0
+          macos-target: '15.0'
           build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Build macOS 26 package
         uses: ./.github/actions/build-macos-release
         with:
-          macos-target: 26.0
+          macos-target: '26.0'
           build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Upload MLX artifacts
         uses: actions/upload-artifact@v7

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
-mlx>=0.21.0
+mlx>=0.31.2
 nanobind==2.12.0

setup.py

@@ -234,7 +234,7 @@ def get_tag(self) -> tuple[str, str, str]:
             "ml_dtypes",
             "numpy>=2",
             "pre-commit",
-            "psutil",
+            "psutil>=7.2",
             "torch>=2.9",
             "typing_extensions",
         ],