diff --git a/slide2vec/runtime/dense_regions.py b/slide2vec/runtime/dense_regions.py index f413297..0b7aead 100644 --- a/slide2vec/runtime/dense_regions.py +++ b/slide2vec/runtime/dense_regions.py @@ -31,6 +31,9 @@ from __future__ import annotations +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager from dataclasses import dataclass from typing import Callable, Iterator, Sequence @@ -44,6 +47,41 @@ from slide2vec.runtime.slide_encode import slide_encode_autocast_ctx +@contextmanager +def _pinned_intraop_threads() -> Iterator[None]: + """Pin torch (and cv2) intra-op threads to 1 for the duration of a dense read stream. + + The pooled ``DataLoader`` path gets this for free — each worker *process* calls + ``torch.set_num_threads(1)`` on start. The threaded dense prefetcher shares one + process, so we pin process-wide while the stream runs and restore on exit. Pinning + collapses the per-ROI dense transform (``get_dense_transform`` + ``pad``) from ~71 ms + at torch's default 32 intra-op threads (which thrash on a loaded box) to ~5 ms at a + single thread. The GPU forward is unaffected — intra-op threads are CPU-only — so the + only visible effect on the encode is a ~1e-6 reduction-order perturbation on a CPU + forward, well within the reproducibility tolerance (see ADR 0002). + """ + prev_torch = torch.get_num_threads() + try: + import cv2 # local import: cv2 is heavy and only needed on the read path + except Exception: # pragma: no cover - cv2 always present in the runtime image + cv2 = None + prev_cv2 = None + if cv2 is not None: + try: + prev_cv2 = cv2.getNumThreads() + except Exception: # pragma: no cover - defensive + prev_cv2 = None + torch.set_num_threads(1) + if cv2 is not None and prev_cv2 is not None: + cv2.setNumThreads(1) + try: + yield + finally: + torch.set_num_threads(prev_torch) + if cv2 is not None and prev_cv2 is not None: + cv2.setNumThreads(prev_cv2) + + def _resolve_output_dtype(output_dtype: "torch.dtype | None", precision: str) -> "torch.dtype": """Resolve the dtype emitted grids are materialized in. @@ -182,6 +220,7 @@ def iter_regions_dense( attention_blocks: tuple[int, ...] = (-1,), attention_include_registers: bool = False, batch_size: int = 1, + num_workers: int | None = None, precision: str = "fp32", output_dtype: "torch.dtype | None" = None, dense_transform: Callable | None = None, @@ -191,8 +230,20 @@ def iter_regions_dense( Yields one ``(d, grid_h, grid_w)`` grid per coordinate, in coordinate order, in the model's compute ``precision`` by default (fp16 runs yield fp16 grids; see ``output_dtype``). Regions are read and encoded one ``batch_size`` chunk at a time, so - resident host memory is bounded by ``batch_size`` rather than by a slide's ROI count - (the loop holds at most one batch of grids resident — no per-slide accumulation). + resident host memory is bounded by ``batch_size`` (serial) / the prefetch window + (``num_workers + batch_size`` in-flight reads, see ``num_workers``) rather than by a + slide's ROI count — no per-slide accumulation. + + Reads and the encoder forward run serially on the main thread by default + (``num_workers is None``): read a chunk → forward → yield, one chunk at a time. This + leaves the GPU idle on CPU JPEG decode. Pass ``num_workers=K`` to read regions through + a ``ThreadPoolExecutor`` of width ``K`` (threads, not processes — cuCIM/hs2p release + the GIL during a region read), double-buffered so the next batch's reads overlap the + current forward, with torch/cv2 intra-op threads pinned to 1 for the read path. Batches + are the *same* fixed ``batch_size`` chunks of ``coordinates`` and are gathered in + coordinate order, so the prefetcher only changes read *timing*, never the forward batch + size ``B`` or its membership (ADR 0002: grids are reproducible to a tolerance; thread + pinning perturbs a CPU forward by ~1e-6, a GPU forward not at all). Injectable core: takes a constructed dense-capable ``model`` (with ``encode_tiles_dense`` / ``encode_tiles_attention`` / ``patch_size`` / @@ -218,6 +269,13 @@ def iter_regions_dense( sliding is internal to extraction. overlap: fractional window overlap in ``[0, 1)`` for the sliding path (ignored when ``window_size is None``); the stride is ``window * (1 - overlap)``. + num_workers: prefetch width. ``None`` (default) reads serially on the main thread + (the legacy path). A positive ``K`` reads regions through a ``K``-thread + ``ThreadPoolExecutor`` that overlaps reads with the forward. Callers size this + from ``ExecutionOptions.num_workers_per_gpu`` (resolved by + ``slide2vec.runtime.cpu_budget.resolve_on_the_fly_num_workers``); it is taken as + a plain int here so the primitive stays injectable/offline-testable. ``< 1`` + is rejected eagerly. output_dtype: torch dtype the grids are materialized in. ``None`` (default) follows the compute ``precision`` — fp16 → fp16, fp32 → fp32, bf16 → fp32 (numpy has no bfloat16). Pass e.g. ``torch.float32`` to force a lossless cache regardless of @@ -231,6 +289,8 @@ def iter_regions_dense( """ if pad_mode not in _PAD_MODES: raise ValueError(f"unsupported pad_mode {pad_mode!r}; expected one of {sorted(_PAD_MODES)}") + if num_workers is not None and int(num_workers) < 1: + raise ValueError(f"num_workers must be a positive int or None, got {num_workers!r}") resolved_output_dtype = _resolve_output_dtype(output_dtype, precision) geometry = compute_dense_geometry(target_size=target_size, patch_size=model.patch_size) if dense_transform is None: @@ -269,33 +329,69 @@ def _read_padded(location: tuple[int, int]) -> torch.Tensor: tensor, geometry, pad_mode=pad_mode, image_pad_value=image_pad_value ) - def _stream() -> Iterator[np.ndarray]: + def _encode_batch(tiles: list[torch.Tensor]) -> Iterator[np.ndarray]: + """Stack one fixed chunk of read tiles, run the forward, yield its grids. + + The batch is exactly the chunk's tiles in coordinate order — identical batch size + ``B`` and membership whether the tiles were read serially or via the prefetcher, so + reordering reads never changes the forward's grouping (acceptance c). + """ + batch = torch.stack(tiles).to(device, non_blocking=True) + # Every batch goes through the one windowed primitive: window_size=None + # short-circuits to a single whole-tile forward (byte-identical to the + # whole-region encode), so there is no separate whole-region branch. + out = encode_dense_sliding( + model, + batch, + geometry=geometry, + window_size=window_size, + overlap=overlap, + encode_fn=encode_fn, + ) + if out.ndim != 4: + raise ValueError( + f"{feature_kind} encode returned a {out.ndim}-D tensor; expected (B, d, gh, gw)." + ) + batch_np = out.detach().to(resolved_output_dtype).cpu().numpy() + for i in range(batch_np.shape[0]): + # Standalone C-contiguous copy: a per-row view would pin the whole + # batch alive (the blended sliding output is contiguous, so a view of + # it would not copy). ``.copy()`` always copies, in C order. + yield batch_np[i].copy() + + def _stream_serial() -> Iterator[np.ndarray]: with torch.inference_mode(), slide_encode_autocast_ctx(device, precision): for start in range(0, len(coords), step): chunk = coords[start : start + step] - batch = torch.stack([_read_padded(loc) for loc in chunk]).to( - device, non_blocking=True - ) - # Every batch goes through the one windowed primitive: window_size=None - # short-circuits to a single whole-tile forward (byte-identical to the - # whole-region encode), so there is no separate whole-region branch. - out = encode_dense_sliding( - model, - batch, - geometry=geometry, - window_size=window_size, - overlap=overlap, - encode_fn=encode_fn, - ) - if out.ndim != 4: - raise ValueError( - f"{feature_kind} encode returned a {out.ndim}-D tensor; expected (B, d, gh, gw)." - ) - batch_np = out.detach().to(resolved_output_dtype).cpu().numpy() - for i in range(batch_np.shape[0]): - # Standalone C-contiguous copy: a per-row view would pin the whole - # batch alive (the blended sliding output is contiguous, so a view of - # it would not copy). ``.copy()`` always copies, in C order. - yield batch_np[i].copy() - - return _stream() + yield from _encode_batch([_read_padded(loc) for loc in chunk]) + + def _stream_prefetch(workers: int) -> Iterator[np.ndarray]: + # Double-buffered read/forward overlap. ``capacity`` is the max in-flight reads: + # ``workers + step`` keeps the whole pool busy AND guarantees the next batch's + # reads are already submitted while the current batch forwards. Reads are gathered + # in submission order (== coordinate order), so batches match the serial chunks. + n = len(coords) + capacity = workers + step + with _pinned_intraop_threads(), torch.inference_mode(), slide_encode_autocast_ctx( + device, precision + ): + with ThreadPoolExecutor(max_workers=workers) as pool: + futures: deque[Future[torch.Tensor]] = deque() + submitted = 0 + + def _fill() -> None: + nonlocal submitted + while submitted < n and len(futures) < capacity: + futures.append(pool.submit(_read_padded, coords[submitted])) + submitted += 1 + + _fill() + for start in range(0, n, step): + count = min(step, n - start) + tiles = [futures.popleft().result() for _ in range(count)] + _fill() # top up so the next batch's reads overlap this forward + yield from _encode_batch(tiles) + + if num_workers is None: + return _stream_serial() + return _stream_prefetch(int(num_workers)) diff --git a/tests/test_dense_regions.py b/tests/test_dense_regions.py index a5438b6..9170f77 100644 --- a/tests/test_dense_regions.py +++ b/tests/test_dense_regions.py @@ -262,3 +262,287 @@ def test_iter_regions_dense_rejects_bfloat16_output_eagerly(): requested_spacing_um=0.5, target_size=target_size, output_dtype=torch.bfloat16, ) assert wsi.calls == [] + + +# --------------------------------------------------------------------------- +# Prefetch path (num_workers): overlap reads with the forward. +# +# ``num_workers=None`` is the legacy serial path (byte-identical, exercised by the +# tests above). ``num_workers=K`` reads regions through a ``ThreadPoolExecutor`` of +# width ``K`` (threads, not processes — the reader releases the GIL), double-buffered +# so the next batch's reads overlap the current forward, with torch/cv2 intra-op +# threads pinned to 1 for the read path. These tests prove the *mechanism* on CPU with +# a fake reader (there is no GPU here to measure the ~2x throughput acceptance target). +# --------------------------------------------------------------------------- + +import threading # noqa: E402 +import time # noqa: E402 +from concurrent.futures import ThreadPoolExecutor # noqa: E402 + +from slide2vec.runtime import dense_regions as _dense_regions # noqa: E402 +from slide2vec.runtime.cpu_budget import resolve_on_the_fly_num_workers # noqa: E402 + + +class _ConcurrentFakeWSI: + """Thread-safe fake reader that records read concurrency and (optionally) gates. + + ``gate`` reads are held on a ``threading.Barrier(gate)`` so they can only proceed + once ``gate`` reads are in flight simultaneously — if the executor is narrower than + ``gate`` the barrier times out and the read raises, so a clean run *proves* at least + ``gate`` reads ran concurrently. ``max_active`` records the peak concurrency reached. + """ + + def __init__(self, *, target_h: int, target_w: int, gate: int = 0, sleep_s: float = 0.0): + self._target_h = target_h + self._target_w = target_w + self._sleep_s = sleep_s + self._lock = threading.Lock() + self.calls: list[tuple] = [] + self.active = 0 + self.max_active = 0 + self._gated_remaining = gate + self._barrier = threading.Barrier(gate) if gate else None + + def read_region_at_spacing(self, location, requested_spacing_um, size, *, tolerance, interpolation): + with self._lock: + self.calls.append((tuple(location), requested_spacing_um, tuple(size), tolerance, interpolation)) + self.active += 1 + self.max_active = max(self.max_active, self.active) + gate_this = self._gated_remaining > 0 + if gate_this: + self._gated_remaining -= 1 + if gate_this and self._barrier is not None: + self._barrier.wait(timeout=10.0) # releases only once `gate` reads coexist + if self._sleep_s: + time.sleep(self._sleep_s) + with self._lock: + self.active -= 1 + width, height = size + x, y = location + rng = np.random.default_rng(abs(hash((int(x), int(y)))) % (2**32)) + return rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8) + + +def _cosine(a: np.ndarray, b: np.ndarray) -> float: + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b))) + + +def test_prefetch_pool_width_matches_num_workers(monkeypatch): + """Acceptance (b): num_workers sets the ThreadPoolExecutor width.""" + enc = _encoder() + target_size = 64 + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + captured: list[int] = [] + + def _spy_executor(*args, **kwargs): + captured.append(int(kwargs.get("max_workers"))) + return ThreadPoolExecutor(*args, **kwargs) + + monkeypatch.setattr(_dense_regions, "ThreadPoolExecutor", _spy_executor) + list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=[(0, 0), (64, 0), (0, 64)], + requested_spacing_um=0.5, target_size=target_size, batch_size=2, num_workers=3, + )) + assert captured == [3] + + +def test_prefetch_width_driven_by_resolve_on_the_fly_num_workers(monkeypatch): + """Acceptance (b): the width comes from the num_workers_per_gpu resolver path.""" + enc = _encoder() + target_size = 64 + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + resolved, _ = resolve_on_the_fly_num_workers(num_cucim_workers=4, num_gpus=1) + assert resolved >= 1 + captured: list[int] = [] + + def _spy_executor(*args, **kwargs): + captured.append(int(kwargs.get("max_workers"))) + return ThreadPoolExecutor(*args, **kwargs) + + monkeypatch.setattr(_dense_regions, "ThreadPoolExecutor", _spy_executor) + list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=[(0, 0)], + requested_spacing_um=0.5, target_size=target_size, num_workers=resolved, + )) + assert captured == [resolved] + + +@pytest.mark.parametrize("num_workers", [2, 3]) +def test_prefetch_reads_run_concurrently(num_workers): + """Acceptance (a)/(b): reads are issued concurrently, not strictly serially. + + A barrier of width ``num_workers`` only releases if that many reads are in flight + at once — a serial reader would time out on it and raise. + """ + enc = _encoder() + target_size = 64 + wsi = _ConcurrentFakeWSI(target_h=target_size, target_w=target_size, gate=num_workers) + coords = [(i * 64, 0) for i in range(num_workers + 2)] + grids = list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=coords, + requested_spacing_um=0.5, target_size=target_size, + batch_size=1, num_workers=num_workers, + )) + assert len(grids) == len(coords) + assert wsi.max_active == num_workers # exactly the configured width ran at once + + +def test_prefetch_reads_overlap_the_forward(): + """Acceptance (a): the next batch's reads are in flight while the forward runs. + + batch_size=1 with a gate of 2 means: for the first forward (of coord 0) to obtain + its input, coord 0's read must complete — which the barrier only permits once a + *second* read (coord 1, a later batch) is also in flight. A serial path would read + coord 0 alone, time out on the barrier, and raise. Completing proves overlap. + """ + enc = _encoder() + target_size = 64 + wsi = _ConcurrentFakeWSI(target_h=target_size, target_w=target_size, gate=2) + coords = [(0, 0), (64, 0), (128, 0)] + gen = iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=coords, + requested_spacing_um=0.5, target_size=target_size, + batch_size=1, num_workers=2, + ) + first = next(gen) # would raise BrokenBarrierError if reads did not overlap + assert first.shape[1:] == (4, 4) + assert len(list(gen)) == 2 + + +def test_prefetch_pins_and_restores_intraop_threads(monkeypatch): + """Thread-pinning: the read path sets torch intra-op threads to 1 and restores them.""" + enc = _encoder() + target_size = 64 + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + before = torch.get_num_threads() + seen: list[int] = [] + real_set = torch.set_num_threads + + def _spy_set(n): + seen.append(int(n)) + return real_set(n) + + monkeypatch.setattr(torch, "set_num_threads", _spy_set) + list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=[(0, 0), (64, 0)], + requested_spacing_um=0.5, target_size=target_size, batch_size=1, num_workers=2, + )) + assert 1 in seen # pinned to a single intra-op thread in the read path + assert torch.get_num_threads() == before # restored afterwards + + +def test_prefetch_does_not_pin_threads_on_serial_path(monkeypatch): + """The legacy serial path (num_workers=None) leaves intra-op threads untouched.""" + enc = _encoder() + target_size = 64 + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + seen: list[int] = [] + real_set = torch.set_num_threads + monkeypatch.setattr(torch, "set_num_threads", lambda n: (seen.append(int(n)), real_set(n))[1]) + list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=[(0, 0)], + requested_spacing_um=0.5, target_size=target_size, + )) + assert seen == [] + + +@pytest.mark.parametrize("batch_size", [1, 2, 3]) +def test_prefetch_preserves_forward_batch_sizes(monkeypatch, batch_size): + """Acceptance (c): prefetch reordering does not change the forward batch sizes B.""" + enc = _encoder() + target_size = 64 + coords = [(i * 64, 0) for i in range(5)] # 5 coords => last batch is a remainder + + def _run(num_workers): + seen: list[int] = [] + real = encode_dense_sliding + + def _spy(model, batch, **kwargs): + seen.append(int(batch.shape[0])) + return real(model, batch, **kwargs) + + monkeypatch.setattr(_dense_regions, "encode_dense_sliding", _spy) + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=coords, + requested_spacing_um=0.5, target_size=target_size, + batch_size=batch_size, num_workers=num_workers, + )) + return seen + + assert _run(None) == _run(4) # identical B sequence, prefetch vs serial + + +@pytest.mark.parametrize("feature_kind", ["patch_features", "cls_attention"]) +def test_prefetch_grids_match_serial_within_tolerance(feature_kind): + """Acceptance (c): prefetched grids match the serial path (cosine >= 1 - 1e-4). + + Byte-identity is not required (nor achievable): pinning intra-op threads to 1 + perturbs the CPU forward's reduction order by ~1e-6. The ADR tolerance applies. + """ + enc = _encoder() + target_size = 64 + coords = [(0, 0), (64, 0), (0, 64), (64, 64), (128, 0)] + + def _grids(num_workers): + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + return list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=coords, + requested_spacing_um=0.5, target_size=target_size, + batch_size=2, num_workers=num_workers, feature_kind=feature_kind, + )) + + serial = _grids(None) + prefetched = _grids(3) + assert len(serial) == len(prefetched) == len(coords) + for s, p in zip(serial, prefetched): + assert s.shape == p.shape + assert _cosine(s, p) >= 1.0 - 1e-4 + + +def test_prefetch_reads_every_coordinate_once(): + """Acceptance (c): the same set of reads happens, one per coordinate.""" + enc = _encoder() + target_size = 64 + coords = [(0, 0), (64, 0), (0, 64), (64, 64), (128, 0)] + wsi = _ConcurrentFakeWSI(target_h=target_size, target_w=target_size) + list(iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=coords, + requested_spacing_um=0.5, target_size=target_size, batch_size=2, num_workers=3, + )) + read_locs = sorted(c[0] for c in wsi.calls) + assert read_locs == sorted(coords) + + +def test_prefetch_validates_num_workers_eagerly(): + """num_workers < 1 raises at the call site, before any region is read.""" + enc = _encoder() + target_size = 64 + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + with pytest.raises(ValueError): + iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=[(0, 0)], + requested_spacing_um=0.5, target_size=target_size, num_workers=0, + ) + assert wsi.calls == [] + + +def test_prefetch_is_lazy_and_reads_nothing_on_build(): + """Building the prefetch generator reads nothing; empty coords yield nothing.""" + enc = _encoder() + target_size = 64 + wsi = _FakeWSI(target_h=target_size, target_w=target_size) + gen = iter_regions_dense( + model=enc, device="cpu", wsi=wsi, coordinates=[(0, 0), (64, 0)], + requested_spacing_um=0.5, target_size=target_size, num_workers=2, + ) + assert wsi.calls == [] # no read until first next() + assert len(list(gen)) == 2 + empty = _FakeWSI(target_h=target_size, target_w=target_size) + assert list(iter_regions_dense( + model=enc, device="cpu", wsi=empty, coordinates=[], + requested_spacing_um=0.5, target_size=target_size, num_workers=2, + )) == [] + assert empty.calls == []