fix(mac): serialise all Apple-MPS torch work behind one device lock

CLAUDE.md

@@ -366,10 +366,17 @@ The hot paths (vector finding, vector orientation mapping) are GPU-accelerated.
 **Avoid per-item Python loops around tiny kernels.** The original coarse seed looped templates × angles in Python (hundreds of thousands of tiny kernel launches) → **289s** for a realistic library. Rewriting it as a polar-histogram angular cross-correlation (one batched FFT, no Python loop) → **1.6s**. When a GPU step is slow, the cause is almost always a Python loop launching small kernels or a blown-up intermediate tensor — not the arithmetic. Reach for FFTs / matmuls / `scatter_add_` over explicit loops, and **chunk the batch dimension** to bound the largest intermediate (e.g. the `(P,T,n_a)` correlation is chunked over patterns) rather than materialising it whole (a full `(P,T,…)` tensor OOMs).
 
 **Windows + torch-CUDA-autograd gotchas (hard-won):**
-- `backward()` segfaults when run off the **main thread** under CUDA on Windows. The GPU orientation fit therefore runs **inline on the main thread** (it's only ~1-2s of compute) with an `on_yield` callback (pumps the event loop / flushes pending work) so the UI stays responsive; it is *not* offloaded to a worker thread. The CPU fallback (numpy/scipy) is thread-safe and *does* run on a worker.
+- `backward()` segfaults when run off the **main thread** under CUDA on Windows, the first time it runs on a thread whose autograd engine isn't initialised. Both OM handlers (`om_run`, `vom_run`) *do* dispatch the fit to a daemon worker via `run_on_worker`, so the two mitigations in the code are load-bearing: `warmup_autograd()` runs one trivial backward on the **dispatch thread** before the worker starts (CUDA-gated; a no-op on MPS/CPU), and the refine loop pins backward to its calling thread (next bullet). The fit takes an `on_yield` callback (pumps the event loop / flushes pending work) so the UI stays responsive. The CPU fallback (numpy/scipy) is thread-safe and needs neither.
 - Pin backward to the calling thread with `torch.autograd.set_multithreading_enabled(False)` around the refine loop.
 - Yield *inside* the step loop (every ~12 steps), not just per anneal stage — otherwise the window freezes for seconds and the progress bar appears stuck. Drive the progress label from the compute's own `progress(done,total)` callback; do not derive % from a lagging live-preview cell count.
 
+**Mac + Apple-MPS is NOT thread-safe — every torch user takes ONE shared device lock (`spyde/device_lock.py`):**
+- Two threads submitting to Metal at once corrupts the command encoder and raises an **uncatchable native SIGSEGV** — no `try`/`except` sees it, the backend process dies, and the user gets the "Analysis backend stopped" dialog. Observed faulting frames: `at::native::relu_mps_` → `MetalShaderLibrary::exec_unary_kernel`, and `at::native::zero_` → `fill_mps_kernel` → `[AGXG13GFamilyComputeContext setComputePipelineState:]`. Reproduces in ~30 lines: 4 threads running one small conv/ReLU net on MPS segfaults (or hangs) within a few dozen iterations; the same loop behind one lock runs clean.
+- **`DEVICE_LOCK` in `spyde/device_lock.py` is THE lock** — a reentrant, process-wide `RLock`. Use `accelerator_lock(device)`, which is a **null context off MPS** (CUDA is thread-safe and its stream concurrency is a deliberate throughput win — never serialise it). Current holders: the neural batch / single-frame preview / calibration / cold model load (`models.infer.load_model`'s `.to(device)` + smoke-test forward), the torch NXCORR + DoG peak finders, and the batched vector-orientation fit.
+- **A lock only works if EVERY participant takes it.** This crash existed because the lock was private to `find_vectors_torch` and only the neural *batch* and NXCORR paths took it — so the live preview (fires on navigator moves), the calibration, the cold model load, and the whole vector-orientation fit submitted unserialised. Adding a new torch call site without the lock silently re-opens it. `test_device_lock.py` pins the contract (shared-object identity + every entry point).
+- **Long GPU work hands the device back at its yield points**, it does not hold the lock end-to-end: `compute_vector_orientation_gpu` wraps the caller's `on_yield` so each yield does `mps_sync()` → release → yield → re-acquire. So a concurrent preview waits one yield window (~12 refine steps), not the whole anneal. **Always `mps_sync()` before releasing** — handing off while kernels are still in flight lets the next thread submit into a live encoder, which is the race itself.
+- The failure is **probabilistic per submission**, so a short stress test is NOT a regression gate: SpyDE's own preview path does heavy GIL-holding numpy work between short MPS bursts, and 6 threads × 40 frames with the lock disabled still completed cleanly here. It bites over a real multi-minute run (thousands of frames + a fit). Trust the lock-contract unit tests, not a crash-or-not stress run.
+
 **Numerical traps that only show on real/strained data** (unit tests on uniform synthetic data won't catch these):
 - *Rotation-branch ambiguity*: a centrosymmetric diffraction pattern is invariant under 180°, so the seed may pick θ≈±180° where an SPD-bounded stretch can't fit → garbage strain. Collapse the seed angle into `(−π/2, π/2]`.
 - *Coarse-σ shrink bias*: at wide Gaussian σ the soft-assign cost is minimised by shrinking the template (spurious negative strain pinned at the cap). Fit a **rigid pose through the coarse stages and only release the strain DOF at the finest σ**, where the true strain is the global minimum.

spyde/actions/find_vectors_neural.py

@@ -108,13 +108,20 @@ def _find_vectors_single_frame_neural(
     replaced by the robust disk-mean frame intensity)."""
     from spyde import models
     from spyde.actions.find_vectors import _disk_mean_intensity
+    from spyde.device_lock import accelerator_lock
 
     f = np.asarray(frame, dtype=np.float32)
     model, device = models.get_model(model_id)
-    pred = models.detect(model, f, device, thresh=float(threshold),
-                         min_distance=int(min_distance),
-                         bg_sigma=(float(bg_sigma) if bg_sigma is not None else None),
-                         spot_diameter=(2.0 * spot_radius) if spot_radius else None)
+    # Serialise the MPS forward against every other torch user in the process.
+    # This single-frame path is the LIVE PREVIEW: it fires from navigator moves on
+    # worker threads, so two previews (or a preview and a batch/orientation fit)
+    # could submit to Metal at once — an uncatchable SIGSEGV that killed the whole
+    # backend. The batch path was already locked; this one silently was not.
+    with accelerator_lock(device):
+        pred = models.detect(model, f, device, thresh=float(threshold),
+                             min_distance=int(min_distance),
+                             bg_sigma=(float(bg_sigma) if bg_sigma is not None else None),
+                             spot_diameter=(2.0 * spot_radius) if spot_radius else None)
     pred = np.asarray(pred, dtype=np.float32).reshape(-1, 3)
     pred = _apply_beamstop(pred, beamstop_mask, f.shape)
 
@@ -205,14 +212,20 @@ def _refine_block(peaks_grid, ny, nx, flat, beamstop_mask):
 def _mps_forward_lock():
     """Whole-device lock that serialises the MPS forward on Mac — ONE forward at a
     time per process (fix 3). Thread-race crashes on MPS come from concurrent
-    forwards hitting a single Metal context; this gives the neural path the same
-    serial discipline the NXCORR torch path has via ``_GPU_LOCK``. Reuses that
-    same lock object so a mixed neural+NXCORR run shares one device-serialisation
-    budget. Off Mac this returns a null context (CUDA keeps its own concurrency)."""
+    forwards hitting a single Metal context.
+
+    This is ``spyde.device_lock.DEVICE_LOCK``, THE process-wide accelerator lock —
+    the same object the NXCORR torch paths, the single-frame preview, the neural
+    calibration and the batched vector-orientation fit take, so every torch user
+    in the process shares one serialisation budget. (It used to be a lock private
+    to ``find_vectors_torch``, and only some of those callers took it — which is
+    exactly how a preview or an orientation fit ended up submitting to Metal
+    concurrently and segfaulting the backend.) Off Mac this returns a null context
+    (CUDA keeps its own concurrency)."""
     if sys.platform != "darwin":
         return None
-    from spyde.actions.find_vectors_torch import _GPU_LOCK
-    return _GPU_LOCK
+    from spyde.device_lock import DEVICE_LOCK
+    return DEVICE_LOCK
 
 
 def _neural_block(b4d, threshold, min_dist, subpixel, beamstop_mask, model_id,
@@ -383,15 +396,19 @@ def calibrate_neural(sample_frames, *, sigma=1.0, model_id=None, spot_radius=Non
     from scipy.ndimage import gaussian_filter
 
     from spyde import models
+    from spyde.device_lock import accelerator_lock
 
     model, device = models.get_model(model_id)
     frames = []
     for f in sample_frames:
         f = np.asarray(f, np.float32)
         # light single-frame smoothing as a stand-in for nav-blur on isolated samples
         frames.append(gaussian_filter(f, 0.6) if sigma else f)
-    cal = models.calibrate(model, frames, device,
-                           spot_diameter=(2.0 * spot_radius) if spot_radius else None)
+    # Calibration runs a sweep of real forwards, so it needs the same device
+    # serialisation as the preview and the batch (see device_lock.py).
+    with accelerator_lock(device):
+        cal = models.calibrate(model, frames, device,
+                               spot_diameter=(2.0 * spot_radius) if spot_radius else None)
     log.info("[find_vectors] neural calibration: bg_sigma=%.1f thresh=%.2f "
              "scale=%.2f conf=%.2f", cal["bg_sigma"], cal["thresh"],
              cal["scale_factor"], cal.get("confidence", float("nan")))

spyde/actions/find_vectors_torch.py

@@ -19,13 +19,18 @@
 """
 from __future__ import annotations
 
-import threading
 from typing import Optional
 
 import numpy as np
 
+from spyde.device_lock import DEVICE_LOCK
+
 _TORCH_DEV = "unset"        # cache: torch.device | None
-_GPU_LOCK = threading.Lock()
+# THE process-wide accelerator lock (spyde.device_lock) — NOT a lock private to
+# this module. Every torch user in the process must serialise against the SAME
+# object or concurrent Metal submission segfaults the backend; see device_lock.py
+# for the crash stacks. Kept under the old name because callers import it.
+_GPU_LOCK = DEVICE_LOCK
 
 
 def torch_gpu_device():

spyde/actions/vector_orientation_gpu.py

@@ -29,6 +29,9 @@
 
 import numpy as np
 
+from spyde.device_lock import (
+    DEVICE_LOCK, accelerator_lock, is_mps, mps_sync,
+)
 from spyde.signals.diffraction_vectors import COL_KX, COL_KY, COL_INTENSITY
 
 log = logging.getLogger(__name__)
@@ -359,9 +362,63 @@ def compute_vector_orientation_gpu(
     t: Optional[int] = None, progress=None, stopped_flag=None,
     shm_name: Optional[str] = None, n_seed_angles: int = 72,
     refine_steps: int = 60, on_yield=None,
+) -> Optional[VectorOrientationResult]:
+    """Public entry point for the batched fit — see ``_compute_vector_orientation_batched``
+    for the fit itself. This wrapper adds Apple-MPS device serialisation.
+
+    On MPS, torch is NOT thread-safe: two threads submitting to the Metal
+    command encoder at once is an uncatchable native SIGSEGV that kills the whole
+    backend process (the user sees "Analysis backend stopped"). This fit runs on a
+    worker thread and can easily overlap the neural spot detector's live preview
+    or a find-vectors batch, so it must hold THE process-wide device lock
+    (``spyde.device_lock``) like every other torch user.
+
+    The lock is handed back at each ``on_yield`` point (the fit already yields
+    every ~12 refine steps to keep the UI responsive), so a concurrent live
+    preview waits for one yield window rather than the whole anneal. The device is
+    synchronised before each hand-off so no kernels are still in flight.
+
+    Off MPS this is a straight pass-through — CUDA is thread-safe and its
+    concurrency is a deliberate throughput win, so that path is unchanged.
+    """
+    fit_kwargs = dict(
+        params=params, t=t, progress=progress, stopped_flag=stopped_flag,
+        shm_name=shm_name, n_seed_angles=n_seed_angles,
+        refine_steps=refine_steps, on_yield=on_yield,
+    )
+    dev = select_device()
+    if not is_mps(dev):
+        return _compute_vector_orientation_batched(vectors, lib, **fit_kwargs)
+
+    def _yield_releasing():
+        # Called from inside the fit while we hold the lock: quiesce Metal, let
+        # any waiting torch user through, then take the device back.
+        mps_sync()
+        DEVICE_LOCK.release()
+        try:
+            if on_yield is not None:
+                on_yield()
+        finally:
+            DEVICE_LOCK.acquire()
+
+    # Installed even when the caller passed no on_yield: the hand-off window is
+    # the whole point, and the fit calls this on a cadence it already chose.
+    fit_kwargs["on_yield"] = _yield_releasing
+    with accelerator_lock(dev):
+        return _compute_vector_orientation_batched(vectors, lib, **fit_kwargs)
+
+
+def _compute_vector_orientation_batched(
+    vectors, lib: TemplateLibrary, params: Optional[dict] = None,
+    t: Optional[int] = None, progress=None, stopped_flag=None,
+    shm_name: Optional[str] = None, n_seed_angles: int = 72,
+    refine_steps: int = 60, on_yield=None,
 ) -> Optional[VectorOrientationResult]:
     """Fit the whole field on the GPU in one batched pass. See module docstring.
 
+    NB: call the public ``compute_vector_orientation_gpu`` wrapper instead — it
+    adds the Apple-MPS device serialisation this function does NOT do.
+
     progress(done, total) is called a handful of times (seed, each anneal
     stage, decode) so the GUI can show coarse progress + paint the live buffer.
 

spyde/device_lock.py

@@ -0,0 +1,93 @@
+"""device_lock.py — THE process-wide accelerator serialisation lock.
+
+Apple-MPS (Metal) crashes when two Python threads submit work to the device at
+the same time: the Metal compute-context / shader-library bookkeeping inside
+``libtorch_cpu`` is not thread-safe, and a concurrent submission corrupts the
+command encoder. The failure is an **uncatchable native SIGSEGV** — no
+``try``/``except`` sees it, the whole backend process dies, and the user gets
+"Analysis backend stopped". Observed faulting frames (macOS 26.4, torch 2.13):
+
+    at::native::relu_mps_ -> MetalShaderLibrary::exec_unary_kernel
+    at::native::zero_ -> fill_mps_kernel -> [AGXG13GFamilyComputeContext
+                                             setComputePipelineState:]
+
+Reproduced standalone in ~30 lines: 4 threads running the same small conv/ReLU
+net on MPS segfaults (or hangs) within a few dozen iterations; the identical
+loop with one shared lock runs clean indefinitely.
+
+**A lock only works if EVERY participant takes it.** SpyDE has several
+independent torch users that can run on worker threads at the same time — the
+neural spot detector (batch, single-frame preview, and calibration), the torch
+NXCORR/DoG peak finders, and the batched vector-orientation fit. They must all
+serialise against ONE lock object, which is what this module owns. Before this
+existed the lock lived in ``find_vectors_torch`` and only the *batch* neural
+path and the NXCORR paths took it, so a live navigator preview or an
+orientation fit ran concurrently with them and took the process down.
+
+Scope: MPS only. On CUDA the driver is thread-safe and concurrent streams are a
+deliberate throughput win (the find_vectors GPU lane runs several submitters);
+serialising there would be a pure regression. So ``accelerator_lock`` is a null
+context off MPS, and the CUDA behaviour is byte-for-byte unchanged.
+
+The lock is **reentrant**: nested acquisitions on one thread are normal here
+(e.g. the neural batch takes the device lock and then the ``_gpu_slots``
+semaphore, and helpers may re-enter), and an ``RLock`` keeps that from
+self-deadlocking.
+"""
+from __future__ import annotations
+
+import contextlib
+import logging
+import threading
+
+log = logging.getLogger(__name__)
+
+# THE lock. Reentrant so nesting on one thread is safe. Module-global, so every
+# importer in this process shares the single object.
+DEVICE_LOCK = threading.RLock()
+
+
+def is_mps(device) -> bool:
+    """True when ``device`` (a torch.device, a string, or None) is Apple-MPS."""
+    if device is None:
+        return False
+    return getattr(device, "type", str(device)) == "mps"
+
+
+def mps_sync() -> None:
+    """Drain the Metal command queue. Called before releasing the lock so the
+    device is quiesced at the hand-off — releasing while kernels are still in
+    flight would let the next thread start submitting into a live encoder, which
+    is the very race the lock exists to prevent. Best-effort: a torch build
+    without ``torch.mps.synchronize`` (or with MPS uninitialised) just no-ops."""
+    try:
+        import torch
+
+        mps = getattr(torch, "mps", None)
+        sync = getattr(mps, "synchronize", None)
+        if sync is not None:
+            sync()
+    except Exception as e:  # noqa: BLE001 - never let a sync failure escape
+        log.debug("mps_sync skipped: %s", e)
+
+
+@contextlib.contextmanager
+def accelerator_lock(device=None, *, sync: bool = True):
+    """Serialise a block of torch work against every other MPS user in-process.
+
+    ``device`` is the torch device the block will run on; when it is not MPS this
+    is a null context (see module docstring — CUDA concurrency is intentional).
+    Pass ``device=None`` to mean "unknown/any", which also skips locking.
+
+    ``sync=False`` skips the drain on release — only for a block that has already
+    synchronised (or read a tensor back to the host, which syncs implicitly).
+    """
+    if not is_mps(device):
+        yield
+        return
+    with DEVICE_LOCK:
+        try:
+            yield
+        finally:
+            if sync:
+                mps_sync()

spyde/models/infer.py

@@ -112,17 +112,26 @@ def load_model(ckpt_path, device=None, arch: dict | None = None):
         levels=arch.get("levels", ck.get("levels", 2)),
     )
     model.load_state_dict(ck["state_dict"])
-    model = model.to(device)
-    model.eval()
-    if device.type == "mps":
-        # Smoke-test the forward ONCE at load: an MPS build missing an op this
-        # net uses should degrade to CPU here, not fail on every chunk.
-        try:
-            with torch.no_grad():
-                model(torch.zeros(1, in_ch, 32, 32, device=device))
-        except Exception:
-            device = torch.device("cpu")
-            model = model.to(device)
+    # Moving the weights to MPS and the smoke-test forward are both real device
+    # submissions, and ``registry.get_model`` deliberately calls this OUTSIDE its
+    # cache lock (so a slow download doesn't block the cache) — so several
+    # find-vectors chunk threads can land here at once on a cold cache. Concurrent
+    # Metal submission is an uncatchable SIGSEGV, so take THE process-wide device
+    # lock for the device-touching part. Null context off MPS. See spyde.device_lock.
+    from spyde.device_lock import accelerator_lock
+
+    with accelerator_lock(device):
+        model = model.to(device)
+        model.eval()
+        if device.type == "mps":
+            # Smoke-test the forward ONCE at load: an MPS build missing an op this
+            # net uses should degrade to CPU here, not fail on every chunk.
+            try:
+                with torch.no_grad():
+                    model(torch.zeros(1, in_ch, 32, 32, device=device))
+            except Exception:
+                device = torch.device("cpu")
+                model = model.to(device)
     return model, device