fix(cluster): 🐛 size work against the SLURM allocation, not the whole node

os.cpu_count() and psutil.virtual_memory().available describe the machine, so
on a shared node every worker-count and memory decision was made against
resources the job never had: a 32-core, 128 GB allocation on a 128-core,
512 GB node sized itself for 128 cores and 512 GB. That mismatch is how a job
walks into an OOM kill while its own accounting says it has room.

cpu_allocation() prefers SLURM_CPUS_PER_TASK, then the process' affinity mask,
then the machine. _get_available_memory() takes the smallest of
SLURM_MEM_PER_NODE, SLURM_MEM_PER_CPU x cpus, the cgroup limit (v1 and v2 —
the one that actually triggers the kill) and the node's free RAM. Both are
now used everywhere a worker count was derived.

convert.py pins dask's threaded scheduler to the allocation; it previously
ran one worker per machine core, the same bug the profile worked around by
raising mem_mb rather than fixing. merge.py sizes its pool from the RAM budget
and logs what it detected, so `seff` output can be read against it —
merge_workers: null also silently left merge_tile_labels capped at 4 while
the profile's comment claimed the full allocation was in use.

The profile gains retries: 2 with attempt-scaled mem_mb, so a transient OOM
resubmits larger instead of failing the run.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: lguerard

src/patchworks/__init__.py

@@ -29,7 +29,13 @@
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as _pkg_version
 
