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:
parent
aa25f99699
commit
3614b23a2c
2 changed files with 120 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue