Make target_chunk_mb actually bound memory on the FITS path

[AlecThomson]

Three places on the dask FITS path made peak memory a function of cube size, so target_chunk_mb was effectively inert there. The symptom was dask workers stalling or being killed on tasks named astype-from-value-read_fits_block-concatenate-.... That astype is dask's own, inserted by dask.array.concatenate because the reader declared the FITS on-disk dtype >f4; it was just the first task to touch the data, so it absorbed the cost and took the blame.

Diagnosis credit: handed over from a session scoped to flint-crew/flint, which could not push here.

1. The reader returned a lazy memmap view

_read_fits_block did np.asarray(memmap_slice), which does not copy. Nothing was read in the task; whoever materialised the block later faulted in most of the cube, and every pending block kept a whole-file mapping alive. Square spatial tiles made it worse: each channel's slice of a tile is cy short runs scattered across the file.

Now: full-width y-band chunks, memmap=False, and hdul[0].section[...] indexed against the on-disk shape (an int for degenerate length-1 axes, so a genuine cy == 1 band survives where np.squeeze would have eaten it). Blocks come back native-endian, which removes the astype layer.

Full traversal of every block of a 703 MiB cube (288 x 800 x 800 f32), 64 MiB block target, one process, peak RSS delta over the post-import baseline. Apple M-series, warm page cache, so these are not directly comparable to the Linux numbers in the original report (which had the same ranking but a wider spread):

block shape	reader	wall	peak RSS
241x241 tiles (before)	memmap	1.33 s	300 MiB
241x241 tiles	.section	8.55 s	233 MiB
72x800 bands	memmap	0.13 s	132 MiB
72x800 bands (after)	.section	0.11 s	161 MiB

Square tiles are what makes .section slow (241 short reads per channel); with full-width bands each channel is one contiguous read and it is no slower than memmap while staying near 2x the block size.

2. Per-channel noise gathered the whole cube

_channel_mad_lazy did cube.rechunk({1: -1, 2: -1}). Axis 0 was already one chunk, so the result was a single chunk holding the entire cube, reached eagerly from rmsynth_3d_from_fits for every noise-based weight_type on Q and U together:

MemoryError: Task 'channel_mad_std_block-rechunk-merge-finalize-...' has
703.12 MiB worth of input dependencies, but worker memory_limit is 476.84 MiB

flint's RMSynthOptions.weight_type defaults to "variance", so every flint run hit this.

Resizing the rechunk does not help: every input tile feeds every output channel plane, so dask pins the whole cube across it either way. An exact per-channel median simply cannot be computed from spatially chunked blocks. New read_fits_cube_channel_chunks reads the transposed chunking instead (whole image planes, chunked along frequency, one contiguous read per task), and rmsynth_3d_from_fits uses it for the noise estimate since it already has the paths. estimate_channel_noise_mad(q, u) still accepts spatially chunked arrays for array-based callers, but now logs a warning naming the size of the gather rather than doing it silently.

3. Output chunks outgrew the input chunks they were sized from

target_chunk_mb sized the float32 input footprint, but an FDF chunk is n_phi long in complex128 against n_freq in float32. Measured at n_freq=96, default phi grid (n_phi=989): a 4 MiB input chunk produced a 162 MiB FDF chunk and a 325 MiB RMSF chunk. This turned out to be the dominant term, not a footnote: it accounted for 6.4 GB of peak RSS on a 96 MiB cube.

Rather than document it, rmsynth_3d now treats the caller's input chunking as the memory budget and shrinks spatial chunks (y only, so blocks stay contiguous) so the complex128 output fits it. Same config after: 7.7 MiB FDF chunks, ~450 MB compute-phase peak, and the write is faster (4.1 s vs 5.6 s) because the tasks fit better. It only ever shrinks: a caller who chunked coarsely on purpose keeps their chunks when the FDF is no larger than the input.

End to end

rmsynth_3d_from_fits, weight_type="variance", target_chunk_mb=4.0, 2 dask threads, peak RSS delta over the pre-compute snapshot:

cube	main	this branch
48 x 512 x 512 (48 MiB)	723 MB	87 MB
48 x 1024 x 1024 (192 MiB)	1524 MB	129 MB

The point is the second column being flat: on main peak memory follows the cube, here it follows the chunk.

Tests

tests/_dask_memory_worker.py built its inputs with da.from_array on an in-memory array and called rmsynth_3d directly, so it never touched read_fits_cube_dask or estimate_channel_noise_mad. That is how both bugs shipped green.

• tests/_fits_memory_worker.py + test_fits_path_memory_scales_with_chunk_size_not_cube_size: drives rmsynth_3d_from_fits on real FITS files with a noise-based weight_type, at two cube sizes and one chunk target. Fails on main (1524 MB vs 723 MB for 4x the pixels).
• Reader equivalence against the old memmap path, including a degenerate leading Stokes axis, a short final band (ny % cy != 0), and cy == 1.
• No astype key in read_fits_cube_dask's graph, and a native dtype.
• Channel-chunked noise matches the whole-cube reduction.
• Output chunks stay within the input chunk budget, and the FDF is unchanged by the rechunking.

132 passed, 1 skipped (from 121 passed, 1 skipped). The suite also drops from ~250 s to ~40 s, because the old memory test was spending most of that time on a pathologically large FDF.

Workaround for anyone pinned to a release

weight_type: uniform avoids problem 2 entirely, plus worker headroom of about one full cube for problem 1.

Not verified

Not smoke tested through flint's process_rmsynth flow. The fixes are covered by the memory test above, but the numbers here come from rm-lite's own entry points, not from a run under a real distributed cluster with a nanny memory limit.