From 321bb126ec85a6510b10ad0499e0bde90e7a2dfc Mon Sep 17 00:00:00 2001 From: Matthieu Chourrout <83714683+chourroutm@users.noreply.github.com> Date: Sat, 18 Jul 2026 13:02:49 -0400 Subject: [PATCH] Release slab memory to the OS between iterations Copying/downsampling processed slabs with steadily growing resident memory: the large numpy copy buffers and tensorstore shard-encode buffers were freed at the Python level but retained in the glibc allocator arena rather than returned to the OS, so RSS ratcheted up slab-by-slab until a memory-capped session was exhausted. - Add _release_memory() (gc.collect + guarded glibc malloc_trim) and call it at the end of _copy_slab and _downsample_block after dropping the slab/shard buffers and tensorstore handles. - Bound the tensorstore cache pool so the downsample read path cannot grow an unbounded in-memory cache. - Build the joblib job list lazily so each slab's dask view is released as it completes instead of all being pinned for the whole call. - Document that peak memory scales with n_processes. --- docs/guide.rst | 12 +++++++ src/stack_to_chunk/_array_helpers.py | 38 ++++++++++++++++++++ src/stack_to_chunk/main.py | 53 +++++++++++++++------------- 3 files changed, 79 insertions(+), 24 deletions(-) diff --git a/docs/guide.rst b/docs/guide.rst index 87d5412..cf6b1e9 100644 --- a/docs/guide.rst +++ b/docs/guide.rst @@ -24,6 +24,18 @@ Third-party multi-threading This allows the ``n_processes`` argument to be respected when set to ``1``, and prevents issues when ``stack_to_chunk`` uses a larger number of parallel processes. +Memory usage +------------ +Each parallel process holds roughly one slab/shard of data in memory at a time, so +peak memory is approximately ``n_processes`` times +:func:`stack_to_chunk.memory_per_slab_process` (a lower bound on the per-process +size). On a memory-capped machine, choose ``n_processes`` so that this product +stays within the available memory. + +``stack-to-chunk`` releases each slab's buffers back to the operating system as it +finishes (via ``malloc_trim`` on glibc systems), so resident memory should plateau +across slabs rather than climb. + Zarr group layout ----------------- The zarr groups produced by ``stack-to-chunk`` contain zarr arrays that are labelled 0, 1, 2, 3... etc. diff --git a/src/stack_to_chunk/_array_helpers.py b/src/stack_to_chunk/_array_helpers.py index c554f1f..0d8b6e8 100644 --- a/src/stack_to_chunk/_array_helpers.py +++ b/src/stack_to_chunk/_array_helpers.py @@ -1,3 +1,6 @@ +import ctypes +import ctypes.util +import gc from collections.abc import Callable from pathlib import Path @@ -10,6 +13,28 @@ from loguru import logger +def _release_memory() -> None: + """ + Return freed memory to the operating system. + + Large transient numpy and tensorstore buffers are freed at the Python level + when a slab-processing function returns, but with the glibc allocator the pages + are retained in the process arena rather than handed back to the OS. Without + this, resident memory ratchets up slab-by-slab. ``malloc_trim`` is glibc-only, + so it is guarded and becomes a no-op elsewhere. + """ + gc.collect() + libc_name = ctypes.util.find_library("c") + if libc_name is None: + return + try: + libc = ctypes.CDLL(libc_name) + malloc_trim = libc.malloc_trim + except (OSError, AttributeError): + return + malloc_trim(0) + + @delayed # type: ignore[misc] def _copy_slab(arr_path: Path, slab: da.Array, zstart: int, zend: int) -> None: """ @@ -37,6 +62,11 @@ def _copy_slab(arr_path: Path, slab: da.Array, zstart: int, zend: int) -> None: arr_zarr[:, :, zstart:zend].write(data).result() logger.info(f"Finished copying z={zstart} -> {zend - 1}") + # Hand the slab buffer back to the OS before returning so resident memory + # does not ratchet up slab-by-slab. + del data, arr_zarr + _release_memory() + @delayed # type: ignore[misc] def _downsample_block( @@ -101,6 +131,10 @@ def _downsample_block( ) arr_out[out_slice].write(data).result() + # Hand the block/shard buffers back to the OS before returning. + del data, arr_in, arr_out + _release_memory() + def _open_with_tensorstore(arr_path: Path) -> ts.TensorStore: return ts.open( @@ -111,5 +145,9 @@ def _open_with_tensorstore(arr_path: Path) -> ts.TensorStore: "path": str(arr_path), }, "open": True, + # Bound the read cache so repeated opens (e.g. in the downsample read + # path) cannot grow an unbounded in-memory cache. + "cache_pool": {"total_bytes_limit": 100_000_000}, + "recheck_cached_data": "open", } ).result() diff --git a/src/stack_to_chunk/main.py b/src/stack_to_chunk/main.py index 329bef2..24f6b5a 100644 --- a/src/stack_to_chunk/main.py +++ b/src/stack_to_chunk/main.py @@ -272,6 +272,13 @@ def add_full_res_data( multiple slabs in parallel using a compute cluster where the job wants to be split into many small individual Python processes. + Notes + ----- + Peak memory is approximately ``n_processes`` times + :func:`memory_per_slab_process` (a lower bound on the per-process size). + On a memory-capped machine, choose ``n_processes`` so that this product + stays within the available memory. + """ assert data.ndim == 3, "Input array is not 3-dimensional" if start_z_idx % self.chunk_size_z != 0: @@ -301,21 +308,22 @@ def add_full_res_data( slab_idxs: list[tuple[int, int]] = [ (z, min(z + self.chunk_size_z, nz)) for z in range(0, nz, self.chunk_size_z) ] - all_args = [ - ( + + logger.info("Starting full resolution copy to zarr...") + blosc_use_threads = blosc.use_threads + blosc.use_threads = 0 + + # Build jobs lazily so each slab's dask view is created and released + # incrementally rather than all being pinned for the whole call. + jobs = ( + _copy_slab( self._path / "0", data[:, :, zmin:zmax], zmin + start_z_idx, zmax + start_z_idx, ) for (zmin, zmax) in slab_idxs - ] - - logger.info("Starting full resolution copy to zarr...") - blosc_use_threads = blosc.use_threads - blosc.use_threads = 0 - - jobs = [_copy_slab(*args) for args in all_args] + ) Parallel(n_jobs=n_processes)(jobs) blosc.use_threads = blosc_use_threads @@ -356,7 +364,8 @@ def add_downsample_level( added. Running this with one process will use about 5/8 the amount of memory of a - single slab/shard. + single slab/shard. Peak memory scales with ``n_processes``, so on a + memory-capped machine choose ``n_processes`` accordingly. """ logger.info(f"Downsampling to level {level} with {n_processes=}") @@ -406,26 +415,22 @@ def add_downsample_level( for z in range(0, sink_arr.shape[2], sink_arr.shards[2]) ] - all_args: list[ - tuple[ - Path, Path, tuple[int, int, int], Callable[[npt.ArrayLike], npt.NDArray] - ] - ] = [ - ( + logger.info(f"Starting downsampling from level {level - 1} > {level}...") + blosc_use_threads = blosc.use_threads + blosc.use_threads = 0 + + # Build jobs lazily so each block's task is created and released + # incrementally rather than all being pinned for the whole call. + jobs = ( + _downsample_block( self._path / str(level - 1), self._path / str(level), idxs, downsample_func, ) for idxs in block_indices - ] - - logger.info(f"Starting downsampling from level {level - 1} > {level}...") - blosc_use_threads = blosc.use_threads - blosc.use_threads = 0 - - jobs = [_downsample_block(*args) for args in all_args] - logger.info(f"Launching {len(jobs)} jobs") + ) + logger.info(f"Launching {len(block_indices)} jobs") Parallel(n_jobs=n_processes, verbose=10)(jobs) self._add_level_metadata(level)