fix(ngmix): release FitModel observations after each metacal fit

[cailmdaley]

fix(ngmix): release FitModel observations after each metacal fit

Problem

The ngmix module's memory climbs steadily as it works through a tile — by roughly 5 MB for every galaxy it fits — and never comes back down until the tile is done. Each finished fit is held in full, so a tile with tens of thousands of galaxies accumulates gigabytes of image data it no longer needs. On a full production tile (~27,000 galaxies) during the Nibi run in #808, this reached ~11 GB per worker; in the controlled test below (3,000 galaxies) it reaches ~6 GB. Both are the same steady climb, just stopped at different points.

The current CFIS configs set SAVE_BATCH = 1000 to periodically write results out and clear them, which caps how high the climb gets — but the module's own default is to never clear until the tile finishes, and even with the cap the memory is pure waste. That per-worker footprint is what limits how many workers fit on a node, so it directly throttles throughput.

Root cause

Each per-object result in the resdict returned by MetacalBootstrapper.go in do_ngmix_metacal is an ngmix FitModel — a dict subclass that retains the full per-epoch metacal observations on the plain attribute .obs (set by FitModel._set_obs), plus the pixel arrays derived from them in ._pixels_list, together ~5 MB per object. ShapePipe accumulates these results and compile_results reads only scalar keys from each (g, g_cov, T, flux, s2n, flags, …) — the retained image arrays are never touched again. They are dead weight held to the end of the tile.

Fix

In shapepipe/modules/ngmix_package/ngmix.py, immediately after the metacal fit: null .obs and ._pixels_list on each FitModel in resdict, and drop the obsdict reference once average_multiepoch_psf has consumed it. Dict-like access to the fit results is untouched. +14 lines.

Evidence / testing

A/B on tile 196-305, objects 0-3000, SAVE_BATCH=1000, single core, RSS sampled every 30 s:

• Baseline (unpatched): RSS ramps 1.1 GB → 6.1-6.5 GB over the run (~4.8 MB/object).
• With fix: RSS flat, peak 0.90 GB over the identical 53-minute run.
• Output catalog: numerically identical to the unpatched baseline — same HDUs, same 2950 rows and columns in each of 1M/1P/2M/2P/NOSHEAR, all data arrays equal (NaN-aware comparison; only timestamp headers differ).

Found during the Nibi test run in #808.

— Claude (Fable) on behalf of Cail.

🤖 Generated with Claude Code

[martinkilbinger]

Another nice and simple fix for a long-standing problem (that I had forgotten about...)