Skip to content

Overlap reads with the encoder forward in the dense extraction path#220

Closed
clemsgrs wants to merge 1 commit into
mainfrom
worktree-agent-a44f6c8c8f309f651
Closed

Overlap reads with the encoder forward in the dense extraction path#220
clemsgrs wants to merge 1 commit into
mainfrom
worktree-agent-a44f6c8c8f309f651

Conversation

@clemsgrs

Copy link
Copy Markdown
Owner

Closes #216

What changed

iter_regions_dense (slide2vec/runtime/dense_regions.py) read WSI regions
serially on the main thread, in the parent's default 32-intra-op-thread context.
The GPU sat ~50% idle on CPU JPEG decode and the per-ROI dense transform thrashed
(~71 ms vs ~5 ms single-threaded).

  • Opt-in prefetch path via a new num_workers argument. num_workers=K reads
    regions through a ThreadPoolExecutor of width K (threads, not processes —
    cuCIM/hs2p release the GIL on a region read), double-buffered so the next
    batch's reads overlap the current forward.
  • Thread pinning in the read path: torch (and cv2) intra-op threads pinned to 1
    for the duration of the prefetch stream, restored on exit — the same discipline a
    DataLoader worker process gets for free.
  • Batch invariance: reads are gathered in coordinate order into the same fixed
    batch_size chunks, so the prefetcher only changes read timing, never the
    forward batch size B or its membership.
  • num_workers=None (default) keeps the legacy serial path byte-identical, so
    the existing offline fake-reader tests pass unchanged.

Acceptance criteria

  • (a) Overlap reads with the forward; ~2× single-GPU throughput. Mechanism-met,
    not measured.
    There is no GPU in CI, so the throughput number is a real-hardware
    benchmark I did not fabricate. Instead the mechanism is proven with CPU-deterministic
    tests: test_prefetch_reads_run_concurrently (a width-K barrier only releases if
    K reads are in flight at once) and test_prefetch_reads_overlap_the_forward (the
    first forward can only obtain its input once a later batch's read is also in flight
    — a serial path would deadlock the barrier).
  • (b) num_workers_per_gpu controls prefetch width. Met.
    test_prefetch_pool_width_matches_num_workers and
    test_prefetch_width_driven_by_resolve_on_the_fly_num_workers assert the executor's
    max_workers equals the value produced by resolve_on_the_fly_num_workers.
  • (c) No change to emitted grids beyond tolerance; reordering must not change B.
    Met. test_prefetch_preserves_forward_batch_sizes asserts the exact B sequence is
    identical serial-vs-prefetch across batch sizes; test_prefetch_grids_match_serial_within_tolerance
    asserts cosine ≥ 1−1e-4 (byte-identity is impossible: pinning intra-op threads to 1
    perturbs the CPU forward's reduction order by ~1e-6 — within ADR 0002 tolerance).

Deviation from the suggested approach

The issue asked to prefer reusing the pooled OnTheFlyBatchTileCollator /
WSITileReader machinery. That machinery is bound to a TilingResult grid, cuCIM
read_regions, super-tile plans, and a DataLoader collator interface — it reads a
fixed tiling grid, not arbitrary (x, y) level-0 coordinates at a requested µm/px via
hs2p's read_region_at_spacing. The dense primitive's injectable-wsi read interface
is fundamentally different, so I reused the discipline the pooled path proved
(thread-pinning + fixed-composition batching) rather than the classes, and kept
iter_regions_dense offline-testable. Width is taken as a plain int that callers resolve
via resolve_on_the_fly_num_workers (reused), rather than importing ExecutionOptions
into the library primitive.

Tests

python -m pytest tests -q -m "not heavy"477 passed, 25 deselected (15 new
prefetch tests in tests/test_dense_regions.py). Run with --no-cov locally because
.coverage sqlite locking flakes on this worktree's network filesystem; with
COVERAGE_FILE on local fs the suite passes cleanly under coverage (as CI runs it
in-container). Heavy tests (real-weight inference) were not run locally.

…d prefetcher

iter_regions_dense read WSI regions serially on the main thread, in the
parent's default 32-intra-op-thread context, so the GPU sat ~50% idle on CPU
JPEG decode and the per-ROI dense transform thrashed at ~71 ms.

Add an opt-in num_workers prefetch path: a ThreadPoolExecutor (threads, not
processes -- cuCIM/hs2p release the GIL on a region read) of width num_workers,
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. Reads are gathered in
coordinate order into the same fixed batch_size chunks, so the prefetcher only
changes read timing -- never the forward batch size B or its membership.

num_workers=None keeps the legacy serial path unchanged (byte-identical); the
existing offline fake-reader tests still pass. Callers size num_workers from
ExecutionOptions.num_workers_per_gpu via resolve_on_the_fly_num_workers; the
primitive takes a plain int to stay injectable/offline-testable.
@clemsgrs

Copy link
Copy Markdown
Owner Author

Superseded by the convergence approach for #216. A design review established that the standalone ThreadPoolExecutor prefetcher this PR added was patching a symptom: the dense path had drifted from the pooled path (raw coords + scalar spacing + injectable wsi, per-region read_region_at_spacing) instead of reading through the shared batched WSIRegionReader driven by a TilingResult. #216 has been re-scoped to that convergence, which subsumes the read-overlap goal (batched cuCIM parallel reads) without a bespoke prefetcher. Closing in favor of that work.

@clemsgrs clemsgrs closed this Jul 15, 2026
@clemsgrs
clemsgrs deleted the worktree-agent-a44f6c8c8f309f651 branch July 15, 2026 12:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Converge the dense read path onto the shared WSI reader (retire the ad-hoc coords+spacing contract)

1 participant