diff --git a/recognization/code/gallery.py b/recognization/code/gallery.py index f1c21bc..c952ed7 100644 --- a/recognization/code/gallery.py +++ b/recognization/code/gallery.py @@ -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. diff --git a/recognization/code/test_gallery.py b/recognization/code/test_gallery.py index 57e0a2c..1154d29 100644 --- a/recognization/code/test_gallery.py +++ b/recognization/code/test_gallery.py @@ -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