Skip to content

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

Description

@clemsgrs

Summary

Make dense extraction read WSI regions through the same batched reader the pooled path uses,
driven by an hs2p TilingResult, and retire iter_regions_dense's ad-hoc
(coordinates, requested_spacing_um, wsi) contract.

Today the dense path (slide2vec/runtime/dense_regions.py) has drifted from the pooled path:
it takes raw coords + a scalar spacing + an injectable wsi, and reads one region at a time through
the high-level per-call wsi.read_region_at_spacing, instead of keeping the TilingResult and
reading through the low-level batched WSIRegionReader (cuCIM read_regions(num_workers=…)).
That drift is the root cause of the original "overlap reads with the forward" problem — the drifted
primitive has no batched parallel read, so the earlier attempt (PR #220, now to be closed) bolted a
standalone ThreadPoolExecutor on top. Convergence removes the need for that entirely.

Supersedes PR #220. This was originally scoped as "add a prefetcher to the dense read"; a
design review (2026-07-15) established that the prefetcher was patching a symptom of the drift.

Why this is the right fix (design review)

  • read_region_at_spacing (hs2p wsi.py) is not a different I/O path — it is plan_spacing_read
    (pure arithmetic: spacing→level) + the same low-level read_region(loc, level, size) the
    pooled reader uses + an optional resize_array(..., "area"). The pooled path resolves spacing→level
    once, at tiling time (TilingResult.read_level / read_tile_size_px / requested_tile_size_px
    / resized); the dense path re-resolves per read.
  • The dense coordinates already come from hs2p annotation sampling (tile_slide(sampling=…, merged); tissue is the binary special case) — the same tiling machinery as the pooled path. soma
    builds a TilingResult and then unpacks it to coords+spacing before calling
    iter_regions_dense. Convergence = stop unpacking; pass the TilingResult through.
  • The offline-injectable-wsi property is a symptom of the drift, not a requirement. The real
    contract is "slide + tiling." Do not preserve fake-wsi injection in the production API for its
    own sake (see Tests below for the offline seam that replaces it).

Measured (virchow2, 512px ROIs): dense extraction is ~78% read time and CPU-JPEG-decode-bound
(not network); pinning intra-op threads to 1 and batching parallel reads are the wins — both of
which the pooled reader already does.

The key simplification

The TilingResult already resolved the plan. So the converged reader does not re-run
plan_spacing_read per read; it simply:

  1. reads read_tile_size_px at read_level (exactly what WSIRegionReader already does), then
  2. area-resizes to requested_tile_size_px iff read_tile_size_px != requested_tile_size_px
    (i.e. when TilingResult.resized), reusing hs2p's own resize_array(..., interpolation="area")
    — the identical call read_region_at_spacing makes.

This mirrors the pooled tar path (hs2p tiling/tar.py: if read_tile_size_px != requested_tile_size_px: resize). Because the level comes from the same tiling planner and the resize reuses the same
resize_array, the converged dense pixels are essentially identical to today's (within the
encode-side ADR-0002 float tolerance).

Scope

  1. Drive iter_regions_dense from a TilingResult (hs2p type — clean layering: slide2vec depends
    on hs2p; soma passes it). Replace (wsi, coordinates, requested_spacing_um, target_size, tolerance)
    with the TilingResult (carries tiles.x/y, read_level, read_tile_size_px,
    requested_tile_size_px, requested_spacing_um, tolerance, image_path, backend, resized)
    • reader knobs (backend override, num_workers). Keep all encode-side params (model, device, window_size, overlap, feature_kind, attention_*, batch_size, precision, output_dtype, pad_mode, image_pad_value, dense_transform).
  2. Read through WSIRegionReader (slide2vec/data/tile_reader.py) — batched cuCIM
    read_regions(num_workers=…). Remove the bespoke serial torch.stack([_read_padded(loc) …]) loop.
    No standalone ThreadPoolExecutor.
  3. Add the parameterized area geometry-resize to the shared read path (read → target size,
    interpolation a parameter, "area" for dense), reusing hs2p resize_array. Applied only when
    read_tile_size_px != requested_tile_size_px. Preserve area — do not switch to bilinear
    (a real, non-tolerance pixel change; see docs/adr/0002).
  4. Keep get_dense_transform (normalization-only) untouched — the encoder still receives the full
    requested_tile_size_px tile; dynamic_img_size handles it. Resize-by-concern: the reader owns
    the geometry resize; the model-input resize stays out of the dense path.
  5. Do not touch the pooled path. Its model-input resize and the silent-override gap are tracked
    in Pooled path silently resizes tiles to the encoder's native input (no warning, no high-res opt-in) #221.

Tests

  • Retire the fake-wsi injection from the production signature. The offline seam moves down:
    monkeypatch slide2vec.data.tile_reader._open_wsi_backend to return a fake backend that serves
    canned region arrays (implements read_regions/read_region). Production iter_regions_dense
    takes only the TilingResult + knobs — no injection param.
  • Rewrite tests/test_dense_regions.py onto that seam. Preserve the existing behavioural coverage:
    grid-per-coordinate order, eager validation, empty coords, output_dtype, bfloat16 rejection,
    streaming one batch at a time, and the batch-invariance property (composition irrelevant; only B
    matters — docs/adr/0002).
  • Add a case exercising the resized path (read_tile_size_px != requested_tile_size_px) so the
    area-resize is covered.

Acceptance

  • Dense reads go through WSIRegionReader (batched cuCIM parallel reads); no bespoke
    ThreadPoolExecutor, no per-region serial read_region_at_spacing loop.
  • iter_regions_dense is driven by a TilingResult; the (coordinates, requested_spacing_um, wsi)
    contract is gone.
  • num_workers (sized via resolve_on_the_fly_num_workers, slide2vec/runtime/cpu_budget.py)
    controls read parallelism.
  • Output preserved: converged dense grids match the pre-convergence output within ADR-0002 tolerance
    (per-grid cosine ≥ 1−1e-4); the area resize reuses hs2p resize_array so the read pixels are
    identical.
  • Full suite green (CPU, no GPU): pytest -m "not heavy".

Downstream (NOT in this PR)

  • soma's call site (soma/dense_slide_extraction.py) must stop unpacking the TilingResult and pass
    it through — tracked in soma#279 (already sequenced after this). This is a breaking slide2vec
    dense-API change (hard replacement; note in release notes); soma stays on the prior slide2vec until
    updated.

References

  • slide2vec/runtime/dense_regions.py (iter_regions_dense, the read_region_at_spacing loop)
  • slide2vec/data/tile_reader.py (WSIRegionReader, _open_wsi_backend, OnTheFlyHierarchicalBatchCollator)
  • hs2p wsi/wsi.py (read_region_at_spacing = plan_spacing_read + read_region + resize_array),
    tiling/result.py (TilingResult), tiling/tar.py (the read_tile != requested_tile resize precedent)
  • slide2vec/runtime/cpu_budget.py (resolve_on_the_fly_num_workers)
  • docs/adr/0002-features-are-reproducible-to-a-tolerance.md

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions