Release slab memory to the OS between iterations

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.

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()

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)