fix(training,calibration): resolve DataLoader workers everywhere, and guard it (#207)

src/leech/calibration.py

@@ -24,7 +24,7 @@
 from torch.utils.data import DataLoader
 
 from leech.chunking import load_chunks
-from leech.dataset import LeechDataset, collate_fn
+from leech.dataset import LeechDataset, collate_fn, resolve_val_dataloader_workers
 from leech.models import get_model
 from leech.models.inference_wrapper import ModelInferenceWrapper
 
@@ -195,7 +195,11 @@ def calibrate_model(
         batch_size=batch_size,
         shuffle=False,
         collate_fn=collate_fn,
-        num_workers=num_workers,
+        # `num_workers=0` means AUTO, not "no workers" -- on CUDA feeding the
+        # forward pass from the main process alone leaves the GPU idle (#205,
+        # #207). The resolver caps by the Slurm cpuset, returns 0 in daemonic
+        # workers, and keeps 0 when the dataset fell back to per-chunk lists.
+        num_workers=resolve_val_dataloader_workers(val_dataset, num_workers, device),
     )
 
     logger.info(f"Collecting logits from {len(val_dataset)} validation samples...")
@@ -545,7 +549,7 @@ def calibrate_model_multiclass(
         batch_size=batch_size,
         shuffle=False,
         collate_fn=collate_fn,
-        num_workers=num_workers,
+        num_workers=resolve_val_dataloader_workers(val_dataset, num_workers, device),
     )
 
     logger.info(f"Collecting logits from {len(val_dataset)} validation samples...")

src/leech/commands/benchmark.py

@@ -60,6 +60,8 @@ def _build_loader_and_model(
         loader_kwargs["persistent_workers"] = True
         loader_kwargs["prefetch_factor"] = prefetch_factor
 
+    # dataloader-workers: unresolved -- the worker count is the independent
+    # variable being benchmarked here, so resolving it defeats the measurement.
     loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, **loader_kwargs)
 
     first = next(iter(loader))

src/leech/dataset.py

@@ -1003,6 +1003,34 @@ def resolve_dataloader_workers(num_workers: int, device: str) -> int:
     return effective
 
 
+def resolve_val_dataloader_workers(val_dataset, num_workers: int, device: str) -> int:
+    """Workers for the VALIDATION loader.
+
+    Same rule as [`resolve_dataloader_workers`], with one exception: a dataset
+    that fell back to per-chunk lists gets 0.
+
+    Validation used to be hardcoded to 0 on the grounds that its `__getitem__`
+    is trivially fast. That does not follow -- collate, pin, host-to-device and
+    the forward pass still serialize onto one core, which cost ~5 minutes of
+    near-idle GPU at every epoch boundary on a 1.18M-chunk val set (issue #207,
+    the same failure as #205 for `eval test`).
+
+    The memory half of the old rationale is real but narrow. `LeechDataset`
+    stacks per-chunk tensors into contiguous buffers precisely so a fork
+    COW-shares them; only the `_try_stack` list fallback makes each worker
+    fault N PyObject headers into private copies and multiply peak RSS. So the
+    exception is scoped to exactly that case rather than applied to every run.
+    """
+    workers = resolve_dataloader_workers(num_workers, device)
+    if workers and getattr(val_dataset, "_signals_tensor", None) is None:
+        logger.info(
+            "Validation dataset is not contiguously stacked; using 0 DataLoader "
+            "workers to avoid multiplying peak RSS across forks"
+        )
+        return 0
+    return workers
+
+
 class SignalDataset(Dataset):
     """Minimal signal-only dataset for ``SignalCNN``.
 

src/leech/training.py

@@ -27,7 +27,12 @@
 
 import leech
 from leech.cli_config import make_console
-from leech.dataset import LeechDataset, collate_fn, resolve_dataloader_workers
+from leech.dataset import (
+    LeechDataset,
+    collate_fn,
+    resolve_dataloader_workers,
+    resolve_val_dataloader_workers,
+)
 from leech.losses import AdversarialHead, FocalBCEWithLogitsLoss, RegressionHead
 from leech.models import get_model
 from leech.models.inference_wrapper import ModelInferenceWrapper
@@ -258,10 +263,14 @@ def train_signal_classifier(
     from leech.models import SignalCNN
 
     model = SignalCNN(num_classes=num_classes, signal_len=signal_len, channels=model_channels)
+    # dataloader-workers: unresolved -- legacy SignalCNN path. SignalDataset is
+    # not a LeechDataset and has no `_signals_tensor`, so the val guard would
+    # force 0 and change behaviour; converting it needs its own measurement.
     train_loader = DataLoader(
         SignalDataset(x_train, y_train), batch_size=batch_size, shuffle=True, collate_fn=collate_fn
     )
     val_loader = (
+        # dataloader-workers: unresolved -- same legacy SignalCNN path.
         DataLoader(SignalDataset(x_val, y_val), batch_size=batch_size, collate_fn=collate_fn)
         if x_val is not None
         else None
@@ -1582,13 +1591,27 @@ def train_model(
 
     val_loader = None
     if val_dataset is not None:
-        # Validation __getitem__ is trivially fast (no augmentation, pre-tensorized
-        # lookups only), so workers add memory overhead without benefit.  Dropping
-        # val workers halves the total process count and avoids OOM-triggered
-        # segfaults on large multiclass datasets.
+        # Validation goes through `resolve_dataloader_workers` like training and
+        # test, EXCEPT when the dataset fell back to per-chunk lists.
+        #
+        # The old comment here said workers "add memory overhead without
+        # benefit" because validation `__getitem__` is trivially fast. The first
+        # half is conditional and the second half is wrong. `__getitem__` being
+        # cheap does not mean one process can saturate a GPU: collate, pin,
+        # host-to-device and the forward pass all serialize onto that core. On a
+        # 1,176,763-chunk binary val set that cost ~5 minutes of near-idle GPU
+        # at every epoch boundary -- ~75 min per 15-epoch run -- which is the
+        # same failure as `eval test` in #205, and #206 fixed only that one.
+        #
+        # The memory half is real but narrower than a blanket 0. `LeechDataset`
+        # stacks into contiguous tensors precisely so a fork COW-shares the
+        # buffers; it is only the `_try_stack` list fallback (inconsistent
+        # per-chunk shapes) where each worker faults N PyObject headers into
+        # private copies and multiplies peak RSS. So keep 0 exactly there, and
+        # let every other case use workers.
         val_loader_kwargs: dict = {
             "collate_fn": collate_fn,
-            "num_workers": 0,
+            "num_workers": resolve_val_dataloader_workers(val_dataset, num_workers, device),
         }
         if device != "cpu":
             val_loader_kwargs["pin_memory"] = True

tests/test_dataloader_workers_guard.py

@@ -0,0 +1,104 @@
+"""`num_workers` must never be set to a literal in this package.
+
+The same bug shipped three times, each invisible to the check that caught the
+one before:
+
+  #205  `eval test` fed a GPU from one process        -> 8% GPU utilisation
+  #206  added `resolve_dataloader_workers`, fixed eval only
+  #207  the in-training VALIDATION loader still set 0 -> ~5 min idle GPU at
+        every epoch boundary, ~75 min per 15-epoch run; `calibration.py` was
+        passing a literal 0 through on CUDA at the same time
+
+`resolve_dataloader_workers`'s docstring claims "Every caller that builds a
+loader goes through this function". Nothing enforced it. This does.
+
+The check is deliberately on the VALUE, not on `DataLoader(...)` calls: #207
+lived in a `val_loader_kwargs` dict that reached the loader via `**kwargs`, and
+an earlier version of this guard that looked at DataLoader call sites missed it
+because the enclosing function resolved a *different* loader. Mutation-tested
+below against exactly that reintroduction.
+"""
+
+import ast
+from pathlib import Path
+
+SRC = Path(__file__).resolve().parents[1] / "src" / "leech"
+
+#: A literal must be annotated at its call site with this marker plus a reason.
+MARKER = "dataloader-workers: unresolved"
+
+RESOLVERS = {"resolve_dataloader_workers", "resolve_val_dataloader_workers"}
+
+
+def _is_resolved(value: ast.AST) -> bool:
+    """A resolver call, or a name/attribute that carries a resolved count."""
+    if isinstance(value, ast.Call):
+        fn = value.func
+        nm = fn.id if isinstance(fn, ast.Name) else getattr(fn, "attr", None)
+        return nm in RESOLVERS
+    if isinstance(value, ast.Name):
+        return "worker" in value.id.lower()
+    if isinstance(value, ast.Attribute):
+        return "worker" in value.attr.lower()
+    return False
+
+
+def _literal_num_workers():
+    """Every place `num_workers` is bound to a constant, with its context."""
+    for path in sorted(SRC.rglob("*.py")):
+        text = path.read_text()
+        lines = text.splitlines()
+        tree = ast.parse(text, filename=str(path))
+        for node in ast.walk(tree):
+            pairs = []
+            if isinstance(node, ast.Dict):
+                for k, v in zip(node.keys, node.values, strict=False):
+                    if isinstance(k, ast.Constant) and k.value == "num_workers":
+                        pairs.append(v)
+            elif isinstance(node, ast.Call):
+                for kw in node.keywords:
+                    if kw.arg == "num_workers":
+                        pairs.append(kw.value)
+            for v in pairs:
+                if _is_resolved(v):
+                    continue
+                if not isinstance(v, ast.Constant):
+                    continue  # an expression we cannot judge; not a literal
+                lo = max(0, getattr(v, "lineno", 1) - 5)
+                hi = getattr(node, "end_lineno", v.lineno) or v.lineno
+                span = "\n".join(lines[lo:hi])
+                yield str(path.relative_to(SRC)), v.lineno, v.value, MARKER in span
+
+
+def test_num_workers_is_never_a_bare_literal():
+    offenders = [
+        f"{rel}:{ln} num_workers={val!r} is a literal, not a resolved count"
+        for rel, ln, val, annotated in _literal_num_workers()
+        if not annotated
+    ]
+    assert not offenders, (
+        "num_workers must come from resolve_dataloader_workers / "
+        "resolve_val_dataloader_workers, or carry the marker "
+        f"{MARKER!r} with a reason:\n  " + "\n  ".join(offenders)
+    )
+
+
+def test_the_guard_rejects_a_literal():
+    """A guard that cannot fail is not a guard."""
+    tree = ast.parse('kw = {"num_workers": 0}')
+    d = next(n for n in ast.walk(tree) if isinstance(n, ast.Dict))
+    assert not _is_resolved(d.values[0])
+
+
+def test_the_guard_accepts_a_resolver_call():
+    tree = ast.parse("DataLoader(ds, num_workers=resolve_val_dataloader_workers(v, n, d))")
+    call = next(n for n in ast.walk(tree) if isinstance(n, ast.Call))
+    kw = next(k for k in call.keywords if k.arg == "num_workers")
+    assert _is_resolved(kw.value)
+
+
+def test_the_guard_accepts_a_resolved_local():
+    tree = ast.parse("DataLoader(ds, num_workers=effective_workers)")
+    call = next(n for n in ast.walk(tree) if isinstance(n, ast.Call))
+    kw = next(k for k in call.keywords if k.arg == "num_workers")
+    assert _is_resolved(kw.value)

tests/test_evaluation.py

@@ -16,7 +16,7 @@
 
 import leech.dataset as dataset
 from leech.constants import AUTO_DATALOADER_WORKERS
-from leech.dataset import resolve_dataloader_workers
+from leech.dataset import resolve_dataloader_workers, resolve_val_dataloader_workers
 from leech.evaluation import _save_scores
 
 READ_IDS = ["read-a", "read-b", "read-c", "read-d"]
@@ -172,3 +172,66 @@ def test_cli_forwards_num_workers(self, tmp_path):
 
         assert result.exit_code == 0, result.output
         assert captured["num_workers"] == 4
+
+
+class _Stacked:
+    """Stands in for a LeechDataset whose chunks stacked into one tensor."""
+
+    _signals_tensor = object()
+
+
+class _ListFallback:
+    """Stands in for the `_try_stack` fallback: per-chunk lists, fork-unsafe."""
+
+    _signals_tensor = None
+
+
+class TestValLoaderWorkers:
+    """The validation loader starved the GPU once per epoch (issue #207).
+
+    #206 routed `eval test` through `resolve_dataloader_workers` but left the
+    in-training validation pass hardcoded to 0, so a 1.18M-chunk val set spent
+    ~5 minutes at near-idle GPU at every epoch boundary -- ~75 minutes across a
+    15-epoch run.
+    """
+
+    def test_stacked_dataset_gets_workers_on_cuda(self, monkeypatch):
+        """The normal case: contiguous buffers COW-share, so workers are safe."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 32)
+
+        assert resolve_val_dataloader_workers(_Stacked(), 0, "cuda") > 0
+
+    def test_val_matches_train_when_stacked(self, monkeypatch):
+        """Validation should not be a special case just for being validation."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 32)
+
+        assert resolve_val_dataloader_workers(_Stacked(), 0, "cuda") == (
+            resolve_dataloader_workers(0, "cuda")
+        )
+
+    def test_list_fallback_stays_serial(self, monkeypatch):
+        """The one real memory case: forking a list of tensors multiplies RSS.
+
+        `LeechDataset` stacks precisely so a fork shares the buffers; when
+        `_try_stack` could not, each worker faults N PyObject headers into
+        private copies. This is the exception the old blanket 0 was protecting.
+        """
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 32)
+
+        assert resolve_val_dataloader_workers(_ListFallback(), 0, "cuda") == 0
+
+    def test_cpu_stays_serial(self):
+        """On CPU workers compete with compute, stacked or not."""
+        assert resolve_val_dataloader_workers(_Stacked(), 0, "cpu") == 0
+
+    def test_explicit_request_honoured_when_stacked(self, monkeypatch):
+        """An explicit N is not capped, matching the train loader's contract."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 2)
+
+        assert resolve_val_dataloader_workers(_Stacked(), 3, "cuda") == 3
+
+    def test_explicit_request_still_loses_to_the_list_fallback(self, monkeypatch):
+        """Memory safety wins over an explicit request -- OOM is not a tradeoff."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 32)
+
+        assert resolve_val_dataloader_workers(_ListFallback(), 3, "cuda") == 0