fix(eval): feed the GPU from workers instead of one core

CHANGELOG.md

@@ -9,6 +9,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- **`leech eval test` fed the GPU from a single core.** The eval DataLoader was
+  built with `num_workers` pinned to 0 and no flag to change it, so collate, the
+  host-to-device copy and the forward pass all ran serially in one process: 8%
+  GPU utilisation on an A5000 over a 7,835,334-chunk test set, against 98% for
+  `model train` on the same corpus and the same card — same dataset class, same
+  collate function, the only difference being that training had workers.
+
+  The rule for sizing a loader now lives in exactly one place,
+  `dataset.resolve_dataloader_workers`, which training carried inline and
+  evaluation did not have at all. Its semantics are training's: `0` means
+  *auto*, auto is 0 on CPU (workers there would compete with the compute) and
+  >0 on CUDA, and a daemonic process — a grid-search `mp.Pool` worker — always
+  gets 0, because it cannot spawn children.
+
+  Auto is now also capped by the CPUs the process may actually run on
+  (`sched_getaffinity`, which respects the Slurm cpuset). Without that, the
+  pipeline's GPU eval rules, which request `cpus_per_task=2`, would have forked
+  8 workers onto 2 cores. An explicit `--num-workers N` is honoured as given.
+
 - **`leech_core`'s version now tracks `leech`'s.** It sat at `0.3.0` from v0.3.1
   to v0.6.4 — ten releases, spanning #176, #185, #187, #188, #192, #195, #200
   and #202 — while the Rust changed underneath it. That is not cosmetic: `uv`
@@ -31,6 +50,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `--num-workers` on `leech eval test`, so the auto default can be overridden
+  where it is wrong (default `0` = auto, as on `model train`).
+
 - `tests/test_rust_version_pairing.py`: asserts the two declared versions agree
   in the source tree, that `rust/pyproject.toml` defers rather than pinning a
   third copy, and that the *installed* extension matches the tree — the last of

docs/reference/cli.md

@@ -504,8 +504,19 @@ leech eval test --model FILE --test-data FILES --output FILE [OPTIONS]
 | `--test-data FILES` | *(required)* | Test data JSON files |
 | `--output FILE` | *(required)* | Output metrics JSON |
 | `--device STR` | auto | `cuda` or `cpu` |
+| `--num-workers INT` | `0` (auto) | DataLoader workers; auto is 8 on GPU (capped by the job's CPU allocation) and 0 on CPU |
 | `--emit-scores PATH` | *(off)* | Also write per-chunk `read_ids`, `labels`, `probs` to an `.npz` |
 
+!!! note "Why `--num-workers` matters on a GPU"
+
+    Collate, the host-to-device copy and the forward pass all run in whichever
+    process owns the loader. With no workers that is one core feeding an
+    accelerator that then waits: a 7.8M-chunk test set evaluated at **8% GPU**
+    on an A5000, while training the same corpus on the same card ran at 98%.
+    The default now resolves to workers on CUDA, capped by the CPUs the job is
+    actually allowed to use, so a 2-core allocation gets 1 worker rather than 8
+    thrashing ones. Raise it when the eval job has cores to spare.
+
 !!! note "Why `--emit-scores` exists"
 
     The metrics JSON is a summary at **one** threshold. The per-chunk scores

src/leech/cli.py

@@ -596,7 +596,7 @@ def merge(
     "--num-workers",
     type=int,
     default=0,
-    help="DataLoader workers (0=auto: 8 for GPU, 0 for CPU)",
+    help="DataLoader workers (0=auto: up to 8 on GPU, capped by CPUs; 0 on CPU)",
 )
 @click.option(
     "--balance-groups/--no-balance-groups",
@@ -1194,7 +1194,7 @@ def fetch(name, model_version, tag, output_dir, repo):
     "--num-workers",
     type=int,
     default=0,
-    help="DataLoader workers (0=auto: 8 for GPU, 0 for CPU)",
+    help="DataLoader workers (0=auto: up to 8 on GPU, capped by CPUs; 0 on CPU)",
 )
 @click.option(
     "--balance-groups/--no-balance-groups",
@@ -1414,6 +1414,12 @@ def eval():
     default=512,
     help="Batch size for evaluation (default: 512)",
 )
+@click.option(
+    "--num-workers",
+    type=int,
+    default=0,
+    help="DataLoader workers (0=auto: up to 8 on GPU, capped by CPUs; 0 on CPU)",
+)
 @click.option(
     "--emit-scores",
     type=click.Path(path_type=Path),
@@ -1424,7 +1430,7 @@ def eval():
         "question needs."
     ),
 )
