From 40b2a93bf901d21ae6f407e9f9cd79e73335b8e6 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Thu, 27 Aug 2026 06:01:49 +0400 Subject: [PATCH 1/2] Fix a deadlock after test succeed due to mx.clear_streams() holding GIL while synchronously waiting for Metal --- .github/actions/test-macos/action.yml | 2 +- .github/actions/test-wheel/action.yml | 7 +-- .github/workflows/build_and_test.yml | 2 +- .github/workflows/release.yml | 2 +- python/src/stream.cpp | 2 + python/tests/mlx_tests.py | 4 +- python/tests/test_zero_copy.py | 62 +++++++++++++++++++++++++++ 7 files changed, 71 insertions(+), 10 deletions(-) diff --git a/.github/actions/test-macos/action.yml b/.github/actions/test-macos/action.yml index aa043fb66c..55e5acc5c5 100644 --- a/.github/actions/test-macos/action.yml +++ b/.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 diff --git a/.github/actions/test-wheel/action.yml b/.github/actions/test-wheel/action.yml index a29e3c1dc7..e060876503 100644 --- a/.github/actions/test-wheel/action.yml +++ b/.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 diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index a4b5ddf8f3..d590d10b12 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -85,7 +85,7 @@ jobs: matrix: macos-target: ['14.0', '15.0', '26.2'] toolkit: ['cpu', 'metal', 'jit'] - runs-on: 'macos-26-xlarge' + runs-on: 'macos-26' needs: check_lint steps: - uses: actions/checkout@v7 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 0445da0361..b53f003ff4 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -141,7 +141,7 @@ jobs: strategy: matrix: python-version: *pyver - runs-on: 'macos-26-xlarge' + runs-on: 'macos-26' env: *build-env steps: - uses: actions/checkout@v7 diff --git a/python/src/stream.cpp b/python/src/stream.cpp index 467518e991..796da8bcd5 100644 --- a/python/src/stream.cpp +++ b/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(s)) { mx::synchronize(); } else { diff --git a/python/tests/mlx_tests.py b/python/tests/mlx_tests.py index 84ccf51375..2b60f46154 100644 --- a/python/tests/mlx_tests.py +++ b/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) diff --git a/python/tests/test_zero_copy.py b/python/tests/test_zero_copy.py index 231e2e45ee..bc0e0af420 100644 --- a/python/tests/test_zero_copy.py +++ b/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(120, 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(120, 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() From 97910768c47bbdb1d211883c37339c13ed948306 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Thu, 27 Aug 2026 09:53:50 +0400 Subject: [PATCH 2/2] use faster machine and pick the right timeout --- .github/workflows/build_and_test.yml | 2 +- .github/workflows/release.yml | 2 +- python/tests/test_zero_copy.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index d590d10b12..a4b5ddf8f3 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -85,7 +85,7 @@ jobs: matrix: macos-target: ['14.0', '15.0', '26.2'] toolkit: ['cpu', 'metal', 'jit'] - runs-on: 'macos-26' + runs-on: 'macos-26-xlarge' needs: check_lint steps: - uses: actions/checkout@v7 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index b53f003ff4..0445da0361 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -141,7 +141,7 @@ jobs: strategy: matrix: python-version: *pyver - runs-on: 'macos-26' + runs-on: 'macos-26-xlarge' env: *build-env steps: - uses: actions/checkout@v7 diff --git a/python/tests/test_zero_copy.py b/python/tests/test_zero_copy.py index bc0e0af420..248e93a766 100644 --- a/python/tests/test_zero_copy.py +++ b/python/tests/test_zero_copy.py @@ -113,7 +113,7 @@ def _submit_work(a): # in Python would never run. def test_synchronize_releases_gil(self): a = self._adopted_source(1024) - faulthandler.dump_traceback_later(120, exit=True) + faulthandler.dump_traceback_later(30, exit=True) try: for _ in range(4): y = self._submit_work(a) @@ -136,7 +136,7 @@ def worker(): except Exception as e: errors.append(e) - faulthandler.dump_traceback_later(120, exit=True) + faulthandler.dump_traceback_later(30, exit=True) try: thread = threading.Thread(target=worker) thread.start()