diff --git a/src/visionset/inference/cache.py b/src/visionset/inference/cache.py index c28c6e2..0b883a9 100644 --- a/src/visionset/inference/cache.py +++ b/src/visionset/inference/cache.py @@ -24,6 +24,7 @@ from __future__ import annotations +import threading from collections import OrderedDict from typing import Final @@ -50,13 +51,20 @@ class BoundedCache[K, V]: - """Bounded, least-recently-used, and deliberately not thread safe. + """Bounded, least-recently-used, and safe to share between threads. - Not thread safe for the same reason ``LocalTransformersProvider`` is not: a - worker process runs one task at a time, and a server handler holding a - model is already serialised by the device it is talking to. A lock here - would buy nothing and would suggest a concurrency this design does not - have. + **It did not used to be, and the assumption behind that was wrong.** This + was written as "deliberately not thread safe" on the reasoning that a worker + process runs one task at a time and a server handler is serialised by the + device it talks to. Neither holds: the suggest route is a plain ``def``, so + FastAPI runs concurrent requests in parallel threadpool threads, and two of + them reaching one cache is the ordinary case rather than an exotic one. + + The lock here protects the container's own integrity — two threads calling + ``move_to_end`` for *different* keys are still mutating one linked list — and + nothing more. It is never held across a computation, so it cannot serialise + two encodes; stopping the same value from being computed twice is a separate + problem with a separate answer, :class:`KeyedLocks` at the caller. ``get`` counts as a use, which is what makes this LRU rather than first-in-first-out: the asset somebody is clicking on repeatedly is the one @@ -69,6 +77,7 @@ def __init__(self, capacity: int) -> None: raise ValueError(f"a cache holding {capacity} things is not a cache") self._capacity = capacity self._held: OrderedDict[K, V] = OrderedDict() + self._lock = threading.Lock() @property def capacity(self) -> int: @@ -77,10 +86,11 @@ def capacity(self) -> int: def get(self, key: K) -> V | None: """What is held under that key, or ``None`` — and a hit is a use.""" - if key not in self._held: - return None - self._held.move_to_end(key) - return self._held[key] + with self._lock: + if key not in self._held: + return None + self._held.move_to_end(key) + return self._held[key] def put(self, key: K, value: V) -> V: """Hold that, evicting the least recently used if the bound is reached. @@ -88,24 +98,66 @@ def put(self, key: K, value: V) -> V: Returns the value, so a caller can write ``return cache.put(k, compute())`` rather than putting and then reading back. """ - if key in self._held: - self._held.move_to_end(key) - self._held[key] = value - while len(self._held) > self._capacity: - self._held.popitem(last=False) + with self._lock: + if key in self._held: + self._held.move_to_end(key) + self._held[key] = value + while len(self._held) > self._capacity: + self._held.popitem(last=False) return value def discard(self, key: K) -> None: """Forget that key if it is held. A no-op if it is not.""" - self._held.pop(key, None) + with self._lock: + self._held.pop(key, None) def clear(self) -> None: """Forget everything.""" - self._held.clear() + with self._lock: + self._held.clear() def __len__(self) -> int: - return len(self._held) + with self._lock: + return len(self._held) def __contains__(self, key: object) -> bool: """Membership **without** counting as a use, for tests that assert eviction.""" - return key in self._held + with self._lock: + return key in self._held + + +class KeyedLocks[K]: + """One lock per key, so that a value is computed once and only once. + + **What a bounded cache alone cannot do.** A cache tells a caller whether a + value is *there*; it has nothing to say about a value that is on its way. Two + threads asking for the same un-cached thing both see nothing and both compute + it, which is how one image came to be encoded twice and one model came to be + loaded four times in a single process. The remedy is to make the second + thread wait for the first rather than duplicate it, and a lock held across + the computation is what waiting means. + + **Per key, because a global lock would be a different bug.** Holding one lock + across every encode would make a click on one asset wait for an unrelated + click on another — turning a duplicated-work problem into a queueing problem + and costing exactly the latency the cache exists to save. Two keys never + contend here; only two callers wanting the same key do. + + The map grows with distinct keys seen in a process and is never pruned. That + is one ``threading.Lock`` — tens of bytes — per asset anybody has ever + suggested on, against the 16 MB the same asset's embedding occupies while it + is cached, so the bound that matters is already elsewhere. + """ + + def __init__(self) -> None: + self._guard = threading.Lock() + self._locks: dict[K, threading.Lock] = {} + + def for_key(self, key: K) -> threading.Lock: + """The lock for that key, made on first ask. + + The guard is held only for the lookup, never for the work the returned + lock goes on to protect. + """ + with self._guard: + return self._locks.setdefault(key, threading.Lock()) diff --git a/src/visionset/inference/providers.py b/src/visionset/inference/providers.py index 63d7982..4cee0cb 100644 --- a/src/visionset/inference/providers.py +++ b/src/visionset/inference/providers.py @@ -41,7 +41,7 @@ from typing import Any, Final from visionset.inference._extra import require -from visionset.inference.cache import DEFAULT_PROVIDER_CAPACITY, BoundedCache +from visionset.inference.cache import DEFAULT_PROVIDER_CAPACITY, BoundedCache, KeyedLocks from visionset.inference.families import ( DETECTOR_FAMILIES, SEGMENTER_FAMILIES, @@ -82,6 +82,7 @@ class ProviderPool: def __init__(self, capacity: int = DEFAULT_PROVIDER_CAPACITY) -> None: self._held: BoundedCache[_Key, Runner] = BoundedCache(capacity) + self._building: KeyedLocks[_Key] = KeyedLocks() self._builds = 0 @property @@ -99,14 +100,25 @@ def get(self, connection: InferenceConnection, *, workspace_root: Path) -> Runne Every refusal ``provider_for`` can raise is raised here too, and raised *before* anything is cached: a connection that is not ready must not leave a half-answer behind for the request that follows its download. + + **Built once even when the first several clicks arrive together.** The + window between finding nothing here and storing what was built is wide — + a build reads a config off disk and the adapter then loads gigabytes of + weights — and four concurrent first clicks went through it four times in + one process. The lock is per connection, so two connections still build + in parallel, and a refusal still caches nothing. """ key = (str(connection.id), connection.updated_at.isoformat()) held = self._held.get(key) if held is not None: return held - built = provider_for(connection, workspace_root=workspace_root) - self._builds += 1 - return self._held.put(key, built) + with self._building.for_key(key): + held = self._held.get(key) + if held is not None: + return held + built = provider_for(connection, workspace_root=workspace_root) + self._builds += 1 + return self._held.put(key, built) def clear(self) -> None: """Drop everything held. What a test does between cases.""" diff --git a/src/visionset/inference/sam_provider.py b/src/visionset/inference/sam_provider.py index a68b6c4..c4162cb 100644 --- a/src/visionset/inference/sam_provider.py +++ b/src/visionset/inference/sam_provider.py @@ -48,7 +48,7 @@ from visionset.inference import _fp16 from visionset.inference._extra import imported -from visionset.inference.cache import DEFAULT_EMBEDDING_CAPACITY, BoundedCache +from visionset.inference.cache import DEFAULT_EMBEDDING_CAPACITY, BoundedCache, KeyedLocks from visionset.kernel.domain import ( AssetSegmentation, PointPrompt, @@ -146,6 +146,7 @@ def __init__( self._embeddings: BoundedCache[UUID, tuple[Any, tuple[int, int]]] = BoundedCache( embedding_capacity ) + self._encoding: KeyedLocks[UUID] = KeyedLocks() self._encodes = 0 @property @@ -266,17 +267,33 @@ def _embedding( content-addressed: the bytes behind an id cannot change, so a hit can never be stale. Editing the *connection* is what would invalidate these, and that replaces the whole provider rather than reaching in here. + + **Once means once even when two clicks arrive together.** The route is a + plain ``def``, so FastAPI answers concurrent suggests in parallel + threadpool threads; a bare check-then-compute let two clicks on the same + un-encoded asset both encode it, which is the most expensive thing this + adapter does. The lock is taken per asset, so a click on another asset + neither waits for this one nor is waited for. + + The cache is read twice on purpose. The first read is the common case and + takes no lock at all; the second is what the loser of a race sees, and + without it the winner's work would be redone by everybody who queued + behind it. """ held = self._embeddings.get(target.asset_id) if held is not None: return held - image = Image.open(BytesIO(target.content)).convert("RGB") - size = (image.height, image.width) - inputs = processor(images=image, return_tensors="pt").to(device) - self._encodes += 1 - return self._embeddings.put( - target.asset_id, (model.get_image_embeddings(inputs["pixel_values"]), size) - ) + with self._encoding.for_key(target.asset_id): + held = self._embeddings.get(target.asset_id) + if held is not None: + return held + image = Image.open(BytesIO(target.content)).convert("RGB") + size = (image.height, image.width) + inputs = processor(images=image, return_tensors="pt").to(device) + self._encodes += 1 + return self._embeddings.put( + target.asset_id, (model.get_image_embeddings(inputs["pixel_values"]), size) + ) def _ready(self) -> tuple[Any, Any, str, bool]: if self._loaded is None: @@ -284,6 +301,29 @@ def _ready(self) -> tuple[Any, Any, str, bool]: return self._loaded def _load(self) -> tuple[Any, Any, str, bool]: + """Processor and model, once per provider. + + **``transformers`` warns here on every load, and the warning is expected.** + The published SAM 2 checkpoints declare ``model_type: sam2_video``, so + loading one into ``Sam2Model`` prints *"You are using a model of type + ``sam2_video`` to instantiate a model of type ``sam2``"*. Why that is the + right class rather than a mistake is argued where the families are + declared, in ``families.py``; what is worth recording at the load site is + that it was **measured** and not assumed. Asking + ``from_pretrained(..., output_loading_info=True)`` for + ``facebook/sam2.1-hiera-base-plus`` reports ``missing_keys: 0``, + ``unexpected_keys: 0``, ``mismatched_keys: 0`` and no errors — every + parameter this class needs came out of the checkpoint and nothing in the + checkpoint went unused, so no weight is left randomly initialised. + + The alternative class does not fit: ``Sam2VideoModel.forward`` takes an + ``inference_session`` and a frame index, which is the video-tracking path + and has no way to answer a point on a single image. + + The warning is therefore left where a reader can see it. Silencing it + would hide the same sentence on the day a checkpoint genuinely does not + match, and that day it is the only warning there is. + """ torch = imported("torch") transformers = imported("transformers") device, half = self._resolved_device(torch) diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index 0ecd209..8332354 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -12,6 +12,7 @@ DEFAULT_EMBEDDING_CAPACITY, DEFAULT_PROVIDER_CAPACITY, BoundedCache, + KeyedLocks, ) @@ -88,3 +89,19 @@ def test_the_shipped_capacities_leave_room_for_the_co_residency_the_design_assum """Two providers is a detector and a segmenter, which is what D1 describes.""" assert DEFAULT_PROVIDER_CAPACITY >= 2 assert DEFAULT_EMBEDDING_CAPACITY >= 2 + + +# --- the single-flight primitive ---------------------------------------------- + + +def test_one_key_answers_with_one_lock() -> None: + """The whole of single-flight: two callers asking about the same thing wait on + the same object. A fresh lock per call would let both compute.""" + locks: KeyedLocks[str] = KeyedLocks() + assert locks.for_key("a") is locks.for_key("a") + + +def test_two_keys_answer_with_different_locks() -> None: + """Per key, so an encode of one asset never queues behind another's.""" + locks: KeyedLocks[str] = KeyedLocks() + assert locks.for_key("a") is not locks.for_key("b") diff --git a/tests/inference/test_provider_concurrency.py b/tests/inference/test_provider_concurrency.py new file mode 100644 index 0000000..5e45efc --- /dev/null +++ b/tests/inference/test_provider_concurrency.py @@ -0,0 +1,260 @@ +"""What concurrent suggests do to the embedding cache and the provider pool. + +The suite's second threaded file, and it follows the first one's rules +(`tests/kernel/test_concurrency.py`): everything sequences on a `threading` +primitive rather than on sleeps, every thread is joined with a timeout and then +asserted dead, and nothing asserts on wall-clock — a concurrency test that hangs +is a concurrency test nobody runs, and one that measures duration fails for +reasons nobody chose. + +The arrangement under test is the one FastAPI actually builds. `suggest_region` +is a plain ``def``, so concurrent requests run in parallel threadpool threads +against one process-wide pool and one embedding cache per provider. Both were +written for a single caller, and the cost of that showed up in production as two +clicks encoding the same image twice and one process loading the model four +times. + +**Overlap is asserted through a barrier, never through timing.** A barrier that +releases proves two encodes were genuinely in flight at once; a barrier that +times out proves they were not. Neither reading depends on how fast the machine +is. +""" + +from __future__ import annotations + +import threading +from pathlib import Path +from typing import Any +from uuid import UUID, uuid4 + +import pytest +from tests.inference.stubs import StubModel, StubProcessor, StubTorch, disc + +from visionset.inference import providers as providers_module +from visionset.inference import sam_provider +from visionset.inference.providers import ProviderPool +from visionset.inference.sam_provider import LocalSamProvider +from visionset.kernel.domain import ( + ConnectionType, + InferenceConnection, + PointPrompt, + PredictionRequest, + PredictionTarget, +) +from visionset.kernel.services import InferenceConnectionService, WorkspaceService + +#: Every wait in this file. Long enough that a loaded runner does not trip it, +#: short enough that a genuine deadlock fails the suite instead of stalling it. +TIMEOUT_SECONDS = 30.0 + +#: How long a barrier that is *expected* to break waits before giving up. Only +#: ever paid on the fixed code path, where exactly one thread arrives at a +#: barrier sized for several — which is the whole point of the assertion. +LONE_ARRIVAL_SECONDS = 0.5 + +PNG = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06" + b"\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00" + b"\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +class GatedModel(StubModel): + """A model whose encode parks at a barrier, so overlap is observable. + + Holding every arriving thread inside the encode is what makes the + unfixed behaviour deterministic: no thread can finish and populate the + cache while another is still deciding whether to encode, so a suite that + sees one encode saw single-flight rather than a lucky schedule. + """ + + def __init__(self, masks: Any, scores: Any, barrier: threading.Barrier) -> None: + super().__init__(masks, scores) + self._barrier = barrier + self.lone_arrivals = 0 + + def get_image_embeddings(self, pixel_values: Any) -> str: + try: + self._barrier.wait(timeout=LONE_ARRIVAL_SECONDS) + except threading.BrokenBarrierError: + # Nobody else came. Under single-flight that is the expected shape. + self.lone_arrivals += 1 + return super().get_image_embeddings(pixel_values) + + +def gated(monkeypatch: pytest.MonkeyPatch, parties: int) -> tuple[LocalSamProvider, GatedModel]: + """A provider whose encode is observable, with everything else as shipped.""" + processor = StubProcessor([disc(20)], [0.9]) + model = GatedModel([disc(20)], [0.9], threading.Barrier(parties)) + provider = LocalSamProvider( + "some/segmenter", + "abc123", + device="cpu", + precision=None, + cache_dir=Path("/nowhere"), + connection_name="local", + ) + monkeypatch.setattr(provider, "_ready", lambda: (processor, model, "cpu", False)) + monkeypatch.setattr(sam_provider, "imported", lambda _: StubTorch()) + return provider, model + + +def click(provider: LocalSamProvider, asset: UUID) -> None: + request = PredictionRequest( + targets=(PredictionTarget(asset_id=asset, content=PNG, media_type="image/png"),), + prompt=PointPrompt(positive=((10.0, 12.0),)), + ) + list(provider.segment(request)) + + +def run_together(work: list[Any]) -> list[BaseException]: + """Start every callable at once, join them all, and hand back what raised. + + A failure inside a thread is returned rather than printed and lost, which is + the first file's rule: a thread that died quietly makes an assertion about + counts pass for the wrong reason. + """ + failures: list[BaseException] = [] + ready = threading.Barrier(len(work)) + + def guarded(task: Any) -> None: + try: + ready.wait(timeout=TIMEOUT_SECONDS) + task() + except BaseException as error: # noqa: BLE001 — re-raised by the caller + failures.append(error) + + threads = [threading.Thread(target=guarded, args=(task,)) for task in work] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=TIMEOUT_SECONDS) + assert not thread.is_alive(), "a thread outlived the timeout" + return failures + + +# --- the embedding cache ------------------------------------------------------ + + +def test_concurrent_clicks_on_one_asset_encode_it_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The stampede, which cost a production process four model loads and a stack. + + Every thread is held inside the encode until the barrier releases or breaks, + so an implementation without single-flight cannot avoid encoding four times. + """ + provider, model = gated(monkeypatch, parties=4) + asset = uuid4() + + failures = run_together([lambda: click(provider, asset)] * 4) + + assert failures == [] + assert provider.encodes == 1, "four clicks on one asset, one encode" + assert model.encodes == 1 + + +def test_every_concurrent_click_still_gets_an_answer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Single-flight means the others wait for the encode, never that they are dropped.""" + provider, _ = gated(monkeypatch, parties=4) + asset = uuid4() + answers: list[int] = [] + + def ask() -> None: + request = PredictionRequest( + targets=(PredictionTarget(asset_id=asset, content=PNG, media_type="image/png"),), + prompt=PointPrompt(positive=((10.0, 12.0),)), + ) + answers.append(len(list(provider.segment(request))[0].segments)) + + failures = run_together([ask] * 4) + + assert failures == [] + assert answers == [1, 1, 1, 1] + + +def test_two_assets_encode_at_the_same_time(monkeypatch: pytest.MonkeyPatch) -> None: + """Per-key, not global: one asset's encode must not block another's. + + The barrier is sized for both threads and is *not* expected to break. If the + two encodes were serialised behind one lock the second would never arrive, + the first would time out, and `lone_arrivals` would record it. + """ + provider, model = gated(monkeypatch, parties=2) + + failures = run_together([lambda: click(provider, uuid4()) for _ in range(2)]) + + assert failures == [] + assert provider.encodes == 2, "two different assets, two encodes" + assert model.lone_arrivals == 0, "they overlapped rather than serialising" + + +# --- the provider pool -------------------------------------------------------- + + +@pytest.fixture() +def connection(tmp_path: Path) -> Any: + workspace = WorkspaceService.init(tmp_path / "ws", name="concurrency") + try: + connections = InferenceConnectionService(workspace) + made = connections.create( + "seg", + connection_type=ConnectionType.LOCAL, + model_id="some/segmenter", + model_revision="abc123", + device="cpu", + precision="fp32", + ) + yield connections.record_weights_ready(made.id) + finally: + workspace.close() + + +class GatedBuilder: + """A stand-in for `provider_for` that parks, so overlapping builds are visible.""" + + def __init__(self, parties: int) -> None: + self._barrier = threading.Barrier(parties) + self.lone_arrivals = 0 + + def __call__(self, connection: InferenceConnection, *, workspace_root: Path) -> object: + try: + self._barrier.wait(timeout=LONE_ARRIVAL_SECONDS) + except threading.BrokenBarrierError: + self.lone_arrivals += 1 + return object() + + +def test_concurrent_first_clicks_build_one_provider( + connection: InferenceConnection, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Four concurrent first clicks loaded the model four times in production.""" + builder = GatedBuilder(parties=4) + monkeypatch.setattr(providers_module, "provider_for", builder) + pool = ProviderPool() + + failures = run_together( + [lambda: pool.get(connection, workspace_root=tmp_path) for _ in range(4)] + ) + + assert failures == [] + assert pool.builds == 1, "one connection, one provider" + assert len(pool) == 1 + + +def test_concurrent_callers_all_receive_the_same_provider( + connection: InferenceConnection, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(providers_module, "provider_for", GatedBuilder(parties=4)) + pool = ProviderPool() + seen: list[object] = [] + + failures = run_together( + [lambda: seen.append(pool.get(connection, workspace_root=tmp_path)) for _ in range(4)] + ) + + assert failures == [] + assert len(seen) == 4 + assert len({id(one) for one in seen}) == 1, "one provider, handed to everybody"