From d1f4d291e87393f17ca599789a11142cb0820e2d Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Wed, 23 Jul 2025 16:54:19 -0700 Subject: [PATCH] Fix uv install and add dev release (#2411) * fix uv install and add dev release * fix docstring * pin cuda deps * cuda release on cpu-only machine --- .circleci/config.yml | 103 ++++++++++++++++++++++++++++++++-- python/scripts/repair_cuda.sh | 1 + python/src/ops.cpp | 5 +- setup.py | 19 ++++--- 4 files changed, 113 insertions(+), 15 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 5b5f61dcf..38c993401 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -7,6 +7,9 @@ parameters: nightly_build: type: boolean default: false + test_release: + type: boolean + default: false jobs: build_documentation: @@ -366,22 +369,27 @@ jobs: type: string default: "" machine: - image: linux-cuda-12:2024.11.1 - resource_class: gpu.nvidia.small.gen2 + image: ubuntu-2204:current + resource_class: large 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 install cuda-toolkit-12-9 sudo apt-get install libblas-dev liblapack-dev liblapacke-dev sudo apt-get install zip - python -m venv env - source env/bin/activate 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 @@ -392,7 +400,6 @@ jobs: - run: name: Upload package command: | - source env/bin/activate twine upload wheelhouse/*.whl - store_artifacts: path: wheelhouse/ @@ -405,6 +412,7 @@ workflows: pattern: "^(?!pull/)[-\\w]+$" value: << pipeline.git.branch >> - not: << pipeline.parameters.nightly_build >> + - not: << pipeline.parameters.test_release >> jobs: - mac_build_and_test: matrix: @@ -418,6 +426,7 @@ workflows: when: and: - not: << pipeline.parameters.nightly_build >> + - not: << pipeline.parameters.test_release >> jobs: - build_release: filters: @@ -601,3 +610,87 @@ workflows: 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"] diff --git a/python/scripts/repair_cuda.sh b/python/scripts/repair_cuda.sh index 155af7dd2..60586a0c2 100644 --- a/python/scripts/repair_cuda.sh +++ b/python/scripts/repair_cuda.sh @@ -4,6 +4,7 @@ auditwheel repair dist/* \ --plat manylinux_2_35_x86_64 \ --exclude libcublas* \ --exclude libnvrtc* \ + --exclude libcuda* \ -w wheel_tmp diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 178572813..5e00d3073 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4022,8 +4022,9 @@ void init_ops(nb::module_& m) { Args: file (file, str): File in which the array is saved. arrays (dict(str, array)): The dictionary of names to arrays to - be saved. metadata (dict(str, str), optional): The dictionary of - metadata to be saved. + be saved. + metadata (dict(str, str), optional): The dictionary of + metadata to be saved. )pbdoc"); m.def( "save_gguf", diff --git a/setup.py b/setup.py index 915426559..f62627749 100644 --- a/setup.py +++ b/setup.py @@ -205,9 +205,6 @@ if __name__ == "__main__": build_macos = platform.system() == "Darwin" build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "") - install_requires = [] - if build_cuda: - install_requires = ["nvidia-cublas-cu12", "nvidia-cuda-nvrtc-cu12"] version = get_version() _setup = partial( @@ -250,6 +247,7 @@ if __name__ == "__main__": "mlx.distributed_config = mlx.distributed_run:distributed_config", ] } + install_requires = [] # Release builds for PyPi are in two stages. # Each stage should be run from a clean build: @@ -269,11 +267,11 @@ if __name__ == "__main__": # - Package name is back-end specific, e.g mlx-metal if build_stage != 2: if build_stage == 1: - if build_macos: - install_requires += [f"mlx-metal=={version}"] - else: - extras["cuda"] = [f"mlx-cuda=={version}"] - extras["cpu"] = [f"mlx-cpu=={version}"] + install_requires.append( + f'mlx-metal=={version}; platform_system == "Darwin"' + ) + extras["cuda"] = [f'mlx-cuda=={version}; platform_system == "Linux"'] + extras["cpu"] = [f'mlx-cpu=={version}; platform_system == "Linux"'] _setup( name="mlx", @@ -288,9 +286,14 @@ if __name__ == "__main__": name = "mlx-metal" elif build_cuda: name = "mlx-cuda" + install_requires += [ + "nvidia-cublas-cu12==12.9.*", + "nvidia-cuda-nvrtc-cu12==12.9.*", + ] else: name = "mlx-cpu" _setup( name=name, packages=["mlx"], + install_requires=install_requires, )