Add a convenience for making local streams in python (#3355)

Author: angeloskath

mlx/compile.cpp

@@ -374,6 +374,10 @@ class CompilerCache {
     cache_.clear();
   }
 
+  bool empty() {
+    return cache_.empty();
+  }
+
  private:
   CompilerCache() {
     // Make sure the allocator is fully
@@ -1192,6 +1196,10 @@ void compile_clear_cache() {
   detail::compiler_cache().clear();
 }
 
+bool compile_cache_empty() {
+  return detail::compiler_cache().empty();
+}
+
 } // namespace detail
 
 std::function<std::vector<array>(const std::vector<array>&)> compile(

mlx/compile_impl.h

@@ -34,6 +34,9 @@ MLX_API void compile_erase(std::uintptr_t fun_id);
 // when called again.
 MLX_API void compile_clear_cache();
 
+// Return true if the cache is empty.
+MLX_API bool compile_cache_empty();
+
 bool compile_available_for_device(const Device& device);
 
 std::tuple<std::vector<array>, std::vector<array>, std::shared_ptr<void>>

python/mlx/utils.py

@@ -1,4 +1,5 @@
 # Copyright © 2023 Apple Inc.
+
 from collections import defaultdict
 from itertools import zip_longest
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union

python/src/random.cpp

@@ -18,10 +18,15 @@ using namespace nb::literals;
 
 class PyKeySequence {
  public:
-  PyKeySequence() {
-    // Destroy state before the python interpreter exits.
-    auto atexit = nb::module_::import_("atexit");
-    atexit.attr("register")(nb::cpp_function([this]() { state_.reset(); }));
+  ~PyKeySequence() {
+    if (state_.has_value()) {
+      nb::gil_scoped_acquire gil;
+      state_.reset();
+    }
+  }
+
+  void reset() {
+    state_.reset();
   }
 
   void seed(uint64_t seed) {
@@ -521,4 +526,10 @@ void init_random(nb::module_& parent_module) {
             array:
               The generated random permutation or randomly permuted input array.
       )pbdoc");
+
+  // Ensure the main thread cleanup will happen before the interpreter goes
+  // away. As a result if the other threads join the main thread we should have
+  // a clean tear-down.
+  auto atexit = nb::module_::import_("atexit");
+  atexit.attr("register")(nb::cpp_function([]() { default_key().reset(); }));
 }

python/src/stream.cpp

@@ -41,6 +41,26 @@ class PyStreamContext {
   mx::StreamContext* _inner;
 };
 
+class PyThreadLocalStream {
+ public:
+  PyThreadLocalStream(mx::Device d) : device(d) {}
+
+  mx::Stream stream() const {
+    thread_local std::unordered_map<const PyThreadLocalStream*, mx::Stream>
+        streams;
+
+    auto it = streams.find(this);
+    if (it == streams.end()) {
+      auto result = streams.emplace(this, mx::new_stream(device));
+      it = result.first;
+    }
+
+    return it->second;
+  }
+
+  mx::Device device;
+};
+
 void init_stream(nb::module_& m) {
   nb::class_<mx::Stream>(
       m,
@@ -49,6 +69,11 @@ void init_stream(nb::module_& m) {
       A stream for running operations on a given device.
       )pbdoc")
       .def_ro("device", &mx::Stream::device)
+      .def(
+          "__init__",
+          [](mx::Stream* s, const PyThreadLocalStream& tls) {
+            return new (s) mx::Stream(tls.stream());
+          })
       .def(
           "__repr__",
           [](const mx::Stream& s) {
@@ -61,7 +86,29 @@ void init_stream(nb::module_& m) {
             s == nb::cast<mx::Stream>(other);
       });
 
+  nb::class_<PyThreadLocalStream>(
+      m,
+      "ThreadLocalStream",
+      R"pbdoc(
+      A stream that will be unique per thread and can be used to run operations on a given device.
+      )pbdoc")
+      .def_ro("device", &PyThreadLocalStream::device)
+      .def(nb::init<mx::Device>())
+      .def(
+          "__repr__",
+          [](const PyThreadLocalStream& s) {
+            std::ostringstream os;
+            os << "ThreadLocalStream(" << s.device << ")";
+            return os.str();
+          })
+      .def("__eq__", [](const PyThreadLocalStream& s, const nb::object& other) {
+        auto s_other = mx::default_stream(mx::default_device());
+        return nb::try_cast<mx::Stream>(other, s_other) &&
+            s_other == s.stream();
+      });
+
   nb::implicitly_convertible<mx::Device::DeviceType, mx::Device>();
+  nb::implicitly_convertible<PyThreadLocalStream, mx::Stream>();
 
   m.def(
       "default_stream",

python/src/transforms.cpp

@@ -1463,12 +1463,22 @@ void init_transforms(nb::module_& m) {
          bool shapeless) {
         // Make sure each thread using mx.compile would clear its compile cache
         // before python interpreter exits.
-        static thread_local auto clear_cache = []() {
-          auto atexit = nb::module_::import_("atexit");
-          atexit.attr("register")(
-              nb::cpp_function(&mx::detail::compile_clear_cache));
-          return true;
+        struct ThreadCleanup {
+          ~ThreadCleanup() {
+            if (!mx::detail::compile_cache_empty()) {
+              nb::gil_scoped_acquire gil;
+              mx::detail::compile_clear_cache();
+            }
+          }
         };
+        static thread_local auto clear_cache = []() {
+          // Ensure it is created
+          mx::detail::compile_clear_cache();
+
+          // Ensure it will be cleaned up
+          return ThreadCleanup{};
+        }();
+
         return mlx_func(
             nb::cpp_function(PyCompiledFun{fun, inputs, outputs, shapeless}),
             fun,
@@ -1542,4 +1552,10 @@ void init_transforms(nb::module_& m) {
           A callable that recomputes intermediate states during gradient
           computation.
       )pbdoc");
+
+  // Ensure the main thread cleanup will happen before the interpreter goes
+  // away. As a result if the other threads join the main thread we should have
+  // a clean tear-down.
+  auto atexit = nb::module_::import_("atexit");
+  atexit.attr("register")(nb::cpp_function(&mx::detail::compile_clear_cache));
 }

python/tests/test_threads.py

@@ -0,0 +1,44 @@
+# Copyright © 2026 Apple Inc.
+
+import threading
+import unittest
+
+import mlx.core as mx
+import mlx_tests
+
+
+class TestReduce(mlx_tests.MLXTestCase):
+    def test_threadlocal_stream(self):
+        test_stream = mx.new_stream(mx.default_device())
+
+        def test_failure():
+            with self.assertRaises(RuntimeError):
+                with mx.stream(test_stream):
+                    x = mx.arange(10)
+                    mx.eval(2 * x)
+
+        t1 = threading.Thread(target=test_failure)
+        t2 = threading.Thread(target=test_failure)
+        t1.start()
+        t2.start()
+        t1.join()
+        t2.join()
+
+        test_stream = mx.ThreadLocalStream(mx.default_device())
+
+        def test_success():
+            with mx.stream(test_stream):
+                x = mx.arange(10)
+                mx.eval(2 * x)
+                self.assertEqual(x.tolist(), list(range(10)))
+
+        t1 = threading.Thread(target=test_success)
+        t2 = threading.Thread(target=test_success)
+        t1.start()
+        t2.start()
+        t1.join()
+        t2.join()
+
+
+if __name__ == "__main__":
+    mlx_tests.MLXTestRunner()