-from ._chunks import auto_overlap, auto_tile_shape, auto_tile_shape_cellpose
+from ._chunks import (
+    auto_overlap,
+    auto_tile_shape,
+    auto_tile_shape_cellpose,
+    cpu_allocation,
+    safe_worker_count,
+)
 from ._cluster import make_local_cluster
 from ._core import tile_process
 from ._distributed import (
@@ -56,6 +62,8 @@
     "auto_overlap",
     "auto_tile_shape",
     "auto_tile_shape_cellpose",
+    "cpu_allocation",
+    "safe_worker_count",
     "load_ome_zarr",
     "estimate_empty_tiles",
     "auto_empty_threshold",

src/patchworks/_chunks.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+from pathlib import Path
 from typing import Any, Sequence, Union
 
 import numpy as np
@@ -71,21 +72,100 @@ def auto_overlap(
 _GPU_MEMORY_FALLBACK = 8 * 1024**3
 
 
+def cpu_allocation() -> int:
+    """Return the number of CPUs this process may actually use.
+
+    ``os.cpu_count()`` reports the machine's cores, which on a shared cluster
+    node is wildly more than a job was granted -- a 4-core allocation on a
+    128-core node would size itself for 128. Prefer what the scheduler says,
+    then the process' CPU affinity mask, and only then the machine.
+
+    Returns
+    -------
+    int
+        Usable CPU count (always >= 1).
+    """
+    for var in ("SLURM_CPUS_PER_TASK", "SLURM_CPUS_ON_NODE"):
+        try:
+            value = int(os.environ[var])
+        except (KeyError, ValueError):
+            continue
+        if value > 0:
+            return value
+    try:
+        return max(1, len(os.sched_getaffinity(0)))
+    except AttributeError:  # not POSIX
+        return max(1, os.cpu_count() or 1)
+
+
+def _cgroup_memory_limit() -> "int | None":
+    """Memory ceiling from the process' cgroup, if one applies.
+
+    SLURM confines jobs with cgroups, so this is the limit that actually gets
+    the process OOM-killed -- unlike the node-wide figure ``psutil`` reports.
+    """
+    for path in (
+        "/sys/fs/cgroup/memory.max",  # cgroup v2
+        "/sys/fs/cgroup/memory/memory.limit_in_bytes",  # cgroup v1
+    ):
+        try:
+            raw = Path(path).read_text().strip()
+        except OSError:
+            continue
+        if raw == "max":
+            continue
+        try:
+            value = int(raw)
+        except ValueError:
+            continue
+        # v1 reports a sentinel near 2**63 when unlimited.
+        if 0 < value < 2**62:
+            return value
+    return None
+
+
 def _get_available_memory() -> int:
-    """Return available system RAM in bytes.
+    """Return the memory this process may actually use, in bytes.
+
+    Takes the **smallest** of everything that can constrain it: the SLURM
+    allocation, the cgroup limit, and the node's free RAM. A node with 512 GB
+    free must never convince a 128 GB job that it has room -- that mismatch is
+    exactly how a job sizes itself into an OOM kill.
 
     Returns
     -------
     int
-        Available memory via ``psutil``, or an 8 GiB fallback if it is not
-        installed.
+        Usable memory in bytes, or an 8 GiB fallback if nothing is knowable.
     """
+    limits = []
+
+    per_node = os.environ.get("SLURM_MEM_PER_NODE")
+    if per_node:
+        try:
+            limits.append(int(per_node) * 1024**2)  # SLURM reports MB
+        except ValueError:
+            pass
+    per_cpu = os.environ.get("SLURM_MEM_PER_CPU")
+    if per_cpu:
+        try:
+            limits.append(int(per_cpu) * 1024**2 * cpu_allocation())
+        except ValueError:
+            pass
+
+    cgroup = _cgroup_memory_limit()
+    if cgroup is not None:
+        limits.append(cgroup)
+
     try:
         import psutil
 
-        return int(psutil.virtual_memory().available)
+        limits.append(int(psutil.virtual_memory().available))
     except Exception:
+        pass
+
+    if not limits:
         return 8 * 1024**3
+    return max(1, min(limits))
 
 
 def safe_worker_count(
@@ -124,7 +204,7 @@ def safe_worker_count(
     int
         Worker-thread count (always >= 1).
     """
-    cpu_cap = max(1, (os.cpu_count() or 1) - 1)
+    cpu_cap = max(1, cpu_allocation() - 1)
     if use_gpu:
         return 1
     avail = _get_available_memory()
@@ -204,7 +284,7 @@ def auto_tile_shape(
     >>> tile
     (128, 512, 512)
     """
-    n_workers = n_workers or os.cpu_count() or 1
+    n_workers = n_workers or cpu_allocation()
     itemsize = np.dtype(dtype).itemsize
     n_spatial = min(3, len(shape))
 
@@ -308,7 +388,7 @@ def auto_tile_shape_cellpose(
     >>> tile
     (1, 2048, 2048)
     """
-    n_workers = n_workers or os.cpu_count() or 1
+    n_workers = n_workers or cpu_allocation()
     itemsize = np.dtype(dtype).itemsize
 
     if use_gpu:

src/patchworks/_cluster.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 import logging
-import os
+
+from ._chunks import cpu_allocation
 
 logger = logging.getLogger(__name__)
 
@@ -95,7 +96,7 @@ def make_local_cluster(
     from dask.distributed import Client, LocalCluster
 
     if n_workers is None:
-        n_workers = 1 if use_gpu else min(8, os.cpu_count() or 1)
+        n_workers = 1 if use_gpu else min(8, cpu_allocation())
 
     cluster = LocalCluster(
         processes=True,

src/patchworks/_core.py

@@ -11,7 +11,7 @@
 import dask.array as da
 import numpy as np
 
-from ._chunks import auto_tile_shape, safe_worker_count
+from ._chunks import auto_tile_shape, cpu_allocation, safe_worker_count
 from ._cluster import _client_is_in_process, _distributed_client
 from ._io import auto_empty_threshold, load_ome_zarr
 from ._merge import zarr_native_merge
@@ -537,7 +537,7 @@ def active_fn(block, block_info=None):
             if max_workers is not None
             else safe_worker_count(_tile_nbytes, use_gpu=use_gpu)
         )
-        _workers = max(1, min(_workers, os.cpu_count() or 1))
+        _workers = max(1, min(_workers, cpu_allocation()))
         logger.info("Staging with %d worker thread(s)", _workers)
         _sched_ctx: Any = _dask.config.set(
             scheduler="threads", num_workers=_workers

src/patchworks/_merge.py

@@ -32,6 +32,8 @@
 import numpy as np
 import zarr
 
+from ._chunks import cpu_allocation
+
 try:
     from tqdm.auto import tqdm as _tqdm
 except ImportError:
@@ -746,7 +748,7 @@ def merge_tile_labels(
     """
     import dask.array as da
 
-    nw = n_workers if n_workers is not None else min(4, os.cpu_count() or 1)
+    nw = n_workers if n_workers is not None else min(4, cpu_allocation())
 
     # -- Stage dask array to zarr if needed --
     stage_path: str | None = None

src/patchworks/_relations.py

@@ -3,14 +3,15 @@
 from __future__ import annotations
 
 import logging
-import os
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Union
 
 import dask.array as da
 import numpy as np
 
+from ._chunks import cpu_allocation
+
 logger = logging.getLogger(__name__)
 
 
@@ -94,7 +95,7 @@ def label_relations(
 
     n_blocks = a.numblocks
     total = int(np.prod(n_blocks))
-    nw = n_workers if n_workers is not None else min(4, os.cpu_count() or 1)
+    nw = n_workers if n_workers is not None else min(4, cpu_allocation())
 
     def _one(flat_idx: int) -> np.ndarray:
         idx = np.unravel_index(flat_idx, n_blocks)

tests/test_allocation.py

@@ -0,0 +1,78 @@
+"""Tests for cluster-allocation-aware CPU and memory detection.
+
+On a shared node ``os.cpu_count()`` and ``psutil.virtual_memory().available``
+describe the *machine*, not the slice of it this job was granted. Sizing work
+against the machine is how a job walks into an OOM kill, so these are the
+checks that would have caught it.
+"""
+
+import numpy as np
+
+from patchworks import cpu_allocation, safe_worker_count
+from patchworks._chunks import _get_available_memory
+
+GIB = 1024**3
+
+
+def test_cpu_allocation_prefers_slurm(monkeypatch):
+    """SLURM's grant wins over the machine's core count."""
+    monkeypatch.setenv("SLURM_CPUS_PER_TASK", "4")
+    assert cpu_allocation() == 4
+
+
+def test_cpu_allocation_ignores_junk_and_falls_back(monkeypatch):
+    """A malformed or absent value falls through to the affinity mask."""
+    monkeypatch.setenv("SLURM_CPUS_PER_TASK", "not-a-number")
+    monkeypatch.delenv("SLURM_CPUS_ON_NODE", raising=False)
+    assert cpu_allocation() >= 1
+
+    monkeypatch.delenv("SLURM_CPUS_PER_TASK", raising=False)
+    assert cpu_allocation() >= 1
+
+
+def test_available_memory_takes_the_smallest_limit(monkeypatch):
+    """The node's free RAM must never override a smaller allocation."""
+    monkeypatch.setenv("SLURM_MEM_PER_NODE", str(16 * 1024))  # 16 GiB, in MB
+    monkeypatch.setattr(
+        "patchworks._chunks._cgroup_memory_limit", lambda: 512 * GIB
+    )
+    assert _get_available_memory() == 16 * GIB
+
+
+def test_available_memory_respects_the_cgroup(monkeypatch):
+    """With no SLURM hint, the cgroup ceiling still bounds the answer."""
+    monkeypatch.delenv("SLURM_MEM_PER_NODE", raising=False)
+    monkeypatch.delenv("SLURM_MEM_PER_CPU", raising=False)
+    monkeypatch.setattr(
+        "patchworks._chunks._cgroup_memory_limit", lambda: 2 * GIB
+    )
+    assert _get_available_memory() <= 2 * GIB
+
+
+def test_mem_per_cpu_scales_with_the_allocation(monkeypatch):
+    """SLURM_MEM_PER_CPU is per core, so it multiplies by the core grant."""
+    monkeypatch.delenv("SLURM_MEM_PER_NODE", raising=False)
+    monkeypatch.setenv("SLURM_MEM_PER_CPU", str(2 * 1024))  # 2 GiB per cpu
+    monkeypatch.setenv("SLURM_CPUS_PER_TASK", "4")
+    monkeypatch.setattr(
+        "patchworks._chunks._cgroup_memory_limit", lambda: 512 * GIB
+    )
+    assert _get_available_memory() == 8 * GIB
+
+
+def test_worker_count_is_bounded_by_the_allocation(monkeypatch):
+    """The merge-style sizing must fit the grant, not the machine.
+
+    This is the concrete failure: a 32 GiB job on a 512 GiB node, with chunks
+    big enough that only a couple fit the grant.
+    """
+    monkeypatch.setenv("SLURM_CPUS_PER_TASK", "32")
+    monkeypatch.setenv("SLURM_MEM_PER_NODE", str(32 * 1024))  # 32 GiB
+    monkeypatch.setattr(
+        "patchworks._chunks._cgroup_memory_limit", lambda: 512 * GIB
+    )
+
+    chunk_nbytes = int(np.prod((16, 1024, 1024))) * 4  # int32 tile ~64 MB
+    n = safe_worker_count(chunk_nbytes * 40, fn_overhead=3)
+    assert n < 32, "must not size itself to the core count when RAM is tighter"
+    assert n >= 1

workflow/profile/slurm/config.yaml

@@ -20,6 +20,10 @@ rerun-incomplete: true
 rerun-triggers: mtime
 printshellcmds: true
 keep-going: true
+# Retry a failed job twice, asking for more memory each time (see the
+# attempt-scaled mem_mb below). A transient OOM or a node hiccup then costs a
+# resubmission instead of failing the run and requiring a manual rerun.
+retries: 2
 
 # Defaults for every rule (CPU jobs).
 default-resources:
@@ -32,30 +36,40 @@ default-resources:
 # Per-rule overrides.
 set-resources:
   convert:
-    # bioio/dask's default threaded scheduler uses every visible core, so
-    # this scales conversion for free. Verify actual scaling with `seff
-    # <jobid>` after a run and dial back if efficiency is low (scicore
-    # fair-share asks you not to sit on idle cores).
-    mem_mb: 64000
+    # convert.py pins dask's threaded scheduler to cpus_per_task, so this
+    # scales conversion without the scheduler quietly running one worker per
+    # *machine* core. Verify actual scaling with `seff <jobid>` after a run
+    # and dial back if efficiency is low (scicore fair-share asks you not to
+    # sit on idle cores).
+    mem_mb: "attempt * 64000"
     cpus_per_task: 32
     runtime: 360
   prepare:
-    mem_mb: 32000
+    # Also builds the max-pooled occupancy map on first use for an image.zarr
+    # that predates it (one bounded streaming pass, then shared by every
+    # config against that store).
+    mem_mb: "attempt * 32000"
     runtime: 120
   segment:
-    # one GPU per tile — this is what spreads Cellpose across GPUs.
+    # one GPU per job — this is what spreads Cellpose across GPUs.
     # Use the plugin's native `gres` resource → emits --gres=gpu:1, which is
     # how scicore binds a CUDA device. (The `gpu:` resource emits --gpus and
     # leaves the job without a device.)
     slurm_partition: "rtx4090"
     gres: "gpu:1"
     qos: "rtx4090-6hours" # scicore: <partition>-<duration> QOS
-    mem_mb: 32000 # plenty — a tile used ~1G
+    # A job now processes `tiles_per_job` tiles sequentially, so both memory
+    # and runtime scale with that setting — raise it there and re-check here.
+    # The old "a tile used ~1G" note predates tile_shape: "auto", which sizes
+    # tiles against the real GPU and makes them far bigger.
+    mem_mb: "attempt * 32000"
     cpus_per_task: 4
     runtime: 360 # 6 hours — must match the QOS, NOT 120 (=2h → killed early)
   merge:
-    # merge_tile_labels now uses n_workers = this allocation (see merge.py /
-    # config.yaml's merge_workers) instead of silently capping at 4 threads.
-    mem_mb: 128000
+    # merge.py sizes its worker pool from the cgroup/SLURM budget rather than
+    # the node's core count, and the pyramid is written zarr-natively (one
+    # source chunk per task) instead of through a dask rechunk — so this no
+    # longer has to cover an unbounded scheduler.
+    mem_mb: "attempt * 64000"
     cpus_per_task: 32
     runtime: 240