fix(eval): feed the GPU from workers instead of one core - #206
Merged
Conversation
`eval test` built its DataLoader with `num_workers` pinned to 0 and offered no way to change it, so collate, the host-to-device copy and the forward pass all ran serially in one Python process: 8% GPU on an A5000 over a 7.8M-chunk test set, while training the same corpus on the same card ran at 98%. The rule for how many workers a loader gets now lives in exactly one place, `dataset.resolve_dataloader_workers`, which training already had inline and evaluation did not have at all. It keeps training's semantics -- 0 means auto, auto is 0 on CPU (workers would compete with the compute) and >0 on CUDA, and a daemonic process (a grid-search pool worker) always gets 0 because it cannot spawn children -- and adds a cap: auto never exceeds the CPUs the process may actually run on, so a GPU job allocated 2 cores gets 1 worker rather than 8 thrashing ones. An explicit `--num-workers N` is honoured as given. `eval test` gains `--num-workers`, threaded through `handle_test` to `evaluate_model`, so the auto default can be overridden from outside. Fixes #205
jayhesselberth
added a commit
that referenced
this pull request
Aug 24, 2026
… guard it (#207) (#209) * fix(training): route the validation loader through the worker resolver (#207) #206 fixed `eval test` by resolving its DataLoader workers, and added `resolve_dataloader_workers` whose docstring says "Every caller that builds a loader goes through this function, so that guard lives in one place". The validation loader inside the training loop did not, and was hardcoded to 0 three lines below the call that resolves the train loader. The cost is once per epoch. On a 1,176,763-chunk binary val set the GPU sat at ~30% for ~5 minutes at every epoch boundary and then recovered -- ~75 minutes of near-idle accelerator across a 15-epoch run, and it scales with val size. That is the same shape as #205: `__getitem__` being cheap does not mean one process can saturate a GPU, because collate, pin, host-to-device and the forward pass still serialize onto that core. The memory half of the old comment is real, but narrower than a blanket 0. `LeechDataset` stacks per-chunk tensors into contiguous buffers *precisely* so a fork COW-shares them (see the note at dataset.py `_try_stack`); only the list fallback, taken when per-chunk shapes are inconsistent, makes each worker fault N PyObject headers into private copies and multiply peak RSS. So the exception is now scoped to exactly that case instead of penalising every run to protect it -- and it wins even over an explicit `--num-workers N`, because OOM is not a throughput tradeoff. The logic lives in `resolve_val_dataloader_workers` beside its sibling rather than inline in `train()`, which keeps the "one place" claim true and makes it testable without standing up a training loop. Six tests, and both halves are mutation-checked: removing the list-fallback guard fails two of them, and reverting to the old always-zero fails three. Other call sites that still bypass the resolver, not touched here: `calibration.py` 193/543 (runs on CUDA, most likely to matter next), `commands/benchmark.py` 63, and the legacy `SignalCNN` path at training.py 261/265. `gridsearch.py` is fine -- its pool workers are daemonic, where the resolver returns 0 anyway. * fix(calibration): resolve loader workers there too, and guard the whole class Both `calibration.py` loaders passed `num_workers=num_workers` with a default of 0, so on CUDA they got literally zero rather than AUTO -- the same starvation as #205 and #207, in the third place. Both now go through `resolve_val_dataloader_workers`: they feed validation datasets, so the list-fallback exception applies to them as well. The more useful half is the guard. `num_workers` may no longer be a bare literal anywhere in the package; it must come from a resolver, from a local whose name says it carries a resolved count, or carry a call-site marker `dataloader-workers: unresolved` with a reason. Two such markers exist: `commands/benchmark.py` (the worker count is the independent variable being benchmarked) and the legacy SignalCNN path in `training.py` (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). The guard took three attempts, and the first two are why it is written this way: 1. A file-scoped allow-list. Worse than nothing: `("training.py", None)` exempted the whole file, so reintroducing the #207 bug PASSED. 2. Inspecting `DataLoader(...)` call sites. Also passed the #207 reintroduction, because that bug lives in a `val_loader_kwargs` dict that reaches the loader via `**kwargs`, and the enclosing function resolves a DIFFERENT loader -- which is precisely the shape #207 had. So the check is on the VALUE wherever it is bound, not on loader construction. Mutation-tested against both real regressions: reintroducing the #207 literal fails it naming training.py, and reintroducing the calibration.py literal fails it naming calibration.py.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #205.
What was wrong
evaluation.pybuilt the eval DataLoader withnum_workersfixed at 0 andeval testexposed no flag to change it. Collate, the host-to-device copy and the forward pass therefore all ran serially in one process — 8% GPU on an A5000 over a 7.8M-chunk test set, against 98% formodel trainon the same corpus and card, same dataset class, same collate function.What changed
One place decides worker counts.
dataset.resolve_dataloader_workers(num_workers, device)holds the rule thattraining.pyhad inline andevaluation.pydidn't have at all. Same semantics as before for training:0means auto, not "no workers" — 0 on CPU (workers would compete with the compute, and__getitem__is trivially fast against pre-tensorized data), >0 on CUDA.mp.Poolworker) always gets 0: it cannot spawn children, so a loader with workers raises there.New: auto is capped by the CPUs the process may actually use (
sched_getaffinity, which respects the Slurm cpuset). A GPU eval job allocated 2 cores now gets 1 worker rather than 8 fighting over them; an explicit--num-workers Nis honoured as given. This matters here becausepipeline/workflow/rules/evaluate.smkrequestscpus_per_task=2for GPU eval — those rules will benefit more if their CPU request goes up, which I left alone since it changes cluster resource requests.eval testgains--num-workers, threaded throughhandle_test→evaluate_model. Workers also getprefetch_factor=4, matching training.Verification
tests/test_evaluation.pycover auto-on-CUDA, serial-on-CPU, the CPU-allocation cap, the one-core floor, explicit requests bypassing the cap, the daemon guard, and--num-workersreachingevaluate_model.eval testat--num-workers 0and--num-workers 2; the metrics JSONs are byte-identical, so ordering and the positional--emit-scoresjoin are unaffected.Not measured on a GPU here — the mechanism is the same one that keeps training at 98% on the reporter's hardware, but the utilisation number is unverified until it runs on the A5000.
🤖 Generated with Claude Code