-def test(model, test_data, output, device, batch_size, emit_scores):
+def test(model, test_data, output, device, batch_size, num_workers, emit_scores):
     """Test a trained model on a holdout test set."""
     from leech.commands.eval import handle_test
 
@@ -1434,6 +1440,7 @@ def test(model, test_data, output, device, batch_size, emit_scores):
         output=output,
         device=device,
         batch_size=batch_size,
+        num_workers=num_workers,
         emit_scores=emit_scores,
     )
 

src/leech/commands/eval.py

@@ -19,6 +19,7 @@ def handle_test(
     output: Path,
     device: str = "cuda",
     batch_size: int = 512,
+    num_workers: int = 0,
     emit_scores: Path | None = None,
 ) -> None:
     """
@@ -30,6 +31,7 @@ def handle_test(
         output: Output metrics file (JSON)
         device: Device for inference
         batch_size: Batch size for evaluation
+        num_workers: DataLoader workers (0 = auto: up to 8 on GPU, 0 on CPU)
         emit_scores: Optional .npz for per-chunk scores
     """
     from leech.evaluation import evaluate_model
@@ -44,6 +46,7 @@ def handle_test(
         output_path=output,
         device=device,
         batch_size=batch_size,
+        num_workers=num_workers,
         emit_scores=emit_scores,
     )
 

src/leech/constants.py

@@ -55,6 +55,9 @@
 DEFAULT_BATCH_SIZE = 128
 DEFAULT_LEARNING_RATE = 0.001
 DEFAULT_EPOCHS = 50
+# Workers a DataLoader gets when the caller asks for 0 ("auto") on a GPU.
+# Resolved by ``dataset.resolve_dataloader_workers``; 0 on CPU.
+AUTO_DATALOADER_WORKERS = 8
 
 # Advanced training defaults
 DEFAULT_WEIGHT_DECAY = 0.0

src/leech/dataset.py

@@ -34,6 +34,7 @@
 """
 
 import logging
+import os
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -43,6 +44,7 @@
 from torch.utils.data import Dataset
 
 from leech.chunking import load_chunks
+from leech.constants import AUTO_DATALOADER_WORKERS
 from leech.features import encode_signal_kmer, sequence_to_int
 from leech.models.inference_wrapper import ModelInferenceWrapper
 
@@ -953,6 +955,54 @@ def collate_fn(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
     return result
 
 
+def _usable_cpus() -> int:
+    """CPUs this process is allowed to run on, not CPUs the machine has."""
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:  # not Linux
+        return os.cpu_count() or 1
+
+
+def resolve_dataloader_workers(num_workers: int, device: str) -> int:
+    """Resolve how many DataLoader workers to actually use.
+
+    ``num_workers=0`` means AUTO here, not "no workers": on CUDA it becomes
+    ``AUTO_DATALOADER_WORKERS``, on CPU it stays 0. Feeding a GPU from the main
+    process serializes collate, host-to-device copy and forward pass onto one
+    core, which is how ``eval test`` sat at 8% GPU on an A5000 (issue #205).
+    On CPU the workers would compete with the compute for the same cores, and
+    ``__getitem__`` is trivially fast against pre-tensorized data, so they only
+    add overhead.
+
+    The daemon check is not an optimization: daemonic processes (a
+    ``multiprocessing.Pool`` worker, as in grid search) cannot spawn children,
+    so a DataLoader with workers raises there. Every caller that builds a
+    loader goes through this function, so that guard lives in one place.
+
+    The auto count is capped by the CPUs this process may actually run on --
+    ``sched_getaffinity``, which respects the Slurm cpuset -- because a GPU job
+    allocated 2 cores would otherwise fork 8 workers onto them and thrash. An
+    explicit request is honoured as given; only "auto" is capped.
+    """
+    import multiprocessing
+
+    is_daemon = multiprocessing.current_process().daemon
+    if is_daemon:
+        effective = 0
+    elif num_workers > 0:
+        effective = num_workers
+    elif device == "cpu":
+        effective = 0
+    else:
+        effective = min(AUTO_DATALOADER_WORKERS, max(1, _usable_cpus() - 1))
+
+    logger.info(
+        f"DataLoader workers: {effective} "
+        f"(requested={num_workers}, daemon={is_daemon}, device={device})"
+    )
+    return effective
+
+
 class SignalDataset(Dataset):
     """Minimal signal-only dataset for ``SignalCNN``.
 

src/leech/evaluation.py

@@ -20,7 +20,7 @@
 )
 from torch.utils.data import DataLoader
 
