Add clear_streams API for cleanup before exit (#3395)

mlx/backend/cuda/eval.cpp

@@ -80,4 +80,8 @@ void synchronize(Stream s) {
   cu::get_command_encoder(s).synchronize();
 }
 
+void clear_streams() {
+  cu::get_command_encoders().clear();
+}
+
 } // namespace mlx::core::gpu

mlx/backend/gpu/eval.h

@@ -15,5 +15,6 @@ void new_stream(Stream stream);
 void eval(array& arr);
 void finalize(Stream s);
 void synchronize(Stream s);
+void clear_streams();
 
 } // namespace mlx::core::gpu

mlx/backend/metal/eval.cpp

@@ -87,4 +87,8 @@ void synchronize(Stream s) {
   metal::get_command_encoder(s).synchronize();
 }
 
+void clear_streams() {
+  metal::get_command_encoders().clear();
+}
+
 } // namespace mlx::core::gpu

mlx/backend/no_gpu/eval.cpp

@@ -23,4 +23,6 @@ void synchronize(Stream) {
   throw std::runtime_error("[gpu::synchronize]  GPU backend is not available");
 }
 
+void clear_streams() {}
+
 } // namespace mlx::core::gpu

mlx/scheduler.cpp

@@ -20,6 +20,10 @@ void synchronize() {
   synchronize(default_stream(default_device()));
 }
 
+void clear_streams() {
+  gpu::clear_streams();
+}
+
 namespace scheduler {
 
 Scheduler::Scheduler() {

mlx/stream.h

@@ -39,4 +39,7 @@ MLX_API void synchronize();
 /* Synchronize with the provided stream. */
 MLX_API void synchronize(Stream);
 
+/* Destroy all streams created in current thread. */
+MLX_API void clear_streams();
+
 } // namespace mlx::core

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)

tests/scheduler_tests.cpp

@@ -48,8 +48,11 @@ TEST_CASE("test default stream in threads") {
   for (size_t i = 0; i < num_threads; ++i) {
     threads.emplace_back([&thread_streams, &mtx]() {
       auto s = default_stream(gpu::is_available() ? Device::gpu : Device::cpu);
-      std::lock_guard lock(mtx);
-      thread_streams.insert(s);
+      {
+        std::lock_guard lock(mtx);
+        thread_streams.insert(s);
+      }
+      clear_streams();
     });
   }
 
@@ -80,12 +83,27 @@ TEST_CASE("test access stream in other thread") {
     } catch (const std::runtime_error&) {
       error_caught = true;
     }
+    clear_streams();
   });
   t.join();
 
   CHECK(error_caught);
 }
 
+TEST_CASE("test new stream in threads") {
+  std::vector<std::thread> threads;
+  for (int i = 0; i < 1; ++i) {
+    threads.emplace_back([]() {
+      auto s = new_stream(default_device());
+      eval(arange(10, s));
+      clear_streams();
+    });
+  }
+  for (auto& t : threads) {
+    t.join();
+  }
+}
+
 TEST_CASE("test get streams") {
   // Initialize default CPU stream before querying
   default_stream(Device::cpu);