From 8f93ca9e5269dd44e6abe424cb69acc7b60d49cb Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Fri, 11 Jul 2025 08:17:58 -0700
Subject: [PATCH] cleanup circle, fix cuda repair

---
 .circleci/config.yml          |  3 ++-
 python/scripts/repair_cuda.sh | 10 +++++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 341655c93..faf6ca5b0 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -336,9 +336,10 @@ jobs:
             pip install typing_extensions
             python setup.py generate_stubs
             << parameters.extra_env >> python -m build --wheel
-            << parameters.extra_env >> MLX_BUILD_COMMON=1 python -m build --wheel
             auditwheel show dist/*
             auditwheel repair dist/* --plat manylinux_2_31_x86_64
+            << parameters.extra_env >> MLX_BUILD_COMMON=1 \
+              python -m build --wheel --outdir wheelhouse
       - run:
           name: Upload packages
           command: |
diff --git a/python/scripts/repair_cuda.sh b/python/scripts/repair_cuda.sh
index 21e6a977a..3584a7a8a 100644
--- a/python/scripts/repair_cuda.sh
+++ b/python/scripts/repair_cuda.sh
@@ -3,15 +3,19 @@
 auditwheel repair dist/* \
   --plat manylinux_2_35_x86_64 \
   --exclude libcublas* \
-  --exclude libnvrtc*
+  --exclude libnvrtc* \
+  -w wheel_tmp
 
-cd wheelhouse
+
+mkdir wheelhouse
+cd wheel_tmp
 repaired_wheel=$(find . -name "*.whl" -print -quit)
 unzip -q "${repaired_wheel}"
+rm "${repaired_wheel}"
 core_so=$(find mlx -name "core*.so" -print -quit)
 rpath=$(patchelf --print-rpath "${core_so}")
 rpath=$rpath:\$ORIGIN/../nvidia/cublas/lib:\$ORIGIN/../nvidia/cuda_nvrtc/lib
 patchelf --force-rpath --set-rpath "$rpath" "$core_so"
 
 # Re-zip the repaired wheel
-zip -r -q "${repaired_wheel}" .
+zip -r -q "../wheelhouse/${repaired_wheel}" .