tngo/product-recognition-demo
1
0
Fork 0
forked from vnguyen/product-recognition-demo
Code Pull requests Activity

fix: serialize gallery inference so Metal cannot be entered twice

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

74
recognization/code/gallery.py
View file

@ -56,6 +56,17 @@ DINOV2_EMBED_DIM = 384
CACHE_VERSION = 1
# One thread at a time runs the gallery backbone. Metal's shader library is
# process-global and is not thread-safe while kernels are still being compiled
# into it, which is what the first forward on a fresh process does; two threads
# reaching it together segfault the whole server. Module-level because the state
# being protected is the process's, not one gallery's.
# ponytail: serializes every embed. Audit C5's batching embedding worker is the
# upgrade path, and a second MPS consumer (--device mps puts YOLO on Metal too)
# would need to take this same lock rather than its own.
_MPS_INFER_LOCK = threading.Lock()
def dinov2_weights_path() -> str:
override = os.environ.get("DORAYAKI_DINOV2_WEIGHTS")
if override:
@ -444,11 +455,12 @@ class ProductGallery:
if not kept:
return None
try:
model = self._load_model()
tensor = torch.from_numpy(np.stack(kept)).to(self.device)
with torch.no_grad():
features = model(tensor)
features = torch.nn.functional.normalize(features, dim=1).cpu().numpy()
with _MPS_INFER_LOCK:
model = self._load_model()
tensor = torch.from_numpy(np.stack(kept)).to(self.device)
with torch.no_grad():
features = model(tensor)
features = torch.nn.functional.normalize(features, dim=1).cpu().numpy()
except Exception:
return None
return rows, features.astype(np.float32, copy=False)
@ -492,6 +504,16 @@ class ProductGallery:
self._cache_write(key, rows, features)
else:
rows, features = cached
# A cache hit skips the embedding forward, so without this the model
# is built and Metal's kernels are compiled by whichever classify
# worker asks first - up to eight racing into one process-wide cache.
# The readiness path loads the gallery before the socket accepts, so
# the warmup is free and no worker meets a cold model.
#
# Best-effort on purpose: a cache hit is served without a usable
# backbone by contract, and _MPS_INFER_LOCK already makes the lazy
# path safe. This only spares the first worker the compile.
self._warm_backbone()
self._entries = [
GalleryEntry(sku_id=sku_id, view=path, embedding=features[index])
for index, (sku_id, path) in enumerate(rows)
@ -508,13 +530,45 @@ class ProductGallery:
import numpy as np
import torch
model = self._load_model()
tensor = torch.from_numpy(np.stack([_prepare(crop, self.image_size)])).to(self.device)
with torch.no_grad():
features = model(tensor)
features = torch.nn.functional.normalize(features, dim=1).cpu().numpy()
with _MPS_INFER_LOCK:
model = self._load_model()
tensor = torch.from_numpy(np.stack([_prepare(crop, self.image_size)])).to(self.device)
with torch.no_grad():
features = model(tensor)
features = torch.nn.functional.normalize(features, dim=1).cpu().numpy()
return features[0]
def _warm_backbone(self) -> None:
"""Build the backbone and compile its kernels, single-threaded.
Called under self._lock from the load path, so the lock order is
self._lock then _MPS_INFER_LOCK - the same order _embed_views takes, and
no path takes them reversed. The forward is a dummy canvas: what is
being warmed is the model and the shader library, not any one crop.
Metal only. A CPU gallery has no shader library to compile and a cache
hit is served without touching the backbone by contract, so warming one
would spend a model load to protect against nothing.
"""
import numpy as np
if self._resolved_device() != "mps":
return
try:
self.embed(np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8))
except Exception:
pass
def _resolved_device(self) -> str:
"""The device the backbone will run on, without building it."""
if self.device is not None:
return self.device
import torch
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
return "mps"
return "cpu"
def _rank_index(self) -> Tuple[List[str], Any, Any]:
"""Segment bounds for the per-SKU maximum, in first-appearance order.

56
recognization/code/test_gallery.py
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