Merge pull request #28 from ausimian/investigate-sigsegv-streams

Fix SIGABRT/SIGSEGV from concurrent mx::eval dispatch

Author: ausimian

PLAN.md

@@ -674,6 +674,31 @@ silent about it and ships a tested pattern.
 **Exit:** both patterns documented; concurrency soak demonstrates
 the streamed path is stable.
 
+**Post-M14 note — eval serialisation:**
+MLX is not thread-safe (ml-explore/mlx#2133). The Metal
+`CommandEncoder` is shared state — concurrent `mx::eval` calls from
+different OS threads crash with `"A command encoder is already
+encoding to this command buffer"` (SIGABRT) or SIGSEGV from corrupted
+encoder state. The M14 soak tests (8 concurrent workers via
+`Task.async_stream`) exposed this by being the first code path to call
+`mx::eval` from multiple dirty-CPU scheduler threads simultaneously.
+
+Fixed by:
+1. **`emily::safe_eval()`**: a mutex-serialised `mx::eval` wrapper in
+   `c_src/emily/tensor.hpp`. All eval callsites route through it.
+   Graph-building ops (regular scheduler) remain lock-free.
+2. **Removed `set_default_stream` from `with_stream/2`**: the NIF
+   mutated MLX thread-local state on BEAM scheduler threads, which
+   is unreliable since BEAM processes migrate between OS threads.
+   The process-dictionary-based stream routing was already correct.
+3. **Hardened `resolve_stream(-1)`**: the -1 fallback now reads the
+   device default directly instead of the (potentially corrupted)
+   thread-local default.
+
+The mutex serialises Metal dispatch at the cost of true concurrent
+GPU execution. See M15.5 (MLX upgrade) for the plan to restore
+concurrency via MLX's native thread-local `CommandEncoder` support.
+
 ### M15 — Native linalg
 
 `lu`, `svd`, `qr`, `cholesky`, `triangular_solve`, `eigh`,
@@ -698,6 +723,37 @@ BinaryBackend-slow. MLX exposes most natively under `mx::linalg::*`.
 **Exit:** all `mx::linalg::*`-backed callbacks pass property suite;
 remaining `via_binary` linalg paths documented with rationale.
 
+### M15.5 — MLX upgrade (build from source)
+
+Emily pins MLX 0.25.1 via pre-built binaries from `cocoa-xu/mlx-build`.
+MLX gained native thread-safety on `main` in April 2026 (thread-local
+`CommandEncoder` ml-explore/mlx#3348, `ThreadLocalStream` C++ API
+ml-explore/mlx#3405), but neither fix is in any release yet (latest:
+0.31.1). Building from source unblocks true concurrent Metal dispatch
+and removes the `safe_eval` mutex introduced in the post-M14 fix.
+
+- **Build MLX from source** (from `main` or 0.32+ when released)
+  instead of fetching the pre-built 0.25.1 tarball. Extend `mix.exs`
+  to support an `MLX_SOURCE` env var pointing at a local MLX build.
+- **Audit API changes** between 0.25.1 and target version. Emily's
+  C++ surface is narrow (core ops, linalg, streams, eval, allocator),
+  but six months of MLX releases may rename or remove functions.
+- **Adopt `ThreadLocalStream` C++ API**: each BEAM dirty-CPU thread
+  gets its own MLX stream automatically, enabling true per-thread
+  Metal command queues without the eval mutex.
+- **Remove `emily::safe_eval` mutex** once native thread-safety is
+  validated — revert to direct `mx::eval` calls.
+
+**Testing**:
+- Amplified stress test (16 workers, 100 iterations) passes without
+  mutex under the new MLX build.
+- Full test suite 20x with zero crashes.
+- Benchmark `to_binary` latency under concurrent load: confirm
+  throughput improves vs. the mutex-serialised path.
+
+**Exit:** concurrent soak tests pass without mutex; MLX build-from-source
+documented; stress test confirms concurrent Metal dispatch is stable.
+
 ### M16 — Mixed-precision training
 
 bf16 activations + f32 master weights + loss scaling is the standard

RELEASE.md

@@ -2,6 +2,16 @@
 
 ## Fixed
 
+- **Fix SIGABRT/SIGSEGV from concurrent `mx::eval` dispatch.** MLX's
+  Metal `CommandEncoder` is not thread-safe (ml-explore/mlx#2133).
+  Concurrent `mx::eval` calls from BEAM dirty-CPU scheduler threads
+  triggered `"A command encoder is already encoding"` assertions or
+  SIGSEGV from corrupted encoder state. Fixed by serialising all
+  `mx::eval` calls through `emily::safe_eval()` (mutex in
+  `c_src/emily/tensor.hpp`). Also removed `set_default_stream` calls
+  from `with_stream/2` — the NIF mutated MLX thread-local state which
+  is unreliable under BEAM process migration. Hardened
+  `resolve_stream(-1)` to avoid reading the thread-local default.
 - Relax MNIST convergence canary threshold from 97% to 96% to eliminate
   stochastic flaps (observed 96.99% on occasional runs). The test is a
   sanity gate, not a performance benchmark.

c_src/emily/tensor.hpp

@@ -15,6 +15,7 @@
 #include <mlx/mlx.h>
 
 #include <cstdint>
+#include <mutex>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -60,12 +61,38 @@ unwrap_all(const std::vector<fine::ResourcePtr<Tensor>> &tensors) {
 }
 
 // Resolve a stream index from Elixir into an mx::Stream.
-// -1 (the sentinel for "no explicit stream") falls through to the
-// thread-local default — backwards-compatible with pre-M14 code paths.
+// -1 (the sentinel for "no explicit stream") returns the default GPU
+// stream (index 0). We intentionally avoid mx::default_stream() here
+// because that reads a thread-local which can be corrupted on BEAM
+// scheduler threads — the BEAM migrates processes between OS threads,
+// so thread-local state is unreliable.
 inline mx::Stream resolve_stream(int64_t stream_index) {
   if (stream_index < 0)
-    return mx::default_stream(mx::default_device());
+    return mx::default_stream(mx::Device(mx::Device::DeviceType::gpu));
   return mx::get_stream(static_cast<int>(stream_index));
 }
 
+// MLX is not thread-safe (ml-explore/mlx#2133). In particular, the
+// Metal CommandEncoder is shared state — concurrent mx::eval calls
+// from different OS threads crash with "A command encoder is already
+// encoding to this command buffer". BEAM dirty-CPU schedulers are a
+// thread pool, so concurrent to_binary / eval NIF calls race.
+//
+// Serialise all mx::eval calls behind a single mutex until MLX gains
+// native thread-safety (expected 0.32+, see ml-explore/mlx#3348).
+inline std::mutex &eval_mutex() {
+  static std::mutex m;
+  return m;
+}
+
+inline void safe_eval(mx::array &a) {
+  std::lock_guard<std::mutex> lock(eval_mutex());
+  mx::eval(a);
+}
+
+inline void safe_eval(std::initializer_list<mx::array> arrays) {
+  std::lock_guard<std::mutex> lock(eval_mutex());
+  mx::eval(arrays);
+}
+
 } // namespace emily

