Ask SLURM for what the job holds when the fit is launched from the login node (#642)

CHANGELOG.md

@@ -216,6 +216,25 @@ All notable changes to PyBNF are documented below. This project adheres to
   says what it has not bisected.
 
 ### Fixed
+- **A multi-machine fit started from the shell `salloc` opens no longer asks SLURM for more
+  processors than it granted (#642, ADR-0125).** On many clusters `salloc` returns a shell on the
+  **login node** while the allocation is held on a compute node, and that shell is exactly where
+  PyBNF is meant to run: it holds the allocation. Started there, the `srun` launcher
+  (`-t slurm-srun`) stopped at once with `srun: error: Unable to create step for job NNNNN: More
+  processors requested than permitted`, and no worker ever started. PyBNF was sizing the run by
+  `$SLURM_CPUS_ON_NODE`, which SLURM sets only inside a job step running on an allocated node, so
+  on the login node it is absent — and the two remaining numbers describe the machine asking, which
+  there is the login node, not in the allocation and usually several times larger. A job granted 20
+  CPUs was therefore sized as though it held 128, and 128 is a request SLURM refuses. The count now
+  falls back to `$SLURM_JOB_CPUS_PER_NODE`, the per-node list SLURM publishes for the **job**, which
+  is set correctly in that shell (its smallest entry, since one number has to be acceptable on every
+  machine in the step). Inside the allocation nothing changes: `$SLURM_CPUS_ON_NODE` is still
+  preferred where SLURM sets it, and the same `srun` command is built. The default path also hands
+  the per-machine counts it already read straight to the command rather than having it read the
+  environment a second time, so a stale `$SLURM_CPUS_ON_NODE` — including one exported by hand as
+  the workaround for this bug — no longer sizes a later run. The SSH launcher (`-t slurm`) reads the
+  same count, so it too stops starting a login node's worth of worker processes on each machine when
+  it is launched from the login node.
 - **When the workers cannot be started, the message says what went wrong and what to try
   instead (#618).** A multi-machine run whose workers failed to start stopped with
   `Failed to start the dask-ssh cluster (dask-ssh exited with code 1)` and, on the cluster this

docs/adr/0124-machines-of-different-sizes-get-one-srun-step-per-size-so-each-runs-a-worker-per-cpu-it-was-granted.md

@@ -114,7 +114,7 @@ its single-step, even-split behavior, and only the default (auto-sized) path bec
   a total evenly and requests `--cpus-per-task` for the per-node share, a share larger than a
   machine smaller than the one PyBNF runs on can be refused by SLURM on that machine. This predates
   #617 and is out of its scope. It is recorded here rather than fixed; if it needs fixing it needs
-  its own design, and a tracking issue, rather than being folded into this change.
+  its own design, and a tracking issue, rather than being folded into this change. That issue is #643.
 
 ## Alternatives considered
 

docs/adr/0125-a-launcher-running-outside-the-allocation-reads-the-jobs-own-per-node-cpu-list-so-the-srun-step-asks-slurm-for-what-it-granted-rather-than-for-the-login-nodes-size.md

@@ -0,0 +1,122 @@
+# A launcher running outside the allocation reads the job's own per-node CPU list, so the `srun` step asks SLURM for what it granted rather than for the login node's size (issue #642)
+
+**Status: Accepted and implemented (2026-08-22).** The `srun` launcher (#614,
+`cluster_type = slurm-srun`) sizes its worker pool by how many CPUs the job was granted on a node,
+and asks SLURM for that many CPUs per task. It worked that number out from `$SLURM_CPUS_ON_NODE`,
+which SLURM sets only inside a job step running on an allocated node. Launched from a shell that
+is not on one -- which is where `salloc` leaves the user on many clusters -- it fell through to the
+size of the machine it was running on, asked for that, and SLURM refused the step. This ADR makes
+the count come from the allocation in that case too.
+
+## The problem
+
+The reporter's cluster documents `salloc` as the way to open an interactive allocation, and
+`salloc` returns a shell **on the login node** while the allocation is held on a separate compute
+node. That is the shell PyBNF is meant to run in: it holds the allocation, so `srun` starts a step
+inside it rather than queuing a new job, which is what `docs/cluster.rst` tells the user to do and
+what `require_slurm_allocation` checks for.
+
+In that shell the launcher failed immediately:
+
+    srun: error: Unable to create step for job NNNNN: More processors requested than permitted
+
+`Cluster.cpus_per_node` prefers `$SLURM_CPUS_ON_NODE`. That variable is a **step** variable: SLURM
+publishes it to a process running on a node the job holds, and not to a login-node shell, where it
+is simply absent. The two remaining sources are both descriptions of the machine asking --
+`dask.system.CPU_COUNT`, and the whole processor count -- and on the login node they describe the
+login node, which is not in the allocation at all and is typically several times larger than what
+the job holds. So a job granted 20 CPUs was sized as though it held 128, and `--cpus-per-task 128`
+is a request SLURM refuses outright. No worker started, and the fit stopped.
+
+`$SLURM_JOB_CPUS_PER_NODE` was already right in that shell, and PyBNF was already reading it:
+`per_node_cpus` (#617) uses it to size each machine separately. It is a **job** variable, set
+wherever the job's environment reaches, so it survives the trip to the login node. The default
+equal-size path read it, confirmed every machine was the same size -- and then threw the counts
+away and had `srun_worker_command` derive the number again from the environment.
+
+This is not a variant of #616. That issue was about a count that described the *machine* instead of
+the *job* while PyBNF was running inside the job. Here PyBNF is not inside the job's nodes at all,
+so there is no local number that can be right, and the only usable answer is the one the scheduler
+published about the allocation.
+
+## The decision
+
+### The allocation's own per-node list is a source of the single count
+
+`cpus_per_node` gains a source between `$SLURM_CPUS_ON_NODE` and the two machine-level numbers: the
+smallest entry of `$SLURM_JOB_CPUS_PER_NODE`. The order matters and is deliberate. Inside the
+allocation, `$SLURM_CPUS_ON_NODE` stays preferred, exactly as #616 decided, so a launch from an
+allocated node reads what it always read and nothing about it changes. The new source is reached
+only when that variable is absent -- which is the case this issue is about -- and there it is
+strictly better than the two below it, which describe a machine outside the allocation.
+
+The **smallest** entry is the one taken, because this is the one-number answer: it sizes a pool
+started on every machine, and it is what a single `srun` step asks for on every machine in it.
+Asking for fewer CPUs than a machine holds costs speed; asking for more than the smallest machine
+holds is refused outright, which is the failure being fixed. So when the entries differ the safe
+direction is down. (Sizing each machine on its own is `per_node_cpus`, which the default `srun`
+path already uses; this single number is the fallback and the SSH launcher's only option.)
+
+Because both launchers size themselves through `cpus_per_node`, this also stops the SSH launcher
+from starting a login node's worth of worker processes on each allocated machine when it is
+launched from the login node.
+
+### The layout hands its counts on rather than having them derived again
+
+`srun_worker_layout` reads the per-node list before it decides which command to build. In the
+equal-size case it now passes that count (and the phrase naming where it came from, for the log)
+into `srun_worker_command`, which uses it verbatim; the command derives a count for itself only
+when no caller supplied one. This is what the issue proposed, and it is worth doing on its own
+terms even with the source list fixed: a number that has already been read from the allocation
+should not be re-read from the environment, where a different variable can answer differently.
+A `$SLURM_CPUS_ON_NODE` left over from an earlier allocation -- including one exported by hand as
+the workaround for this issue -- no longer decides the size of a later run.
+
+The argument list is otherwise untouched, so an allocation launched from inside itself builds the
+same `srun` command it built before, which the existing regression test still pins.
+
+### `parallel_count` keeps its single-count cap
+
+Setting `parallel_count` still builds one step with an even split, and its CPU request is still
+capped by the single `cpus_per_node` number so that a deliberately oversubscribed count starts all
+the workers the user asked for without the step being refused. What changed is what that number
+is: from the login node it now describes the allocation, so the cap does its job instead of being
+measured against a machine that is not in the run. Making the cap per-machine is the limitation
+ADR-0124 recorded as needing its own design; it is now tracked as #643 and stays there rather than
+being folded in here.
+
+## Consequences
+
+* **The reported failure cannot happen from a count PyBNF chose.** A launch from the login-node
+  shell `salloc` opens asks SLURM for what the job holds. The workaround the reporter found --
+  exporting `SLURM_CPUS_ON_NODE` by hand -- is no longer needed, and no longer has an effect on a
+  later allocation if it is left set.
+* **A launch from inside the allocation is unchanged.** `$SLURM_CPUS_ON_NODE` is still preferred
+  where SLURM sets it, so the command built there, and the number in the log, are what they were.
+* **The SSH launcher is fixed too, for free.** It reads the same single count, so a `-t slurm` run
+  started from a login node no longer sizes each remote machine by the login node's processors.
+* **The log still names the source.** The new phrase names `$SLURM_JOB_CPUS_PER_NODE`, so a user
+  who sees an unexpected count can still trace which number PyBNF believed and where it came from.
+* **What was verified where.** The constructed `srun` argument list is the oracle, as it was for
+  #614 and #617: a login-node environment (no `$SLURM_CPUS_ON_NODE`, a job list of 20, machine-level
+  numbers of 128) is stood up in the test and the command must ask for 20, on the default path and
+  on the `parallel_count` path. That SLURM then accepts the step is what the reporter's cluster
+  verifies.
+
+## Alternatives considered
+
+* **Only route the equal-size path through the per-node counts** (the issue's suggested fix, taken
+  on its own). Rejected as half of the fix: it leaves `parallel_count` and the SSH launcher reading
+  the login node, and leaves the fallback inside `per_node_cpus` -- reached when SLURM's list cannot
+  be lined up with the machines -- reading it too. It is kept as the *other* half, because handing
+  on a count already read is better than re-deriving it however good the source list is.
+* **Prefer `$SLURM_JOB_CPUS_PER_NODE` over `$SLURM_CPUS_ON_NODE` everywhere.** Rejected: on a
+  mixed allocation launched from an allocated node, the step variable describes *this* machine and
+  the job list has to be reduced to one number to compete with it. Reordering would change what
+  every existing SLURM run reads in order to fix a case where the step variable is not set at all.
+* **Refuse to run when no allocation-derived count is available.** Rejected: PyBNF has a usable
+  answer in every case that reaches the machine-level sources -- a single-machine allocation of a
+  whole node, a run outside SLURM entirely -- and refusing would break runs that work today.
+* **Detect the login node and warn.** Rejected as the wrong shape: there is no reliable test for
+  "this machine is not in the allocation" that is better than simply reading the number SLURM
+  published, and once the right number is read there is nothing to warn about.

docs/cluster.rst

@@ -93,7 +93,7 @@ Pass ``-t slurm-srun`` instead (or set ``cluster_type = slurm-srun``) to start t
 2. runs ``srun`` to start one Dask worker process group on each node of the allocation, each reading that file; and
 3. connects through that same file, waiting until at least one worker has registered before the fit starts.
 
-Run PyBNF from the shell that holds the allocation: the one ``salloc`` opened, or your ``sbatch`` script. A separate login into one of the allocated nodes does not inherit the allocation, and ``srun`` would then queue a new job rather than start the workers; PyBNF refuses to start in that case instead of appearing to hang. For the same reason, PyBNF should be the only job step running in the allocation, since a concurrent second ``srun`` can leave the workers waiting for resources.
+Run PyBNF from the shell that holds the allocation: the one ``salloc`` opened, or your ``sbatch`` script. That shell does not have to be *on* one of the allocated machines -- on many clusters ``salloc`` leaves you on the login node while the allocation is held elsewhere, which is fine; what matters is that the shell holds the allocation. A separate login into one of the allocated nodes does not inherit the allocation, and ``srun`` would then queue a new job rather than start the workers; PyBNF refuses to start in that case instead of appearing to hang. For the same reason, PyBNF should be the only job step running in the allocation, since a concurrent second ``srun`` can leave the workers waiting for resources.
 
 An example batch script -- ``examples/tcr/tcr_batch.sh`` with a single word changed::
 
@@ -123,13 +123,14 @@ The two launchers differ in what happens when the machines in one allocation are
 * The **srun** launcher (``-t slurm-srun``) sizes each machine on its own, one worker per CPU that machine was granted. When the machines differ in size it starts one ``srun`` job step per distinct size, so a run on two 40-processor machines and one 96-processor machine starts 40 workers on each of the first two and 96 on the third. The per-machine arrangement is written to the log at the start of the run.
 * The **SSH** launcher (``-t slurm``) uses one worker count for every machine, because ``dask ssh`` takes only a single count for all hosts. The count comes from the node PyBNF is running on, so on a mixed allocation it is right for that machine and may be too high or too low for the others.
 
-The count for a machine comes from the first of these that is available:
+The **srun** launcher's per-machine counts come from ``$SLURM_JOB_CPUS_PER_NODE``, the per-node list SLURM publishes for the job. Where that list cannot be lined up with the machines -- and always for the **SSH** launcher, which has only one count to give -- the count comes from the first of these that is available:
 
 * ``$SLURM_CPUS_ON_NODE``, which is what SLURM granted the job on a node;
+* the smallest entry in ``$SLURM_JOB_CPUS_PER_NODE``. SLURM sets the variable above only inside a job step running on an allocated node, so it is empty when PyBNF is launched from somewhere else -- on many clusters ``salloc`` opens its shell on the login node while the allocation is held on a compute node. This one is set correctly there, and the two numbers below are not: they describe the login node, which is not in the allocation and is usually much larger than what the job holds. The smallest entry is used because one number has to serve every machine, and asking SLURM for more CPUs than the smallest machine holds is refused outright;
 * the CPU count dask derives for this process, which is the machine's processors narrowed by CPU affinity and by any cgroup CPU quota -- the same number a single-machine PyBNF run sizes itself by; or
 * the machine's whole processor count, which is correct only when nothing is limiting the job.
 
-Which number was used, and which of the three it came from, is written to the log at the start of the run, so an unexpected worker count can be traced to the number PyBNF believed.
+Which number was used, and which of these it came from, is written to the log at the start of the run, so an unexpected worker count can be traced to the number PyBNF believed.
 
 Setting ``parallel_count`` overrides all of this with a total number of worker processes over all nodes, divided evenly among them; the log then names ``parallel_count`` as the source. Nodes of different sizes still get equal shares, on either launcher.
 

docs/config_keys.rst

@@ -972,7 +972,7 @@ Parallel Computing
 
   Each parallel job runs in its own **single-threaded worker process**, whether or not this key is set: the simulation backends hold process-wide state that is not thread-safe, so PyBNF never places two concurrently running jobs in one process. This key therefore sets a process count, not a thread count. Lowering it is the way to reduce the memory a run uses, since each worker process holds its own copy of the models.
 
-  Default: Use all available cores -- one single-threaded worker per core. Locally, the core count comes from Dask, which honors CPU affinity and cgroup quotas (so a run confined to 4 cores gets 4 workers, not the host's full count). On a cluster, with either ``cluster_type``, it is the number of CPUs the job was granted on a node (``$SLURM_CPUS_ON_NODE``), falling back to the affinity- and cgroup-aware count Dask derives and then to the machine's whole processor count; PyBNF logs the number it used and which of the three it came from. See :ref:`How many workers run on each node <workercount>`.
+  Default: Use all available cores -- one single-threaded worker per core. Locally, the core count comes from Dask, which honors CPU affinity and cgroup quotas (so a run confined to 4 cores gets 4 workers, not the host's full count). On a cluster, with either ``cluster_type``, it is the number of CPUs the job was granted on a node (``$SLURM_CPUS_ON_NODE``, or the job's own per-node list ``$SLURM_JOB_CPUS_PER_NODE`` when PyBNF is launched from outside the allocation, as it is from the login-node shell ``salloc`` opens on many clusters), falling back to the affinity- and cgroup-aware count Dask derives and then to the machine's whole processor count; PyBNF logs the number it used and which of these it came from. See :ref:`How many workers run on each node <workercount>`.
 
   Example: