Add clear_streams API for cleanup before exit (#3395)

python/src/stream.cpp

@@ -133,6 +133,10 @@ void init_stream(nb::module_& m) {
       &mx::new_stream,
       "device"_a,
       R"pbdoc(Make a new stream on the given device.)pbdoc");
+  m.def(
+      "clear_streams",
+      &mx::clear_streams,
+      R"pbdoc(Destroy all streams created in current thread.)pbdoc");
 
   nb::class_<PyStreamContext>(m, "StreamContext", R"pbdoc(
         A context manager for setting the current device and stream.

python/tests/mlx_tests.py

@@ -9,6 +9,7 @@ os.environ["MLX_ENABLE_TF32"] = "0"
 os.environ["MLX_ENABLE_CACHE_THRASHING_CHECK"] = "0"
 
 import platform
+import sys
 import unittest
 from typing import Any, Callable, List, Tuple, Union
 
@@ -18,6 +19,8 @@ import numpy as np
 
 class MLXTestRunner(unittest.TestProgram):
     def __init__(self, *args, **kwargs):
+        # Do not exit in runTests
+        kwargs["exit"] = False
         super().__init__(*args, **kwargs)
 
     def createTests(self, *args, **kwargs):
@@ -51,6 +54,11 @@ class MLXTestRunner(unittest.TestProgram):
         filter_and_add(self.test)
         self.test = filtered_suite
 
+    def runTests(self):
+        super().runTests()
+        mx.clear_streams()
+        sys.exit(0 if self.result.wasSuccessful() else 1)
+
 
 class MLXTestCase(unittest.TestCase):
     @property

python/tests/test_threads.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx_tests
 
 
-class TestReduce(mlx_tests.MLXTestCase):
+class TestThreads(mlx_tests.MLXTestCase):
     def test_threadlocal_stream(self):
         test_stream = mx.new_stream(mx.default_device())
 
@@ -16,6 +16,7 @@ class TestReduce(mlx_tests.MLXTestCase):
                 with mx.stream(test_stream):
                     x = mx.arange(10)
                     mx.eval(2 * x)
+            mx.clear_streams()
 
         t1 = threading.Thread(target=test_failure)
         t2 = threading.Thread(target=test_failure)
@@ -31,6 +32,7 @@ class TestReduce(mlx_tests.MLXTestCase):
                 x = mx.arange(10)
                 mx.eval(2 * x)
                 self.assertEqual(x.tolist(), list(range(10)))
+            mx.clear_streams()
 
         t1 = threading.Thread(target=test_success)
         t2 = threading.Thread(target=test_success)