Notes on building mlx with uv

[zcbenz]

I'm tired of fighting python's package managers on different platforms, so I tried uv, it is currently not usable as the main frontend for mlx (explained below), but it can surely ease things a lot when bootstrapping on a new machine.

Here is the receipt I'm using, commands assumed running under mlx's root directory.

Setup env and install deps:

uv venv && uv pip install "pre-commit" "numpy" "setuptools<80" "nanobind==2.4.0"
uv run --no-project pre-commit install

Optional env for faster compilation, assuming you have ninja-build and ccache installed:

export CMAKE_BUILD_PARALLEL_LEVEL=16
export CMAKE_ARGS="-GNinja -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache"

If you are building CUDA backend:

export CMAKE_ARGS="-GNinja -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache -DMLX_BUILD_CUDA=ON -DMLX_FAST_COMPILE=ON -DMLX_CUDA_ARCHITECTURES=native"

First-time build which sets up editable install:

uv run --no-project setup.py develop

Incremental builds after:

uv run --no-project setup.py build_ext --inplace

Run python tests:

uv run --no-project python -m unittest discover python/tests

Run transformer example:

uv run --no-project python ~/codes/mlx-examples/transformer_lm/main.py --gpu

[zcbenz]

There are some pitfalls with uv at the time of writing.

The first thing is we can't use pip's editable install, i.e. uv pip install -e ".[dev]". In uv the native extension is always built under a newly created temporary directory, which means every time we run the command, the project will be rebuilt from scratch. Passing --no-build-isolation makes no difference.

So we are using setup.py develop with uv, and have to limit the version of setuptools.

The other thing is we are passing --no-project to uv, because there is no [project] table in pyproject.toml, which is actually good because otherwise it would opt into PEP 517 build isolation and add another layer of complexity (see also discussion in pytorch/pytorch#152276).