-from leech.dataset import LeechDataset, collate_fn
+from leech.dataset import LeechDataset, collate_fn, resolve_dataloader_workers
 from leech.metrics import compute_metrics, print_metrics, save_metrics
 from leech.model_loading import load_model_from_checkpoint
 from leech.models.inference_wrapper import ModelInferenceWrapper
@@ -77,6 +77,7 @@ def evaluate_model(
     kmer_len: int | None = None,
     batch_size: int = 512,
     device: str = "cuda",
+    num_workers: int = 0,
     emit_scores: Path | None = None,
 ) -> dict:
     """
@@ -90,6 +91,8 @@ def evaluate_model(
         kmer_len: K-mer length (if None, read from model config)
         batch_size: Batch size for evaluation
         device: Device for inference
+        num_workers: DataLoader workers; 0 means auto (see
+            ``resolve_dataloader_workers``) -- up to 8 on CUDA, 0 on CPU
         emit_scores: If set, also write per-chunk scores to this .npz. A
             confusion matrix is a summary at ONE threshold; the scores behind
             it answer questions the summary cannot -- per-group error
@@ -160,9 +163,15 @@ def evaluate_model(
         dwell_template_table=dwell_template_table,
     )
 
-    loader_kwargs: dict = {"num_workers": 0}
+    # Collate, the host-to-device copy and the forward pass run serially in
+    # whichever process owns the loader, so a worker-less loader on a GPU means
+    # one core feeding an accelerator that then waits (issue #205).
+    effective_workers = resolve_dataloader_workers(num_workers, device)
+    loader_kwargs: dict = {"num_workers": effective_workers}
     if device != "cpu":
         loader_kwargs["pin_memory"] = True
+    if effective_workers > 0:
+        loader_kwargs["prefetch_factor"] = 4
     test_loader = DataLoader(
         test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn, **loader_kwargs
     )

src/leech/training.py

@@ -27,7 +27,7 @@
 
 import leech
 from leech.cli_config import make_console
-from leech.dataset import LeechDataset, collate_fn
+from leech.dataset import LeechDataset, collate_fn, resolve_dataloader_workers
 from leech.losses import AdversarialHead, FocalBCEWithLogitsLoss, RegressionHead
 from leech.models import get_model
 from leech.models.inference_wrapper import ModelInferenceWrapper
@@ -1487,25 +1487,7 @@ def train_model(
         )
 
     # Create data loaders
-    # Daemon processes (e.g. multiprocessing pool workers) cannot spawn children,
-    # so num_workers must be 0. On CPU, workers compete for CPU time with training;
-    # with pre-tensorized data __getitem__ is trivially fast, so workers add overhead.
-    import multiprocessing
-
-    is_daemon = multiprocessing.current_process().daemon
-    if is_daemon:
-        effective_workers = 0
-    elif num_workers > 0:
-        effective_workers = num_workers
-    elif device == "cpu":
-        effective_workers = 0
-    else:
-        effective_workers = 8  # auto default for CUDA
-
-    logger.info(
-        f"DataLoader workers: {effective_workers} "
-        f"(requested={num_workers}, daemon={is_daemon}, device={device})"
-    )
+    effective_workers = resolve_dataloader_workers(num_workers, device)
 
     # Seed the DataLoader generator from the run seed so shuffle order and each
     # worker's base seed are reproducible regardless of prior global-RNG use.

tests/test_evaluation.py

@@ -3,11 +3,20 @@
 `_save_scores` is where the per-chunk scores are joined back to read ids, and
 the join is positional -- so these tests are mostly about the failure mode that
 join has, not about the happy path.
+
+The rest covers how the eval DataLoader is sized: a worker-less loader on a GPU
+is what left issue #205 running at 8% utilisation.
 """
 
