Fix deadlock caused by mx.clear_streams() holding GIL

.github/actions/test-macos/action.yml

@@ -47,7 +47,7 @@ runs:
           echo "::endgroup::"
 
           echo "::group::Run Python tests"
-          python -m unittest discover -v python/tests
+          uv run python/tests/run.py -v
           echo "::endgroup::"
 
           if ${{ inputs.toolkit != 'cpu' }} ; then

.github/actions/test-wheel/action.yml

@@ -46,9 +46,4 @@ runs:
           echo "No matching backend wheel to install"
           exit 1
         fi
-        if ${{ runner.os == 'macOS' }} ; then
-          # FIXME: run.py does not quit in macOS CI.
-          python -m unittest discover -v python/tests
-        else
-          uv run python/tests/run.py -v --failfast
-        fi
+        uv run python/tests/run.py -v --failfast

python/src/stream.cpp

@@ -140,6 +140,7 @@ void init_stream(nb::module_& m) {
       "clear_streams",
       []() {
         reset_random_state();
+        nb::gil_scoped_release nogil;
         mx::clear_streams();
       },
       R"pbdoc(Destroy all streams created in current thread.)pbdoc");
@@ -190,6 +191,7 @@ void init_stream(nb::module_& m) {
   m.def(
       "synchronize",
       [](mx::StreamOrDevice s) {
+        nb::gil_scoped_release nogil;
         if (std::holds_alternative<std::monostate>(s)) {
           mx::synchronize();
         } else {

python/tests/mlx_tests.py

@@ -15,7 +15,9 @@ def __init__(self, *args, **kwargs):
         # Do not exit in runTests
         kwargs["exit"] = False
         super().__init__(*args, **kwargs)
-        # Do cleanup before exiting
+
+    def runTests(self):
+        super().runTests()
         mx.clear_streams()
         sys.exit(0 if self.result.wasSuccessful() else 1)
 

python/tests/test_zero_copy.py

@@ -1,6 +1,8 @@
 # Copyright © 2024 Apple Inc.
 
+import faulthandler
 import gc
+import threading
 import unittest
 
 import mlx.core as mx
@@ -83,6 +85,66 @@ def test_adopt_in_loop_not_recycled(self):
             mx.eval(r)
         self.assertTrue(True)  # reaching here without crashing is the assertion
 
+    def _adopted_source(self, n):
+        # A square source buffer that mx.asarray can adopt, else a skip.
+        if not mx.metal.is_available():
+            self.skipTest("copy=False requires Metal")
+        a = np.zeros((n, n), dtype=np.float32)
+        if a.ctypes.data % 16384 != 0:
+            self.skipTest("source buffer not page-aligned; adopt path not taken")
+        return a
+
+    @staticmethod
+    def _submit_work(a):
+        # The adopted array is an input of the first matmul and is dropped here,
+        # so only the completion handler of the command buffer that is in flight
+        # keeps its Python owner. The caller must keep the output until the end.
+        n = a.shape[0]
+        w = mx.ones((n, n))
+        y = mx.asarray(a, copy=False) @ w
+        for _ in range(4):
+            y = y @ w
+        mx.async_eval(y)
+        return y
+
+    # A stream callback can free an adopted buffer, which takes the GIL. A call
+    # that waits for a stream must release the GIL, else the two deadlock.
+    # faulthandler reports such a deadlock, its timer is a C thread. A watchdog
+    # in Python would never run.
+    def test_synchronize_releases_gil(self):
+        a = self._adopted_source(1024)
+        faulthandler.dump_traceback_later(30, exit=True)
+        try:
+            for _ in range(4):
+                y = self._submit_work(a)
+                mx.synchronize()
+                del y
+        finally:
+            faulthandler.cancel_dump_traceback_later()
+
+    def test_clear_streams_releases_gil(self):
+        a = self._adopted_source(1024)
+        errors = []
+
+        # clear_streams destroys the streams of the calling thread, so it runs
+        # in a thread that ends right after.
+        def worker():
+            try:
+                y = self._submit_work(a)
+                mx.clear_streams()
+                del y
+            except Exception as e:
+                errors.append(e)
+
+        faulthandler.dump_traceback_later(30, exit=True)
+        try:
+            thread = threading.Thread(target=worker)
+            thread.start()
+            thread.join()
+        finally:
+            faulthandler.cancel_dump_traceback_later()
+        self.assertEqual(errors, [])
+
 
 if __name__ == "__main__":
     unittest.main()