[alxndrkalinin]

What

Predict now reads each (position, timepoint) volume from zarr once and cuts its depth windows from memory (SlidingWindowDataset(cache_volume=...), set by HCSDataModule whenever no plane selection is given). HCSPredictionWriter assembles each volume in memory and writes it once, instead of a zarr read-modify-write per window. write_input keeps the per-window writes.

Why

Stride-1 depth windows read every plane once per covering window and re-wrote every overlapping plane per window. On one H200 (60 lite A549 volumes), FNet3D-VSCyto3DAug end-to-end predict was 76% outside the network.

Measured

• Speed (one H200, same job, 60 volumes, s/vol):
  • FNet3D: wall 15.22 → 8.77, loop 14.20 → 7.98; forward unchanged at 3.38.
  • UNeXt2-3D: wall 13.04 → 5.11.
• Bitwise identity on real data: FNet3D and UNeXt2 stores on the lite A549 mock leg (12 positions, 60 volumes) match the old path byte for byte, max |Δ| 0:
  • on an A40 and on an H200, before cleanup;
  • on the A40 again at this branch's HEAD.
  • The same compare flags A40 vs H200 stores (12/12 differ), so it can detect a difference.
• Tests: 699 passed (viscy-utils, viscy-data, test_submit_benchmark_job; the test_combined_ddp spawn tests also SIGSEGV on the unmodified base in this environment and were deselected).
  • New tests compare the buffered writer against a per-window reference writer: blend, center, target_2d and z_planes, batch 1/3, partial volumes flushed at predict end, and a float16 overwrite of a float32 store.
  • They also compare cached-volume windows against per-window reads.
  • The dtype test fails without its fix.

Commits

• e062ca3: the read-once path.
• f289fc1 (/simplify): drop the predict_read_once flag; buffer every write mode except write_input.
• cdab0a1 (/simplify): buffer in the dtype of an existing output array.
• cd5ef24 (/simplify): read the cached volume with per-channel basic slices (oindex peaks at ~7× the output).

Not in scope (follow-ups)

Each would further cut the ~4.6 s/vol of remaining host work:

• float64 blend on the GPU with one D2H per volume (bit-identical in a spot check);
• normalizing each plane once;
• background volume writes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01APoGhoKfkCho9Z1VMS5jsG

[Copilot]

🟡 Changes recommended

Sparse buffered writes can erase unselected planes, and cache replacement unnecessarily retains two full source volumes.

2 open issues

• Flush only selected planes for noncontiguous plane selections · New
• Release cached volume before loading the replacement · New

What changed in this PR

This PR reduces prediction I/O by caching source volumes and assembling output volumes in memory instead of repeatedly reading and writing overlapping depth windows.

Changes:

• Cache source volumes for single-worker-process prediction without plane selection (num_workers=0).
• Buffer prediction writes while retaining per-window writes for write_input.
• Add equivalence tests covering write modes, partial volumes, overwrites, dtype handling, and cached reads.

File	Description
packages/​viscy-utils/​tests/​test_prediction_writer.py	Adds buffered-write and cached-read regression tests.
packages/​viscy-utils/​src/​viscy_utils/​callbacks/​prediction_writer.py	Assembles and flushes prediction volumes in memory.
packages/​viscy-data/​src/​viscy_data/​sliding_window.py	Adds optional volume caching for window reads.
packages/​viscy-data/​src/​viscy_data/​hcs.py	Enables caching for eligible prediction loaders.

_{Review effort: Balanced}