+from types import SimpleNamespace
+from unittest import mock
+
 import numpy as np
 import pytest
 
+import leech.dataset as dataset
+from leech.constants import AUTO_DATALOADER_WORKERS
+from leech.dataset import resolve_dataloader_workers
 from leech.evaluation import _save_scores
 
 READ_IDS = ["read-a", "read-b", "read-c", "read-d"]
@@ -77,3 +86,89 @@ def test_creates_parent_directory(self, tmp_path):
         _save_scores(out, _test_npz(tmp_path), LABELS, PROBS)
 
         assert out.exists()
+
+
+class TestDataLoaderWorkers:
+    """Eval must feed the GPU from more than one process (issue #205).
+
+    ``eval test`` built its loader with ``num_workers`` pinned to 0, so collate,
+    the host-to-device copy and the forward pass all ran serially in one Python
+    process: 8% GPU utilisation on an A5000 while training on the same corpus
+    and hardware ran at 98%.
+    """
+
+    def test_cuda_auto_gets_workers(self, monkeypatch):
+        """0 means auto, and auto on a GPU is not zero."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 32)
+
+        assert AUTO_DATALOADER_WORKERS > 0
+        assert resolve_dataloader_workers(0, "cuda") == AUTO_DATALOADER_WORKERS
+
+    def test_cpu_auto_stays_serial(self):
+        """On CPU the workers would compete with the compute for the same cores."""
+        assert resolve_dataloader_workers(0, "cpu") == 0
+
+    def test_auto_fits_the_cpu_allocation(self, monkeypatch):
+        """A GPU job given 2 cores must not fork 8 workers onto them."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 2)
+
+        assert resolve_dataloader_workers(0, "cuda") == 1
+
+    def test_auto_keeps_one_worker_on_a_single_core(self, monkeypatch):
+        """Even one worker decouples collate and the H2D copy from the forward pass."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 1)
+
+        assert resolve_dataloader_workers(0, "cuda") == 1
+
+    def test_explicit_request_wins(self, monkeypatch):
+        """Only auto is capped; a caller who asks for N gets N."""
+        monkeypatch.setattr(dataset, "_usable_cpus", lambda: 2)
+
+        assert resolve_dataloader_workers(3, "cuda") == 3
+        assert resolve_dataloader_workers(3, "cpu") == 3
+
+    def test_daemon_forces_zero(self, monkeypatch):
+        """A pool worker (grid search) cannot spawn children; a loader with
+        workers raises there, whatever the caller asked for."""
+        import multiprocessing
+
+        monkeypatch.setattr(
+            multiprocessing, "current_process", lambda: SimpleNamespace(daemon=True)
+        )
+
+        assert resolve_dataloader_workers(8, "cuda") == 0
+
+    def test_cli_forwards_num_workers(self, tmp_path):
+        """--num-workers reaches evaluate_model, so a caller can fix this from
+        outside even where the auto default is wrong."""
+        from click.testing import CliRunner
+
+        import leech.evaluation as evaluation
+        from leech.cli import cli
+
+        captured: dict = {}
+        model = tmp_path / "model.pt"
+        model.touch()
+        test_data = tmp_path / "test.npz"
+        test_data.touch()
+
+        with mock.patch.object(evaluation, "evaluate_model", captured.update):
+            result = CliRunner().invoke(
+                cli,
+                [
+                    "eval",
+                    "test",
+                    "--model",
+                    str(model),
+                    "--test-data",
+                    str(test_data),
+                    "--output",
+                    str(tmp_path / "metrics.json"),
+                    "--num-workers",
+                    "4",
+                ],
+                catch_exceptions=False,
+            )
+
+        assert result.exit_code == 0, result.output
+        assert captured["num_workers"] == 4