c_src/emily_nif.cpp

@@ -79,7 +79,7 @@ FINE_NIF(from_binary, 0);
 fine::Term to_binary(ErlNifEnv *env, fine::ResourcePtr<Tensor> tensor, int64_t s) {
   auto stream = emily::resolve_stream(s);
   auto materialized = mx::contiguous(tensor->array, false, stream);
-  mx::eval(materialized);
+  emily::safe_eval(materialized);
 
   // Defensive: mx::contiguous is supposed to give a row-contiguous
   // layout. If it ever doesn't, we'd be aliasing a strided buffer and
@@ -113,7 +113,7 @@ FINE_NIF(dtype, 0);
 // eval/1 — force evaluation of the lazy graph rooted at this tensor.
 // Dirty CPU: waits for MLX to finish.
 fine::Ok<> eval(ErlNifEnv *, fine::ResourcePtr<Tensor> tensor) {
-  mx::eval(tensor->array);
+  emily::safe_eval(tensor->array);
   return fine::Ok<>{};
 }
 FINE_NIF(eval, ERL_NIF_DIRTY_JOB_CPU_BOUND);

c_src/stream.cpp

@@ -31,8 +31,11 @@ int64_t new_stream(ErlNifEnv *, fine::Atom device_atom) {
 FINE_NIF(new_stream, 0);
 
 // set_default_stream/1 — set the thread-local default stream.
-// Belt-and-suspenders for Emily.Stream.with_stream/2; every op NIF
-// also receives the stream index explicitly.
+//
+// WARNING: BEAM processes migrate between OS threads, so thread-local
+// state is unreliable. with_stream/2 no longer calls this — it routes
+// streams via the process dictionary instead. This NIF is retained for
+// advanced use cases but should be avoided in normal code.
 fine::Ok<> set_default_stream(ErlNifEnv *, int64_t stream_index) {
   mx::set_default_stream(mx::get_stream(static_cast<int>(stream_index)));
   return fine::Ok<>{};

lib/emily/stream.ex

@@ -23,10 +23,9 @@ defmodule Emily.Stream do
   `with_stream/2` stores the stream index in the process dictionary.
   `Emily.Backend` reads it via `Process.get(:emily_stream, -1)` and
   passes it as an explicit argument to every NIF call. Each NIF
-  resolves the index to an `mx::Stream` (or falls back to the
-  thread-local default when the index is -1). This avoids the
-  thread-local race that would occur if we relied solely on
-  `set_default_stream` — BEAM processes can migrate between OS
+  resolves the index to an `mx::Stream` (or falls back to the default
+  GPU stream when the index is -1). This is process-dictionary based,
+  not thread-local, because BEAM processes migrate between OS
   scheduler threads between NIF calls.
 
   ## Concurrent serving patterns
@@ -73,18 +72,13 @@ defmodule Emily.Stream do
   Execute `fun` with the given stream as the default for MLX ops.
 
   Stores the stream index in the process dictionary so that
-  `Emily.Backend` passes it to every NIF call. Also sets the
-  thread-local default stream as a belt-and-suspenders measure for
-  code that calls `Emily.Native` directly.
-
-  The previous stream (if any) is restored in an `after` block, so
-  nesting is safe.
+  `Emily.Backend` passes it to every NIF call. The previous stream
+  (if any) is restored in an `after` block, so nesting is safe.
   """
   @spec with_stream(t(), (-> result)) :: result when result: var
   def with_stream(%__MODULE__{index: index}, fun) when is_function(fun, 0) do
     prev = Process.get(:emily_stream)
     Process.put(:emily_stream, index)
-    Emily.Native.set_default_stream(index)
 
     try do
       fun.()
@@ -93,8 +87,6 @@ defmodule Emily.Stream do
         nil -> Process.delete(:emily_stream)
         idx -> Process.put(:emily_stream, idx)
       end
-
-      if prev, do: Emily.Native.set_default_stream(prev)
     end
   end