diff --git a/CHANGELOG.md b/CHANGELOG.md index 435f848c..6b1f0920 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/adr/0124-machines-of-different-sizes-get-one-srun-step-per-size-so-each-runs-a-worker-per-cpu-it-was-granted.md b/docs/adr/0124-machines-of-different-sizes-get-one-srun-step-per-size-so-each-runs-a-worker-per-cpu-it-was-granted.md index a9723655..7f923393 100644 --- a/docs/adr/0124-machines-of-different-sizes-get-one-srun-step-per-size-so-each-runs-a-worker-per-cpu-it-was-granted.md +++ b/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 diff --git a/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 b/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 new file mode 100644 index 00000000..8ead69f0 --- /dev/null +++ b/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. diff --git a/docs/cluster.rst b/docs/cluster.rst index f46d2f42..145da6a7 100644 --- a/docs/cluster.rst +++ b/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. diff --git a/docs/config_keys.rst b/docs/config_keys.rst index 6db0f0a9..4bf629ac 100644 --- a/docs/config_keys.rst +++ b/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 `. + 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 `. Example: diff --git a/pybnf/cluster.py b/pybnf/cluster.py index 5fcfaf33..757a4524 100644 --- a/pybnf/cluster.py +++ b/pybnf/cluster.py @@ -689,20 +689,32 @@ def cpus_per_node(): for time rather than a fit that runs faster. The two numbers agree only when whole nodes were allocated, which is why the defect stayed hidden. - Three sources are consulted, best first: + Four sources are consulted, best first: * ``$SLURM_CPUS_ON_NODE`` -- what the allocation granted on a node. Preferred - because it is the only one that describes the *allocation* rather than the process - doing the asking, so it is still the right number for a worker the SSH launcher - starts on some other machine. (When nodes differ in size it describes this node; - per-node counts are issue #617.) + because it describes the *allocation* rather than the process doing the asking, + so it is still the right number for a worker the SSH launcher starts on some + other machine. (When nodes differ in size it describes this node; per-node counts + are issue #617.) + * ``$SLURM_JOB_CPUS_PER_NODE`` -- the same question answered for the job as a whole, + as a per-node list, reduced here to its smallest entry. SLURM sets the variable + above only inside a step running on an allocated node, so it is *empty* when the + launching process is not on one: ``salloc`` opens its shell on the login node on + many clusters, and the launcher then runs there while the allocation is held + elsewhere (#642). This variable is set correctly in that shell, and the two + machine-level numbers below are not -- they describe a login node that is not in + the allocation at all, and is typically far larger than what the job holds. The + *smallest* entry is taken because this single number both sizes a pool started on + every machine and is what an srun step asks for on every machine in it: asking for + less than a machine holds costs speed, while asking for more than the smallest + machine holds is refused outright. * ``dask.system.CPU_COUNT`` -- what the operating system will let this process run on: the machine's processors narrowed by CPU affinity and by any cgroup CPU quota. This is the number a local run already sizes itself by, and it is the right one whenever the job is confined on the machine PyBNF is running on but no scheduler published a count. * ``multiprocessing.cpu_count()`` -- the whole machine, correct only when nothing is - limiting the job at all, and reached only if neither number above is usable. + limiting the job at all, and reached only if no number above is usable. The srun launcher does not merely count workers with this: it also asks SLURM for that many CPUs per task, and a request larger than the allocation is refused @@ -714,6 +726,10 @@ def cpus_per_node(): granted = os.environ.get('SLURM_CPUS_ON_NODE', '').strip() if granted.isdigit() and int(granted) > 0: return int(granted), 'what SLURM granted the job ($SLURM_CPUS_ON_NODE)' + allocation = expand_cpus_per_node(os.environ.get('SLURM_JOB_CPUS_PER_NODE', '')) + if allocation and min(allocation) > 0: + return min(allocation), ('the smallest machine in the allocation ' + '($SLURM_JOB_CPUS_PER_NODE)') if DASK_CPU_COUNT > 0: return DASK_CPU_COUNT, ("this process's CPU affinity and cgroup limits " '(dask.system.CPU_COUNT)') @@ -771,10 +787,19 @@ def dask_scheduler_command(scheduler_file): return [*DASK_CLI, 'scheduler', '--scheduler-file', scheduler_file] @staticmethod - def srun_worker_command(scheduler_file, node_count, parallel_count=None): + def srun_worker_command(scheduler_file, node_count, parallel_count=None, + granted=None, source=None): """ Build the srun invocation that starts one ``dask worker`` process group per node. + ``granted`` is how many CPUs the job holds on a node, which both sizes the default + worker pool and is what this step asks SLURM for. The caller passes it when it + already knows -- :meth:`srun_worker_layout` reads the allocation's own per-node list + before it decides which command to build, and that list is right even when the + launching process is not on an allocated node, which is where re-deriving the number + here went wrong (#642). It is derived from :meth:`cpus_per_node` only when no caller + supplied it. + :param scheduler_file: Path of the scheduler file the workers should read :type scheduler_file: str :param node_count: Number of nodes in the allocation @@ -782,10 +807,17 @@ def srun_worker_command(scheduler_file, node_count, parallel_count=None): :param parallel_count: Total number of worker processes over all nodes, or None for one per granted CPU :type parallel_count: int or None + :param granted: CPUs the job holds on a node, or None to work it out here + :type granted: int or None + :param source: Phrase naming where ``granted`` came from, for the log + :type source: str or None :return: the srun argument list :rtype: list """ - granted, source = Cluster.cpus_per_node() + if granted is None: + granted, source = Cluster.cpus_per_node() + elif source is None: + source = 'the allocation' if parallel_count is None: n_per_node = granted else: @@ -944,14 +976,22 @@ def srun_worker_layout(scheduler_file, out_dir, node_names, parallel_count): main_log = os.path.join(out_dir, SRUN_WORKER_LOG) if parallel_count is not None: # The explicit override is left exactly as it was: one srun, an even split over all - # nodes. Making it per-machine is deliberately out of scope (#617). + # nodes, its CPU request capped by the single count cpus_per_node decides. Capping + # it per-machine instead is ADR-0124's deferred question, not this one's; what #642 + # changed is that the single count now describes the allocation rather than the + # machine PyBNF happens to be launching from. cmd = Cluster.srun_worker_command(scheduler_file, node_count, parallel_count) per_node = int(cmd[cmd.index('--nworkers') + 1]) return [cmd], [main_log], per_node * node_count counts, source = Cluster.per_node_cpus(node_names) if len(set(counts)) <= 1: # Every machine the same size (the norm): the current single command, unchanged. - cmd = Cluster.srun_worker_command(scheduler_file, node_count, None) + # The count is handed on rather than worked out again, because the number in hand + # came from the allocation's own per-node list and so is right wherever this + # process is running, while re-deriving it asks a variable SLURM leaves empty off + # an allocated node and falls through to the size of the login node (#642). + cmd = Cluster.srun_worker_command(scheduler_file, node_count, None, + granted=counts[0], source=source) per_node = int(cmd[cmd.index('--nworkers') + 1]) return [cmd], [main_log], per_node * node_count # Machines of different sizes: one srun step per distinct size, each machine in the diff --git a/tests/test_cluster.py b/tests/test_cluster.py index cbcf9604..45171c11 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -1368,14 +1368,17 @@ def test_scheduler_file_chooses_where_it_is_written(self): class TestCpusPerNode: """The one place either launcher decides how many workers a node gets (#616). - The three sources are given three different numbers throughout, so each test + The four sources are given four different numbers throughout, so each test pins *which* one was consulted rather than merely a plausible count. """ - def _sources(self, monkeypatch, granted=None, affinity=6, cpu=64): + def _sources(self, monkeypatch, granted=None, allocation=None, affinity=6, cpu=64): monkeypatch.delenv('SLURM_CPUS_ON_NODE', raising=False) + monkeypatch.delenv('SLURM_JOB_CPUS_PER_NODE', raising=False) if granted is not None: monkeypatch.setenv('SLURM_CPUS_ON_NODE', str(granted)) + if allocation is not None: + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', str(allocation)) monkeypatch.setattr(cluster, 'DASK_CPU_COUNT', affinity) monkeypatch.setattr(cluster, 'cpu_count', lambda: cpu) @@ -1415,6 +1418,46 @@ def test_a_granted_count_never_reports_the_machine_count(self, monkeypatch): self._sources(monkeypatch, granted=4, affinity=4, cpu=128) assert cluster.Cluster.cpus_per_node()[0] == 4 + def test_the_jobs_own_per_node_list_is_read_when_the_step_variable_is_empty(self, monkeypatch): + """#642: SLURM sets $SLURM_CPUS_ON_NODE only inside a step running on an + allocated node, so it is empty in the shell ``salloc`` opens on a login node -- + while $SLURM_JOB_CPUS_PER_NODE, which answers the same question for the job as + a whole, is set correctly there. Reading it keeps the count describing the + allocation (20) instead of falling through to the machine-level numbers, which + describe a login node that is not in the allocation at all (128).""" + self._sources(monkeypatch, granted=None, allocation='20', affinity=128, cpu=96) + count, source = cluster.Cluster.cpus_per_node() + assert count == 20 + assert 'SLURM_JOB_CPUS_PER_NODE' in source + + def test_a_mixed_allocation_reduces_to_its_smallest_machine(self, monkeypatch): + """One number has to serve every machine here: it sizes a pool started on all of + them, and an srun step asks SLURM for it on all of them. Asking for fewer CPUs + than a machine holds costs speed, while asking for more than the smallest machine + holds is refused outright, so the safe direction is the smallest grant. (Sizing + each machine on its own is per_node_cpus, which the default srun path uses.)""" + self._sources(monkeypatch, granted=None, allocation='40(x2),96', affinity=128, cpu=96) + assert cluster.Cluster.cpus_per_node()[0] == 40 + + def test_the_per_step_count_still_wins_when_both_are_set(self, monkeypatch): + """Inside the allocation both are set, and the per-step count is the one that + describes *this* node. The job-level list is a fallback for a launching process + that is not on an allocated node, and does not change what a launch from inside + the allocation reads.""" + self._sources(monkeypatch, granted=12, allocation='40(x2),96', affinity=6, cpu=64) + count, source = cluster.Cluster.cpus_per_node() + assert count == 12 + assert 'SLURM_CPUS_ON_NODE' in source + + @pytest.mark.parametrize('value', ['', 'nonsense', '0']) + def test_an_unusable_per_node_list_falls_through_to_the_local_numbers(self, monkeypatch, value): + """A list this code cannot read, or one that says no CPUs at all, is treated as + absent rather than guessed at -- the local numbers are still better than nothing.""" + self._sources(monkeypatch, granted=None, allocation=value, affinity=6, cpu=64) + count, source = cluster.Cluster.cpus_per_node() + assert count == 6 + assert 'affinity' in source + class TestExpandCpusPerNode: """Reading $SLURM_JOB_CPUS_PER_NODE into one count per machine (#617). SLURM @@ -1474,6 +1517,18 @@ def test_unparseable_variable_falls_back(self, monkeypatch): counts, _ = cluster.Cluster.per_node_cpus(['n1', 'n2']) assert counts == [8, 8] + def test_the_fallback_off_an_allocated_node_still_reads_the_allocation(self, monkeypatch): + """The two fixes compose. A list that cannot be lined up with the machines is not + used to size them individually, but its values still describe machines in *this* + job, so the single count they fall back to is the smallest of them -- not the size + of the login node the launcher happens to be running on (#642).""" + monkeypatch.delenv('SLURM_CPUS_ON_NODE', raising=False) # not on an allocated node + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '20,40') # two, but three nodes + monkeypatch.setattr(cluster, 'DASK_CPU_COUNT', 128) + counts, source = cluster.Cluster.per_node_cpus(['n1', 'n2', 'n3']) + assert counts == [20, 20, 20] + assert 'SLURM_JOB_CPUS_PER_NODE' in source + def test_the_fallback_is_logged_so_a_mixed_cluster_user_is_told(self, monkeypatch, caplog): """A user on a mixed cluster is expecting each machine to be sized on its own, so falling back to one size for all is worth a warning that says why.""" @@ -1518,6 +1573,7 @@ class TestSrunWorkerCommand: def _patch(self, monkeypatch, granted=8): monkeypatch.setenv('SLURM_CPUS_ON_NODE', str(granted)) + monkeypatch.delenv('SLURM_JOB_CPUS_PER_NODE', raising=False) def test_default_is_one_worker_per_granted_cpu(self, monkeypatch): """parallel_count unset ⇒ one single-threaded worker per CPU the job holds @@ -1583,6 +1639,31 @@ def test_scheduler_file_is_one_literal_argument(self, monkeypatch): cmd = cluster.Cluster.srun_worker_command('/tmp/a b$(whoami).json', 1) assert cmd[cmd.index('--scheduler-file') + 1] == '/tmp/a b$(whoami).json' + def test_a_caller_supplied_count_is_used_instead_of_the_environment(self, monkeypatch, caplog): + """#642: a caller that already knows what the allocation granted passes it in, and + that number is what sizes the pool and what SLURM is asked for -- the environment + is not consulted at all. srun_worker_layout is that caller: it reads the + allocation's own per-node list, which is right wherever this process is running, + before it decides which command to build. The phrase naming where the count came + from is the caller's too, since the log is how an unexpected count is traced.""" + self._patch(monkeypatch, granted=128) # must not be consulted + with caplog.at_level('INFO'): + cmd = cluster.Cluster.srun_worker_command( + '/s.json', 2, None, granted=20, + source='what SLURM granted each machine ($SLURM_JOB_CPUS_PER_NODE)') + assert cmd[cmd.index('--nworkers') + 1] == '20' + assert cmd[cmd.index('--cpus-per-task') + 1] == '20' + assert any('SLURM_JOB_CPUS_PER_NODE' in r.message for r in caplog.records) + + def test_a_caller_supplied_count_still_caps_an_oversubscribed_request(self, monkeypatch): + """The cap is the same rule as before, measured against the caller's count: the + workers a user explicitly asked for all start, while the CPU request stays inside + what the job holds so SLURM does not refuse the step.""" + self._patch(monkeypatch, granted=128) # must not be consulted + cmd = cluster.Cluster.srun_worker_command('/s.json', 1, parallel_count=64, granted=20) + assert cmd[cmd.index('--nworkers') + 1] == '64' + assert cmd[cmd.index('--cpus-per-task') + 1] == '20' + class _SchedulerSpy: """A fake ``Popen`` for the srun launcher: records each command, and optionally @@ -1951,6 +2032,78 @@ def test_an_explicit_parallel_count_is_still_one_step_and_an_even_split(self, mo assert expected == 12 assert worker_logs == [str(tmp_path / 'dask_workers.log')] + def _login_node(self, monkeypatch, allocation, cores=128): + """The environment of the shell ``salloc`` opens on a login node (#642): SLURM + publishes the job's per-node CPU list, but not the per-step count, which it sets + only inside a step running on an allocated node. Both machine-level numbers + describe the login node, which is large and is not in the allocation.""" + monkeypatch.delenv('SLURM_CPUS_ON_NODE', raising=False) + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', allocation) + monkeypatch.setattr(cluster, 'DASK_CPU_COUNT', cores) + monkeypatch.setattr(cluster, 'cpu_count', lambda: cores) + + def test_a_launch_from_the_login_node_asks_for_what_the_job_holds(self, monkeypatch, tmp_path): + """#642, the reported failure: on many clusters the shell ``salloc`` opens runs on + the login node while the allocation is held on a compute node. The launcher used to + fall through to the login node's own processor count and ask SLURM for that many + CPUs per task, which SLURM refuses outright -- "Unable to create step ...: More + processors requested than permitted" -- so no worker ever started. The oracle is the + argv: 20 workers and 20 CPUs, from the allocation, not 128 from the login node.""" + sched_file = tmp_path / 'dask_scheduler.json' + spy = _SchedulerSpy(str(sched_file), write_after=1) + monkeypatch.setattr(cluster, 'Popen', spy) + monkeypatch.setattr(cluster.time, 'sleep', spy.sleep) + self._login_node(monkeypatch, '20', cores=128) + + _, worker_procs, expected, _ = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1'], parallel_count=None) + + (_, _), (srun_cmd, _) = spy.calls # scheduler + exactly one srun + assert worker_procs == [spy.procs[1]] + assert srun_cmd[srun_cmd.index('--cpus-per-task') + 1] == '20' + assert srun_cmd[srun_cmd.index('--nworkers') + 1] == '20' + assert expected == 20 + + def test_an_explicit_parallel_count_is_capped_by_the_job_not_by_the_login_node(self, monkeypatch, tmp_path): + """The CPU request is capped so that a deliberately oversubscribed parallel_count + still runs: every worker the user asked for starts, while the request stays inside + what the job holds, because SLURM refuses a larger one. From the login node that cap + was measured against the login node's processors and so did nothing -- 64 workers + asked for 64 CPUs of a 20-CPU allocation and the step was refused.""" + sched_file = tmp_path / 'dask_scheduler.json' + spy = _SchedulerSpy(str(sched_file), write_after=1) + monkeypatch.setattr(cluster, 'Popen', spy) + monkeypatch.setattr(cluster.time, 'sleep', spy.sleep) + self._login_node(monkeypatch, '20', cores=128) + + _, _, expected, _ = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1'], parallel_count=64) + + (_, _), (srun_cmd, _) = spy.calls + assert srun_cmd[srun_cmd.index('--nworkers') + 1] == '64' # what the user asked for + assert srun_cmd[srun_cmd.index('--cpus-per-task') + 1] == '20' + assert expected == 64 + + def test_the_count_the_layout_read_is_the_count_the_command_uses(self, monkeypatch, tmp_path): + """The per-machine list PyBNF already read is handed to the command rather than the + environment being read a second time. So a $SLURM_CPUS_ON_NODE that does not describe + this job -- exported by hand as the #642 workaround, or left over from an earlier + allocation -- no longer sizes the run: the list SLURM published for this job does.""" + sched_file = tmp_path / 'dask_scheduler.json' + spy = _SchedulerSpy(str(sched_file), write_after=1) + monkeypatch.setattr(cluster, 'Popen', spy) + monkeypatch.setattr(cluster.time, 'sleep', spy.sleep) + monkeypatch.setenv('SLURM_CPUS_ON_NODE', '128') # stale, not this job + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '20(x2)') # this job, two machines + + _, _, expected, _ = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2'], parallel_count=None) + + (_, _), (srun_cmd, _) = spy.calls + assert srun_cmd[srun_cmd.index('--cpus-per-task') + 1] == '20' + assert srun_cmd[srun_cmd.index('--nworkers') + 1] == '20' + assert expected == 40 + def test_an_unusable_cpus_per_node_falls_back_to_one_uniform_step(self, monkeypatch, tmp_path, caplog): """If SLURM did not publish a per-machine list PyBNF can line up with the node names, it cannot size each machine, so it does the safe thing the launcher did before #617: