diff --git a/CLAUDE.md b/CLAUDE.md index 267bd240..c16ea7af 100644 --- a/CLAUDE.md +++ b/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. diff --git a/spyde/actions/find_vectors_neural.py b/spyde/actions/find_vectors_neural.py index 71fc0265..db970b01 100644 --- a/spyde/actions/find_vectors_neural.py +++ b/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,6 +396,7 @@ 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 = [] @@ -390,8 +404,11 @@ def calibrate_neural(sample_frames, *, sigma=1.0, model_id=None, spot_radius=Non 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"))) diff --git a/spyde/actions/find_vectors_torch.py b/spyde/actions/find_vectors_torch.py index 83a2f5ed..010d9975 100644 --- a/spyde/actions/find_vectors_torch.py +++ b/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(): diff --git a/spyde/actions/vector_orientation_gpu.py b/spyde/actions/vector_orientation_gpu.py index b2944168..6db86602 100644 --- a/spyde/actions/vector_orientation_gpu.py +++ b/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. diff --git a/spyde/device_lock.py b/spyde/device_lock.py new file mode 100644 index 00000000..d76b2ac6 --- /dev/null +++ b/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() diff --git a/spyde/models/infer.py b/spyde/models/infer.py index b4acd9e0..d8983a03 100644 --- a/spyde/models/infer.py +++ b/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 diff --git a/spyde/tests/migrated/test_device_lock.py b/spyde/tests/migrated/test_device_lock.py new file mode 100644 index 00000000..eecd3876 --- /dev/null +++ b/spyde/tests/migrated/test_device_lock.py @@ -0,0 +1,279 @@ +"""Pins the Apple-MPS device-serialisation contract. + +Concurrent Metal submission from two threads is an uncatchable native SIGSEGV +that kills the whole backend process ("Analysis backend stopped"). The guard is +ONE process-wide lock — and a lock only works if EVERY torch user takes it. + +The bug these tests exist for: the lock lived privately in ``find_vectors_torch`` +and only the neural BATCH and the NXCORR paths took it, so the single-frame live +preview, the neural calibration, the cold model load, and the batched +vector-orientation fit all submitted to Metal unserialised. An orientation-mapping +run overlapping a preview took the backend down with: + + at::native::relu_mps_ -> MetalShaderLibrary::exec_unary_kernel + at::native::zero_ -> fill_mps_kernel -> setComputePipelineState: + +These tests are Qt-free, GPU-free and fast: they use a FAKE device whose ``.type`` +is ``"mps"`` so the locking contract is exercised on any machine (CI included) +without touching Metal. +""" +from __future__ import annotations + +import threading + +import numpy as np +import pytest + +from spyde import device_lock +from spyde.device_lock import DEVICE_LOCK, accelerator_lock, is_mps + + +class _FakeDev: + """Stands in for ``torch.device('mps')`` — only ``.type`` is consulted.""" + + def __init__(self, type_="mps"): + self.type = type_ + + def __str__(self): + return self.type + + +def _lock_is_held_by_me() -> bool: + """True when THIS thread already owns DEVICE_LOCK. ``RLock.acquire`` from the + owning thread always succeeds (it's reentrant), so probe from another thread.""" + got = [] + + def probe(): + got.append(DEVICE_LOCK.acquire(blocking=False)) + if got[-1]: + DEVICE_LOCK.release() + + t = threading.Thread(target=probe) + t.start() + t.join() + return not got[0] + + +class TestSharedLockIdentity: + """Every torch user must serialise against the SAME object.""" + + def test_find_vectors_torch_reuses_the_shared_lock(self): + from spyde.actions import find_vectors_torch + + assert find_vectors_torch._GPU_LOCK is DEVICE_LOCK + + def test_mps_forward_lock_is_the_shared_lock(self): + import sys + + from spyde.actions.find_vectors_neural import _mps_forward_lock + + if sys.platform == "darwin": + assert _mps_forward_lock() is DEVICE_LOCK + else: + # Off Mac the neural path deliberately keeps CUDA concurrency. + assert _mps_forward_lock() is None + + def test_lock_is_reentrant(self): + """Nesting on one thread is normal (batch takes the device lock and then + the _gpu_slots semaphore); a plain Lock would self-deadlock.""" + with accelerator_lock(_FakeDev()): + with accelerator_lock(_FakeDev()): + assert _lock_is_held_by_me() + + +class TestAcceleratorLock: + def test_noop_off_mps(self): + """CUDA/CPU must NOT be serialised — concurrent CUDA streams are a + deliberate throughput win, so locking there would be a pure regression.""" + for dev in (None, _FakeDev("cpu"), _FakeDev("cuda")): + with accelerator_lock(dev): + assert not _lock_is_held_by_me(), f"{dev} should not lock" + + def test_holds_lock_on_mps(self): + with accelerator_lock(_FakeDev()): + assert _lock_is_held_by_me() + assert not _lock_is_held_by_me() + + def test_released_on_exception(self): + with pytest.raises(RuntimeError): + with accelerator_lock(_FakeDev()): + raise RuntimeError("boom") + assert not _lock_is_held_by_me() + + def test_syncs_before_release(self, monkeypatch): + """The device must be quiesced at the hand-off: releasing while kernels + are in flight lets the next thread submit into a live encoder.""" + calls = [] + monkeypatch.setattr(device_lock, "mps_sync", lambda: calls.append("sync")) + with accelerator_lock(_FakeDev()): + pass + assert calls == ["sync"] + + def test_serialises_two_threads(self): + """The actual invariant: two threads never inside the block at once.""" + overlap = [] + inside = [0] + counter_lock = threading.Lock() + + def work(): + for _ in range(50): + with accelerator_lock(_FakeDev()): + with counter_lock: + inside[0] += 1 + if inside[0] > 1: + overlap.append(inside[0]) + with counter_lock: + inside[0] -= 1 + + ts = [threading.Thread(target=work) for _ in range(4)] + for t in ts: + t.start() + for t in ts: + t.join() + assert overlap == [], f"concurrent entry: {overlap}" + + +class TestOrientationFitTakesLock: + """The vector-orientation fit runs on a worker thread and previously took no + lock at all — the ``zero_``/``fill_mps_kernel`` crash.""" + + def test_fit_runs_under_the_lock(self, monkeypatch): + from spyde.actions import vector_orientation_gpu as vog + + held = [] + monkeypatch.setattr(vog, "select_device", lambda: _FakeDev()) + monkeypatch.setattr( + vog, "_compute_vector_orientation_batched", + lambda *a, **k: held.append(_lock_is_held_by_me()) or "result") + + assert vog.compute_vector_orientation_gpu(None, None) == "result" + assert held == [True] + assert not _lock_is_held_by_me(), "lock leaked after the fit" + + def test_yield_hands_the_device_back(self, monkeypatch): + """The fit yields every ~12 refine steps; each yield must release the + device so a concurrent preview waits one yield window, not the whole + anneal — and must re-acquire it afterwards.""" + from spyde.actions import vector_orientation_gpu as vog + + observed = {} + + def fake_fit(*a, **k): + # k["on_yield"] is the wrapper's releasing shim, which calls the + # caller's on_yield in the middle of the hand-off window. + assert _lock_is_held_by_me() + k["on_yield"]() + observed["held_after_yield"] = _lock_is_held_by_me() + return "ok" + + user_calls = [] + monkeypatch.setattr(vog, "select_device", lambda: _FakeDev()) + monkeypatch.setattr(vog, "_compute_vector_orientation_batched", fake_fit) + vog.compute_vector_orientation_gpu( + None, None, on_yield=lambda: user_calls.append( + _lock_is_held_by_me())) + + assert user_calls == [False], "device not released during the yield" + assert observed["held_after_yield"] is True, "device not re-acquired" + + def test_no_deadlock_against_concurrent_users(self, monkeypatch): + """Adding a lock risks deadlock, and the fit's release/re-acquire around + each yield is the delicate part. Run the real yield protocol against + preview-style acquirers on other threads and require completion.""" + from spyde.actions import vector_orientation_gpu as vog + + monkeypatch.setattr(vog, "select_device", lambda: _FakeDev()) + + def fake_fit(*a, **k): + for _ in range(40): + k["on_yield"]() # release -> hand off -> re-acquire + return "ok" + + monkeypatch.setattr(vog, "_compute_vector_orientation_batched", fake_fit) + done, errors = [], [] + + def fit_thread(): + try: + done.append(vog.compute_vector_orientation_gpu(None, None)) + except Exception as e: # noqa: BLE001 + errors.append(repr(e)) + + def preview_thread(): + try: + for _ in range(60): + with accelerator_lock(_FakeDev()): + pass + done.append("preview") + except Exception as e: # noqa: BLE001 + errors.append(repr(e)) + + ts = [threading.Thread(target=fit_thread)] + [ + threading.Thread(target=preview_thread) for _ in range(3)] + for t in ts: + t.start() + for t in ts: + t.join(timeout=30) + + assert not [t for t in ts if t.is_alive()], "deadlock: thread never finished" + assert errors == [], errors + assert len(done) == 4 + assert not _lock_is_held_by_me(), "lock leaked" + + def test_off_mps_is_a_passthrough(self, monkeypatch): + """CUDA path must be byte-for-byte unchanged: no lock, original on_yield + passed straight through.""" + from spyde.actions import vector_orientation_gpu as vog + + seen = {} + sentinel = object() + + def fake_fit(*a, **k): + seen["on_yield"] = k.get("on_yield") + seen["held"] = _lock_is_held_by_me() + return "cuda" + + monkeypatch.setattr(vog, "select_device", lambda: _FakeDev("cuda")) + monkeypatch.setattr(vog, "_compute_vector_orientation_batched", fake_fit) + vog.compute_vector_orientation_gpu(None, None, on_yield=sentinel) + + assert seen["on_yield"] is sentinel + assert seen["held"] is False + + +class TestNeuralPathsTakeLock: + """The single-frame preview and the calibration were the unlocked neural + entry points (the two concurrent ``relu_mps_`` threads in the crash).""" + + def test_single_frame_preview_locks(self, monkeypatch): + from spyde import models + from spyde.actions import find_vectors_neural as fvn + + held = [] + monkeypatch.setattr(models, "get_model", + lambda mid=None: (object(), _FakeDev())) + + def fake_detect(model, f, device, **kw): + held.append(_lock_is_held_by_me()) + return np.zeros((0, 3), dtype=np.float32) + + monkeypatch.setattr(models, "detect", fake_detect) + fvn._find_vectors_single_frame_neural( + np.zeros((32, 32), dtype=np.float32), threshold=0.5, min_distance=3) + assert held == [True], "preview forward ran unserialised on MPS" + + def test_calibration_locks(self, monkeypatch): + from spyde import models + from spyde.actions import find_vectors_neural as fvn + + held = [] + monkeypatch.setattr(models, "get_model", + lambda mid=None: (object(), _FakeDev())) + + def fake_calibrate(model, frames, device, **kw): + held.append(_lock_is_held_by_me()) + return {"bg_sigma": 12.0, "thresh": 0.5, "scale_factor": 1.0, + "confidence": 0.9} + + monkeypatch.setattr(models, "calibrate", fake_calibrate) + fvn.calibrate_neural([np.zeros((32, 32), dtype=np.float32)]) + assert held == [True], "calibration forwards ran unserialised on MPS"