diff --git a/applications/dynaclr/scripts/profiling/README.md b/applications/dynaclr/scripts/profiling/README.md new file mode 100644 index 000000000..b1b44a730 --- /dev/null +++ b/applications/dynaclr/scripts/profiling/README.md @@ -0,0 +1,40 @@ +# DynaCLR I/O profiling scripts + +Scripts that validate data-loading performance on VAST/NFS for the DynaCLR +contrastive training pipeline. + +## Current scripts + +### `benchmark_recheck_cached_data.py` + +Measures the effect of `TensorStoreConfig.recheck_cached_data` on NFS read +latency for the DynaCLR contrastive read pattern. Exercises the iohub +tensorstore implementation directly (no training stack involved) so it can +be run **before** the dynaclr datamodule is ported to iohub 0.3.x. + +**Prerequisite.** Requires an iohub build with the upstream +`recheck_cached_data` knob on `TensorStoreConfig`. Until that lands, either +install iohub from the feature branch locally, or skip this script. + +Run: + +``` +uv run python applications/dynaclr/scripts/profiling/benchmark_recheck_cached_data.py +``` + +Output is a markdown table comparing median/p95 batch latency, patches/s, +and MiB/s across three configurations (`none`, `"open"`, `false`). Run +twice back-to-back and compare: if the `none` vs `"open"` gap shrinks on +the second run, the Linux NFS client page cache is masking the +per-chunk revalidation cost on this node. + +## Planned follow-ups (after iohub 0.3.x merge into dynadtw) + +- **Dataset-level A/B** — same configurations, but driven through + `MultiExperimentDataModule` + `MultiExperimentTripletDataset` so we + exercise `_get_position`/`_get_tensorstore`/`_slice_patches` and the + `ts.stack(...).read().result()` batched read path exactly as training + does. +- **SLURM DDP A/B** — 200-step fastdev runs with Lightning's + `SimpleProfiler`, comparing `data_time`/`batch_time` and GETATTR/s + from `nfsiostat` across ranks. diff --git a/applications/dynaclr/scripts/profiling/benchmark_boc2d_real.py b/applications/dynaclr/scripts/profiling/benchmark_boc2d_real.py new file mode 100644 index 000000000..206ddcf2e --- /dev/null +++ b/applications/dynaclr/scripts/profiling/benchmark_boc2d_real.py @@ -0,0 +1,264 @@ +"""Production-config DataLoader benchmark + batch-composition sanity check. + +Exercises the real +``DynaCLR-2D-MIP-BagOfChannels.yml`` training settings against the +committed v2 parquet to measure end-to-end DataLoader throughput and +verify that batch-grouping/stratification actually do what the config +says. + +Two parts +--------- + +**1. Composition check** — forces ``batch_group_by="marker"`` and checks +the first 20 batches: + +- every batch contains exactly one marker (single-marker batches), +- different batches surface different markers (proves the grouping is + shuffled across the epoch, not stuck on one value). + +**2. Throughput A/B** — runs the production config +(``batch_size=256``, ``channels_per_sample=1``, ``stratify_by=[perturbation, marker]``, +``num_workers=2``) under two ``recheck_cached_data`` settings: + +- ``None`` — TensorStore driver default. +- ``"open"`` — validate at open only (our merge's default). + +Reports median/p95 per-iter latency, iter/s, samples/s for each leg. +Because this runs on the real VAST-resident parquet with 7k+ FOVs, the +FOV-open amortisation is representative of real training. + +Usage +----- + uv run python applications/dynaclr/scripts/profiling/benchmark_boc2d_real.py +""" + +from __future__ import annotations + +import statistics +import time +from dataclasses import dataclass + +import numpy as np + +from dynaclr.data.datamodule import MultiExperimentDataModule +from viscy_transforms import ( + BatchedChannelWiseZReductiond, + BatchedRandSpatialCropd, + NormalizeSampled, +) + +CELL_INDEX_PARQUET = "/hpc/projects/organelle_phenotyping/models/collections/DynaCLR-2D-MIP-BagOfChannels-v2.parquet" + +BATCH_SIZE = 256 +NUM_WORKERS = 2 +WARMUP_BATCHES = 10 +N_BATCHES = 60 +SEED = 42 + +Z_WINDOW = 1 +Z_EXTRACTION_WINDOW = 20 +Z_FOCUS_OFFSET = 0.3 +YX_PATCH_SIZE = (256, 256) +FINAL_YX_PATCH_SIZE = (160, 160) + +COMPOSITION_BATCHES = 20 + +RECHECK_LEGS: list[tuple[str, str | bool | None]] = [ + ("None (driver default)", None), + ("open (our default)", "open"), +] + + +@dataclass +class LegResult: + """Timing outcome for one recheck_cached_data leg on the real parquet.""" + + label: str + iter_latencies_s: list[float] + total_s: float + + @property + def median_ms(self) -> float: + """Return median per-iter latency in milliseconds.""" + return statistics.median(self.iter_latencies_s) * 1000.0 + + @property + def p95_ms(self) -> float: + """Return p95 per-iter latency in milliseconds.""" + return float(np.percentile(self.iter_latencies_s, 95)) * 1000.0 + + @property + def iter_per_s(self) -> float: + """Return sustained iterations per second.""" + return len(self.iter_latencies_s) / self.total_s + + @property + def samples_per_s(self) -> float: + """Return sustained samples per second.""" + return self.iter_per_s * BATCH_SIZE + + +def _build_production_dm( + recheck_cached_data: str | bool | None, + batch_group_by: str | list[str] | None = None, + stratify_by: list[str] | None = None, + num_workers: int = NUM_WORKERS, +) -> MultiExperimentDataModule: + """Build a DataModule matching the production 2D-MIP-BoC training recipe.""" + normalizations = [ + NormalizeSampled( + keys=["channel_0"], + level="timepoint_statistics", + subtrahend="mean", + divisor="std", + ), + ] + augmentations = [ + BatchedRandSpatialCropd(keys=["channel_0"], roi_size=(10, 192, 192)), + BatchedChannelWiseZReductiond(keys=["channel_0"], allow_missing_keys=True), + ] + dm = MultiExperimentDataModule( + cell_index_path=CELL_INDEX_PARQUET, + focus_channel="Phase3D", + reference_pixel_size_xy_um=0.1494, + z_window=Z_WINDOW, + z_extraction_window=Z_EXTRACTION_WINDOW, + z_focus_offset=Z_FOCUS_OFFSET, + yx_patch_size=YX_PATCH_SIZE, + final_yx_patch_size=FINAL_YX_PATCH_SIZE, + channels_per_sample=1, + positive_cell_source="lookup", + positive_match_columns=["lineage_id"], + positive_channel_source="same", + tau_range=(0.5, 2.0), + tau_decay_rate=2.0, + batch_group_by=batch_group_by, + stratify_by=stratify_by if stratify_by is not None else ["perturbation", "marker"], + split_ratio=0.8, + batch_size=BATCH_SIZE, + num_workers=num_workers, + seed=SEED, + normalizations=normalizations, + augmentations=augmentations, + ) + dm.tensorstore_config = dm.tensorstore_config.model_copy(update={"recheck_cached_data": recheck_cached_data}) + return dm + + +def _composition_check() -> None: + """Verify batch_group_by='marker' yields single-marker, shuffled batches.""" + print("=" * 72) + print("Composition check: batch_group_by='marker'") + print("=" * 72) + + dm = _build_production_dm( + recheck_cached_data="open", + batch_group_by="marker", + stratify_by=None, + num_workers=0, + ) + dm.setup("fit") + loader = dm.train_dataloader() + it = iter(loader) + + markers_by_batch: list[set[str]] = [] + for i in range(COMPOSITION_BATCHES): + batch = next(it) + metas = batch["anchor_meta"] + batch_markers = {m["marker"] for m in metas} + markers_by_batch.append(batch_markers) + print(f" batch {i:>2}: {len(batch_markers)} unique markers → {sorted(batch_markers)[:4]}") + + non_singleton = [i for i, ms in enumerate(markers_by_batch) if len(ms) != 1] + if non_singleton: + print(f"\n FAIL: {len(non_singleton)} of {COMPOSITION_BATCHES} batches had >1 marker") + print(f" offending batches: {non_singleton}") + raise AssertionError("batch_group_by='marker' did not produce single-marker batches") + + unique_markers_seen = set().union(*markers_by_batch) + print(f"\n PASS: all {COMPOSITION_BATCHES} batches are single-marker") + print(f" markers touched across the {COMPOSITION_BATCHES} batches: {len(unique_markers_seen)}") + print(f" → {sorted(unique_markers_seen)}") + + if len(unique_markers_seen) < 2: + print("\n WARNING: only 1 marker touched across all batches — epoch may be stuck on one group") + else: + print(" → grouping is shuffled across markers (good)") + + del it + del loader + + +def _run_throughput_leg(label: str, recheck_cached_data: str | bool | None) -> LegResult: + """Run one throughput leg on the production config.""" + print(f"\n-- Throughput leg: recheck_cached_data = {label} --") + dm = _build_production_dm( + recheck_cached_data=recheck_cached_data, + batch_group_by=None, + stratify_by=["perturbation", "marker"], + num_workers=NUM_WORKERS, + ) + dm.setup("fit") + loader = dm.train_dataloader() + it = iter(loader) + + for _ in range(WARMUP_BATCHES): + _ = next(it) + + latencies_s: list[float] = [] + t_total = time.perf_counter() + t_prev = time.perf_counter() + for _ in range(N_BATCHES): + _ = next(it) + t_now = time.perf_counter() + latencies_s.append(t_now - t_prev) + t_prev = t_now + total_s = time.perf_counter() - t_total + + del it + del loader + + result = LegResult(label=label, iter_latencies_s=latencies_s, total_s=total_s) + print( + f" median {result.median_ms:.1f} ms | p95 {result.p95_ms:.1f} ms | " + f"{result.iter_per_s:.2f} iter/s | {result.samples_per_s:.1f} samples/s" + ) + return result + + +def _print_markdown(results: list[LegResult]) -> None: + """Emit a markdown-formatted throughput table.""" + print() + print("## Throughput (real 2D-MIP-BoC v2 parquet)") + print() + print(f"- Parquet: `{CELL_INDEX_PARQUET.split('/')[-1]}`") + print(f"- Batch size: {BATCH_SIZE}, num_workers: {NUM_WORKERS}") + print(f"- Warmup: {WARMUP_BATCHES} batches; timed: {N_BATCHES} batches") + print(f"- Z_extraction={Z_EXTRACTION_WINDOW}, YX={YX_PATCH_SIZE}, final_YX={FINAL_YX_PATCH_SIZE}") + print("- channels_per_sample=1, stratify_by=[perturbation, marker]") + print() + print("| recheck_cached_data | median ms | p95 ms | iter/s | samples/s |") + print("|---|---:|---:|---:|---:|") + for r in results: + print(f"| {r.label} | {r.median_ms:.1f} | {r.p95_ms:.1f} | {r.iter_per_s:.2f} | {r.samples_per_s:.1f} |") + print() + + +def main() -> None: + """Run composition check, then the throughput A/B, and print a summary.""" + _composition_check() + + print() + print("=" * 72) + print("Throughput A/B: production config, real parquet") + print("=" * 72) + + results: list[LegResult] = [] + for label, value in RECHECK_LEGS: + results.append(_run_throughput_leg(label, value)) + + _print_markdown(results) + + +if __name__ == "__main__": + main() diff --git a/applications/dynaclr/scripts/profiling/benchmark_dataloader_recheck.py b/applications/dynaclr/scripts/profiling/benchmark_dataloader_recheck.py new file mode 100644 index 000000000..bd8626fad --- /dev/null +++ b/applications/dynaclr/scripts/profiling/benchmark_dataloader_recheck.py @@ -0,0 +1,198 @@ +"""Full-pipeline A/B benchmark for TensorStoreConfig.recheck_cached_data. + +Drives :class:`dynaclr.data.datamodule.MultiExperimentDataModule` +end-to-end — ``__getitems__`` + ``collate_fn=lambda x:x`` + +PyTorch DataLoader with ``num_workers`` forked workers — to measure the +effect of ``recheck_cached_data`` on sustained training-loader +throughput, the only number that actually matters for GPU utilization. + +Three legs are compared against the same parquet, in the same process, +with the same FOVs and the same seed so sampling is deterministic: + +- ``"open"`` — validate at open only, trust cache thereafter (our + expected production setting). +- ``None`` — driver default, revalidate cached chunk metadata every + read (one stat/GETATTR per chunk per read on NFS). +- ``False`` — never revalidate (included for completeness). + +Per leg the script: + +1. Constructs a fresh ``MultiExperimentDataModule``, forcibly overriding + ``self.tensorstore_config.recheck_cached_data`` after ``__init__`` so + every Plate opens with the configured setting. +2. Runs ``setup("fit")`` once. +3. Warms the DataLoader with ``WARMUP_BATCHES`` batches (discarded). +4. Times ``N_BATCHES`` steady-state batches by wall-clocking the + iterator yield interval — this is what the training loop sees. +5. Reports median/p95 iteration time and steady-state iter/s. + +Because we use forked DataLoader workers, each config opens its own +Plates inside the worker after fork — matching real DDP training. + +Usage +----- + uv run python applications/dynaclr/scripts/profiling/benchmark_dataloader_recheck.py + +Requires: + +- iohub with ``recheck_cached_data`` on ``TensorStoreConfig`` + (czbiohub-sf/iohub#406 or later). +- A parquet whose ``store_path`` entries are readable on this node. +""" + +from __future__ import annotations + +import statistics +import time +from dataclasses import dataclass + +import numpy as np + +from dynaclr.data.datamodule import MultiExperimentDataModule + +CELL_INDEX_PARQUET = "/home/eduardo.hirata/repos/viscy/applications/dynaclr/configs/cell_index/benchmark_2exp.parquet" + +BATCH_SIZE = 32 +NUM_WORKERS = 4 +WARMUP_BATCHES = 10 +N_BATCHES = 100 +SEED = 42 + +Z_WINDOW = 8 +YX_PATCH_SIZE = (192, 192) +FINAL_YX_PATCH_SIZE = (160, 160) + +LEGS: list[tuple[str, str | bool | None]] = [ + ("open (recommended)", "open"), + ("None (driver default)", None), + ("False (never revalidate)", False), +] + + +@dataclass +class LegResult: + """Timing outcome for one recheck_cached_data leg.""" + + label: str + iter_latencies_s: list[float] + total_s: float + + @property + def median_ms(self) -> float: + """Return the median inter-batch iteration time in milliseconds.""" + return statistics.median(self.iter_latencies_s) * 1000.0 + + @property + def p95_ms(self) -> float: + """Return the p95 inter-batch iteration time in milliseconds.""" + return float(np.percentile(self.iter_latencies_s, 95)) * 1000.0 + + @property + def iter_per_s(self) -> float: + """Return steady-state iterations per second.""" + return len(self.iter_latencies_s) / self.total_s + + @property + def samples_per_s(self) -> float: + """Return steady-state samples per second.""" + return self.iter_per_s * BATCH_SIZE + + +def _build_datamodule(recheck_cached_data: str | bool | None) -> MultiExperimentDataModule: + """Construct a DataModule and force the recheck_cached_data leg onto its config.""" + dm = MultiExperimentDataModule( + cell_index_path=CELL_INDEX_PARQUET, + z_window=Z_WINDOW, + yx_patch_size=YX_PATCH_SIZE, + final_yx_patch_size=FINAL_YX_PATCH_SIZE, + channels_per_sample=None, + positive_cell_source="lookup", + positive_match_columns=["lineage_id"], + tau_range=(0.5, 2.0), + tau_decay_rate=2.0, + stratify_by=None, + split_ratio=0.8, + batch_size=BATCH_SIZE, + num_workers=NUM_WORKERS, + seed=SEED, + normalizations=[], + augmentations=[], + ) + # The datamodule sets recheck_cached_data="open" by default; override + # it here so every leg can dial the knob independently without editing + # the production code path. + dm.tensorstore_config = dm.tensorstore_config.model_copy(update={"recheck_cached_data": recheck_cached_data}) + return dm + + +def _run_leg(label: str, recheck_cached_data: str | bool | None) -> LegResult: + """Run one A/B leg and return a populated LegResult.""" + print(f"\n-- Leg: recheck_cached_data = {label} --") + dm = _build_datamodule(recheck_cached_data) + dm.setup("fit") + loader = dm.train_dataloader() + + it = iter(loader) + + # Warmup — discard. Forks workers, populates each worker's plate/ts + # caches, amortises Python import cost in the forked child. + for _ in range(WARMUP_BATCHES): + _ = next(it) + + # Steady-state timing. We measure the inter-batch yield interval, + # which is exactly what the training loop observes. + latencies_s: list[float] = [] + t_total = time.perf_counter() + t_prev = time.perf_counter() + for _ in range(N_BATCHES): + _ = next(it) + t_now = time.perf_counter() + latencies_s.append(t_now - t_prev) + t_prev = t_now + total_s = time.perf_counter() - t_total + + # Release workers before the next leg so forked processes do not + # pile up and compete for file descriptors. + del it + del loader + + result = LegResult(label=label, iter_latencies_s=latencies_s, total_s=total_s) + print( + f" median {result.median_ms:.1f} ms | p95 {result.p95_ms:.1f} ms | " + f"{result.iter_per_s:.2f} iter/s | {result.samples_per_s:.1f} samples/s" + ) + return result + + +def _print_markdown(results: list[LegResult]) -> None: + """Emit a markdown-formatted summary for the PR / Confluence.""" + print() + print("## Results (dataloader-level A/B)") + print() + print(f"- Parquet: `{CELL_INDEX_PARQUET.split('/')[-1]}`") + print(f"- Batch size: {BATCH_SIZE}, num_workers: {NUM_WORKERS}") + print(f"- Warmup: {WARMUP_BATCHES} batches; timed: {N_BATCHES} batches") + print(f"- Z={Z_WINDOW}, YX={YX_PATCH_SIZE}, final_YX={FINAL_YX_PATCH_SIZE}") + print() + print("| recheck_cached_data | median ms | p95 ms | iter/s | samples/s |") + print("|---|---:|---:|---:|---:|") + for r in results: + print(f"| {r.label} | {r.median_ms:.1f} | {r.p95_ms:.1f} | {r.iter_per_s:.2f} | {r.samples_per_s:.1f} |") + print() + + +def main() -> None: + """Run all three legs and print a combined markdown summary.""" + print("=" * 72) + print("Dataloader-level recheck_cached_data benchmark — MultiExperimentDataModule") + print("=" * 72) + + results: list[LegResult] = [] + for label, value in LEGS: + results.append(_run_leg(label, value)) + + _print_markdown(results) + + +if __name__ == "__main__": + main() diff --git a/applications/dynaclr/scripts/profiling/benchmark_dataloader_workers_sweep.py b/applications/dynaclr/scripts/profiling/benchmark_dataloader_workers_sweep.py new file mode 100644 index 000000000..3a650e155 --- /dev/null +++ b/applications/dynaclr/scripts/profiling/benchmark_dataloader_workers_sweep.py @@ -0,0 +1,183 @@ +"""Sweep num_workers × recheck_cached_data for the DynaCLR dataloader. + +Purpose +------- + +The first pass A/B (``benchmark_dataloader_recheck.py``) showed a counter- +intuitive result on ``MultiExperimentDataModule.train_dataloader()`` with +``num_workers=4``: ``recheck_cached_data="open"`` was slower than the +driver default. The raw ``ts.stack`` benchmark showed the opposite. Most +likely the p95 tails were dominated by first-touch FOV opens while the +ThreadDataLoader prefetch buffer masked them differently per leg. + +This sweep pins down the cause by running every ``recheck_cached_data`` +value across several ``num_workers`` settings with generous warmup, so we +can tell: + +- Does the ordering flip between ``num_workers=0`` (no fork, no thread + buffer) and ``num_workers>0`` (forked workers)? +- Is the ``"open"`` penalty paid only on cold FOV opens? If yes, longer + warmup should close the gap. +- Does the ``p95`` converge once steady-state is reached? + +Usage +----- + uv run python applications/dynaclr/scripts/profiling/benchmark_dataloader_workers_sweep.py +""" + +from __future__ import annotations + +import statistics +import time +from dataclasses import dataclass + +import numpy as np + +from dynaclr.data.datamodule import MultiExperimentDataModule + +CELL_INDEX_PARQUET = "/home/eduardo.hirata/repos/viscy/applications/dynaclr/configs/cell_index/benchmark_2exp.parquet" + +BATCH_SIZE = 32 +WARMUP_BATCHES = 30 +N_BATCHES = 150 +SEED = 42 + +Z_WINDOW = 8 +YX_PATCH_SIZE = (192, 192) +FINAL_YX_PATCH_SIZE = (160, 160) + +WORKER_COUNTS: list[int] = [0, 1, 4] +RECHECK_VALUES: list[tuple[str, str | bool | None]] = [ + ("None", None), + ("open", "open"), + ("False", False), +] + + +@dataclass +class SweepResult: + """One cell of the ``num_workers`` × ``recheck_cached_data`` grid.""" + + num_workers: int + recheck_label: str + iter_latencies_s: list[float] + total_s: float + + @property + def median_ms(self) -> float: + """Return median per-iter latency in milliseconds.""" + return statistics.median(self.iter_latencies_s) * 1000.0 + + @property + def p95_ms(self) -> float: + """Return p95 per-iter latency in milliseconds.""" + return float(np.percentile(self.iter_latencies_s, 95)) * 1000.0 + + @property + def iter_per_s(self) -> float: + """Return sustained iterations per second across timed batches.""" + return len(self.iter_latencies_s) / self.total_s + + @property + def samples_per_s(self) -> float: + """Return sustained samples per second (iter/s × batch).""" + return self.iter_per_s * BATCH_SIZE + + +def _build(num_workers: int, recheck_cached_data: str | bool | None) -> MultiExperimentDataModule: + """Build one datamodule with forced num_workers and recheck_cached_data.""" + dm = MultiExperimentDataModule( + cell_index_path=CELL_INDEX_PARQUET, + z_window=Z_WINDOW, + yx_patch_size=YX_PATCH_SIZE, + final_yx_patch_size=FINAL_YX_PATCH_SIZE, + channels_per_sample=None, + positive_cell_source="lookup", + positive_match_columns=["lineage_id"], + tau_range=(0.5, 2.0), + tau_decay_rate=2.0, + stratify_by=None, + split_ratio=0.8, + batch_size=BATCH_SIZE, + num_workers=num_workers, + seed=SEED, + normalizations=[], + augmentations=[], + ) + dm.tensorstore_config = dm.tensorstore_config.model_copy(update={"recheck_cached_data": recheck_cached_data}) + return dm + + +def _run_cell(num_workers: int, label: str, recheck_cached_data: str | bool | None) -> SweepResult: + """Run one cell of the sweep.""" + print(f"\n-- num_workers={num_workers}, recheck_cached_data={label} --") + dm = _build(num_workers, recheck_cached_data) + dm.setup("fit") + loader = dm.train_dataloader() + it = iter(loader) + + for _ in range(WARMUP_BATCHES): + _ = next(it) + + latencies_s: list[float] = [] + t_total = time.perf_counter() + t_prev = time.perf_counter() + for _ in range(N_BATCHES): + _ = next(it) + t_now = time.perf_counter() + latencies_s.append(t_now - t_prev) + t_prev = t_now + total_s = time.perf_counter() - t_total + + del it + del loader + + result = SweepResult( + num_workers=num_workers, + recheck_label=label, + iter_latencies_s=latencies_s, + total_s=total_s, + ) + print( + f" median {result.median_ms:.1f} ms | p95 {result.p95_ms:.1f} ms | " + f"{result.iter_per_s:.2f} iter/s | {result.samples_per_s:.1f} samples/s" + ) + return result + + +def _print_markdown(results: list[SweepResult]) -> None: + """Emit a markdown-formatted sweep table for the PR / Confluence.""" + print() + print("## Sweep results") + print() + print(f"- Parquet: `{CELL_INDEX_PARQUET.split('/')[-1]}`") + print(f"- Batch size: {BATCH_SIZE}, warmup: {WARMUP_BATCHES}, timed: {N_BATCHES}") + print(f"- Z={Z_WINDOW}, YX={YX_PATCH_SIZE}, final_YX={FINAL_YX_PATCH_SIZE}") + print() + print("| num_workers | recheck | median ms | p95 ms | iter/s | samples/s |") + print("|---:|---|---:|---:|---:|---:|") + for r in results: + print( + f"| {r.num_workers} | {r.recheck_label} | " + f"{r.median_ms:.1f} | {r.p95_ms:.1f} | " + f"{r.iter_per_s:.2f} | {r.samples_per_s:.1f} |" + ) + print() + + +def main() -> None: + """Run the full sweep and print a combined markdown summary.""" + print("=" * 72) + print("num_workers × recheck_cached_data sweep — MultiExperimentDataModule") + print("=" * 72) + + results: list[SweepResult] = [] + for nw in WORKER_COUNTS: + for label, value in RECHECK_VALUES: + results.append(_run_cell(nw, label, value)) + + _print_markdown(results) + + +if __name__ == "__main__": + main() diff --git a/applications/dynaclr/scripts/profiling/benchmark_recheck_cached_data.py b/applications/dynaclr/scripts/profiling/benchmark_recheck_cached_data.py new file mode 100644 index 000000000..98f503a11 --- /dev/null +++ b/applications/dynaclr/scripts/profiling/benchmark_recheck_cached_data.py @@ -0,0 +1,215 @@ +"""Measure the impact of ``TensorStoreConfig.recheck_cached_data`` on NFS reads. + +Single-process raw ``ts.stack(...).read().result()`` loop against a +2-experiment parquet for three TensorStoreConfig settings: + +- ``none`` — driver default, revalidate on every read (one stat/GETATTR + per chunk per read). +- ``open`` — validate only at open time, trust the cache thereafter. +- ``false`` — never revalidate. + +The loop issues ``N_BATCHES`` batches of stacked 3D crops sampled from +random FOVs, reports median/p95 read latency and sustained patches/s. +For the DataLoader-driven end-to-end view see +``benchmark_dataloader_workers_sweep.py``. + +Usage +----- + uv run python applications/dynaclr/scripts/profiling/benchmark_recheck_cached_data.py +""" + +from __future__ import annotations + +import statistics +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np +import pandas as pd +import tensorstore as ts +from iohub import open_ome_zarr +from iohub.core.config import TensorStoreConfig + +CELL_INDEX_PARQUET = "/home/eduardo.hirata/repos/viscy/applications/dynaclr/configs/cell_index/benchmark_2exp.parquet" + +BATCH_SIZE = 32 +N_BATCHES = 50 +PATCH_Z = 8 +PATCH_YX = (192, 192) +SEED = 0 + +DATA_COPY_CONCURRENCY = 16 +FILE_IO_CONCURRENCY = 64 +CACHE_POOL_BYTES: int | None = None + +CONFIGS: list[tuple[str, dict[str, Any]]] = [ + ("none (driver default)", {}), + ("open", {"recheck_cached_data": "open"}), + ("false", {"recheck_cached_data": False}), +] + + +@dataclass +class Result: + """Timing results for one ``recheck_cached_data`` configuration.""" + + label: str + batch_latencies_ms: list[float] + total_bytes: int + total_s: float + + @property + def median_ms(self) -> float: + """Return the median per-batch read latency in milliseconds.""" + return statistics.median(self.batch_latencies_ms) + + @property + def p95_ms(self) -> float: + """Return the p95 per-batch read latency in milliseconds.""" + return float(np.percentile(self.batch_latencies_ms, 95)) + + @property + def patches_per_s(self) -> float: + """Return the sustained patch-read throughput.""" + return BATCH_SIZE * len(self.batch_latencies_ms) / self.total_s + + @property + def mib_per_s(self) -> float: + """Return the sustained read throughput in MiB/s.""" + return (self.total_bytes / (1024 * 1024)) / self.total_s + + +def _load_fov_index() -> pd.DataFrame: + """Return unique (store_path, well, fov, shape) rows from the benchmark parquet.""" + df = pd.read_parquet(CELL_INDEX_PARQUET) + unique = df[["store_path", "well", "fov", "C_shape", "Z_shape", "Y_shape", "X_shape"]].drop_duplicates( + subset=["store_path", "well", "fov"] + ) + return unique.reset_index(drop=True) + + +def _open_stores(fov_df: pd.DataFrame, ts_config: TensorStoreConfig) -> dict[str, Any]: + """Open each unique zarr store once with the given TensorStoreConfig.""" + store_paths = fov_df["store_path"].drop_duplicates().tolist() + plates: dict[str, Any] = {} + for sp in store_paths: + plates[sp] = open_ome_zarr( + sp, + mode="r", + implementation="tensorstore", + implementation_config=ts_config, + ) + return plates + + +def _sample_patches( + fov_df: pd.DataFrame, + plates: dict[str, Any], + batch_size: int, + rng: np.random.Generator, +) -> tuple[list[ts.TensorStore], int]: + """Pick ``batch_size`` random (fov, z, y, x) crops and return lazy slices + byte count. + + Returns a list of tensorstore lazy slices (one per crop) plus the + total number of bytes the resulting stacked read will pull. + """ + lazies: list[ts.TensorStore] = [] + total_bytes = 0 + rows = fov_df.sample(n=batch_size, replace=True, random_state=rng.integers(0, 2**31 - 1)) + for _, row in rows.iterrows(): + plate = plates[row["store_path"]] + position_path = f"{row['well']}/{row['fov']}" + arr = plate[position_path]["0"].native + z_start = int(rng.integers(0, max(1, row["Z_shape"] - PATCH_Z + 1))) + y_start = int(rng.integers(0, max(1, row["Y_shape"] - PATCH_YX[0] + 1))) + x_start = int(rng.integers(0, max(1, row["X_shape"] - PATCH_YX[1] + 1))) + lazy = arr[ + 0, # t=0 — keep indexing simple; timepoint is not what we're benchmarking + :, + z_start : z_start + PATCH_Z, + y_start : y_start + PATCH_YX[0], + x_start : x_start + PATCH_YX[1], + ] + lazies.append(lazy) + total_bytes += PATCH_Z * PATCH_YX[0] * PATCH_YX[1] * row["C_shape"] * 4 # assume float32 + return lazies, total_bytes + + +def _run_one_config(label: str, extra_cfg: dict[str, Any], fov_df: pd.DataFrame) -> Result: + """Run the read-loop benchmark for one recheck_cached_data setting.""" + ts_config = TensorStoreConfig( + data_copy_concurrency=DATA_COPY_CONCURRENCY, + file_io_concurrency=FILE_IO_CONCURRENCY, + cache_pool_bytes=CACHE_POOL_BYTES, + **extra_cfg, + ) + plates = _open_stores(fov_df, ts_config) + + def _translate_all(lazies: list[ts.TensorStore]) -> list[ts.TensorStore]: + """Translate each lazy slice to origin so ts.stack can combine them.""" + return [p.translate_to[0] for p in lazies] # noqa: PD013 + + rng_warm = np.random.default_rng(SEED) + warm_lazies, _ = _sample_patches(fov_df, plates, BATCH_SIZE, rng_warm) + _ = ts.stack(_translate_all(warm_lazies)).read().result() + + rng = np.random.default_rng(SEED + 1) + latencies_ms: list[float] = [] + total_bytes = 0 + t_total = time.perf_counter() + for _ in range(N_BATCHES): + lazies, batch_bytes = _sample_patches(fov_df, plates, BATCH_SIZE, rng) + t0 = time.perf_counter() + _ = ts.stack(_translate_all(lazies)).read().result() + latencies_ms.append((time.perf_counter() - t0) * 1000.0) + total_bytes += batch_bytes + total_s = time.perf_counter() - t_total + + for plate in plates.values(): + plate.close() + + return Result(label=label, batch_latencies_ms=latencies_ms, total_bytes=total_bytes, total_s=total_s) + + +def _print_markdown_table(results: list[Result]) -> None: + """Print a markdown-formatted results table suitable for Confluence/PR pasting.""" + print() + print("## Results") + print() + print(f"- Parquet: `{CELL_INDEX_PARQUET.split('/')[-1]}`") + print(f"- Batch size: {BATCH_SIZE}, N batches: {N_BATCHES}") + print(f"- Patch shape: (C, Z={PATCH_Z}, Y={PATCH_YX[0]}, X={PATCH_YX[1]})") + print(f"- data_copy_concurrency={DATA_COPY_CONCURRENCY}, file_io_concurrency={FILE_IO_CONCURRENCY}") + print() + print("| recheck_cached_data | median ms | p95 ms | patches/s | MiB/s | total s |") + print("|---|---:|---:|---:|---:|---:|") + for r in results: + print( + f"| {r.label} | {r.median_ms:.1f} | {r.p95_ms:.1f} | " + f"{r.patches_per_s:.1f} | {r.mib_per_s:.1f} | {r.total_s:.2f} |" + ) + print() + + +def main() -> None: + """Run the three configurations back-to-back and print a markdown summary.""" + print("=" * 72) + print("recheck_cached_data benchmark — DynaCLR contrastive read pattern on VAST") + print("=" * 72) + + fov_df = _load_fov_index() + print(f"Loaded {len(fov_df)} unique FOVs across {fov_df['store_path'].nunique()} stores") + + results: list[Result] = [] + for label, extra_cfg in CONFIGS: + print(f"\n-- Running: recheck_cached_data = {label} --") + r = _run_one_config(label, extra_cfg, fov_df) + print(f" median {r.median_ms:.1f} ms | p95 {r.p95_ms:.1f} ms | {r.patches_per_s:.1f} patches/s") + results.append(r) + + _print_markdown_table(results) + + +if __name__ == "__main__": + main() diff --git a/applications/dynaclr/src/dynaclr/data/datamodule.py b/applications/dynaclr/src/dynaclr/data/datamodule.py index 08826d236..7d559ecb4 100644 --- a/applications/dynaclr/src/dynaclr/data/datamodule.py +++ b/applications/dynaclr/src/dynaclr/data/datamodule.py @@ -599,8 +599,26 @@ def _materialize_strings(df: pd.DataFrame) -> pd.DataFrame: # Dataloaders # ------------------------------------------------------------------ + def _ddp_topology(self) -> tuple[int, int]: + """Return ``(num_replicas, rank)`` for the current trainer. + + Lightning's auto-wrap hook only passes ``world_size``/``rank`` to + ``sampler``, not ``batch_sampler``. With ``use_distributed_sampler: + false`` and a batch sampler, the datamodule must read them from the + trainer itself and forward them; otherwise every rank iterates the + full sequence and yields identical batches. + + Returns ``(1, 0)`` when no trainer is attached (e.g. bare + dataloader construction in tests). + """ + trainer = getattr(self, "trainer", None) + if trainer is None: + return 1, 0 + return trainer.world_size, trainer.global_rank + def train_dataloader(self) -> ThreadDataLoader: """Return training data loader with FlexibleBatchSampler.""" + num_replicas, rank = self._ddp_topology() sampler = FlexibleBatchSampler( valid_anchors=self.train_dataset.index.valid_anchors, batch_size=self.batch_size, @@ -611,6 +629,8 @@ def train_dataloader(self) -> ThreadDataLoader: temporal_enrichment=self.temporal_enrichment, temporal_window_hours=self.temporal_window_hours, temporal_global_fraction=self.temporal_global_fraction, + num_replicas=num_replicas, + rank=rank, seed=self.seed, ) return ThreadDataLoader( @@ -643,6 +663,7 @@ def val_dataloader(self) -> ThreadDataLoader | None: """ if self.val_dataset is None: return None + num_replicas, rank = self._ddp_topology() sampler = FlexibleBatchSampler( valid_anchors=self.val_dataset.index.valid_anchors, batch_size=self.batch_size, @@ -651,6 +672,8 @@ def val_dataloader(self) -> ThreadDataLoader | None: group_weights=self.group_weights, stratify_by=self.stratify_by, temporal_enrichment=False, + num_replicas=num_replicas, + rank=rank, seed=self.seed, ) return ThreadDataLoader( diff --git a/applications/dynaclr/src/dynaclr/data/dataset.py b/applications/dynaclr/src/dynaclr/data/dataset.py index 9c9af4bf7..61034c6f8 100644 --- a/applications/dynaclr/src/dynaclr/data/dataset.py +++ b/applications/dynaclr/src/dynaclr/data/dataset.py @@ -918,4 +918,12 @@ def _slice_patches( rescaled = [] for i in range(len(patches)): rescaled.append(_rescale_patch(read_tensors[i], scales[i], targets[i])) + channel_counts = {t.shape[0] for t in rescaled} + if len(channel_counts) > 1: + raise RuntimeError( + f"Batch mixes samples with different channel counts: {sorted(channel_counts)}. " + "This happens with channels_per_sample=None across experiments that have " + "different channel counts. Set channels_per_sample=1 (bag-of-channels) " + "or channels_per_sample=[...] (fixed channel list)." + ) return torch.stack(rescaled), norms diff --git a/applications/dynaclr/tests/test_datamodule.py b/applications/dynaclr/tests/test_datamodule.py index e4b7b5815..eca373557 100644 --- a/applications/dynaclr/tests/test_datamodule.py +++ b/applications/dynaclr/tests/test_datamodule.py @@ -223,6 +223,33 @@ def test_train_dataloader_uses_flexible_batch_sampler(self, two_experiments): assert sampler.temporal_enrichment is False +class TestTrainDataloaderWiresDDPTopology: + """train_dataloader must forward Trainer world_size/rank to the sampler.""" + + def test_reads_world_size_and_rank_from_trainer(self, two_experiments): + from types import SimpleNamespace + + from dynaclr.data.datamodule import MultiExperimentDataModule + + parquet_path, _ = two_experiments + dm = MultiExperimentDataModule( + cell_index_path=str(parquet_path), + z_window=1, + yx_patch_size=_YX_PATCH, + final_yx_patch_size=_FINAL_YX_PATCH, + val_experiments=["exp_b"], + tau_range=(0.5, 2.0), + batch_size=8, + batch_group_by="experiment", + stratify_by="perturbation", + temporal_enrichment=False, + ) + dm.setup("fit") + dm.__dict__["trainer"] = SimpleNamespace(world_size=4, global_rank=2) + sampler = dm.train_dataloader().batch_sampler + assert (sampler.num_replicas, sampler.rank) == (4, 2) + + class TestValDataloaderNoBatchSampler: """Validation should be deterministic without FlexibleBatchSampler.""" diff --git a/applications/dynaclr/tests/test_dataset.py b/applications/dynaclr/tests/test_dataset.py index 5369c2898..d33860998 100644 --- a/applications/dynaclr/tests/test_dataset.py +++ b/applications/dynaclr/tests/test_dataset.py @@ -419,6 +419,70 @@ def test_int_gt1_raises(self, single_experiment_index): ) +class TestMixedChannelCountErrors: + """``channels_per_sample=None`` on a parquet whose experiments have different + channel counts must raise a clear error instead of a cryptic torch.stack + failure deep in a dataloader thread.""" + + def test_raises_when_experiments_have_different_channel_counts(self, tmp_path, _make_tracks_csv, hcs_dims): + from dynaclr.data.dataset import MultiExperimentTripletDataset + from dynaclr.data.experiment import ExperimentRegistry + from dynaclr.data.index import MultiExperimentIndex + from viscy_data.collection import ChannelEntry, Collection, ExperimentEntry + + # exp_a: 2 channels; exp_b: 1 channel. + zarr_a, tracks_a = _create_zarr_and_tracks( + tmp_path, + name="exp_a", + channel_names=["Phase", "GFP"], + wells=[("A", "1")], + hcs_dims=hcs_dims, + _make_tracks_csv=_make_tracks_csv, + ) + zarr_b, tracks_b = _create_zarr_and_tracks( + tmp_path, + name="exp_b", + channel_names=["Phase"], + wells=[("A", "1")], + hcs_dims=hcs_dims, + _make_tracks_csv=_make_tracks_csv, + ) + registry = ExperimentRegistry( + collection=Collection( + name="test", + experiments=[ + ExperimentEntry( + name="exp_a", + data_path=str(zarr_a), + tracks_path=str(tracks_a), + channels=[ChannelEntry(name="Phase", marker="Phase"), ChannelEntry(name="GFP", marker="GFP")], + channel_names=["Phase", "GFP"], + perturbation_wells={"c": ["A/1"]}, + interval_minutes=30.0, + ), + ExperimentEntry( + name="exp_b", + data_path=str(zarr_b), + tracks_path=str(tracks_b), + channels=[ChannelEntry(name="Phase", marker="Phase")], + channel_names=["Phase"], + perturbation_wells={"c": ["A/1"]}, + interval_minutes=30.0, + ), + ], + ), + z_window=1, + ) + index = MultiExperimentIndex(registry=registry, yx_patch_size=_YX_PATCH, tau_range_hours=(0.5, 2.0)) + ds = MultiExperimentTripletDataset(index=index, fit=True, channels_per_sample=None) + + va = index.valid_anchors + idx_a = int(va.index[va["experiment"] == "exp_a"][0]) + idx_b = int(va.index[va["experiment"] == "exp_b"][0]) + with pytest.raises(RuntimeError, match="different channel counts"): + ds.__getitems__([idx_a, idx_b]) + + class TestDatasetLength: """Test dataset length matches valid_anchors.""" diff --git a/packages/viscy-data/src/viscy_data/sampler.py b/packages/viscy-data/src/viscy_data/sampler.py index 982b01e6b..68534b059 100644 --- a/packages/viscy-data/src/viscy_data/sampler.py +++ b/packages/viscy-data/src/viscy_data/sampler.py @@ -282,8 +282,26 @@ def __iter__(self) -> Iterator[list[int]]: ``limit_train_batches`` interacts with this: Lightning stops pulling from the generator after its cap, so we never pay for the unused suffix of the epoch. + + The epoch counter auto-advances at the start of each iteration + so that the next ``__iter__`` call reseeds the RNG with a fresh + ``seed + epoch`` and yields a different batch sequence. Advancing + at the start (not the end) is robust against early generator + termination from ``limit_train_batches``: Lightning stops pulling + after its cap and garbage-collects the generator, which would + skip any end-of-iter bookkeeping. + + PyTorch Lightning does not call ``set_epoch`` on custom + ``batch_sampler`` instances (``use_distributed_sampler: false`` + with a batch sampler means Lightning's auto-wrap skips us), so + we self-advance. ``set_epoch`` still works if a caller wants + deterministic resume from a specific epoch — call it before the + iteration and the advance will take the resumed epoch as its + starting point. """ - rng = np.random.default_rng(self.seed + self.epoch) + seed_offset = self.epoch + self.epoch += 1 + rng = np.random.default_rng(self.seed + seed_offset) total_batches = len(self.valid_anchors) // self.batch_size rank = self.rank replicas = self.num_replicas diff --git a/packages/viscy-data/tests/test_sampler.py b/packages/viscy-data/tests/test_sampler.py index 20571033d..202cb2937 100644 --- a/packages/viscy-data/tests/test_sampler.py +++ b/packages/viscy-data/tests/test_sampler.py @@ -181,6 +181,107 @@ def test_batch_group_by_none_allows_mixing(self, two_experiment_anchors: pd.Data assert any_mixed, "With batch_group_by=None, at least one batch should mix experiments" +# --------------------------------------------------------------------------- +# Marker-aware batching (bag-of-channels regime) +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def multi_marker_anchors() -> pd.DataFrame: + """DataFrame with 1 experiment, 4 markers, 2 conditions, 320 rows total. + + Represents the bag-of-channels regime where each row is one (cell, + timepoint, channel) observation and ``marker`` identifies which + channel/protein the patch came from. + """ + rng = np.random.default_rng(7) + rows = [] + for marker in ["Phase3D", "TOMM20", "SEC61B", "Brightfield"]: + for cond in ["infected", "uninfected"]: + for i in range(40): + rows.append( + { + "experiment": "exp_boc", + "condition": cond, + "marker": marker, + "hours_post_perturbation": rng.uniform(0, 24), + "global_track_id": f"{marker}_{cond}_{i}", + "t": rng.integers(0, 20), + } + ) + df = pd.DataFrame(rows) + return df.reset_index(drop=True) + + +class TestMarkerAware: + """batch_group_by="marker" produces single-marker batches shuffled across markers. + + This is the bag-of-channels training regime — the config asks for one + marker per batch so contrastive pairs stay within the same channel, + while different batches traverse the full marker pool across an + epoch. + """ + + def test_every_batch_is_single_marker(self, multi_marker_anchors: pd.DataFrame): + """Every batch must contain rows from exactly one marker.""" + sampler = FlexibleBatchSampler( + valid_anchors=multi_marker_anchors, + batch_size=16, + batch_group_by="marker", + stratify_by=None, + leaky=0.0, + seed=42, + ) + batches = list(sampler) + assert batches, "Sampler should yield batches" + for batch in batches: + markers = multi_marker_anchors.iloc[batch]["marker"].unique() + assert len(markers) == 1, f"batch_group_by='marker' batch has {len(markers)} markers: {markers}" + + def test_all_markers_appear_across_epoch(self, multi_marker_anchors: pd.DataFrame): + """Across one epoch every marker surfaces in at least one batch.""" + sampler = FlexibleBatchSampler( + valid_anchors=multi_marker_anchors, + batch_size=16, + batch_group_by="marker", + stratify_by=None, + leaky=0.0, + seed=42, + ) + seen: set[str] = set() + for batch in sampler: + seen.update(multi_marker_anchors.iloc[batch]["marker"].unique()) + expected = {"Phase3D", "TOMM20", "SEC61B", "Brightfield"} + assert seen == expected, f"Not all markers surfaced in one epoch: {seen} vs {expected}" + + def test_batches_shuffled_across_markers(self, multi_marker_anchors: pd.DataFrame): + """Consecutive batches should not all be the same marker — the sampler + must interleave marker groups rather than drain them sequentially. + + We require at least half of the marker-to-marker batch transitions + to be a change (pathological samplers that yield all Phase3D + batches first, then all TOMM20, etc., would get a change-ratio + close to ``1/num_batches`` which this threshold catches). + """ + sampler = FlexibleBatchSampler( + valid_anchors=multi_marker_anchors, + batch_size=16, + batch_group_by="marker", + stratify_by=None, + leaky=0.0, + seed=42, + ) + per_batch_marker: list[str] = [] + for batch in sampler: + per_batch_marker.append(multi_marker_anchors.iloc[batch]["marker"].iloc[0]) + transitions = [a != b for a, b in zip(per_batch_marker[:-1], per_batch_marker[1:], strict=False)] + change_ratio = sum(transitions) / len(transitions) + assert change_ratio >= 0.5, ( + f"Only {change_ratio:.1%} of consecutive batches changed marker; " + "sampler appears to drain groups sequentially instead of shuffling" + ) + + # --------------------------------------------------------------------------- # Stratified sampling (SAMP-02) # --------------------------------------------------------------------------- @@ -391,6 +492,22 @@ def test_set_epoch_same_epoch_same_result(self, two_experiment_anchors: pd.DataF batches_b = list(sampler) assert batches_a == batches_b + def test_iter_auto_advances_epoch(self, two_experiment_anchors: pd.DataFrame): + """Consecutive iterations must yield different sequences without set_epoch. + + PL does not call ``set_epoch`` on ``batch_sampler`` instances, so the + sampler must self-advance. Regression guard for the frozen-dataset bug. + """ + sampler = FlexibleBatchSampler( + valid_anchors=two_experiment_anchors, + batch_size=8, + batch_group_by="experiment", + stratify_by=None, + leaky=0.0, + seed=42, + ) + assert list(sampler) != list(sampler) + # --------------------------------------------------------------------------- # __len__ and __iter__ protocol