Unified worker sizing, resilient bulk OSM/NLD downloads, faster branch processing

MANIFEST.in

@@ -5,3 +5,9 @@ prune data
 prune docs
 prune workflows
 exclude uv.lock
+
+# dist/ is committed to git, so setuptools-scm's file finder would otherwise
+# pack every previously released wheel and tarball into the new sdist -- each
+# build roughly doubling the one before it.
+prune dist
+global-exclude *.whl *.tar.gz

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fimbox"
-version = "0.1.11"
+version = "0.1.12"
 description = "A modular open-source testbed framework to standardize Flood Inundation Mapping simulations and evaluation with custom datasets and hydrologic parameters with reproducible workflows."
 authors = [
   { name = "Surface Dynamics Modeling Lab (SDML), The University of Alabama" },

src/fimbox/_dask.py

@@ -18,9 +18,11 @@
 
 The new defaults:
 
-1. ``n_workers`` is the minimum of CPU count and (free RAM GB / 6),
-   floored at 2. On a 32 GB / 8-core box that gives 5 workers with
-   ~6 GB each — enough for a HUC8 branch with margin.
+1. ``n_workers`` is the minimum of CPU count and (RAM GB / 8), floored
+   at 2. On a 64 GB / 12-core box that gives 8 workers with ~8 GB each
+   — enough for a HUC8 branch with margin. A caller-supplied count is
+   capped by the same ceiling (see ``fimbox._workers``), so
+   ``n_workers=64`` on a laptop lands on what the laptop can feed.
 2. Dask's "pause-at-80% / terminate-at-95%" governor is disabled.
    Branches run with stable RSS once the rasters are loaded; the
    transient peaks were causing Dask to pause workers, which made the
@@ -46,6 +48,8 @@
 import threading
 from typing import Optional
 
+from ._workers import resolve_workers, system_ram_gb
+
 log = logging.getLogger(__name__)
 
 _lock = threading.Lock()
@@ -62,26 +66,22 @@
 
 def _system_ram_gb() -> float:
     """Best-effort total RAM in GB. Returns 8.0 if it can't be detected."""
-    try:
-        import psutil  # type: ignore
+    return system_ram_gb()
 
-        return psutil.virtual_memory().total / (1024**3)
-    except Exception:
-        pass
-    # POSIX fallback (Linux + macOS): sysconf
-    try:
-        return (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (1024**3)
-    except (ValueError, AttributeError, OSError):
-        return 8.0
 
+def _ram_per_branch_gb() -> float:
+    return float(
+        os.environ.get("FIMBOX_DASK_RAM_PER_BRANCH", _DEFAULT_RAM_PER_BRANCH_GB)
+    )
 
-def _resolve_n_workers() -> int:
-    """Pick worker count that fits comfortably in RAM.
 
-    Explicit env override wins. Otherwise: min(cpu_count, RAM_GB /
-    RAM_PER_BRANCH_GB), floored at 2 so very small machines still get
-    some parallelism, capped at cpu_count so we don't oversubscribe
-    cores.
+def _resolve_n_workers(requested: Optional[int] = None) -> int:
+    """Worker count that fits comfortably in RAM.
+
+    ``FIMBOX_DASK_WORKERS`` still wins outright — it's the escape hatch for
+    people who know their box. Otherwise ``requested`` is honoured only as far
+    as CPU and RAM allow: asking for 64 workers on a 12-core laptop gets you 12
+    fighting over 8 GB slices, which is slower than the 8 the machine can feed.
     """
     env = os.environ.get("FIMBOX_DASK_WORKERS")
     if env:
@@ -92,23 +92,11 @@ def _resolve_n_workers() -> int:
         except ValueError:
             log.warning("FIMBOX_DASK_WORKERS=%r is not an int; auto-sizing", env)
 
-    cpu = max(1, os.cpu_count() or 1)
-    ram_gb = _system_ram_gb()
-
-    per_branch = float(
-        os.environ.get("FIMBOX_DASK_RAM_PER_BRANCH", _DEFAULT_RAM_PER_BRANCH_GB)
-    )
-    by_ram = max(2, int(ram_gb // per_branch))
-    n = min(cpu, by_ram)
-    log.info(
-        "Dask worker sizing: cpu_count=%d, system_ram=%.1f GB, "
-        "ram_per_branch=%.1f GB => n_workers=%d",
-        cpu,
-        ram_gb,
-        per_branch,
-        n,
+    return resolve_workers(
+        requested,
+        ram_per_worker_gb=_ram_per_branch_gb(),
+        label="Dask worker sizing",
     )
-    return n
 
 
 def _resolve_memory_limit():
@@ -153,7 +141,8 @@ def get_client(n_workers: Optional[int] = None):
 
     Pass ``n_workers`` to override the env/CPU default on first call
     only — once the cluster exists, later calls return the same client
-    regardless of the argument.
+    regardless of the argument. The request is bounded by what the machine
+    can feed; see :func:`_resolve_n_workers`.
     """
     global _client, _cluster
 
@@ -168,7 +157,7 @@ def get_client(n_workers: Optional[int] = None):
 
         from distributed import Client, LocalCluster
 
-        n = n_workers if n_workers else _resolve_n_workers()
+        n = _resolve_n_workers(n_workers)
         mem = _resolve_memory_limit()
 
         log.info(

src/fimbox/_workers.py

@@ -0,0 +1,113 @@
+"""
+One place that answers "how many workers can this machine actually run?".
+
+Callers hand over whatever the user gave them — nothing, ``None``, ``0``, or a
+number picked optimistically — and get back a count the box can sustain:
+
+* ``None`` / ``0`` / negative -> auto-size to the machine (the common case)
+* ``1``                       -> serial, always honoured; it's how you debug
+* anything bigger             -> honoured up to what CPU + RAM allow, then clamped
+* never more workers than there are branches to chew through
+
+RAM is the binding constraint, not core count. Every branch worker is its own
+process holding its own rasters and tables, so 16 workers on a 16 GB laptop buy
+swap thrash instead of speed — auto-sizing is ``min(cpu, RAM / per_worker)``.
+The per-worker budget differs by workload, hence the three constants below.
+
+Advanced overrides:
+- ``FIMBOX_WORKERS``        — pin the count for every pool (still capped by task count)
+- ``FIMBOX_RAM_PER_WORKER`` — GB to budget per worker when auto-sizing
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional
+
+log = logging.getLogger(__name__)
+
+# Per-worker RAM budgets, by how raster-hungry the step is. Branch processing
+# peaks hardest (AGREE DEM conditioning on a HUC8 at 10m), inundation loads a
+# HAND + catchment raster pair, and the SRC/calibration steps are pandas tables.
+RAM_PER_WORKER_BRANCH_GB = 8.0
+RAM_PER_WORKER_FIM_GB = 4.0
+RAM_PER_WORKER_TABLE_GB = 2.0
+
+
+def system_ram_gb() -> float:
+    """Best-effort total RAM in GB. Returns 8.0 if it can't be detected."""
+    try:
+        import psutil  # type: ignore
+
+        return psutil.virtual_memory().total / (1024**3)
+    except Exception:
+        pass
+    # POSIX fallback (Linux + macOS): sysconf
+    try:
+        return (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (1024**3)
+    except (ValueError, AttributeError, OSError):
+        return 8.0
+
+
+def _cpu_count() -> int:
+    # process_cpu_count honours cgroup/affinity limits (containers, SLURM);
+    # cpu_count reports the whole host and would oversubscribe a cpu-limited job.
+    getter = getattr(os, "process_cpu_count", None)
+    return max(1, (getter() if getter else os.cpu_count()) or 1)
+
+
+def max_supported_workers(ram_per_worker_gb: float = RAM_PER_WORKER_FIM_GB) -> int:
+    """Most workers this machine can feed without going to swap.
+
+    Floored at 2 so a modest laptop still gets some parallelism, then capped at
+    the core count so we never oversubscribe CPUs.
+    """
+    per_worker = float(os.environ.get("FIMBOX_RAM_PER_WORKER", ram_per_worker_gb))
+    by_ram = max(2, int(system_ram_gb() // max(per_worker, 0.5)))
+    return min(_cpu_count(), by_ram)
+
+
+def resolve_workers(
+    requested: Optional[int],
+    *,
+    n_tasks: Optional[int] = None,
+    ram_per_worker_gb: float = RAM_PER_WORKER_FIM_GB,
+    label: str = "workers",
+) -> int:
+    """Turn a user-supplied worker count into one this machine can honour.
+
+    ``1`` is passed straight through — an explicit request for serial is a
+    debugging tool, not a number to second-guess. Everything else is bounded by
+    :func:`max_supported_workers` and by ``n_tasks``, since idle processes still
+    cost a spawn and a full package re-import.
+    """
+    env = os.environ.get("FIMBOX_WORKERS")
+    if env:
+        try:
+            pinned = int(env)
+            if pinned > 0:
+                requested = pinned
+        except ValueError:
+            log.warning("FIMBOX_WORKERS=%r is not an int; ignoring", env)
+
+    if requested == 1:
+        return 1
+
+    ceiling = max_supported_workers(ram_per_worker_gb)
+
+    if requested is None or requested <= 0:
+        n = ceiling
+        why = f"auto-sized (cpu={_cpu_count()}, ram={system_ram_gb():.0f} GB)"
+    elif requested > ceiling:
+        n = ceiling
+        why = f"requested {requested}, clamped to what this machine supports"
+    else:
+        n = requested
+        why = ""  # honouring the given number needs no explanation
+
+    if n_tasks is not None:
+        n = max(1, min(n, n_tasks))
+
+    log.info("%s: %d worker(s)%s", label, n, f" — {why}" if why else "")
+    return n

src/fimbox/fimgeneration/pipeline.py

@@ -35,6 +35,7 @@
 
 import pandas as pd
 
+from .._workers import RAM_PER_WORKER_FIM_GB, resolve_workers
 from ..logging_utils import WATERSHED_DIR_NAME, aoi_root
 from .inundator import InundationResult, Inundator
 from .mosaic import BranchMosaic, MosaicResult
@@ -107,10 +108,10 @@ class FimGenerator:
     # Whether to mosaic per-branch outputs into AOI-level rasters at the end. Default True.
     mosaic: bool = True
 
-    # Parallel workers for the per-branch loop. 1 = serial. When use_dask
-    # is enabled n_workers is ignored — the shared Dask LocalCluster sizes
-    # itself from the machine.
-    n_workers: int = 1
+    # Parallel workers for the per-branch loop. None/0 = size to the machine,
+    # 1 = serial, anything above what CPU + RAM support is clamped down. When
+    # use_dask is enabled this is a ceiling on the shared LocalCluster instead.
+    n_workers: Optional[int] = None
 
     # When True, dispatch the branch loop through the process-wide Dask
     # LocalCluster (auto-sized to the machine). When None, Dask is used
@@ -177,10 +178,13 @@ def run(self) -> FimGenerationResult:
             )
             use_dask = False
 
-        # n_workers=None means "auto": let ProcessPoolExecutor size to the
-        # machine (only an explicit <=1 runs the serial loop). Guarding here
-        # keeps the comparison from raising on None.
-        serial = self.n_workers is not None and self.n_workers <= 1
+        n_workers = resolve_workers(
+            self.n_workers,
+            n_tasks=len(bids),
+            ram_per_worker_gb=RAM_PER_WORKER_FIM_GB,
+            label="FimGenerator",
+        )
+        serial = n_workers <= 1
 
         log.info(
             f"FimGenerator: AOI={self.aoi_dir.name} branches={len(bids)} "
@@ -189,7 +193,9 @@ def run(self) -> FimGenerationResult:
         )
 
         if use_dask:
-            results = self._run_with_dask(bids, branch_root, forecast_df, tmp_dir)
+            results = self._run_with_dask(
+                bids, branch_root, forecast_df, tmp_dir, n_workers
+            )
         elif serial:
             results = [
                 _run_one_branch(
@@ -205,7 +211,7 @@ def run(self) -> FimGenerationResult:
             ]
         else:
             results = []
-            with ProcessPoolExecutor(max_workers=self.n_workers) as pool:
+            with ProcessPoolExecutor(max_workers=n_workers) as pool:
                 fut_to_bid = {
                     pool.submit(
                         _run_one_branch,
@@ -282,12 +288,15 @@ def _run_with_dask(
         branch_root: Path,
         forecast_df: pd.DataFrame,
         tmp_dir: Path,
+        n_workers: int,
     ) -> list[InundationResult]:
         # Submit every branch to the shared LocalCluster. retries=0
         # mirrors the process_branches setup: an OOMed branch should
         # surface immediately as 'failed' so siblings finish instead
         # of cascading into a KilledWorker storm.
-        client = _get_dask_client()
+        # n_workers only bites if this call is what creates the cluster —
+        # a pool the preprocessing run already warmed up is reused as-is.
+        client = _get_dask_client(n_workers)
         log.info(
             "FimGenerator (dask): %d branches -> %d workers (dashboard %s)",
             len(bids),
@@ -471,7 +480,7 @@ class generateFIM:
 
     aoi_dir: PathLike
     feature_id_csv: Optional[PathLike] = None
-    n_workers: int = 4
+    n_workers: Optional[int] = None  # None -> size to the machine
     int16_mode: bool = True
     depth: bool = False
 
@@ -616,7 +625,12 @@ def generate(self, discharge_csvs: Sequence[PathLike]) -> list[FimGenerationResu
         help="Step 2: a single discharge CSV. When omitted, every CSV under "
         "<aoi_dir>/discharge-inputs/ is processed.",
     )
-    parser.add_argument("--workers", type=int, default=4)
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=None,
+        help="Parallel branch workers. Omit to size to the machine; 1 for serial.",
+    )
     parser.add_argument(
         "--int16",
         action="store_true",

src/fimbox/preprocessing/calculate_branch/_d8.py

@@ -0,0 +1,53 @@
+"""D8 flow-direction plumbing shared by the raster traversals.
+
+Both the label propagation in :mod:`gage_catchments` and the flow accumulation in
+:mod:`flowacc_dem` need the same thing first: for every cell, where does its water
+go. One table, one builder, so the two can't drift apart.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+# WBT D8 pointer encoding: powers-of-2 --> (row_offset, col_offset)
+# 64=N  128=NE  1=E  2=SE  4=S  8=SW  16=W  32=NW
+D8_OFFSETS: dict[int, tuple[int, int]] = {
+    1: (0, 1),
+    2: (1, 1),
+    4: (1, 0),
+    8: (1, -1),
+    16: (0, -1),
+    32: (-1, -1),
+    64: (-1, 0),
+    128: (-1, 1),
+}
+
+
+def downstream_index(d8: np.ndarray) -> np.ndarray:
+    """Flat index of each cell's downstream neighbour; itself where flow stops.
+
+    Nodata must already be zeroed by the caller. Uses shifted slices rather than
+    per-cell row/col arithmetic, so a HUC8 costs a few views instead of several
+    whole-grid index copies.
+    """
+    rows, cols = d8.shape
+    codes = d8.astype(np.int32, copy=False)
+    flat_base = np.arange(rows * cols, dtype=np.int32).reshape(rows, cols)
+    ds = flat_base.copy()
+
+    for code, (dr, dc) in D8_OFFSETS.items():
+        sel = codes == code
+        if not sel.any():
+            continue
+        # Only the window whose neighbour lands on the grid can flow; cells
+        # pointing off the edge keep their self-loop.
+        r0, r1 = max(0, -dr), rows - max(0, dr)
+        c0, c1 = max(0, -dc), cols - max(0, dc)
+        if r0 >= r1 or c0 >= c1:
+            continue
+        np.copyto(
+            ds[r0:r1, c0:c1],
+            flat_base[r0 + dr : r1 + dr, c0 + dc : c1 + dc],
+            where=sel[r0:r1, c0:c1],
+        )
+    return ds.ravel()