Set install rpath of python bindings with cmake (#2934)

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

@@ -22,7 +22,7 @@ runs:
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w
 
-        auditwheel repair dist/* \
+        auditwheel repair dist/mlx_cuda*.whl \
           --plat manylinux_2_35_${{ inputs.arch }} \
           --exclude libcublas* \
           --exclude libcuda* \

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

@@ -24,13 +24,18 @@ runs:
         pip install -e ".[dev]" -v
         pip install typing_extensions
         python setup.py generate_stubs
+
     - name: Build Python package
       shell: bash
       run: |
         pip install auditwheel patchelf build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
-        bash python/scripts/repair_linux.sh ${{ inputs.arch }}
+        auditwheel repair dist/mlx-*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libmlx.so* \
+          --only-plat
+
     - name: Build backend package
       if: ${{ inputs.build-backend }}
       shell: bash

python/scripts/repair_linux.sh

@@ -1,20 +0,0 @@
-#!/bin/bash
-
-auditwheel repair dist/* \
-  --plat manylinux_2_35_${1} \
-  --only-plat \
-  --exclude libmlx* \
-  -w wheel_tmp
-
-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="\$ORIGIN/lib"
-patchelf --force-rpath --set-rpath "$rpath" "$core_so"
-python ../python/scripts/repair_record.py ${core_so}
-
-# Re-zip the repaired wheel
-zip -r -q "../wheelhouse/${repaired_wheel}" .

python/scripts/repair_record.py

@@ -1,33 +0,0 @@
-import base64
-import glob
-import hashlib
-import sys
-
-filename = sys.argv[1]
-
-
-# Compute the new hash and size
-def urlsafe_b64encode(data: bytes) -> bytes:
-    return base64.urlsafe_b64encode(data).rstrip(b"=")
-
-
-hasher = hashlib.sha256()
-with open(filename, "rb") as f:
-    data = f.read()
-    hasher.update(data)
-hash_str = urlsafe_b64encode(hasher.digest()).decode("ascii")
-size = len(data)
-
-# Update the record file
-record_file = glob.glob("*/RECORD")[0]
-with open(record_file, "r") as f:
-    lines = [l.split(",") for l in f.readlines()]
-
-for l in lines:
-    if filename == l[0]:
-        l[1] = hash_str
-        l[2] = f"{size}\n"
-
-with open(record_file, "w") as f:
-    for l in lines:
-        f.write(",".join(l))

python/src/CMakeLists.txt

@@ -55,8 +55,10 @@ target_link_libraries(core PRIVATE mlx)
 
 if(BUILD_SHARED_LIBS)
   if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
-    target_link_options(core PRIVATE -Wl,-rpath,@loader_path/lib)
+    set_target_properties(core PROPERTIES INSTALL_RPATH "@loader_path/lib")
   else()
-    target_link_options(core PRIVATE -Wl,-rpath,\$ORIGIN/lib)
+    set_target_properties(core PROPERTIES INSTALL_RPATH "\$ORIGIN/lib")
   endif()
+  # Do not add build dir to rpath.
+  set_target_properties(core PROPERTIES BUILD_WITH_INSTALL_RPATH ON)
 endif()