feat(nextflow): self-heal torn-shard/IO failures on per-position steps + default to non-preemptible cpu partition

[tayllatheodoro]

Summary

Fixes the recurrent zarr torn-shard corruption that kept aborting mantis-v2 per-position steps and looping on -resume (root-cause writeup: #286). Two complementary changes:

1. Generalized per-position self-heal

checksum_heal(output_zarr, position, cmd) in nextflow/modules/common.nf, wired into run_flat_field, run_deskew, run_apply_inv_tf, run_virtual_stain.

A per-position task killed mid-write (SLURM preemption / timeout / OOM, or a transient storage EIO) leaves a torn/partial zarr v3 shard. These steps write partial chunks via read-modify-write, so the next attempt (Nextflow retry or -resume) reads the torn shard back and the zarrs codec aborts with a non-signal exit (1) that the global errorStrategy would terminate on. The same root cause surfaces as several different messages depending on the codec:
the checksum is invalid / encoded shard is smaller than the expected size / blosc encoded value is invalid / …

Rather than enumerate every message, the wrapper heals on any non-zero exit: because a per-position task always fully recomputes its position from the input, it is always safe to clear that position's chunk data (the c/ dirs, preserving the zarr.json scaffold from --init) and run the command once more. If the retry also fails, its exit status propagates (genuine bugs are not masked — they just cost one extra attempt). No-op on success; scoped to a single position group so concurrent tasks never collide. This also gives transient EIO a second chance (the EIO-retry idea from #286).

2. Route cpu steps off the preempted partition

nextflow/nextflow.config: withLabel: 'cpu' now queues on the non-preemptible cpu partition instead of preempted. Preemptive reclamation mid-write was the primary source of torn shards; running on cpu removes it. gpu is unchanged. Switching back to preempted remains a one-line change for cost-sensitive runs.

Commit history (evolution)

• feat: self-heal checksum is invalid
• fix: also encoded shard is smaller than the expected size
• fix: default cpu steps to non-preemptible cpu partition
• fix: also blosc encoded value is invalid
• refactor: generalize to clean-and-retry on any failure (supersedes the string matching)

Verification

• nextflow lint clean on all changed files.
• Heal logic tested under bash -ue (Nextflow's shell): success is a no-op; a torn shard (incl. blosc) is cleaned + retried → success; a genuine bug is cleaned + retried once then propagates; transient EIO gets a retry.
• Validated in production on 2026_07_14_A549_MAP4_ZIKV: healed checksum / encoded-shard / blosc positions across flat_field & deskew; on cpu there are zero preemption (exit-143) kills, and runs pass flat_field → deskew → reconstruct.

Refs #286, #287, czbiohub-sf/iohub#415, czbiohub-sf/iohub#442, zarrs/zarrs#420.

[tayllatheodoro]

Errors encountered & decision log

Full record of the errors that drove this PR and why each change was made. Root-cause tracker: #286.

Errors encountered (in the order they surfaced)

#	Error	Where / trigger	Root cause	Resolution
1	RuntimeError: the checksum is invalid	flat_field write (RMW), exit 1	CRC32C shard-index mismatch reading a torn shard left by a killed write	self-heal: clean chunks + retry
2	exit 143 (SIGTERM)	per-position tasks on preempted	SLURM preemption/timeout reclaims job mid-write → torn shard (the source)	move cpu steps off preempted
3	RuntimeError: The encoded shard is smaller than the expected size of its index.	flat_field write (RMW), exit 1	truncated shard (same torn-shard root cause, different codec message)	broadened heal pattern
4	RuntimeError: blosc encoded value is invalid	flat_field write (RMW), exit 1	corrupt blosc inner chunk (same root cause, 3rd message)	broadened heal → then generalized
5	RuntimeError: Input/output error (os error 5) (EIO)	flat_field write, exit 1	transient storage (Lustre) error mid-write; also seeded torn shards in in-flight neighbours when it aborted the session	generalized heal gives it a retry
6	exit 140 (SIGUSR2)	C/4/000000, C/4/000002	near-walltime (chronically slow positions); retried via errorStrategy (time scales by attempt)	noted; bump flat_field walltime if it persists
7	KeyError: 'median' → RuntimeError: applying transform NormalizeSampled	virtual-stain, exit 1	predict.yml used level: timepoint_statistics (nested per-timepoint), but biahub's in-process VS path (#267) doesn't pre-resolve the timepoint as viscy's dataloader does → not corruption	predict.yml → level: fov_statistics (unblock); code gap tracked

Key decisions (and why)

1. Self-heal at the Nextflow layer, all per-position steps — chosen over a biahub Python change: contained to nextflow/modules/, fixes both the retry and -resume cases, and works for every per-position CLI without touching core write logic.
2. Detect, then generalize — started by matching specific torn-shard strings (checksum → encoded-shard → blosc). After the 3rd string it was clearly whack-a-mole, so generalized to clean-and-retry on any non-zero exit. Safe because a per-position task always fully recomputes from input, so clearing its chunks is never destructive; a genuine bug just fails the retry and propagates (not masked). Also covers transient EIO for free.
3. preempted → cpu partition — removes the source of new torn shards (mid-write preemption). Verified: zero exit-143 after the switch, vs 36 before. Kept as a one-line, reversible knob for cost-sensitive runs; filed Per-position steps hardcoded to preempted partition and not user-configurable #287 to make it configurable.
4. Virtual-stain normalization — timepoint_statistics is offered as a config option but unsupported by biahub's in-process VS path; quick unblock via fov_statistics, code fix (pre-resolve timepoint) left as a tracked gap.
5. Issues filed — Recurrent zarr torn-shard corruption on preemptible per-position steps (mantis-v2) #286 (parent root-cause), Per-position steps hardcoded to preempted partition and not user-configurable #287 (partition config), create_empty_plate shard/chunk geometry forces read-modify-write, making torn shards recur iohub#442 (shard geometry forces RMW), all cross-linked to RuntimeError: the checksum is invalid when writing zarr chunks via zarrs codec on Lustre iohub#415 and Add storage adapter that writes a tmp file then atomically moves on success zarrs/zarrs#420.

Process missteps & corrections (for transparency)

• Editing the live repo checkout invalidated resume cache. The pipeline resolves code from this repo at runtime (--biahub_project); editing its working tree changed run_flat_field task hashes, so -resume re-ran all positions and the live run picked up un-merged branch code. Corrected by restoring main and being explicit about which branch runs use.
• Orphaned preempted jobs after a cancel — a cancel left previous-run deskew jobs running unmanaged, risking two runs writing the same store. Cancelled the orphans before continuing.
• Started from scratch once, to clear accumulated torn shards + cache confusion from many cancel/restart cycles.