Distributed multi-GPU dense feature extraction over slide ROIs

[clemsgrs]

Summary

Adds first-class distributed extraction for dense (d, gh, gw) feature grids over caller-supplied WSI regions. ROIs are split across visible GPUs; each rank writes its final grids directly, avoiding a parent gather for artifacts that are roughly 1000× larger than pooled embeddings.

This is the slide2vec prerequisite for soma #279. Scope is Python-only: no CLI, YAML, or manifest changes.

Architecture

1. Parent orchestration flattens SlideRegions, filters completed ROIs before sharding, resolves each slide's spacing-aware read plan once, and dispatches in-process for num_gpus=1 or through torchrun for multiple GPUs.
2. plan_dense_shards performs deterministic, contiguous ROI-level np.array_split, balanced to ±1 ROI. No batch alignment is needed because features depend on batch size, not batch composition (ADR 0002).
3. run_dense_shard groups consecutive ROIs by slide, reuses iter_regions_dense for batched reads and encoding, and persists one payload plus sidecar per ROI.
4. dense_worker reads RANK, WORLD_SIZE, and LOCAL_RANK from torchrun, loads coordinates from NPZ, and runs its shard without NCCL or process-group initialization.

Ranks persist dense artifacts directly; nobody gathers grids in the parent (ADR 0001).

Public interface and artifacts

from slide2vec import DenseOptions, DenseRegionArtifact, ExecutionOptions, SlideRegions

artifacts: list[DenseRegionArtifact] = model.embed_regions_dense(
    regions,
    dense=DenseOptions(...),
    execution=ExecutionOptions(...),
)

Artifact layout:

dense_embeddings/[<class>/]<sample_id>/<x>_<y>.pt
dense_embeddings/[<class>/]<sample_id>/<x>_<y>.meta.json

Sidecars contain extraction geometry and encoding parameters owned by slide2vec; downstream supervision metadata remains downstream.

Reliability and safety

• Payloads and sidecars are published atomically, with the sidecar last as the completion marker.
• Resume filtering happens before sharding; payloads without sidecars are re-encoded.
• Output, request, and slide paths are absolute before workers change cwd.
• Sample IDs and annotations are validated as path components, including Windows drive-relative paths, and resolved artifact paths must remain within output_dir even through existing symlinks.
• Both single- and multi-GPU paths honor num_workers_per_gpu.

Verification

Behavior	Proof
Exact, balanced, deterministic partition	test_plan_dense_shards_*
N ROIs → N payloads + sidecars	test_run_dense_shard_writes_payload_and_sidecar_per_region
Multi-rank equivalence within cosine tolerance	test_multi_rank_matches_single_rank
Resume and crash recovery	test_run_dense_shard_skips_regions_with_existing_sidecar, test_run_dense_shard_reencodes_payload_missing_its_sidecar, test_write_dense_region_does_not_publish_interrupted_sidecar
Single-GPU in-process execution and worker count	test_embed_regions_dense_num_gpus_one_runs_in_process
Torchrun fan-out, NPZ coordinates, and absolute worker paths	test_embed_regions_dense_num_gpus_gt_one_launches_dense_worker, test_dense_worker_encodes_only_its_rank_shard, test_region_specs_round_trip_through_request
Artifact containment and public exports	test_region_dense_paths_rejects_*, test_package_root_exports_api

Test result

python -m pytest --no-cov -q -m "not heavy": 496 passed, 1 skipped, 25 deselected.

Final structured review: clean, no actionable findings.

Near-linear speedup remains a manual multi-GPU acceptance check because CI has no GPU.

Closes #217