From ebf46b8f92a903674fc240237f20e201f3ae6b55 Mon Sep 17 00:00:00 2001 From: Michael Trefry Date: Sat, 18 Jul 2026 16:55:00 -0400 Subject: [PATCH] =?UTF-8?q?sc-12959:=20thread-safe-eval=20=E2=80=94=20eval?= =?UTF-8?q?=20lock,=20synced=20encoder=20maps,=20process-global=20default?= =?UTF-8?q?=20stream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the three hazards the sc-12937 thread-shared-streams patch left open, making the fork's threading model coherent: streams visible everywhere, all evaluation serialized, O(devices) streams. An audit of the inference repo confirmed both risks are real: MLX eval safety was discipline-only (!Send traits + docs saying "the worker serializes jobs per thread" — nothing enforced it, and two engine instances on two threads could race a shared command encoder into a Metal SIGABRT), and the worker's spawn_blocking-per-job contract leaked one CommandEncoder + MTLCommandQueue per dead blocking-pool thread. thread-safe-eval.patch (5th in the apply chain, required=true): 1. detail::eval_mutex() — process-global recursive lock held for eval_impl and synchronize(Stream). Every encode path funnels through these two, so concurrent cross-thread eval serializes instead of aborting. eval()'s host-side completion wait stays outside the lock; deadlock-free (completion handlers / cpu StreamThreads never take it; lock order eval -> default-stream-storage -> all_streams -> encoder-map has no reverse edge). 2. shared_mutex on the global encoder maps (metal + cpu): shared lock on the fallback read (element reference bound inside the locked scope — references survive rehash, iterators do not), exclusive lock on registration/clear. Closes the read-during-rehash UB. 3. Process-global default stream (was per-thread): stream/encoder count is O(devices), killing the thread-churn leak structurally. Safe because of (1). ThreadLocalStream keeps per-thread semantics. Two new regression tests (test_cross_thread_stream.rs, now 6): test_default_stream_shared_across_threads (red without this patch) and test_concurrent_eval_from_many_threads (8 threads x 25 evals contending on the shared default stream). Full suite: 972 passed / 0 failed (--test-threads=1); fmt + clippy clean; all 5 patches apply through a from-scratch FetchContent rebuild; adversarial review findings fixed (iterator-deref-after-unlock, comment scope, mutex lifetime hardening). Co-Authored-By: Claude Fable 5 --- mlx-sys/build.rs | 19 ++ mlx-sys/patches/thread-safe-eval.patch | 281 ++++++++++++++++++++ mlx-tests/tests/test_cross_thread_stream.rs | 48 ++++ 3 files changed, 348 insertions(+) create mode 100644 mlx-sys/patches/thread-safe-eval.patch diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs index aa83b1c28..bdb409b33 100644 --- a/mlx-sys/build.rs +++ b/mlx-sys/build.rs @@ -166,6 +166,24 @@ fn prepare_mlx_c_source() -> PathBuf { // new_thread_unsafe_stream, #3578), restoring cross-thread stream // visibility. Concurrent eval on the SAME stream from two threads // remains the caller's responsibility (unchanged from upstream). + // - thread-safe-eval.patch : sc-12959 — closes the three + // hazards the sc-12937 patch left open, making the fork's threading + // model coherent ("streams visible everywhere, all eval serialized, + // O(devices) streams"): + // (1) process-global recursive eval lock (detail::eval_mutex) taken by + // eval_impl and synchronize — concurrent cross-thread eval now + // serializes instead of racing a shared command encoder into a + // Metal SIGABRT. eval()'s host-side completion wait stays OUTSIDE + // the lock (synchronize(Stream) drains under it by design). This + // enforces the pmetal consumer contract (single eval thread, + // batch-dimension concurrency) that was previously discipline-only. + // (2) shared_mutex on the global encoder maps (metal + cpu) — the + // fallback read in get_command_encoder no longer races another + // thread's first-stream registration (unordered_map rehash UB). + // (3) process-global default stream (was per-thread) — a dying worker + // thread (tokio spawn_blocking churn) no longer strands one + // CommandEncoder + MTLCommandQueue per thread in the global map; + // stream count is O(devices). Safe because of (1). // // The sc-2714 (dense GEMM) and sc-2770 (fast SDPA) NAX-dispatch gate patches were // REMOVED in sc-2772. Their root cause was never the dispatch or the MLX version: the @@ -202,6 +220,7 @@ fn prepare_mlx_c_source() -> PathBuf { ("patches/command-buffer-recoverable.patch", true), ("patches/pad-copy-int64.patch", true), ("patches/thread-shared-streams.patch", true), + ("patches/thread-safe-eval.patch", true), ]; // sc-12780 idempotency guard: CMake FetchContent may re-run PATCH_COMMAND against an // mlx-src that is ALREADY patched (e.g. an incremental rebuild that does not re-fetch). diff --git a/mlx-sys/patches/thread-safe-eval.patch b/mlx-sys/patches/thread-safe-eval.patch new file mode 100644 index 000000000..11f1ee68f --- /dev/null +++ b/mlx-sys/patches/thread-safe-eval.patch @@ -0,0 +1,281 @@ +diff --git a/mlx/backend/cpu/encoder.cpp b/mlx/backend/cpu/encoder.cpp +index bb2c06ea..627c7a7c 100644 +--- a/mlx/backend/cpu/encoder.cpp ++++ b/mlx/backend/cpu/encoder.cpp +@@ -10,13 +10,16 @@ CommandEncoder& get_command_encoder(Stream s) { + auto& encoders = get_command_encoders(); + auto it = encoders.find(s.index); + if (it == encoders.end()) { ++ // pmetal (sc-12959): see the matching comment in metal/device.cpp. ++ std::shared_lock lock(get_global_command_encoders_mutex()); + auto& global_encoders = get_global_command_encoders(); +- it = global_encoders.find(s.index); +- if (it == global_encoders.end()) { ++ auto git = global_encoders.find(s.index); ++ if (git == global_encoders.end()) { + throw std::runtime_error( + fmt::format( + "There is no Stream(cpu, {}) in current thread.", s.index)); + } ++ return git->second; + } + return it->second; + } +@@ -31,4 +34,10 @@ std::unordered_map& get_global_command_encoders() { + return encoders; + } + ++std::shared_mutex& get_global_command_encoders_mutex() { ++ // Leaked like the scheduler singleton (see metal/device.cpp). ++ static auto* mtx = new std::shared_mutex(); ++ return *mtx; ++} ++ + } // namespace mlx::core::cpu +diff --git a/mlx/backend/cpu/encoder.h b/mlx/backend/cpu/encoder.h +index cd015623..b4080375 100644 +--- a/mlx/backend/cpu/encoder.h ++++ b/mlx/backend/cpu/encoder.h +@@ -2,6 +2,7 @@ + + #pragma once + ++#include + #include + + #include "mlx/array.h" +@@ -66,5 +67,7 @@ MLX_API CommandEncoder& get_command_encoder(Stream s); + + std::unordered_map& get_command_encoders(); + std::unordered_map& get_global_command_encoders(); ++// pmetal (sc-12959): see the matching comment in backend/metal/device.h. ++std::shared_mutex& get_global_command_encoders_mutex(); + + } // namespace mlx::core::cpu +diff --git a/mlx/backend/cpu/eval.cpp b/mlx/backend/cpu/eval.cpp +index 7d13fd2f..a76350f5 100644 +--- a/mlx/backend/cpu/eval.cpp ++++ b/mlx/backend/cpu/eval.cpp +@@ -16,6 +16,8 @@ void new_stream(Stream s) { + } + + void new_thread_unsafe_stream(Stream s) { ++ // pmetal (sc-12959): see the matching comment in metal/eval.cpp. ++ std::unique_lock lock(get_global_command_encoders_mutex()); + auto& encoders = get_global_command_encoders(); + encoders.try_emplace(s.index, s); + } +diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp +index 2f4b007b..00d7a1f5 100644 +--- a/mlx/backend/metal/device.cpp ++++ b/mlx/backend/metal/device.cpp +@@ -957,13 +957,19 @@ CommandEncoder& get_command_encoder(Stream s) { + auto& encoders = get_command_encoders(); + auto it = encoders.find(s.index); + if (it == encoders.end()) { ++ // pmetal (sc-12959): the global map is written by any thread's first ++ // stream creation (sc-12937), so the fallback read takes a shared lock — ++ // an unguarded find during rehash is UB. Bind the element REFERENCE ++ // inside the locked scope: references survive rehash, iterators do not. ++ std::shared_lock lock(get_global_command_encoders_mutex()); + auto& global_encoders = get_global_command_encoders(); +- it = global_encoders.find(s.index); +- if (it == global_encoders.end()) { ++ auto git = global_encoders.find(s.index); ++ if (git == global_encoders.end()) { + throw std::runtime_error( + fmt::format( + "There is no Stream(gpu, {}) in current thread.", s.index)); + } ++ return git->second; + } + return it->second; + } +@@ -978,6 +984,13 @@ std::unordered_map& get_global_command_encoders() { + return encoders; + } + ++std::shared_mutex& get_global_command_encoders_mutex() { ++ // Leaked like the scheduler singleton: worker threads may still look up ++ // encoders during static destruction. ++ static auto* mtx = new std::shared_mutex(); ++ return *mtx; ++} ++ + NS::SharedPtr new_scoped_memory_pool() { + return NS::TransferPtr(NS::AutoreleasePool::alloc()->init()); + } +diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h +index 871e95cc..02d5caa2 100644 +--- a/mlx/backend/metal/device.h ++++ b/mlx/backend/metal/device.h +@@ -251,6 +251,11 @@ MLX_API CommandEncoder& get_command_encoder(Stream s); + + std::unordered_map& get_command_encoders(); + std::unordered_map& get_global_command_encoders(); ++// pmetal (sc-12959): guards get_global_command_encoders() — readers take a ++// shared lock, writers (stream registration / clear) an exclusive one. ++// unordered_map element references remain valid after unlock (rehash never ++// invalidates references), so returning an encoder reference is safe. ++std::shared_mutex& get_global_command_encoders_mutex(); + NS::SharedPtr new_scoped_memory_pool(); + + bool is_nax_available(); +diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp +index 3d0c56ef..9e3e18cf 100644 +--- a/mlx/backend/metal/eval.cpp ++++ b/mlx/backend/metal/eval.cpp +@@ -98,6 +98,9 @@ void new_stream(Stream s) { + + void new_thread_unsafe_stream(Stream s) { + assert(s.device == Device::gpu); ++ // pmetal (sc-12959): exclusive lock — registration may race another ++ // thread's fallback read in get_command_encoder. ++ std::unique_lock lock(metal::get_global_command_encoders_mutex()); + auto& encoders = metal::get_global_command_encoders(); + auto& d = metal::device(s.device); + encoders.try_emplace(s.index, d, s.index, d.residency_set()); +@@ -160,6 +163,8 @@ void synchronize(Stream s) { + void clear_streams() { + metal::get_command_encoders().clear(); + if (is_main_thread()) { ++ // pmetal (sc-12959): exclusive lock, same reason as registration. ++ std::unique_lock lock(metal::get_global_command_encoders_mutex()); + metal::get_global_command_encoders().clear(); + } + } +diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp +index 7507917f..01e471f7 100644 +--- a/mlx/scheduler.cpp ++++ b/mlx/scheduler.cpp +@@ -3,11 +3,16 @@ + #include "mlx/scheduler.h" + #include "mlx/backend/cpu/eval.h" + #include "mlx/backend/gpu/eval.h" ++#include "mlx/transforms_impl.h" + #include "mlx/utils.h" + + namespace mlx::core { + + void synchronize(Stream s) { ++ // pmetal (sc-12959): CommandEncoder::synchronize() mutates encoder state ++ // (end_encoding + commit); take the process-global eval lock so it cannot ++ // race an in-flight eval encoding on the same encoder. ++ std::lock_guard eval_lk(detail::eval_mutex()); + if (s.device == mlx::core::Device::cpu) { + auto p = std::make_shared>(); + std::future f = p->get_future(); +diff --git a/mlx/stream.cpp b/mlx/stream.cpp +index b78ee67d..987639bd 100644 +--- a/mlx/stream.cpp ++++ b/mlx/stream.cpp +@@ -8,6 +8,7 @@ + + #include + #include ++#include + #include + #include + +@@ -15,9 +16,25 @@ namespace mlx::core { + + namespace { + ++auto& default_stream_storage_mutex() { ++ // Leaked like the scheduler singleton: worker threads may still resolve ++ // the default stream during static destruction. ++ static auto* mtx = new std::mutex(); ++ return *mtx; ++} ++ + auto& default_stream_storage(Device d) { +- // Each device has its own default stream in each thread. +- static thread_local auto default_streams = []() { ++ // pmetal (sc-12959): the default stream is PROCESS-GLOBAL (upstream: one ++ // per thread). Per-thread defaults meant every OS thread that touched MLX ++ // lazily created its own stream whose CommandEncoder + MTLCommandQueue now ++ // live in the global encoder map for process lifetime (sc-12937) — a real ++ // leak under the pmetal worker's spawn_blocking-per-job pattern, where ++ // blocking-pool threads die on idle. One shared default stream per device ++ // caps stream/encoder count at O(devices) instead of O(threads). Safe ++ // because all evaluation is serialized by detail::eval_mutex() (see ++ // transforms_impl.h). Callers must hold default_stream_storage_mutex(). ++ // ThreadLocalStream (new_thread_local_stream) keeps per-thread semantics. ++ static auto default_streams = []() { + std::array>, 2> streams; + streams[static_cast(Device::cpu)].resize(cpu::device_count()); + streams[static_cast(Device::gpu)].resize(gpu::device_count()); +@@ -43,6 +60,7 @@ Stream default_stream(Device d) { + throw std::invalid_argument( + "[default_stream] Cannot get gpu stream without gpu backend."); + } ++ std::lock_guard lock(default_stream_storage_mutex()); + auto& s = default_stream_storage(d); + if (!s.has_value()) { + s = new_stream(d.type); +@@ -55,6 +73,7 @@ void set_default_stream(Stream s) { + throw std::invalid_argument( + "[set_default_stream] Cannot set gpu stream without gpu backend."); + } ++ std::lock_guard lock(default_stream_storage_mutex()); + default_stream_storage(s.device) = s; + } + +diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp +index 9a820733..a8585c95 100644 +--- a/mlx/transforms.cpp ++++ b/mlx/transforms.cpp +@@ -77,7 +77,19 @@ int& detail::InExportTracing::counter() { + } + thread_local int detail::RetainGraph::tracing_counter{0}; + ++namespace detail { ++std::recursive_mutex& eval_mutex() { ++ // Leaked like the scheduler singleton: worker threads may still encode ++ // during static destruction. ++ static auto* mtx = new std::recursive_mutex(); ++ return *mtx; ++} ++} // namespace detail ++ + array eval_impl(std::vector outputs, bool async) { ++ // pmetal (sc-12959): serialize all evaluation process-wide (see ++ // transforms_impl.h). Scope ends before the caller's host-side wait(). ++ std::lock_guard eval_lk(detail::eval_mutex()); + std::deque tape; + + // Make an effort to choose a good output stream +diff --git a/mlx/transforms_impl.h b/mlx/transforms_impl.h +index 4fc78a1a..03a3194e 100644 +--- a/mlx/transforms_impl.h ++++ b/mlx/transforms_impl.h +@@ -2,10 +2,29 @@ + + #pragma once + ++#include ++ + #include "mlx/api.h" + + namespace mlx::core::detail { + ++// pmetal (sc-12959): process-global evaluation lock. MLX evaluation encodes ++// GPU work on per-stream command encoders and is not safe to run from two ++// threads at once; with the sc-12937 thread-shared-streams patch every stream ++// (including the process-global default stream) is visible from every thread, ++// so concurrent eval_impl calls could race the same encoder into a Metal ++// assertion (SIGABRT). The pmetal consumer contract is "single eval thread / ++// batch-dimension concurrency" (inference gen-core runtime.rs, schedule.rs); ++// this lock turns that documented discipline into an enforced invariant. ++// Recursive so an unexpected nested entry deadlocks nothing. eval()'s and ++// array::wait()'s host-side completion waits happen OUTSIDE the lock, so a ++// thread waiting on finished-encoding work does not starve other evals. ++// synchronize(Stream) intentionally drains the stream UNDER the lock, and ++// eval_impl can host-block under it in scheduler::wait_for_one (memory ++// backpressure) — both are safe: Metal completion handlers and cpu ++// StreamThread tasks never acquire this mutex. ++MLX_API std::recursive_mutex& eval_mutex(); ++ + MLX_API std::pair, std::vector> vmap_trace( + const std::function(const std::vector&)>& fun, + const std::vector& inputs, diff --git a/mlx-tests/tests/test_cross_thread_stream.rs b/mlx-tests/tests/test_cross_thread_stream.rs index 8f807ac50..f519d524f 100644 --- a/mlx-tests/tests/test_cross_thread_stream.rs +++ b/mlx-tests/tests/test_cross_thread_stream.rs @@ -116,3 +116,51 @@ fn test_evaluated_array_used_from_other_thread() { assert_eq!(out, 14.0); } + +/// sc-12959: the default stream is process-global — every thread sees the +/// same stream for the default device (upstream: one per thread). This is +/// what caps encoder/queue count at O(devices) instead of O(threads that +/// ever touched MLX). RED without the thread-safe-eval patch (per-thread +/// defaults produce distinct stream indices). +#[test] +fn test_default_stream_shared_across_threads() { + let _guard = SERIAL.lock().unwrap_or_else(|e| e.into_inner()); + + let here = mlx_rs::Stream::new().get_index().unwrap(); + let there = std::thread::spawn(|| mlx_rs::Stream::new().get_index().unwrap()) + .join() + .unwrap(); + assert_eq!( + here, there, + "default stream must be process-global (sc-12959), got {here} vs {there}" + ); +} + +/// sc-12959: concurrent evaluation from many threads must be safe — the +/// process-global eval lock serializes encoding, so this must neither SIGABRT +/// ("A command encoder is already encoding") nor corrupt results. With the +/// process-global default stream every thread contends on the SAME stream, +/// making this the worst case. Deterministic assertions per thread. +#[test] +fn test_concurrent_eval_from_many_threads() { + let _guard = SERIAL.lock().unwrap_or_else(|e| e.into_inner()); + + let handles: Vec<_> = (0..8) + .map(|i| { + std::thread::spawn(move || { + let scale = (i + 1) as f32; + for _ in 0..25 { + let a = Array::from_slice(&[scale, scale, scale, scale], &[2, 2]); + let b = ones::(&[2, 2]).unwrap(); + let sum = a.add(&b).unwrap().sum(None).unwrap(); + sum.eval().unwrap(); + assert_eq!(sum.item::(), 4.0 * (scale + 1.0)); + } + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } +}