diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b1f0920..8c6abce1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,7 +23,8 @@ All notable changes to PyBNF are documented below. This project adheres to machines, it falls back to sizing every machine the same, with a warning saying so. This applies only to the `srun` launcher. The SSH launcher (`-t slurm`) still uses one count for all machines, because `dask ssh` takes only one. Setting `parallel_count` still splits that total evenly across - the machines on either launcher. + same-size machines on either launcher; on machines of different sizes the `srun` launcher now + splits it in proportion to machine size (#643, see below). - **A multi-machine fit can start its workers without logging in anywhere, so clusters that use host-based or Kerberos SSH can run one at all (#614, ADR-0122).** PyBNF had exactly one way to run across several machines: `dask-ssh`, which logs in to every node with **paramiko** @@ -216,6 +217,20 @@ All notable changes to PyBNF are documented below. This project adheres to says what it has not bisected. ### Fixed +- **Setting `parallel_count` on a cluster whose machines are not all the same size no longer stops + the `srun` launcher from starting any worker (#643, ADR-0126).** `parallel_count` gives a total + number of workers over all machines. The `srun` launcher (`-t slurm-srun`) split that total + evenly and asked SLURM for the per-machine share on every machine in one job step. On a mixed + allocation the even share can be more than a smaller machine was granted, so SLURM refused the + step with `srun: error: Unable to create step for job NNNNN: More processors requested than + permitted`, and no worker started. For example, on a 96-CPU and a 40-CPU machine with + `parallel_count = 136` the even share is 68 per machine, and 68 is more than the 40-CPU machine + holds. The total is now split in proportion to each machine's granted CPUs, one job step per + distinct size, the same way the automatic (unset `parallel_count`) sizing already works, so each + step asks only for what its machines hold and none is refused. On the 96-CPU and 40-CPU example + that is 96 workers on the larger machine and 40 on the smaller. An allocation whose machines are + all the same size is unchanged: `parallel_count` is still split evenly in one step. The SSH + launcher (`-t slurm`) is unchanged, because `dask ssh` takes one worker count for all machines. - **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 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 7f923393..36e285cb 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 @@ -87,6 +87,10 @@ launcher's own readiness check is unchanged. ### `parallel_count` is left as an even split +**Superseded by issue #643 / ADR-0126:** the override is now split in proportion to machine size on +a mixed allocation, one `srun` step per size, so it no longer asks a smaller machine for the larger +machines' share. The rest of this section records why #617 left it alone at the time. + Setting `parallel_count` overrides the automatic sizing with a total number of workers split evenly across the machines, exactly as the SSH launcher has always done. This change does not make that override per-machine. A user who names a total is asking for that many workers, and dividing an @@ -114,7 +118,8 @@ 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. That issue is #643. + its own design, and a tracking issue, rather than being folded into this change. That issue is + #643, now fixed: the override is split in proportion to machine size, one step per size (ADR-0126). ## Alternatives considered diff --git a/docs/adr/0126-an-explicit-parallel-count-is-split-across-a-mixed-allocation-in-proportion-to-machine-size-one-srun-step-per-size-so-no-step-asks-a-machine-for-more-cpus-than-it-holds.md b/docs/adr/0126-an-explicit-parallel-count-is-split-across-a-mixed-allocation-in-proportion-to-machine-size-one-srun-step-per-size-so-no-step-asks-a-machine-for-more-cpus-than-it-holds.md new file mode 100644 index 00000000..0fecdb83 --- /dev/null +++ b/docs/adr/0126-an-explicit-parallel-count-is-split-across-a-mixed-allocation-in-proportion-to-machine-size-one-srun-step-per-size-so-no-step-asks-a-machine-for-more-cpus-than-it-holds.md @@ -0,0 +1,119 @@ +# An explicit `parallel_count` is split across a mixed allocation in proportion to machine size, one `srun` step per size, so no step asks a machine for more CPUs than it holds (issue #643) + +**Status: Accepted and implemented (2026-08-22).** Setting `parallel_count` gives the `srun` +launcher (#614, `cluster_type = slurm-srun`) a total number of workers to start over all machines. +The launcher split that total evenly and asked SLURM, in a single job step, for the per-machine +share on every machine. On an allocation whose machines are not all the same size, the even share +can be more than a smaller machine was granted, and SLURM refuses the whole step. This ADR splits +the total in proportion to each machine's size and starts one step per size, the way the automatic +sizing already works, so each step asks only for what its machines hold. It supersedes the decision +ADR-0124 recorded to leave this path an even split. + +## The problem + +`parallel_count` is a total, and the `srun` launcher divided it evenly: +`n_per_node = ceil(parallel_count / node_count)`. One `srun` step then asked SLURM for that many +CPUs per task with `--cpus-per-task`. That request has to be satisfiable on *every* machine the +step runs on, but it was capped against a single CPU count rather than against the smallest machine +in the allocation. On a mixed allocation the even share can exceed what the smaller machines were +granted, and SLURM refuses the step: + + srun: error: Unable to create step for job NNNNN: More processors requested than permitted + +On a 96-CPU machine and a 40-CPU machine with `parallel_count = 136`, the even share is 68 per +machine. Launched from the 96-CPU machine the cap is 96, so the step asks for 68 CPUs on both +machines, and the 40-CPU machine cannot supply them. No worker starts, and the fit stops. + +The automatic path does not have this problem. #617 (ADR-0124) gave it one `srun` step per distinct +machine size, each asking only for what those machines hold. Only the explicit-override path was +left as an even split. #642 (ADR-0125) did not change it either: it changed which single number the +even split's cap is measured against, so that the number describes the allocation rather than the +machine PyBNF was launched from, which is a different bug. + +ADR-0124 recorded this deliberately, under "`parallel_count` is left as an even split", and asked +for it to be tracked separately rather than folded into #617, because a fix changes what an +existing config key means on a mixed allocation and wanted its own decision. This is that decision. + +## The decision + +### On a mixed allocation, split the total in proportion to machine size + +The automatic path already groups the machines by how many CPUs each was granted and gives each +group its own `srun` step, so a smaller machine and a larger one are never in the same step and +never bound to the same `--cpus-per-task`. `parallel_count` now uses the same grouping. For a size +group whose machines were each granted `cpus` CPUs, with the allocation holding `total_cpus` across +all machines: + + workers per machine in the group = max(1, ceil(cpus * parallel_count / total_cpus)) + CPUs the step asks for = max(1, min(workers per machine, cpus)) + +A group holds `cpus / total_cpus` of the allocation's processors, so it gets that share of the +requested total. Each step's CPU request is capped at its own group's machine size, so SLURM +accepts every step; an oversubscribed `parallel_count` still starts every worker asked for while +the request stays satisfiable, the same contract the single-step path already had. + +On the 96-CPU and 40-CPU example with `parallel_count = 136`, which is exactly what the allocation +holds, this is one worker per CPU: 96 workers asking for 96 CPUs on the larger machine, 40 asking +for 40 on the smaller. Both steps are accepted, and both machines are fully and correctly bound. + +### `ceil`, and a floor of one, so the homogeneous case is the even split exactly + +The rounding is the same `ceil` the even split has always used, and this is not a coincidence: when +every machine is the same size, `cpus / total_cpus` is `1 / node_count`, so the per-machine count is +`ceil(parallel_count / node_count)`, the number the even split produces. The per-size split is the +generalization of the even split, not a different rule bolted on beside it, and a homogeneous +allocation is left running the single even-split step it ran before, byte for byte. Because `ceil` +rounds each machine's share up, the realized total can sit just above the request, exactly as the +even split's total already can (`ceil(5 / 3) * 3 = 6`). The floor of one guarantees every machine in +the allocation runs at least one worker even when its proportional share rounds to nothing. + +### Only the mixed case changes; the layout reads the sizes once and branches + +`srun_worker_layout` now reads the allocation's own per-node CPU list first, for both the automatic +and the explicit paths, and asks a single question: are the machines all one size? If they are, it +builds the one step it always built, with the count coming from `parallel_count` or from the grant +as before. If they are not, it builds one step per size, with each machine's worker count coming +from the grant (automatic) or from the proportional split (explicit). Reading the list first also +means the explicit path hands the count it read straight to the step rather than the step reading +the environment a second time, which is the property #642 relied on: a `$SLURM_CPUS_ON_NODE` that +does not describe this job no longer sizes the run. + +## Consequences + +* **The reported failure cannot happen from a count PyBNF chose.** Every step asks only for what + its own machines hold, so `parallel_count` on a mixed allocation no longer produces a request + SLURM refuses. That SLURM then places the concurrent per-size steps as intended on a real + heterogeneous allocation is what the reporter's cluster verifies, as it did for #614 and #617. +* **What `parallel_count` means on a mixed allocation changed.** It was an even split; it is now a + size-weighted split. This is the change ADR-0124 declined to make in passing. On a homogeneous + allocation, which is the common case, nothing changed. The realized per-machine arrangement and + the total are written to the log, one line per size, so an unexpected count can be traced. +* **A mixed allocation with `parallel_count` writes more than one worker log**, the same way the + automatic mixed path does (`dask_workers.log`, `dask_workers_2.log`, and so on), because + concurrent steps writing one file would interleave and truncate each other. +* **The common case is byte-for-byte unchanged.** A homogeneous allocation, with or without + `parallel_count`, builds exactly the single `srun` command it built before, which a regression + test pins. The automatic mixed path (#617) is unchanged: the per-group command builder defaults + its worker count to one per granted CPU, so a caller that does not pass a count gets the old argv. +* **What was verified where.** The constructed `srun` argument list is the oracle, as it was for + #614, #617, and #642. Tests stand up a mixed allocation (`40(x2),96`, and the reporter's `96,40`) + with `parallel_count` set and require each step to ask only for what its machines hold, on the + ordinary and the oversubscribed cases, and require the homogeneous case to stay one even-split + step. + +## Alternatives considered + +* **Cap the single even-split step at the smallest machine.** The step would be accepted, but a + task on a larger machine would fork the even share while `--cpus-per-task` asked only for the + smallest machine's count, and under `task/cgroup` binding that task is confined to the CPUs it + asked for (ADR-0122 documents this confinement as the reason the count is requested at all). The + larger machines would run partly serialized and the smaller machine oversubscribed, so the fit + would run slowly for no visible reason. Rejected: it trades a loud failure for a quiet one. +* **Distribute the total by a largest-remainder method so it sums to `parallel_count` exactly.** + Rejected as more machinery than the case needs. It would also have to give two same-size machines + different counts to spend the last unit, which breaks "one step per size" and would need a step + per machine. `ceil` keeps same-size machines identical, which is what lets them share one step, + and the even split already does not hit the total exactly, so matching it precisely buys nothing. +* **Leave the override an even split and only document the limitation.** This is what ADR-0124 did, + on the ground that the fix wanted its own decision. That decision is this ADR; leaving it undone + keeps a config key that fails outright on the mixed allocations #617 exists to support. diff --git a/docs/cluster.rst b/docs/cluster.rst index 145da6a7..84468228 100644 --- a/docs/cluster.rst +++ b/docs/cluster.rst @@ -132,7 +132,7 @@ The **srun** launcher's per-machine counts come from ``$SLURM_JOB_CPUS_PER_NODE` 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. +Setting ``parallel_count`` overrides all of this with a total number of worker processes over all nodes; the log then names ``parallel_count`` as the source. When the nodes are all the same size the total is divided evenly among them, on either launcher. When they differ in size, the SSH launcher still divides it evenly, but the ``slurm-srun`` launcher divides it in proportion to each machine's granted CPUs, so no machine is asked for more CPUs than it holds (a step that asked for more would be refused by SLURM). .. _sizing: diff --git a/pybnf/cluster.py b/pybnf/cluster.py index 757a4524..7da8ccda 100644 --- a/pybnf/cluster.py +++ b/pybnf/cluster.py @@ -821,9 +821,11 @@ def srun_worker_command(scheduler_file, node_count, parallel_count=None, if parallel_count is None: n_per_node = granted else: - # Same arithmetic as the SSH launcher: parallel_count is a total over all nodes. - # Per-node counts on nodes of different sizes are issue #617, which needs a - # per-node command either launcher can express -- ``dask ssh`` cannot. + # Same arithmetic as the SSH launcher: parallel_count is a total over all nodes, + # split evenly. This single-step command is used only when the machines are all + # one size; a mixed allocation is split per size by srun_worker_layout instead + # (#617 for the automatic count, #643 for parallel_count), because one step's + # --cpus-per-task cannot ask each machine for a different number. n_per_node = max(1, int(np.ceil(parallel_count / node_count))) # One task per node, each of which forks n_per_node single-threaded workers. The # task has to be given CPUs for all of them: with task/cgroup binding, a task that @@ -854,37 +856,49 @@ def srun_worker_command(scheduler_file, node_count, parallel_count=None, '--nthreads', '1'] @staticmethod - def srun_worker_command_for_group(scheduler_file, nodes, cpus): + def srun_worker_command_for_group(scheduler_file, nodes, cpus, n_workers=None): """ Build the srun invocation for one group of machines that were all granted the same - number of CPUs, starting one worker per granted CPU on each (#617). - - This is what the default path uses when the allocation holds machines of different - sizes. A single srun step cannot start different numbers of workers on different - machines -- ``--cpus-per-task`` is one value for the whole step, and under task/cgroup - binding a task that under-asked for CPUs is confined to them -- so each distinct size - is its own step, named by ``--nodelist``. The homogeneous case does not come here; it - stays the single :meth:`srun_worker_command`. + number of CPUs (#617). + + This is what PyBNF uses when the allocation holds machines of different sizes. A single + srun step cannot start different numbers of workers on different machines -- + ``--cpus-per-task`` is one value for the whole step, and under task/cgroup binding a + task that under-asked for CPUs is confined to them -- so each distinct size is its own + step, named by ``--nodelist``. The homogeneous case does not come here; it stays the + single :meth:`srun_worker_command`. + + ``n_workers`` is how many workers to start on each machine in the group. The default + (auto sizing) is one per granted CPU, so a caller that leaves it ``None`` gets that. + An explicit ``parallel_count`` split in proportion to machine size passes the group's + share instead (#643); the CPU request is then capped at what the machine holds, for the + same task/cgroup reason :meth:`srun_worker_command` caps its own, so an oversubscribed + count still starts every worker rather than being refused by SLURM. :param scheduler_file: Path of the scheduler file the workers should read :type scheduler_file: str :param nodes: The machines in this group, all granted the same CPU count :type nodes: list - :param cpus: CPUs granted on each machine in the group, and so workers to start there + :param cpus: CPUs granted on each machine in the group :type cpus: int + :param n_workers: Workers to start on each machine, or None for one per granted CPU + :type n_workers: int or None :return: the srun argument list for this group :rtype: list """ + if n_workers is None: + n_workers = cpus + cpus_per_task = max(1, min(n_workers, cpus)) return ['srun', # Name the exact machines this step runs on, so the steps for the different # sizes land on disjoint machines and can run at the same time. '--nodelist', ','.join(nodes), '--nodes', str(len(nodes)), '--ntasks', str(len(nodes)), - '--ntasks-per-node', '1', '--cpus-per-task', str(cpus), + '--ntasks-per-node', '1', '--cpus-per-task', str(cpus_per_task), '--label', *DASK_CLI, 'worker', '--scheduler-file', scheduler_file, - '--nworkers', str(cpus), + '--nworkers', str(n_workers), '--nthreads', '1'] @staticmethod @@ -952,14 +966,17 @@ def setup_srun_cluster(scheduler_file, out_dir, node_names, parallel_count=None) def srun_worker_layout(scheduler_file, out_dir, node_names, parallel_count): """ Work out the srun command(s) that start the workers, the log each writes, and how many - workers to expect in total (#617). - - There is one command in every case except the one this issue is about: a default - (auto-sized) run on machines of different sizes, which becomes one command per distinct - size so each machine can be given a worker per CPU it holds. Everything else -- an - explicit ``parallel_count``, and a default run where every machine is the same size -- - stays the single :meth:`srun_worker_command` it was before, so the common case is - unchanged. + workers to expect in total (#617, #643). + + A homogeneous allocation is one srun step, whether the count comes from + ``parallel_count`` or from what each machine was granted, so the common case is + unchanged. An allocation of different-sized machines becomes one step per distinct size, + because a single step's ``--cpus-per-task`` cannot ask each machine for a different + number without asking the smaller ones for more than they hold. On that mixed path the + auto (unset ``parallel_count``) run gives each machine one worker per CPU it was granted, + and an explicit ``parallel_count`` is split across the machines in proportion to their + size (#643): the even split a single step would use asks the smaller machines for the + larger machines' share, which SLURM refuses. :param scheduler_file: Path of the scheduler file the workers should read :type scheduler_file: str @@ -974,43 +991,51 @@ def srun_worker_layout(scheduler_file, out_dir, node_names, parallel_count): """ node_count = len(node_names) 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, 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 + # Read the allocation's own per-node list first: both paths need it to tell a + # homogeneous allocation from a mixed one, and the count in hand 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). 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. - # 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, + # Every machine the same size (the norm): one srun step. An even split when + # parallel_count is set, one worker per granted CPU otherwise -- both unchanged. + cmd = Cluster.srun_worker_command(scheduler_file, node_count, parallel_count, 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 - # step given one worker per CPU it was granted. Group by size, keeping first-seen order. + # Machines of different sizes: one srun step per distinct size. Group by size, keeping + # first-seen order. + total_cpus = sum(counts) groups = {} for name, cpus in zip(node_names, counts): groups.setdefault(cpus, []).append(name) - logger.info('Machines of different sizes in this allocation; sizing each by %s. ' - 'Starting one srun step per size:' % source) + if parallel_count is None: + logger.info('Machines of different sizes in this allocation; sizing each by %s. ' + 'Starting one srun step per size:' % source) + else: + logger.info('Machines of different sizes in this allocation (sized by %s); splitting ' + 'the requested %i worker(s) in proportion to machine size. Starting one ' + 'srun step per size:' % (source, parallel_count)) commands, logs, expected_total = [], [], 0 for index, (cpus, nodes) in enumerate(groups.items()): - commands.append(Cluster.srun_worker_command_for_group(scheduler_file, nodes, cpus)) + if parallel_count is None: + # Auto (#617): one worker per CPU the machine was granted. + n_per_machine = cpus + else: + # This size group holds cpus/total_cpus of the allocation's processors, so it + # gets that share of the requested total. ceil, and a floor of one, match the + # single-step even split (ceil(parallel_count/node_count)), which this is the + # per-size generalization of: on a homogeneous allocation the two agree exactly. + n_per_machine = max(1, int(np.ceil(cpus * parallel_count / total_cpus))) + commands.append(Cluster.srun_worker_command_for_group( + scheduler_file, nodes, cpus, n_per_machine)) # The first step keeps the usual log name; the rest are numbered, so the concurrent # steps do not truncate and interleave one another's output. logs.append(main_log if index == 0 else os.path.join(out_dir, 'dask_workers_%i.log' % (index + 1))) - expected_total += cpus * len(nodes) - logger.info(' %i worker(s) on %i machine(s) (%s)' - % (cpus, len(nodes), ', '.join(nodes))) + expected_total += n_per_machine * len(nodes) + logger.info(' %i worker(s) on %i machine(s) of %i CPU(s) (%s)' + % (n_per_machine, len(nodes), cpus, ', '.join(nodes))) return commands, logs, expected_total @staticmethod diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 45171c11..dbb37891 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -2012,16 +2012,92 @@ def test_the_per_machine_arrangement_is_logged(self, monkeypatch, tmp_path, capl assert 'n1' in text and 'n2' in text and '40' in text assert 'n3' in text and '96' in text - def test_an_explicit_parallel_count_is_still_one_step_and_an_even_split(self, monkeypatch, tmp_path): - """A user who sets parallel_count is asking for a specific number of workers, split - evenly across the machines the way the SSH launcher has always done. That path is - left alone by #617: one srun step, the count taken from parallel_count rather than - from what each machine was granted, even on a mixed allocation.""" + def test_an_explicit_parallel_count_on_a_mixed_allocation_is_split_by_machine_size(self, monkeypatch, tmp_path): + """#643: on machines of different sizes, parallel_count is split in proportion to what + each machine holds, one srun step per size, rather than split evenly in one step. The + even split would ask the smaller machines for the larger machines' share, which SLURM + refuses. Here 12 workers over 40,40,96 (176 CPUs) is ceil(40*12/176)=3 on each of the + two 40-CPU machines and ceil(96*12/176)=7 on the 96-CPU one, each step asking only for + what its machines hold.""" 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_JOB_CPUS_PER_NODE', '40(x2),96') # ignored: the user chose + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '40(x2),96') + + scheduler_proc, worker_procs, expected, worker_logs = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2', 'n3'], parallel_count=12) + + (_, _), (first_cmd, _), (second_cmd, _) = spy.calls # scheduler + one step per size + assert first_cmd[first_cmd.index('--nodelist') + 1] == 'n1,n2' + assert first_cmd[first_cmd.index('--nworkers') + 1] == '3' + assert first_cmd[first_cmd.index('--cpus-per-task') + 1] == '3' + assert second_cmd[second_cmd.index('--nodelist') + 1] == 'n3' + assert second_cmd[second_cmd.index('--nworkers') + 1] == '7' + assert second_cmd[second_cmd.index('--cpus-per-task') + 1] == '7' + assert worker_procs == [spy.procs[1], spy.procs[2]] + assert expected == 3 * 2 + 7 # every worker across both steps + assert worker_logs == [str(tmp_path / 'dask_workers.log'), + str(tmp_path / 'dask_workers_2.log')] + + def test_the_reported_parallel_count_failure_no_longer_asks_a_machine_for_too_much(self, monkeypatch, tmp_path): + """The exact failure #643 reports: 96-CPU and 40-CPU machines with parallel_count=136. + The old even split was 68 per machine, and a single step asked SLURM for 68 CPUs on the + 40-CPU machine, which it cannot supply -- "More processors requested than permitted". + 136 is exactly what the allocation holds (96+40), so the proportional split is one + worker per CPU: 96 on the 96-CPU machine and 40 on the 40-CPU one, each step asking only + for its machine's own size, so no step is 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) + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '96,40') + + _, _, expected, _ = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['big', 'small'], parallel_count=136) + + (_, _), (first_cmd, _), (second_cmd, _) = spy.calls + assert first_cmd[first_cmd.index('--nodelist') + 1] == 'big' + assert first_cmd[first_cmd.index('--nworkers') + 1] == '96' + assert first_cmd[first_cmd.index('--cpus-per-task') + 1] == '96' # == 96, accepted + assert second_cmd[second_cmd.index('--nodelist') + 1] == 'small' + assert second_cmd[second_cmd.index('--nworkers') + 1] == '40' + assert second_cmd[second_cmd.index('--cpus-per-task') + 1] == '40' # == 40, accepted + assert expected == 96 + 40 + + def test_an_oversubscribed_parallel_count_on_a_mixed_allocation_stays_bound(self, monkeypatch, tmp_path): + """A parallel_count above the total the allocation holds still starts every worker the + user asked for, split by machine size, while each step's CPU request stays capped at + what its machines hold so SLURM accepts it. 352 workers over 40,40,96 (176 CPUs) is + twice the CPUs, so ceil(40*352/176)=80 workers on each 40-CPU machine and + ceil(96*352/176)=192 on the 96-CPU one, but the requests are 40 and 96 -- the machines' + own sizes -- not 80 and 192.""" + 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_JOB_CPUS_PER_NODE', '40(x2),96') + + _, _, expected, _ = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2', 'n3'], parallel_count=352) + + (_, _), (first_cmd, _), (second_cmd, _) = spy.calls + assert first_cmd[first_cmd.index('--nworkers') + 1] == '80' + assert first_cmd[first_cmd.index('--cpus-per-task') + 1] == '40' # capped at the machine + assert second_cmd[second_cmd.index('--nworkers') + 1] == '192' + assert second_cmd[second_cmd.index('--cpus-per-task') + 1] == '96' # capped at the machine + assert expected == 80 * 2 + 192 + + def test_an_explicit_parallel_count_on_same_size_machines_is_still_one_even_step(self, monkeypatch, tmp_path): + """When every machine is the same size, parallel_count keeps its single-step even split + (#643 changes only the mixed case). 12 workers over three same-size machines is + ceil(12/3)=4 per machine in one step, exactly as before, so the common case a user runs + is untouched by the per-size split.""" + 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_JOB_CPUS_PER_NODE', '18(x3)') scheduler_proc, worker_procs, expected, worker_logs = cluster.Cluster.setup_srun_cluster( str(sched_file), str(tmp_path), ['n1', 'n2', 'n3'], parallel_count=12)