Skip to content

fix(cluster): size the worker pool by what the job was granted, not by the machine (#616) - #632

Merged
wshlavacek merged 1 commit into
mainfrom
fix/616-worker-count-from-allocation
Aug 21, 2026
Merged

fix(cluster): size the worker pool by what the job was granted, not by the machine (#616)#632
wshlavacek merged 1 commit into
mainfrom
fix/616-worker-count-from-allocation

Conversation

@wshlavacek

Copy link
Copy Markdown
Collaborator

Closes #616.

The defect

PyBNF decided how many worker processes to start on each node by calling multiprocessing.cpu_count(), which reports every processor the machine has whatever the job scheduler granted it. On the cluster where this was measured, a job that asked for 4 CPUs was told the node had 128: PyBNF would have started one worker per processor and overshot the job's real capacity 32-fold.

Every worker is a separate process, so that multiplies memory use and leaves the workers competing for the same four CPUs — a fit that runs slower than it would have on the share it was given, or that runs out of memory. The defect could hide because a job that asks for whole nodes gets the right answer by coincidence: there the two numbers are equal.

The fix

Cluster.cpus_per_node is now the one place either launcher decides this, and it returns the count together with a phrase naming where it came from. The precedence:

  1. $SLURM_CPUS_ON_NODE — what the allocation granted. Preferred because it is the only one of the three that describes the allocation rather than the process asking, so it is still the right number for a worker the SSH launcher starts on some other machine.
  2. dask.system.CPU_COUNT — the machine's processors narrowed by CPU affinity and by any cgroup quota, i.e. what the operating system will actually permit here. This is the number a single-machine run already sizes itself by.
  3. multiprocessing.cpu_count() — the whole machine, correct only when nothing is limiting the job at all.

Both launchers log the count, the node count and the source, so a user who sees an unexpected number of workers can trace it to the number PyBNF believed:

INFO Starting 4 worker process(es) on each of 1 node(s), one per CPU, from what SLURM granted the job ($SLURM_CPUS_ON_NODE)

Setting parallel_count still overrides all of it, and its branch now names that key as the source rather than logging a bare "Manually setting N workers per node". -t slurm-srun, which already read $SLURM_CPUS_ON_NODE, is unchanged apart from the added provenance in its log line.

Tests

The tests give the three sources three different numbers, so each one pins which source PyBNF consulted rather than merely a plausible count. The reported case — 4 granted of a 128-processor node — is asserted directly. All four new setup_cluster assertions were confirmed to go red against the old code and green against the new.

Docs

docs/cluster.rst gains a "How many workers run on each node" section covering both launchers; the parallel_count key entry and the manual-Dask paragraph no longer name multiprocessing.cpu_count() as the cluster default. ADR-0089 and ADR-0122 both recorded the old split — 0089's parenthetical justification, that the number "is a remote node's core count anyway", was the defect itself — and now carry superseding notes.

Verification

  • The real SSH bring-up run against an unreachable host under SLURM_CPUS_ON_NODE=4: dask's CLI accepts the constructed argv (--nthreads 1 --nworkers 4) and the process is still running after the bring-up wait.
  • The srun launcher re-exercised end to end with stand-ins for srun/scontrol on PATH — a real scheduler, two real workers, a real task through the client, no orphans after teardown.
  • Full suite green: 4472 passed, 23 skipped. ruff clean. Docs build clean under -W --keep-going.

Not in scope

Per-node counts on nodes of different sizes (#617) — $SLURM_CPUS_ON_NODE describes the node PyBNF runs on, and dask ssh takes one count for all hosts.

…y the machine (#616)

PyBNF decided how many worker processes to start on each node by calling
`multiprocessing.cpu_count()`, which reports every processor the machine has
whatever the job scheduler granted it. On the cluster where this was measured, a
job that asked for 4 CPUs was told the node had 128: PyBNF would have started one
worker per processor and overshot the job's real capacity 32-fold.

Every worker is a separate process, so that multiplies memory use and leaves the
workers competing for the same four CPUs -- a fit that runs slower than it would
have on the share it was given, or that runs out of memory. The defect could hide
because a job that asks for *whole* nodes gets the right answer by coincidence:
there the two numbers are equal.

`Cluster.cpus_per_node` is now the one place either launcher decides this, and it
returns the count together with a phrase naming where it came from. The precedence
is:

  1. `$SLURM_CPUS_ON_NODE` -- what the allocation granted. Preferred because it is
     the only one of the three that describes the *allocation* rather than the
     process asking, so it is still the right number for a worker the SSH launcher
     starts on some other machine.
  2. `dask.system.CPU_COUNT` -- the machine's processors narrowed by CPU affinity
     and by any cgroup quota, i.e. what the operating system will actually permit
     here. This is the number a single-machine run already sizes itself by.
  3. `multiprocessing.cpu_count()` -- the whole machine, correct only when nothing
     is limiting the job at all.

Both launchers log the count, the node count and the source, so a user who sees an
unexpected number of workers can trace it to the number PyBNF believed. Setting
`parallel_count` still overrides all of it, and its branch now names that key as
the source rather than logging a bare "Manually setting N workers per node".
`-t slurm-srun`, which already read `$SLURM_CPUS_ON_NODE`, is unchanged apart from
the added provenance in its log line.

The tests give the three sources three different numbers, so each one pins *which*
source PyBNF consulted rather than merely a plausible count; the reported case (4
granted of a 128-processor node) is asserted directly, and all four new
setup_cluster assertions go red against the old code.

ADR-0089 and ADR-0122 both recorded the old split -- 0089's parenthesis that a
remote node's core count "is a remote node's core count anyway" was the defect
itself -- and now carry superseding notes.

Verified by running the real SSH bring-up against an unreachable host under
`SLURM_CPUS_ON_NODE=4`: dask's CLI accepts the constructed argv (`--nthreads 1
--nworkers 4`) and the process is still running after the bring-up wait. The srun
launcher was re-exercised end to end with stand-ins for srun/scontrol on PATH --
a real scheduler, two real workers, a real task through the client, no orphans
after teardown. Full suite green (4472 passed, 23 skipped); docs build clean
under -W --keep-going.
@wshlavacek
wshlavacek merged commit 95cce90 into main Aug 21, 2026
9 checks passed
@wshlavacek
wshlavacek deleted the fix/616-worker-count-from-allocation branch August 21, 2026 19:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

The worker count is taken from the whole machine rather than from what the job was granted

1 participant