Overlap reads with the encoder forward in the dense extraction path#220
Closed
clemsgrs wants to merge 1 commit into
Closed
Overlap reads with the encoder forward in the dense extraction path#220clemsgrs wants to merge 1 commit into
clemsgrs wants to merge 1 commit into
Conversation
…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.
Owner
Author
|
Superseded by the convergence approach for #216. A design review established that the standalone |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #216
What changed
iter_regions_dense(slide2vec/runtime/dense_regions.py) read WSI regionsserially 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).
num_workersargument.num_workers=Kreadsregions through a
ThreadPoolExecutorof widthK(threads, not processes —cuCIM/hs2p release the GIL on a region read), double-buffered so the next
batch's reads overlap the current forward.
for the duration of the prefetch stream, restored on exit — the same discipline a
DataLoaderworker process gets for free.batch_sizechunks, so the prefetcher only changes read timing, never theforward batch size
Bor its membership.num_workers=None(default) keeps the legacy serial path byte-identical, sothe existing offline fake-reader tests pass unchanged.
Acceptance criteria
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-Kbarrier only releases ifKreads are in flight at once) andtest_prefetch_reads_overlap_the_forward(thefirst forward can only obtain its input once a later batch's read is also in flight
— a serial path would deadlock the barrier).
num_workers_per_gpucontrols prefetch width. Met.test_prefetch_pool_width_matches_num_workersandtest_prefetch_width_driven_by_resolve_on_the_fly_num_workersassert the executor'smax_workersequals the value produced byresolve_on_the_fly_num_workers.B.Met.
test_prefetch_preserves_forward_batch_sizesasserts the exactBsequence isidentical serial-vs-prefetch across batch sizes;
test_prefetch_grids_match_serial_within_toleranceasserts 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/WSITileReadermachinery. That machinery is bound to aTilingResultgrid, cuCIMread_regions, super-tile plans, and aDataLoadercollator interface — it reads afixed tiling grid, not arbitrary
(x, y)level-0 coordinates at a requested µm/px viahs2p's
read_region_at_spacing. The dense primitive's injectable-wsiread interfaceis fundamentally different, so I reused the discipline the pooled path proved
(thread-pinning + fixed-composition batching) rather than the classes, and kept
iter_regions_denseoffline-testable. Width is taken as a plain int that callers resolvevia
resolve_on_the_fly_num_workers(reused), rather than importingExecutionOptionsinto the library primitive.
Tests
python -m pytest tests -q -m "not heavy"→ 477 passed, 25 deselected (15 newprefetch tests in
tests/test_dense_regions.py). Run with--no-covlocally because.coveragesqlite locking flakes on this worktree's network filesystem; withCOVERAGE_FILEon local fs the suite passes cleanly under coverage (as CI runs itin-container). Heavy tests (real-weight inference) were not run locally.