fix: serialize gallery inference so Metal cannot be entered twice

The server segfaulted a few seconds into every live scan that used the crop
lane. The crash report put the fault in libtorch's Metal shader library
(`MetalShaderLibrary::exec_unary_kernel` inside an unordered_map insert),
reached from `copy_cast_kernel_mps` under `clone`/`reshape`.

Reproduced deterministically: eight threads calling `embed()` on a fresh
process crash it, while the same eight threads after one warm forward are
fine. Metal's shader library is process-global and is not thread-safe while
kernels are still being compiled into it, which is exactly what the first
forward does. The detector already serializes its predict for this reason;
the gallery did not.

Two changes:

* A module-level `_MPS_INFER_LOCK` covers the whole MPS region of a forward -
  model load, `.to(device)`, the forward, and the normalize that follows it.
  Normalize dispatches its own Metal kernels, so leaving it outside the lock
  left the same first-compile race open. Module-level because the state being
  protected is the process's, not one gallery's.
* A cache hit skips the embedding forward, so the model used to be built and
  the kernels compiled by whichever classify worker asked first - up to eight
  racing into one process-wide cache. The load path now warms the backbone
  single-threaded under the existing lock, on the cache-hit branch too. The
  readiness path loads the gallery before the socket accepts, so the warmup
  is free. It is Metal-only and best-effort: a CPU gallery has no shader
  library to compile, and a cache hit is served without a usable backbone by
  contract.

The regression test records overlap rather than timing: a fake backbone that
sleeps and counts concurrent entries, run from four threads, asserting the
count never exceeds one. It fails on the previous gallery.py and needs no MPS.
This commit is contained in:
kiyo9w 2026-09-28 17:22:55 +09:00
parent aa25f99699
commit 3614b23a2c
2 changed files with 120 additions and 10 deletions

View file

@ -532,6 +532,62 @@ class HeadPinTests(unittest.TestCase):
self.assertIsNone(head_class_set(["sku_a", "sku_a", ""]))
class ConcurrentInferenceTests(unittest.TestCase):
"""The backbone is never entered by two threads at once.
Metal's shader library is process-global and is not thread-safe while
kernels are still being compiled into it, so two threads reaching the first
forward together segfault the whole server. This is a contract test, not a
timing test: it records overlap rather than measuring speed.
"""
def test_two_threads_never_enter_the_backbone_together(self):
import time
import threading
import tempfile
with tempfile.TemporaryDirectory() as directory:
a_store(directory, names=("sku_a", "sku_b"))
gallery = ProductGallery(directory, ["sku_a", "sku_b"], device="cpu", image_size=32)
inside = 0
overlapped = []
gate = threading.Lock()
class SlowBackbone:
def __call__(self, tensor):
nonlocal inside
with gate:
inside += 1
if inside > 1:
overlapped.append(inside)
time.sleep(0.01)
with gate:
inside -= 1
return tensor.mean(dim=(2, 3)).repeat(1, gallery_module.DINOV2_EMBED_DIM // 3)
gallery._load_model = lambda: SlowBackbone()
crop = np.zeros((24, 32, 3), dtype=np.uint8)
errors = []
def work():
try:
for _ in range(6):
gallery.embed(crop)
except Exception as error: # pragma: no cover - reported below
errors.append(repr(error))
threads = [threading.Thread(target=work) for _ in range(4)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
self.assertEqual(errors, [])
self.assertEqual(overlapped, [], "two threads ran the backbone at the same time")
class TrainerPinTests(unittest.TestCase):
def test_head_training_loads_the_pinned_file_and_keeps_preprocess(self):
import finetune_identity