mlx/.circleci/config.yml

version: 2.1

orbs:
  apple: ml-explore/pr-approval@0.1.0

parameters:
  nightly_build:
    type: boolean
    default: false
  test_release:
    type: boolean
    default: false

jobs:
  build_documentation:
    parameters:
      upload-docs:
        type: boolean
        default: false
    macos:
      xcode: "16.2.0"
    resource_class: m2pro.medium
    steps:
      - checkout
      - run:
          name: Install
          command: |
            brew install python@3.9
            brew install doxygen
            python3.9 -m venv env
            source env/bin/activate
            pip install --upgrade pip
            pip install --upgrade cmake
            pip install -r docs/requirements.txt
            pip install . -v
      - when:
          condition:
            not: << parameters.upload-docs >>
          steps:
            - run:
               name: Build documentation
               command: |
                 source env/bin/activate
                 cd docs && doxygen && make html O=-W
      - when:
          condition: << parameters.upload-docs >>
          steps:
            - add_ssh_keys:
                fingerprints:
                  - "SHA256:OhcVVMovbT0pkgMeiVRyxMnjV9R2t+hKBsNcuxq9h+0"
            - run:
               name: Upload documentation
               command: |
                 source env/bin/activate
                 git config user.email "mlx@group.apple.com"
                 git config user.name "CircleCI Docs"
                 git checkout gh-pages
                 git rebase main
                 cd docs
                 git rm -rf build/html
                 doxygen && make html O=-W
                 git add -f build/html
                 git commit -m "rebase"
                 git push -f origin gh-pages

  linux_build_and_test:
    machine:
      image: ubuntu-2204:current
      resource_class: large
    steps:
      - checkout
      - run:
          name: Run style checks
          command: |
            pip install pre-commit
            pre-commit run --all
            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
      - run:
          name: Install dependencies
          command: |
            export DEBIAN_FRONTEND=noninteractive
            export NEEDRESTART_MODE=a
            sudo apt-get update
            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
            sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
            curl -LsSf https://astral.sh/uv/install.sh | sh
      - run:
          name: Install Python package
          command: |
            uv venv
            uv pip install cmake
            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
              uv pip install -e ".[dev]" -v
      - run:
          name: Generate package stubs
          command: |
            uv pip install typing_extensions
            uv run --no-project setup.py generate_stubs
      - run:
          name: Run Python tests
          command: |
            source .venv/bin/activate
            python -m unittest discover python/tests -v
            mpirun --bind-to none -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 "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
      - run:
          name: Build CPP only
          command: |
            source .venv/bin/activate
            mkdir -p build && cd build
            cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
            make -j `nproc`
      - run:
          name: Run CPP tests
          command: ./build/tests/tests

  mac_build_and_test:
    parameters:
      xcode_version:
        type: string
        default: "16.2.0"
      macosx_deployment_target:
        type: string
        default: ""
    macos:
      xcode: << parameters.xcode_version >>
    environment:
      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
    resource_class: m2pro.medium
    steps:
      - checkout
      - run:
          name: Install dependencies
          command: |
            HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALL_CLEANUP=1 \
              brew install openmpi uv
      - run:
          name: Install Python package
          command: |
            uv venv --python 3.9
            uv pip install \
              nanobind==2.4.0 \
              cmake \
              numpy \
              torch \
              tensorflow \
              unittest-xml-reporting
            DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
              uv pip install -e . -v
      - run:
          name: Generate package stubs
          command: |
            uv pip install typing_extensions
            uv run --no-project setup.py generate_stubs
      - run:
          name: Run Python tests
          command: |
            source .venv/bin/activate
            LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
            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/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
      - run:
          name: Build example extension
          command: |
            source .venv/bin/activate
            cd examples/extensions
            uv pip install -r requirements.txt
            uv run --no-project setup.py build_ext --inplace
            uv run --no-project python test.py
      - store_test_results:
          path: test-results
      - run:
          name: Build CPP only
          command: |
            source .venv/bin/activate
            mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
      - run:
          name: Run CPP tests
          command: |
            DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 ./build/tests/tests
      - run:
          name: Build small binary
          command: |
            source .venv/bin/activate
            cd build/
            cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
              -DBUILD_SHARED_LIBS=ON \
              -DMLX_BUILD_CPU=OFF \
              -DMLX_BUILD_SAFETENSORS=OFF \
              -DMLX_BUILD_GGUF=OFF \
              -DMLX_METAL_JIT=ON
            make -j `sysctl -n hw.ncpu`
      - run:
          name: Run Python tests with JIT
          command: |
            CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
              uv pip install -e .
            LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
              METAL_DEBUG_ERROR_MODE=0 \
              uv run --no-project python -m xmlrunner discover \
                -v python/tests \
                -o test-results/gpu_jit

  cuda_build_and_test:
    parameters:
      image_date:
        type: string
        default: "2023.11.1"
    machine:
      image: "linux-cuda-12:<< parameters.image_date >>"
      resource_class: gpu.nvidia.small.gen2
    steps:
      - checkout
      - restore_cache:
          keys:
            - cuda-<< parameters.image_date >>-{{ arch }}-
      - run:
          name: Install dependencies
          command: |
            sudo apt-get update
            sudo apt-get install libcudnn9-dev-cuda-12
            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
            sudo apt-get install libnccl2 libnccl-dev
            curl -sL https://github.com/ccache/ccache/releases/download/v4.11.3/ccache-4.11.3-linux-x86_64.tar.xz | tar xJf -
            sudo mv ccache-4.11.3-linux-x86_64/ccache /usr/bin/ccache
            rm -rf ccache-4.11.3-linux-x86_64
            curl -LsSf https://astral.sh/uv/install.sh | sh
      - run:
          name: Install Python package
          command: |
            uv venv
            uv pip install cmake
            DEBUG=1 CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
              uv pip install -e ".[dev]" -v
      - run:
          name: Run Python tests
          command: |
            source .venv/bin/activate
            LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
            LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
      - run:
          name: Build CPP only
          command: |
            source .venv/bin/activate
            cmake . -B build \
              -DMLX_BUILD_CUDA=ON \
              -DCMAKE_CUDA_COMPILER=`which nvcc` \
              -DCMAKE_BUILD_TYPE=DEBUG
            cmake --build build -j `nproc`
      - run:
          name: Run CPP tests
          command: ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
      - run:
          name: CCache report
          command: |
            ccache --show-stats
            ccache --zero-stats
            ccache --max-size 400MB
            ccache --cleanup
      - save_cache:
          key: cuda-<< parameters.image_date >>-{{ arch }}-{{ epoch }}
          paths:
            - /home/circleci/.cache/ccache

  build_release:
    parameters:
      python_version:
        type: string
        default: "3.9"
      xcode_version:
        type: string
        default: "16.2.0"
      build_env:
        type: string
        default: ""
      macosx_deployment_target:
        type: string
        default: ""
    macos:
      xcode: << parameters.xcode_version >>
    resource_class: m2pro.medium
    environment:
      MACOSX_DEPLOYMENT_TARGET: << parameters.macosx_deployment_target >>
    steps:
      - checkout
      - run:
          name: Install dependencies
          command: |
            brew install python@<< parameters.python_version >>
            brew install openmpi
            python<< parameters.python_version >> -m venv env
            source env/bin/activate
            pip install --upgrade pip
            pip install --upgrade cmake
            pip install nanobind==2.4.0
            pip install --upgrade setuptools
            pip install numpy
            pip install twine
            pip install build
      - run:
          name: Install Python package
          command: |
            source env/bin/activate
            env -u MACOSX_DEPLOYMENT_TARGET DEV_RELEASE=1 \
              pip install . -v
      - run:
          name: Generate package stubs
          command: |
            source env/bin/activate
            pip install typing_extensions
            python setup.py generate_stubs
      - run:
          name: Build Python package
          command: |
            source env/bin/activate
            python setup.py clean --all
            << parameters.build_env >> MLX_BUILD_STAGE=1 python -m build -w
      - when:
          condition:
            equal: ["3.9", << parameters.python_version >>]
          steps:
            - run:
                name: Build common package
                command: |
                  source env/bin/activate
                  python setup.py clean --all
                  << parameters.build_env >> MLX_BUILD_STAGE=2 python -m build -w
      - when:
          condition: << parameters.build_env >>
          steps:
            - run:
                name: Upload package
                command: |
                  source env/bin/activate
                  twine upload dist/*
      - store_artifacts:
          path: dist/

  build_linux_release:
    parameters:
      python_version:
        type: string
        default: "3.9"
      build_env:
        type: string
        default: ""
    machine:
      image: ubuntu-2204:current
      resource_class: large
    steps:
      - checkout
      - run:
          name: Build wheel
          command: |
            PYTHON=python<< parameters.python_version >>
            export DEBIAN_FRONTEND=noninteractive
            export NEEDRESTART_MODE=a
            sudo apt-get update
            TZ=Etc/UTC sudo apt-get -y install tzdata
            sudo add-apt-repository -y ppa:deadsnakes/ppa
            sudo apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
            sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
            $PYTHON -m venv env
            source env/bin/activate
            pip install --upgrade pip
            pip install --upgrade cmake
            pip install auditwheel
            pip install patchelf
            pip install build
            pip install twine
            << parameters.build_env >> pip install ".[dev]" -v
            pip install typing_extensions
            python setup.py generate_stubs
            python setup.py clean --all
            MLX_BUILD_STAGE=1 << parameters.build_env >> python -m build -w
            bash python/scripts/repair_linux.sh
      - when:
          condition:
            equal: ["3.9", << parameters.python_version >>]
          steps:
            - run:
                name: Build common package
                command: |
                  source env/bin/activate
                  python setup.py clean --all
                  << parameters.build_env >> MLX_BUILD_STAGE=2 \
                    python -m build -w
                  auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_x86_64
      - when:
          condition: << parameters.build_env >>
          steps:
            - run:
                name: Upload packages
                command: |
                  source env/bin/activate
                  twine upload wheelhouse/*.whl
      - store_artifacts:
          path: wheelhouse/

  build_cuda_release:
    parameters:
      build_env:
        type: string
        default: ""
    machine:
      image: ubuntu-2204:current
      resource_class: xlarge
    steps:
      - checkout
      - run:
          name: Build wheel
          command: |
            export DEBIAN_FRONTEND=noninteractive
            export NEEDRESTART_MODE=a
            wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/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 cuda-toolkit-12-9 libcudnn9-dev-cuda-12
            sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
            sudo apt-get install zip
            pip install auditwheel
            pip install patchelf
            pip install build
            pip install twine
            export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
            export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
            << parameters.build_env >> MLX_BUILD_STAGE=2 \
              CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
              python -m build -w
            bash python/scripts/repair_cuda.sh
      - when:
          condition: << parameters.build_env >>
          steps:
            - run:
                name: Upload package
                command: |
                  twine upload wheelhouse/*.whl
      - store_artifacts:
          path: wheelhouse/

workflows:
  build_and_test:
    when:
      and:
        - matches:
            pattern: "^(?!pull/)[-\\w]+$"
            value: << pipeline.git.branch >>
        - not: << pipeline.parameters.nightly_build >>
        - not: << pipeline.parameters.test_release >>
    jobs:
      - mac_build_and_test:
          matrix:
            parameters:
              macosx_deployment_target: ["13.5", "14.0"]
      - linux_build_and_test
      - cuda_build_and_test:
          matrix:
            parameters:
              image_date: ["2023.11.1", "2025.05.1"]
      - build_documentation 

  build_pypi_release:
    when:
      and:
        - not: << pipeline.parameters.nightly_build >>
        - not: << pipeline.parameters.test_release >>
    jobs:
      - build_release:
          filters:
            tags:
              only: /^v.*/
            branches:
              ignore: /.*/
          matrix:
            parameters:
              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
              macosx_deployment_target: ["13.5", "14.0", "15.0"]
              build_env: ["PYPI_RELEASE=1"]
              xcode_version: ["16.2.0", "15.0.0"]
            exclude:
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.9"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.10"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.11"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.12"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.13"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.9"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.10"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.11"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.12"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.13"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.9"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.10"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.11"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.12"
                build_env: "PYPI_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.13"
                build_env: "PYPI_RELEASE=1"
      - build_documentation:
          filters:
            tags:
              only: /^v.*/
            branches:
              ignore: /.*/
          upload-docs: true
      - build_linux_release:
          filters:
            tags:
              only: /^v.*/
            branches:
              ignore: /.*/
          matrix:
            parameters:
              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
              build_env: ["PYPI_RELEASE=1"]
      - build_cuda_release:
          filters:
            tags:
              only: /^v.*/
            branches:
              ignore: /.*/
          matrix:
            parameters:
              build_env: ["PYPI_RELEASE=1"]

  prb:
    when:
      matches:
        pattern: "^pull/\\d+(/head)?$"
        value: << pipeline.git.branch >>
    jobs:
      - hold:
          type: approval
      - apple/authenticate:
          context: pr-approval
      - mac_build_and_test:
          requires: [ hold ]
          matrix:
            parameters:
              macosx_deployment_target: ["13.5", "14.0"]
      - linux_build_and_test:
          requires: [ hold ]
      - cuda_build_and_test:
          requires: [ hold ]
          matrix:
            parameters:
              image_date: ["2023.11.1", "2025.05.1"]
  nightly_build:
    when:
      and:
        - equal: [ main, << pipeline.git.branch >> ]
        - << pipeline.parameters.nightly_build >>
    jobs:
      - build_release:
          matrix:
            parameters:
              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
              macosx_deployment_target: ["13.5", "14.0", "15.0"]
              xcode_version: ["16.2.0", "15.0.0"]
            exclude:
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.9"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.10"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.11"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.12"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.13"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.9"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.10"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.11"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.12"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.13"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.9"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.10"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.11"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.12"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.13"
      - build_linux_release:
          matrix:
            parameters:
              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
      - build_cuda_release

  build_dev_release:
    when:
      and:
        - equal: [ main, << pipeline.git.branch >> ]
        - << pipeline.parameters.test_release >>
    jobs:
      - build_release:
          matrix:
            parameters:
              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
              macosx_deployment_target: ["13.5", "14.0", "15.0"]
              build_env: ["DEV_RELEASE=1"]
              xcode_version: ["16.2.0", "15.0.0"]
            exclude:
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.9"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.10"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.11"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.12"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "13.5"
                xcode_version: "16.2.0"
                python_version: "3.13"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.9"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.10"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.11"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.12"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "14.0"
                xcode_version: "15.0.0"
                python_version: "3.13"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.9"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.10"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.11"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.12"
                build_env: "DEV_RELEASE=1"
              - macosx_deployment_target: "15.0"
                xcode_version: "15.0.0"
                python_version: "3.13"
                build_env: "DEV_RELEASE=1"
      - build_linux_release:
          matrix:
            parameters:
              python_version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
              build_env: ["DEV_RELEASE=1"]
      - build_cuda_release:
          matrix:
            parameters:
              build_env: ["DEV_RELEASE=1"]