[akhanf]

Addresses #215. Requires zarrnii 0.22.0, which includes khanlab/zarrnii#271.

Cause

I measured the steps of n4_biasfield.py one at a time on F1A1Te4 at levels 2 and 1, with 16 threads per worker and 18 GB per worker:

• Applying the bias field (main cause): ZarrNii's map_blocks back end built a float64 coordinate grid for each block, about 52 bytes per voxel. That is 0.8 GB per 256³ block for each of a worker's 16 threads. At level 1 this reproduced the KilledWorker on _fused_block seen on the first cluster attempt.
• Writing the pyramid: ome-zarr computes all levels in one dask graph, and memory crept up during the write.
• Pickled task size (level 0 only): each apply task carried its own copy of the low-resolution field, the mask and the input's dask graph, about 8 MB per task. At level 0's 18,480 chunks the queued tasks filled the workers. This showed up only in the first level-0 cluster run below.
• The .ims temp zarr: not a cause. It peaked at 3 GB at level 1, and /tmp on the nodes is local disk.

khanlab/zarrnii#271 fixes all three, with separable interpolation, a pyramid written one level at a time, and the low-resolution arrays passed into the graph once.

level	image	ZarrNii 0.21.6	with khanlab/zarrnii#271
2	4.75 GB	48–56 GB, workers at 15–17 GB	15 GB, 4.4 GB per worker
1	38 GB	KilledWorker	20 GB, 6.0 GB per worker

On the cluster with this PR's request (128 threads, 96 GB), for the apply + pyramid stage:

subject, level	image	before	with khanlab/zarrnii#271
F1A1Te4, 1	38 GB	first attempt killed at 147 GB	53 s, 29.8 GB, 5.2 GB per worker
M1A1Te3, 0	563 GB	333 GB	355 s, 33.2 GB, 5.4 GB per worker

Both succeeded on the first attempt. Including the .ims conversion, the jobs took 4 min 23 s and 41 min 49 s. Slurm's MaxRSS (65 GB and 98 GB) also counts page cache from writing to the node's local disk, and the jobs were not memory-limited.

That is about 0.4 GB per thread for 256³ chunks, regardless of image size. At level 2, outputs match 0.21.6 to within float32 precision. The only exception was 2 of 2.4 billion voxels at the mask's 0.5 threshold.

Changes

• n4_biasfield memory: a fixed 96 GB, multiplied by the attempt, instead of scaled_mem_mb(128000, 500). dask splits the job's memory across its workers, so a 16-thread .ims worker gets 12 GB. It peaked at 5.4 GB at level 0, well below dask's 80% pause threshold. For the 563 GB subject at level 0 this goes from 409 GB to 96 GB.
• .ims temp zarr: written to resources.tmpdir, with only level 0, since nothing reads its lower levels.

Before merging

• Cluster runs with the new request: M1A1Te3 at level 0 and F1A1Te4 at level 1, results above.
• Release ZarrNii with Keep apply_scaled_processing(map_blocks) and pyramid writes from growing with image size zarrnii#271 (0.22.0) and raise the zarrnii pin and pixi.lock. The lock also moves ome-zarr from 0.18.0 to 0.21.0, adds transformnd and array-api-compat, updates ome-zarr-models to 1.8.1, and moves pydantic back to 2.12.5. With the new lock, n4_biasfield.py ran on F1A1Te4 at level 3, and its output kept the input's shape, voxel size, units and affine.

Notes

• Since the lock moved to ome-zarr 0.18 (0b45135), OME-Zarr outputs written through ZarrNii have had no voxel size, origin or units (OME-Zarr writes lose voxel size, origin and units with ome-zarr-py >= 0.18 zarrnii#272, write_image() silently ignores coordinate_transformations (and axis units) since 0.18.0 ome/ome-zarr-py#710). zarrnii 0.22.0 fixes this, so this PR fixes it for new outputs. OME-Zarr outputs written since 0b45135 still have the wrong geometry and need to be regenerated.
• OME-Zarr inputs use 2 threads per worker, so at 96 GB each of the 64 workers gets 1.5 GB. That is enough for input chunks up to 256³, but larger input chunks are untested.
• The pyramid change applies to every rule that writes through to_ome_zarr. Rules whose requests were raised because workers paused (compute_filtered_regionprops, signed_distance_transform) could be re-benchmarked after the release.

🤖 Generated with Claude Code