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/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..248e93a766 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(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()