From ca6525d4c45582b664aa752c173f7e4828158543 Mon Sep 17 00:00:00 2001 From: Bill Hlavacek Date: Fri, 21 Aug 2026 18:12:10 -0600 Subject: [PATCH] Size each machine by what SLURM granted it on mixed clusters (#617) When PyBNF started workers across several machines it sent the same worker count to every one. On a cluster whose machines differ in size, that single count is too many workers for a small machine, where they compete for its processors, and too few for a large one, which sits partly idle. The srun launcher (cluster_type = slurm-srun) now works out each machine's count from what SLURM granted it and starts that many there. A single srun job step binds every machine in it to the same number of CPUs, so a mixed allocation is started as one job step per distinct machine size, each on its own machines, which SLURM runs at the same time. 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, and each job step writes its own worker log so their output does not interleave. An allocation whose machines are all the same size still runs as one job step, exactly as before. If SLURM does not report a per-machine list PyBNF can line up with its machines, it falls back to sizing every machine the same and warns that it did. Readiness now waits for every worker the steps should produce, not just the first, so a job step whose placement failed is caught rather than masked by another that succeeded. This applies only to the srun launcher. The SSH launcher (-t slurm) is unchanged, because dask ssh takes only one worker count for all hosts. Setting parallel_count still splits that total evenly across the machines on either launcher. Adds ADR-0124, a CHANGELOG entry, docs, and tests. --- CHANGELOG.md | 18 + ...ch-runs-a-worker-per-cpu-it-was-granted.md | 136 +++++++ docs/cluster.rst | 13 +- pybnf/cluster.py | 287 ++++++++++++--- tests/test_cluster.py | 338 +++++++++++++++--- 5 files changed, 694 insertions(+), 98 deletions(-) create mode 100644 docs/adr/0124-machines-of-different-sizes-get-one-srun-step-per-size-so-each-runs-a-worker-per-cpu-it-was-granted.md diff --git a/CHANGELOG.md b/CHANGELOG.md index b2042dba..435f848c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,24 @@ All notable changes to PyBNF are documented below. This project adheres to ## [Unreleased] ### Added +- **On a cluster whose machines are not all the same size, each machine now runs a worker per + CPU it was granted rather than one count for all of them (#617, ADR-0124).** When PyBNF started + workers across several machines it sent the same worker count to every one. On a cluster whose + machines differ in size, that single count is too many workers for a small machine, where they + compete for its processors, and too few for a large one, which sits partly idle. The `srun` + launcher (`-t slurm-srun`) now works out each machine's count from what SLURM granted it and + starts that many there. A single worker count cannot express this, because one `srun` job step + binds every machine in it to the same number of CPUs, so PyBNF starts one job step per distinct + machine size, each on its own machines, which SLURM runs at the same time. 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 arrangement is written to the log, and each job step writes its own worker + log (`dask_workers.log`, `dask_workers_2.log` and so on) so their output does not interleave. An + allocation whose machines are all the same size, which is the common case, still runs as one job + step exactly as before. If SLURM does not report a per-machine list PyBNF can line up with its + 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. - **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** 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 new file mode 100644 index 00000000..a9723655 --- /dev/null +++ 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 @@ -0,0 +1,136 @@ +# Machines of different sizes get one srun step per size, so each machine runs one worker per CPU it was granted rather than a count borrowed from another machine (issue #617) + +**Status: Accepted and implemented (2026-08-21).** When PyBNF started workers across several +machines it sent the same worker count to every one. On a cluster whose machines differ in size, +that single count is wrong for some of them. This ADR changes the `srun` launcher (#614, +`cluster_type = slurm-srun`) to size each machine by what SLURM granted it, starting one `srun` +job step per distinct machine size. The SSH launcher is unchanged, and so is every allocation +whose machines are all the same size. + +## The problem + +`Cluster.cpus_per_node` returns one number, and both launchers used to start that many workers on +every machine. The number describes the machine PyBNF is running on. On a cluster whose machines +are all the same size that is right for all of them, which is the case the launcher was first built +and tested against. + +The reporter's cluster is not that case. It has a single queue whose machines differ in processor +count by more than a factor of two, and separate requests land on unequal machines. One count for +all of them is then too many workers for a small machine, where the extra workers compete for its +processors, and too few for a large machine, which runs fewer workers than it has processors and +sits partly idle. Neither failure stops the run or prints anything wrong. The fit is simply slower +than the allocation could have made it, in a way nothing in the log accounts for. + +This could not be fixed until a launcher existed that starts each machine's workers separately. +`dask ssh` takes one worker count for all hosts and has no way to say a different number per host, +which is why this fix cannot be built on the SSH launcher. The `srun` launcher, added in #614, runs +`srun` itself and can run it more than once. ADR-0122 recorded this as the reason that launcher +would be where #617 was fixed: "`srun` ... can be invoked per node group, so this launcher is what +that fix will be built on." + +## The decision + +### One `srun` job step per distinct machine size + +A single `srun` step cannot express different worker counts on different machines. `--cpus-per-task` +is one value for the whole step, and it is not decoration: under `task/cgroup` binding a task that +asks for fewer CPUs than the workers it forks is confined to the CPUs it asked for, which quietly +serializes the node (ADR-0122 documents this as the reason the count is requested at all). So a step +that asked for the small machine's CPU count would throttle the large machines in the same step, and +a step that asked for the large machine's count would be refused on the small ones, which do not have +that many. + +The machines are therefore grouped by how many CPUs each was granted, and each group gets its own +`srun` step with its own `--cpus-per-task` and `--nworkers`, naming its machines with `--nodelist`. +A run on two 40-processor machines and one 96-processor machine becomes two steps: 40 workers on each +of the first two, 96 on the third. The steps run on disjoint machines, which SLURM allows to run at +the same time, so this does not serialize the bring-up. + +Grouping by size, rather than one step per machine, is what keeps the common case unchanged. An +allocation whose machines are all one size is a single group, so it runs as exactly one `srun` step, +the same command the launcher built before this change. Only a genuinely mixed allocation starts +more than one step, and then only as many as there are distinct sizes. + +### Where the per-machine counts come from + +`Cluster.per_node_cpus` reads `$SLURM_JOB_CPUS_PER_NODE`, which SLURM publishes as a compressed +per-node list (for example `40(x2),96` for two 40-processor machines and one 96-processor machine) +in the same order as the node list `scontrol show hostname` returns. Expanding it gives one count +per machine, lined up with the names PyBNF already read. + +This is the granted allocation, not the size of the machine, for the same reason the single-count +path reads `$SLURM_CPUS_ON_NODE`: the count is not only how many workers to start, it is how many +CPUs the step asks SLURM for, and a number taken from the whole machine would be refused. The +existing per-worker cap (`--cpus-per-task` no larger than what the job holds) is unchanged. + +### The fallback is the old behavior, said out loud + +If `$SLURM_JOB_CPUS_PER_NODE` is unset, does not parse, or does not have exactly one entry per +machine in the allocation, PyBNF cannot line a count up with each machine. Rather than guess, it +falls back to the single `cpus_per_node` count for every machine, which is what the launcher did +before this change, and logs a warning naming which of those three reasons applied. A user on a +mixed cluster is expecting each machine to be sized on its own, so the one case where that did not +happen is worth a line in the log rather than silent. + +### Readiness now waits for every worker, not just the first + +The `srun` launcher used to treat one registered worker as the readiness signal. With more than one +step that is not enough: a second step whose placement failed would go unnoticed as long as the +first step's workers registered. The wait now counts all of the workers the steps should produce +and returns when that many have registered, and it watches every step's process while it waits, so +a step that dies is reported with that step's own log rather than masked by another that succeeded. + +This is a deliberate strengthening of the readiness bar on this path, not only for the multi-step +case. It also closes on the `srun` launcher the same gap #200 describes for the SSH launcher, where +a bring-up that produced fewer workers than asked for was accepted as long as one arrived. The SSH +launcher's own readiness check is unchanged. + +### `parallel_count` is left as an even split + +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 +explicit total by machine size is a different decision from sizing an automatic run, with its own +question of what the total means on machines that cannot hold an equal share. So the override keeps +its single-step, even-split behavior, and only the default (auto-sized) path became per-machine. + +## Consequences + +* **A mixed-size allocation is used as fully as the launcher can use it.** Each machine runs a + worker per CPU it was granted, so small machines are not oversubscribed and large machines are + not left idle. The per-machine arrangement is in the log, so an unexpected count can be traced. +* **The common case is byte-for-byte unchanged.** An allocation of same-size machines, and any run + with `parallel_count` set, builds exactly the single `srun` command it did before, which a + regression test pins against the pre-change argv. +* **A mixed allocation writes more than one worker log.** Each step writes its own + (`dask_workers.log`, `dask_workers_2.log`, and so on), because concurrent steps writing one file + would interleave and truncate each other. The readiness error names the logs it read. +* **What was verified where.** The commands PyBNF builds for the grouped case, the grouping itself, + the fallback, and the stronger readiness count were exercised on one machine with stand-ins for + `srun` and `scontrol`, in the same way #614 was: the constructed command string is the oracle. + That SLURM places concurrent per-group steps as intended on a real heterogeneous allocation is + what the reporter's cluster verifies, exactly as it did for #614. +* **A known limitation on the `parallel_count` path is left in place.** Because that override splits + 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. + +## Alternatives considered + +* **One `srun` step per machine.** Rejected: it is more steps than the problem needs, and it makes + even a same-size allocation, the common case, stop building the single command that is already + tested and known good. Grouping by size gives each distinct size its own binding while leaving the + homogeneous case as one step. +* **Make `parallel_count` per-machine too.** Rejected as out of scope. Sizing an automatic run by + what each machine holds is a clear rule; splitting a user's explicit total by machine size raises + a separate question of what the total means when machines cannot take equal shares, and the SSH + launcher's even split is the behavior users of `parallel_count` have today. +* **Keep the single-worker readiness check.** Rejected: with more than one step, a step whose + placement failed would be masked by another that succeeded, turning a half-empty cluster into a + fit that runs slowly for no visible reason. Waiting for the full count is what makes the multi-step + bring-up safe. +* **Fail rather than fall back when `$SLURM_JOB_CPUS_PER_NODE` cannot be lined up with the machines.** + Rejected: the launcher has a correct, if less precise, thing to do (size every machine the same, as + it did before), so refusing to run would be worse than doing that and saying so. The warning makes + the degraded case visible without stopping the fit. diff --git a/docs/cluster.rst b/docs/cluster.rst index 3bebf260..f46d2f42 100644 --- a/docs/cluster.rst +++ b/docs/cluster.rst @@ -109,16 +109,21 @@ An example batch script -- ``examples/tcr/tcr_batch.sh`` with a single word chan pybnf -c tcr-ss.conf -t slurm-srun -o -Two log files are written to the output directory: ``dask_scheduler.log`` and ``dask_workers.log``. The second is where ``srun`` reports anything that went wrong with placing the workers, and PyBNF quotes from it in the error message if no worker ever registers. +Two log files are written to the output directory: ``dask_scheduler.log`` and ``dask_workers.log``. The second is where ``srun`` reports anything that went wrong with placing the workers, and PyBNF quotes from it in the error message if no worker ever registers. When the allocation holds machines of more than one size, each size runs as its own ``srun`` job step and writes its own worker log (``dask_workers.log``, ``dask_workers_2.log`` and so on), so their output does not interleave. .. _workercount: How many workers run on each node ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -By default -- with either launcher -- each node runs one single-threaded worker process per CPU **the job was granted** on that node. That is the number your allocation asked for, not the number of processors the machine happens to have: a job given 4 CPUs of a 128-processor node runs 4 workers on it, not 128. Each worker is a separate process, so a pool sized to the machine rather than to the job would multiply memory use and leave the workers competing for the same few CPUs. +By default, each node runs one single-threaded worker process per CPU **the job was granted** on that node. That is the number your allocation asked for, not the number of processors the machine happens to have: a job given 4 CPUs of a 128-processor node runs 4 workers on it, not 128. Each worker is a separate process, so a pool sized to the machine rather than to the job would multiply memory use and leave the workers competing for the same few CPUs. -PyBNF takes that number from the first of these that is available: +The two launchers differ in what happens when the machines in one allocation are not all the same size: + +* 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: * ``$SLURM_CPUS_ON_NODE``, which is what SLURM granted the job on a node; * 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 @@ -126,7 +131,7 @@ PyBNF takes that number from the first of these that is available: 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. -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. +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. .. _sizing: diff --git a/pybnf/cluster.py b/pybnf/cluster.py index 962342e2..5fcfaf33 100644 --- a/pybnf/cluster.py +++ b/pybnf/cluster.py @@ -178,6 +178,40 @@ def uses_srun(cluster_type): return SRUN_CLUSTER_TYPE_RE.fullmatch(cluster_type.strip()) is not None +# One compressed CPU-count group as SLURM writes it in $SLURM_JOB_CPUS_PER_NODE: a number, +# optionally followed by "(xN)" meaning N nodes in a row were granted that many CPUs. So +# "40(x2),96" is three nodes granted 40, 40 and 96. The whole variable, once expanded, lines +# up one-to-one with the node list `scontrol show hostname` returns, which is how the srun +# launcher learns how many workers each machine should run (#617). +CPUS_PER_NODE_GROUP_RE = re.compile(r'\s*(\d+)(?:\(x(\d+)\))?\s*') + + +def expand_cpus_per_node(spec): + """ + Expand a ``$SLURM_JOB_CPUS_PER_NODE`` string into one CPU count per node (#617). + + SLURM records the per-node counts in a run-length form, e.g. ``40(x2),96`` for three + nodes granted 40, 40 and 96 CPUs. This turns that into ``[40, 40, 96]``, in the same + order as the node list, so each machine's granted count can be matched to its name. + + :param spec: The value of ``$SLURM_JOB_CPUS_PER_NODE``, or None + :type spec: str or None + :return: one CPU count per node, or None if the text is empty or does not parse + :rtype: list or None + """ + if not spec or not spec.strip(): + return None + counts = [] + for group in spec.strip().split(','): + match = CPUS_PER_NODE_GROUP_RE.fullmatch(group) + if not match: + return None + count = int(match.group(1)) + repeats = int(match.group(2)) if match.group(2) else 1 + counts.extend([count] * repeats) + return counts + + class Cluster: """ Class handling the setup and teardown of the dask Client used to submit simulation jobs @@ -208,7 +242,13 @@ def __init__(self, config, log_prefix, debug, log_level_name): # either attaches to an existing scheduler or lets dask ssh start one (#614). self._scheduler_proc = None self._own_scheduler_file = None - self._srun_worker_log = None + # srun launcher readiness and teardown (#614, #617). The workers are one or more srun + # processes -- more than one only when the allocation holds machines of different + # sizes, which take one job step each -- so this is a list. The expected count is the + # total number of workers over all of them, and the logs are what each step wrote. + self._srun_worker_procs = [] + self._srun_expected_workers = None + self._srun_worker_logs = [] # SSH launcher readiness (#398). When PyBNF starts the cluster with dask ssh, these # hold what the readiness wait needs: the file dask ssh's output was captured to (so a # failed login can be quoted), how many workers it should bring up, and the directory @@ -225,9 +265,10 @@ def __init__(self, config, log_prefix, debug, log_level_name): # Find the name of the scheduler node, and a list of all available nodes (node_string), depending on what # cluster options are set if uses_srun(config.config['cluster_type']): - # srun launcher (#614): no node names are needed to *reach* the nodes -- SLURM - # places the workers itself -- but the node count is, because it is what the - # per-node worker arithmetic divides by. + # srun launcher (#614): SLURM places the workers itself, so the node names are not + # needed to *reach* the nodes -- but they are needed to size each machine, because + # the default worker count on a machine is one per CPU that machine was granted, + # and machines in one allocation can differ in size (#617). scheduler_node = None node_string = None if config.config['scheduler_node'] or config.config['worker_nodes']: @@ -238,10 +279,10 @@ def __init__(self, config, log_prefix, debug, log_level_name): scheduler_file = self.srun_scheduler_file(config) self._own_scheduler_file = scheduler_file out_dir = config.config['output_dir'] - self._srun_worker_log = os.path.join(out_dir, SRUN_WORKER_LOG) dummy, srun_nodes = self.read_node_names(config) - self._scheduler_proc, self._dask_proc = self.setup_srun_cluster( - scheduler_file, out_dir, len(srun_nodes.split()), config.config['parallel_count']) + (self._scheduler_proc, self._srun_worker_procs, self._srun_expected_workers, + self._srun_worker_logs) = self.setup_srun_cluster( + scheduler_file, out_dir, srun_nodes.split(), config.config['parallel_count']) elif config.config['scheduler_file']: # Scheduler node will be read in from scheduler file stored on shared file system node_string = None @@ -275,7 +316,8 @@ def __init__(self, config, log_prefix, debug, log_level_name): try: self.client = Client(scheduler_file=scheduler_file) self.local = False - self.wait_for_srun_workers(self.client, self._dask_proc, self._srun_worker_log) + self.wait_for_srun_workers(self.client, self._srun_worker_procs, + self._srun_expected_workers, self._srun_worker_logs) except Exception: self.stop_own_processes() raise @@ -677,6 +719,41 @@ def cpus_per_node(): '(dask.system.CPU_COUNT)') return cpu_count(), "this machine's whole processor count (multiprocessing.cpu_count)" + @staticmethod + def per_node_cpus(node_names): + """ + How many CPUs the job was granted on each machine, one number per node (#617). + + The srun launcher's default is one worker per granted CPU, and machines in one + allocation can differ in size, so this returns a count for each machine rather than + the single number :meth:`cpus_per_node` gives. It reads ``$SLURM_JOB_CPUS_PER_NODE``, + which lists the per-node counts in the same order as ``node_names``. If that variable + is missing, does not parse, or does not have one entry per node, per-machine sizing is + not available, so this falls back to the single :meth:`cpus_per_node` count for every + machine -- the behaviour before this change, where every machine was sized the same -- + and says so, since a user on a mixed cluster is expecting each machine to be sized on + its own. + + :param node_names: The machines in the allocation, in the order SLURM lists them + :type node_names: list + :return: a CPU count for each machine, and a phrase naming where the counts came from + :rtype: tuple + """ + spec = os.environ.get('SLURM_JOB_CPUS_PER_NODE', '') + counts = expand_cpus_per_node(spec) + if counts is not None and len(counts) == len(node_names): + return counts, 'what SLURM granted each machine ($SLURM_JOB_CPUS_PER_NODE)' + granted, source = Cluster.cpus_per_node() + if not spec.strip(): + reason = '$SLURM_JOB_CPUS_PER_NODE is not set' + elif counts is None: + reason = 'could not read $SLURM_JOB_CPUS_PER_NODE (%r)' % spec + else: + reason = ('$SLURM_JOB_CPUS_PER_NODE lists %i machine(s) but the allocation has %i' + % (len(counts), len(node_names))) + logger.warning('Sizing every machine the same because %s; using %s.' % (reason, source)) + return [granted] * len(node_names), source + @staticmethod def dask_scheduler_command(scheduler_file): """ @@ -745,24 +822,61 @@ def srun_worker_command(scheduler_file, node_count, parallel_count=None): '--nthreads', '1'] @staticmethod - def setup_srun_cluster(scheduler_file, out_dir, node_count, parallel_count=None): + def srun_worker_command_for_group(scheduler_file, nodes, cpus): + """ + 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`. + + :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 + :type cpus: int + :return: the srun argument list for this group + :rtype: list + """ + 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), + '--label', + *DASK_CLI, 'worker', + '--scheduler-file', scheduler_file, + '--nworkers', str(cpus), + '--nthreads', '1'] + + @staticmethod + def setup_srun_cluster(scheduler_file, out_dir, node_names, parallel_count=None): """ Start a dask scheduler here and a set of dask workers with srun, with no SSH login. The scheduler runs as an ordinary subprocess of this process, on this node, and is - told to write ``scheduler_file``; the workers are one srun task per node, each - reading that file. Nothing authenticates anywhere: SLURM already granted the - allocation, which is the whole point of the launcher (#614). + told to write ``scheduler_file``; the workers are started with srun, each reading that + file. Nothing authenticates anywhere: SLURM already granted the allocation, which is + the whole point of the launcher (#614). With no ``parallel_count`` set, each machine + runs one worker per CPU it was granted, and an allocation of different-sized machines + is brought up as one srun step per distinct size (#617). :param scheduler_file: Path the scheduler should write its connection information to :type scheduler_file: str :param out_dir: Directory for the scheduler and worker logs :type out_dir: str - :param node_count: Number of nodes in the allocation - :type node_count: int + :param node_names: The machines in the allocation, in the order SLURM lists them + :type node_names: list :param parallel_count: Total number of worker processes over all nodes, or None :type parallel_count: int or None - :return: the scheduler process and the srun process + :return: the scheduler process, the list of srun worker processes, the total number of + workers to wait for, and the log file each srun step is writing :rtype: tuple """ # Both of these would otherwise fail inside dask, or as a bare OSError from opening @@ -794,11 +908,70 @@ def setup_srun_cluster(scheduler_file, out_dir, node_count, parallel_count=None) raise logger.info('The dask scheduler is listening at %s' % address) - worker_log = os.path.join(out_dir, SRUN_WORKER_LOG) - srun_cmd = Cluster.srun_worker_command(scheduler_file, node_count, parallel_count) - logger.info('Starting dask workers with srun, logging to %s' % worker_log) - srun_proc = Cluster.popen_logged(srun_cmd, worker_log) - return scheduler_proc, srun_proc + commands, worker_logs, expected_total = Cluster.srun_worker_layout( + scheduler_file, out_dir, node_names, parallel_count) + worker_procs = [] + for srun_cmd, worker_log in zip(commands, worker_logs): + logger.info('Starting dask workers with srun, logging to %s' % worker_log) + worker_procs.append(Cluster.popen_logged(srun_cmd, worker_log)) + return scheduler_proc, worker_procs, expected_total, worker_logs + + @staticmethod + 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. + + :param scheduler_file: Path of the scheduler file the workers should read + :type scheduler_file: str + :param out_dir: Directory the worker logs are written in + :type out_dir: str + :param node_names: The machines in the allocation, in the order SLURM lists them + :type node_names: list + :param parallel_count: Total number of worker processes over all nodes, or None + :type parallel_count: int or None + :return: the srun command list(s), the matching log path(s), and the total worker count + :rtype: tuple + """ + 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. Making it per-machine is deliberately out of scope (#617). + 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) + 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. + 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) + 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)) + # 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))) + return commands, logs, expected_total @staticmethod def popen_logged(cmd, log_path): @@ -892,7 +1065,7 @@ def wait_for_scheduler_file(scheduler_file, scheduler_proc, scheduler_log, 'See %s for details.' % (scheduler_file, timeout, scheduler_log)) @staticmethod - def _poll_for_workers(client, worker_proc, expected, timeout, poll): + def _poll_for_workers(client, worker_procs, expected, timeout, poll): """ Poll the scheduler until enough workers have registered, watching the process (#398). @@ -911,7 +1084,10 @@ def _poll_for_workers(client, worker_proc, expected, timeout, poll): for staying open was that the SSH login could not be tested on the cluster available). :param client: The connected dask Client - :param worker_proc: The process that is bringing the workers up (srun, or dask ssh) + :param worker_procs: The process bringing the workers up (dask ssh), or the list of + them (the srun launcher runs more than one when machines differ in size, #617). + Any one of them exiting is reported as ``'exited'``, since it means a group of + workers failed to start. :param expected: Number of registered workers that counts as ready :type expected: int :param timeout: Seconds to wait before giving up @@ -922,11 +1098,13 @@ def _poll_for_workers(client, worker_proc, expected, timeout, poll): of workers that had registered, and the process return code (only set on ``'exited'``) :rtype: tuple """ + procs = worker_procs if isinstance(worker_procs, (list, tuple)) else [worker_procs] n_workers = 0 for _ in range(max(1, int(timeout / poll))): - returncode = worker_proc.poll() - if returncode is not None: - return 'exited', n_workers, returncode + for worker_proc in procs: + returncode = worker_proc.poll() + if returncode is not None: + return 'exited', n_workers, returncode try: n_workers = len(client.scheduler_info()['workers']) except Exception: @@ -937,44 +1115,59 @@ def _poll_for_workers(client, worker_proc, expected, timeout, poll): return 'timeout', n_workers, None @staticmethod - def wait_for_srun_workers(client, srun_proc, worker_log, + def wait_for_srun_workers(client, worker_procs, expected, worker_logs, timeout=SRUN_WORKER_TIMEOUT, poll=READINESS_POLL_INTERVAL): """ - Wait until at least one srun-launched worker has registered with the scheduler. + Wait until all of the srun-launched workers have registered with the scheduler. The waiting itself is :meth:`_poll_for_workers`, the loop both launchers share; this - adds what is specific to srun -- that one worker is enough to call the cluster up, and - what to say, quoting srun's own log, when srun dies or no worker ever arrives. + adds what is specific to srun -- the count to wait for, and what to say, quoting srun's + own log(s), when an srun step dies or the workers never all arrive. + + The full ``expected`` count is required rather than just one worker (#200, #617). On a + default run over machines of different sizes there is one srun step per size, and a + step that never places its workers -- a queued job step, a request larger than that + part of the allocation -- would otherwise be masked by another step that did place its + own. Connecting to our own scheduler always succeeds, so nothing else would report it, + and the fit would quietly run on fewer machines than were reserved. :param client: The connected dask Client - :param srun_proc: The running srun process - :param worker_log: Path of the srun output log, quoted if the workers never arrive - :type worker_log: str + :param worker_procs: The running srun process(es), one per machine size + :type worker_procs: list + :param expected: Total number of workers that should register across all the steps + :type expected: int + :param worker_logs: The srun output log(s), quoted if the workers never all arrive + :type worker_logs: list :param timeout: Seconds to wait before giving up :type timeout: float :param poll: Seconds between checks :type poll: float :return: the number of workers that had registered :rtype: int - :raises PybnfError: if srun exits, or no worker registers in time + :raises PybnfError: if an srun step exits, or the workers do not all register in time """ outcome, n_workers, returncode = Cluster._poll_for_workers( - client, srun_proc, expected=1, timeout=timeout, poll=poll) + client, worker_procs, expected=expected, timeout=timeout, poll=poll) if outcome == 'ready': - logger.info('%i dask worker process(es) registered with the scheduler' % n_workers) + logger.info('%i of %i dask worker process(es) registered with the scheduler' + % (n_workers, expected)) return n_workers + details = '\n'.join(t for t in (Cluster.log_tail(log) for log in worker_logs) if t) + where = ', '.join(worker_logs) if outcome == 'exited': - details = Cluster.log_tail(worker_log) - logger.error('srun exited with code %s before any worker started. Log:\n%s' + logger.error('An srun step exited with code %s before all workers started. Log:\n%s' % (returncode, details)) - raise PybnfError('srun exited with code %s before any dask worker started. %s' + raise PybnfError('srun exited with code %s before all the dask workers started. %s' % (returncode, ('Details:\n%s' % details) if details - else 'See %s for details.' % worker_log)) - details = Cluster.log_tail(worker_log) - logger.error('No dask worker registered within %s s. Log:\n%s' % (timeout, details)) - raise PybnfError('No dask worker started by srun registered with the scheduler within ' - '%s s.%s' % (timeout, ('\nLog:\n%s' % details) if details else ''), - hint=['Read %s: srun reports there what it is waiting for.' % worker_log, + else 'See %s for details.' % where)) + logger.error('Only %i of %i dask worker(s) registered within %s s. Log:\n%s' + % (n_workers, expected, timeout, details)) + raise PybnfError('Only %i of the %i expected dask worker process(es) started by srun ' + 'registered with the scheduler within %s s.%s' + % (n_workers, expected, timeout, + ('\nLog:\n%s' % details) if details else ''), + hint=['Read the srun log(s) (%s): srun reports there what it is waiting ' + 'for.' % where, 'A message about job step creation means another step already holds ' 'the allocation; run PyBNF as the only job step.']) @@ -1072,6 +1265,12 @@ def stop_own_processes(self): if self._dask_proc: self.stop_process(self._dask_proc, 'worker launcher subprocess') self._dask_proc = None + # The srun launcher's workers, one process per machine size (#617). Stopped before the + # scheduler, for the same reason the SSH worker launcher is: terminating them signals + # workers that would otherwise be left talking to a scheduler that is already gone. + for srun_proc in getattr(self, '_srun_worker_procs', None) or []: + self.stop_process(srun_proc, 'srun worker launcher subprocess') + self._srun_worker_procs = [] if self._scheduler_proc: self.stop_process(self._scheduler_proc, 'dask scheduler subprocess') self._scheduler_proc = None diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 5d686013..cbcf9604 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -489,6 +489,20 @@ def test_srun_still_takes_every_option_the_launcher_passes(self, monkeypatch): assert not missing, ('`%s` no longer mentions %s, which PyBNF passes it' % (srun_argv[0], ', '.join(missing))) + def test_srun_still_takes_every_option_the_group_builder_passes(self, monkeypatch): + """The heterogeneous path (#617) hands srun a few more options than the single + step does: ``--nodelist``, ``--nodes``, ``--ntasks`` and ``--ntasks-per-node``, + which name exactly the machines one size group runs on. Held to the installed + srun the same way, so a rename in any of them is caught on a cluster rather than + in a fit.""" + srun_argv, _ = _split_at_dask_cli( + cluster.Cluster.srun_worker_command_for_group( + '/shared/dask_scheduler.json', ['n1', 'n2'], 40)) + help_text = _help_text((srun_argv[0],)) + missing = [opt for opt in _options_in(srun_argv) if not _help_mentions(help_text, opt)] + assert not missing, ('`%s` no longer mentions %s, which PyBNF passes it' + % (srun_argv[0], ', '.join(missing))) + def test_the_node_list_is_read_with_a_program_that_is_installed(self, monkeypatch): """``scontrol`` is the other SLURM program PyBNF runs, and the one whose absence would stop a cluster fit first: the node list is read before @@ -933,8 +947,9 @@ def __init__(self): self.setup_calls = [] self.setup_proc = None # the dask ssh process fake_setup handed back self.ssh_wait_calls = [] # (client, dask_proc, expected, output_file, out_dir) - self.srun_setup_calls = [] # (scheduler_file, out_dir, node_count, parallel_count) - self.srun_wait_calls = [] # (client, srun_proc, worker_log) + self.srun_setup_calls = [] # (scheduler_file, out_dir, node_names, parallel_count) + self.srun_worker_procs = None # the srun worker procs fake_setup_srun handed back + self.srun_wait_calls = [] # (client, worker_procs, expected, worker_logs) def Client(self, *args, **kwargs): self.client_calls.append((args, kwargs)) @@ -968,9 +983,13 @@ def fake_setup(node_string, out_dir, parallel_count): rec.setup_proc = _ProcStub() return rec.setup_proc, 2, None - def fake_setup_srun(scheduler_file, out_dir, node_count, parallel_count): - rec.srun_setup_calls.append((scheduler_file, out_dir, node_count, parallel_count)) - return _ProcStub(), _ProcStub() + def fake_setup_srun(scheduler_file, out_dir, node_names, parallel_count): + rec.srun_setup_calls.append((scheduler_file, out_dir, node_names, parallel_count)) + # scheduler proc, the list of srun worker procs, the total worker count to wait for, + # and the log each srun step writes. wait_for_srun_workers, which consumes them, is + # faked below, so the count and logs here are placeholders. + rec.srun_worker_procs = [_ProcStub()] + return _ProcStub(), rec.srun_worker_procs, 1, [os.path.join(out_dir, 'dask_workers.log')] def fake_ssh_wait(client, dask_proc, expected, output_file, out_dir, **kwargs): rec.ssh_wait_calls.append((client, dask_proc, expected, output_file, out_dir)) @@ -978,11 +997,11 @@ def fake_ssh_wait(client, dask_proc, expected, output_file, out_dir, **kwargs): raise ssh_raises return expected - def fake_wait(client, srun_proc, worker_log, **kwargs): - rec.srun_wait_calls.append((client, srun_proc, worker_log)) + def fake_wait(client, worker_procs, expected, worker_logs, **kwargs): + rec.srun_wait_calls.append((client, worker_procs, expected, worker_logs)) if srun_raises: raise srun_raises - return 1 + return expected monkeypatch.setattr(cluster.Cluster, 'read_node_names', staticmethod(fake_read)) monkeypatch.setattr(cluster.Cluster, 'setup_cluster', staticmethod(fake_setup)) @@ -1397,6 +1416,104 @@ def test_a_granted_count_never_reports_the_machine_count(self, monkeypatch): assert cluster.Cluster.cpus_per_node()[0] == 4 +class TestExpandCpusPerNode: + """Reading $SLURM_JOB_CPUS_PER_NODE into one count per machine (#617). SLURM + writes the counts run-length encoded, in the same order as the node list.""" + + @pytest.mark.parametrize('spec, expected', [ + ('40(x2),96', [40, 40, 96]), # the mixed-size case this issue is about + ('8', [8]), # a single machine + ('8(x3)', [8, 8, 8]), # one size, several machines + ('4(x2),8,16(x2)', [4, 4, 8, 16, 16]), + (' 40(x2) , 96 ', [40, 40, 96]), # whitespace around the groups is tolerated + ]) + def test_expands_the_run_length_form(self, spec, expected): + assert cluster.expand_cpus_per_node(spec) == expected + + @pytest.mark.parametrize('spec', [None, '', ' ', 'many', '4(x)', '(x2)', '4,', 'a,b']) + def test_unusable_text_is_none(self, spec): + """Anything that is not the run-length form -- unset, empty, or a value this + code does not recognize -- returns None, so the caller falls back to sizing + every machine the same rather than acting on a misread.""" + assert cluster.expand_cpus_per_node(spec) is None + + +class TestPerNodeCpus: + """How many CPUs each machine was granted, one number per node (#617).""" + + def test_reads_the_per_machine_counts_in_node_order(self, monkeypatch): + """The preferred source is $SLURM_JOB_CPUS_PER_NODE, which lines up with the + node list; each machine is sized by what it was granted.""" + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '40(x2),96') + counts, source = cluster.Cluster.per_node_cpus(['n1', 'n2', 'n3']) + assert counts == [40, 40, 96] + assert 'SLURM_JOB_CPUS_PER_NODE' in source + + def test_unset_variable_sizes_every_machine_the_same(self, monkeypatch): + """With the per-machine variable unset, there is nothing to size each machine + by, so every machine gets the single cpus_per_node count -- the behaviour + before this change. Here that count comes from $SLURM_CPUS_ON_NODE.""" + monkeypatch.delenv('SLURM_JOB_CPUS_PER_NODE', raising=False) + monkeypatch.setenv('SLURM_CPUS_ON_NODE', '8') + counts, source = cluster.Cluster.per_node_cpus(['n1', 'n2', 'n3']) + assert counts == [8, 8, 8] + assert 'SLURM_CPUS_ON_NODE' in source # the fallback source, named + + def test_a_length_mismatch_falls_back_rather_than_misaligning(self, monkeypatch): + """A per-machine list that does not have one entry per node cannot be trusted + to line up, so it is not used: every machine is sized the same instead.""" + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '40,96') # two, but three nodes + monkeypatch.setenv('SLURM_CPUS_ON_NODE', '8') + counts, _ = cluster.Cluster.per_node_cpus(['n1', 'n2', 'n3']) + assert counts == [8, 8, 8] + + def test_unparseable_variable_falls_back(self, monkeypatch): + """A value this code cannot read is treated as unavailable, not guessed at.""" + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', 'nonsense') + monkeypatch.setenv('SLURM_CPUS_ON_NODE', '8') + counts, _ = cluster.Cluster.per_node_cpus(['n1', 'n2']) + assert counts == [8, 8] + + 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.""" + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '40,96') + monkeypatch.setenv('SLURM_CPUS_ON_NODE', '8') + with caplog.at_level('WARNING'): + cluster.Cluster.per_node_cpus(['n1', 'n2', 'n3']) + assert any('SLURM_JOB_CPUS_PER_NODE' in r.message for r in caplog.records) + + +class TestSrunWorkerCommandForGroup: + """The srun command for one group of same-sized machines (#617). One worker per + granted CPU on each machine, the machines named so concurrent groups stay + disjoint, and the CPU request tracking the worker count for the same cgroup + reason srun_worker_command has.""" + + def test_names_the_machines_and_sizes_them_by_their_grant(self): + cmd = cluster.Cluster.srun_worker_command_for_group('/s.json', ['n1', 'n2'], 40) + assert cmd == ['srun', '--nodelist', 'n1,n2', + '--nodes', '2', '--ntasks', '2', '--ntasks-per-node', '1', + '--cpus-per-task', '40', '--label', + *DASK, 'worker', '--scheduler-file', '/s.json', + '--nworkers', '40', '--nthreads', '1'] + + def test_a_single_machine_group(self): + cmd = cluster.Cluster.srun_worker_command_for_group('/s.json', ['big'], 96) + assert cmd[cmd.index('--nodelist') + 1] == 'big' + assert cmd[cmd.index('--nodes') + 1] == '1' + assert cmd[cmd.index('--nworkers') + 1] == '96' + assert cmd[cmd.index('--cpus-per-task') + 1] == '96' + + def test_workers_are_single_threaded(self): + cmd = cluster.Cluster.srun_worker_command_for_group('/s.json', ['n1'], 4) + assert cmd[cmd.index('--nthreads') + 1] == '1' + + def test_scheduler_file_is_one_literal_argument(self): + cmd = cluster.Cluster.srun_worker_command_for_group('/tmp/a b$(whoami).json', ['n1'], 2) + assert cmd[cmd.index('--scheduler-file') + 1] == '/tmp/a b$(whoami).json' + + class TestSrunWorkerCommand: def _patch(self, monkeypatch, granted=8): @@ -1559,27 +1676,31 @@ def test_a_file_that_never_appears_times_out_naming_the_log(self, monkeypatch, t class TestWaitForSrunWorkers: - def test_returns_once_a_worker_registers(self, monkeypatch): - """The readiness signal for the workers is a worker registering with the - scheduler -- not srun having been launched, which says nothing.""" + def test_returns_once_all_the_workers_register(self, monkeypatch): + """The readiness signal for the workers is all of them registering with the + scheduler -- not srun having been launched, which says nothing. Two expected, + two connected, so the wait returns two.""" monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) n = cluster.Cluster.wait_for_srun_workers( - _ClientStub(workers=('tcp://n1:1', 'tcp://n2:1')), _FakeDaskProc(), '/log') + _ClientStub(workers=('tcp://n1:1', 'tcp://n2:1')), [_FakeDaskProc()], + expected=2, worker_logs=['/log']) assert n == 2 - def test_waits_while_the_cluster_is_still_empty(self, monkeypatch): - """A scheduler with no workers yet is not an error: the poll continues - until the workers arrive.""" - client = _ClientStub(workers=()) + def test_waits_while_the_cluster_is_still_filling_up(self, monkeypatch): + """A scheduler short of the expected count is not ready and not an error: the + poll continues until the last worker arrives (#200, #617). Here one of two is + up, then the second arrives, and only then does the wait return.""" + client = _ClientStub(workers=('tcp://n1:1',)) polls = [] def fake_sleep(_seconds): polls.append(1) if len(polls) == 3: - client._workers = {'tcp://n1:1': {}} + client._workers = dict.fromkeys(('tcp://n1:1', 'tcp://n2:1'), {}) monkeypatch.setattr(cluster.time, 'sleep', fake_sleep) - assert cluster.Cluster.wait_for_srun_workers(client, _FakeDaskProc(), '/log') == 1 + assert cluster.Cluster.wait_for_srun_workers( + client, [_FakeDaskProc()], expected=2, worker_logs=['/log']) == 2 assert len(polls) == 3 def test_a_transient_scheduler_error_is_not_fatal(self, monkeypatch): @@ -1597,32 +1718,36 @@ def scheduler_info(self): return {'workers': {'tcp://n1:1': {}}} monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) - assert cluster.Cluster.wait_for_srun_workers(_FlakyClient(), _FakeDaskProc(), '/log') == 1 - - def test_srun_exiting_early_is_reported_with_its_log(self, monkeypatch, tmp_path): - """srun exiting before any worker registered means the placement failed -- - a bad flag, a request larger than the allocation. Reported at once, quoting - srun's own message rather than waiting out the timeout.""" + assert cluster.Cluster.wait_for_srun_workers( + _FlakyClient(), [_FakeDaskProc()], expected=1, worker_logs=['/log']) == 1 + + def test_any_srun_step_exiting_early_is_reported_with_its_log(self, monkeypatch, tmp_path): + """An srun step exiting before its workers registered means the placement + failed -- a bad flag, a request larger than that part of the allocation. With + machines of different sizes there is one step per size, so any of them exiting + is caught, reported at once and quoting srun's own message.""" log = tmp_path / 'workers.log' log.write_text('srun: error: Unable to allocate resources') monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) with pytest.raises(printing.PybnfError) as exc: cluster.Cluster.wait_for_srun_workers( - _ClientStub(workers=()), _FakeDaskProc(returncode=1), str(log)) + _ClientStub(workers=()), [_FakeDaskProc(returncode=1)], + expected=1, worker_logs=[str(log)]) assert 'code 1' in str(exc.value) assert 'Unable to allocate resources' in str(exc.value) - def test_no_worker_in_time_names_the_log_and_the_step_hazard(self, monkeypatch, tmp_path): - """srun still running with no worker registered is the shape of a queued - job step: connecting to our own scheduler succeeded, so nothing else would - report it, and the fit would submit jobs no one takes. The message points - at srun's own log and at the likely cause.""" + def test_too_few_in_time_names_the_logs_and_the_step_hazard(self, monkeypatch, tmp_path): + """srun still running with the workers short of the count is the shape of a + queued job step: connecting to our own scheduler succeeded, so nothing else + would report it, and the fit would run on fewer machines than reserved. The + message points at srun's own log(s) and at the likely cause.""" log = tmp_path / 'workers.log' log.write_text('srun: Job step creation temporarily disabled, retrying') monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) with pytest.raises(printing.PybnfError) as exc: cluster.Cluster.wait_for_srun_workers( - _ClientStub(workers=()), _FakeDaskProc(), str(log), timeout=1.) + _ClientStub(workers=()), [_FakeDaskProc()], expected=2, + worker_logs=[str(log)], timeout=1.) assert 'srun: Job step creation' in str(exc.value) # srun's own words, in the log assert 'workers.log' in exc.value.message # ... and where to read more assert 'job step' in exc.value.message.lower() @@ -1660,6 +1785,17 @@ def test_timeout_when_too_few_arrive_in_time(self, monkeypatch): expected=3, timeout=1., poll=0.25) assert (outcome, n, rc) == ('timeout', 1, None) + def test_exited_when_any_of_several_processes_is_gone(self, monkeypatch): + """The srun launcher passes a list of processes, one per machine size. If any one of + them exits the placement failed, so the loop reports 'exited' with that process's + code even though the others are still running.""" + monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + alive, dead = _FakeDaskProc(), _FakeDaskProc(returncode=2) + outcome, n, rc = cluster.Cluster._poll_for_workers( + _ClientStub(workers=()), [alive, dead], + expected=4, timeout=5., poll=0.25) + assert (outcome, rc) == ('exited', 2) + class TestWaitForSSHWorkers: """The SSH launcher's startup readiness check (#398), driven directly. The dead-process @@ -1721,9 +1857,10 @@ def test_starts_the_scheduler_here_then_the_workers_with_srun(self, monkeypatch, monkeypatch.setattr(cluster, 'Popen', spy) monkeypatch.setattr(cluster.time, 'sleep', spy.sleep) monkeypatch.setenv('SLURM_CPUS_ON_NODE', '4') + monkeypatch.delenv('SLURM_JOB_CPUS_PER_NODE', raising=False) # same size everywhere - scheduler_proc, srun_proc = cluster.Cluster.setup_srun_cluster( - str(sched_file), str(tmp_path), 2, parallel_count=None) + scheduler_proc, worker_procs, expected, worker_logs = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2'], parallel_count=None) (sched_cmd, sched_kwargs), (srun_cmd, srun_kwargs) = spy.calls assert sched_cmd == [*DASK, 'scheduler', '--scheduler-file', str(sched_file)] @@ -1736,7 +1873,107 @@ def test_starts_the_scheduler_here_then_the_workers_with_srun(self, monkeypatch, assert hasattr(kwargs['stdout'], 'write') assert (tmp_path / 'dask_scheduler.log').exists() assert (tmp_path / 'dask_workers.log').exists() - assert (scheduler_proc, srun_proc) == (spy.procs[0], spy.procs[1]) + # One srun step (same size everywhere), and the worker total is what the readiness + # wait needs: four workers on each of two machines. + assert (scheduler_proc, worker_procs) == (spy.procs[0], [spy.procs[1]]) + assert expected == 8 + assert worker_logs == [str(tmp_path / 'dask_workers.log')] + + def test_machines_of_different_sizes_each_get_their_own_srun_step(self, monkeypatch, tmp_path): + """#617: a mixed-size allocation cannot be placed by one srun step, because + --cpus-per-task is a single number for the whole step and under cgroup binding a + task that under-requests is confined to what it asked for. So each distinct size + gets its own step: two 40-processor machines in one step at 40 workers each, the + 96-processor machine in a second step at 96. The steps run on disjoint machines, + which SLURM allows concurrently, and each writes its own log so their output does + not interleave.""" + 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') + + scheduler_proc, worker_procs, expected, worker_logs = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2', 'n3'], parallel_count=None) + + (sched_cmd, _), (first_cmd, _), (second_cmd, _) = spy.calls + assert sched_cmd == [*DASK, 'scheduler', '--scheduler-file', str(sched_file)] + # First step: the two same-size machines, 40 workers each. + assert first_cmd[first_cmd.index('--nodelist') + 1] == 'n1,n2' + assert first_cmd[first_cmd.index('--cpus-per-task') + 1] == '40' + assert first_cmd[first_cmd.index('--nworkers') + 1] == '40' + # Second step: the larger machine on its own, 96 workers. + assert second_cmd[second_cmd.index('--nodelist') + 1] == 'n3' + assert second_cmd[second_cmd.index('--cpus-per-task') + 1] == '96' + assert second_cmd[second_cmd.index('--nworkers') + 1] == '96' + assert worker_procs == [spy.procs[1], spy.procs[2]] + assert expected == 40 * 2 + 96 # every worker across both steps + assert worker_logs == [str(tmp_path / 'dask_workers.log'), + str(tmp_path / 'dask_workers_2.log')] + assert (tmp_path / 'dask_workers.log').exists() + assert (tmp_path / 'dask_workers_2.log').exists() + + def test_the_per_machine_arrangement_is_logged(self, monkeypatch, tmp_path, caplog): + """#617 asks for the arrangement to be recorded, so a user can see which machines + got how many workers. One line per distinct size, naming the machines and their + count.""" + 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') + + with caplog.at_level('INFO'): + cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2', 'n3'], parallel_count=None) + + text = '\n'.join(r.message for r in caplog.records) + 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.""" + 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 + + scheduler_proc, worker_procs, expected, worker_logs = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2', 'n3'], parallel_count=12) + + (_, _), (srun_cmd, _) = spy.calls # scheduler + exactly one srun + assert worker_procs == [spy.procs[1]] + assert srun_cmd[srun_cmd.index('--nworkers') + 1] == '4' # 12 across three machines + assert expected == 12 + assert worker_logs == [str(tmp_path / 'dask_workers.log')] + + 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: + one step, the same count everywhere, with a warning that per-machine sizing was not + available.""" + 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', '4') + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', 'garbage') # will not parse + + with caplog.at_level('WARNING'): + scheduler_proc, worker_procs, expected, worker_logs = cluster.Cluster.setup_srun_cluster( + str(sched_file), str(tmp_path), ['n1', 'n2'], parallel_count=None) + + (_, _), (srun_cmd, _) = spy.calls # scheduler + exactly one srun + assert worker_procs == [spy.procs[1]] + assert srun_cmd[srun_cmd.index('--nworkers') + 1] == '4' + assert expected == 8 + assert worker_logs == [str(tmp_path / 'dask_workers.log')] + assert any('every machine the same' in r.message.lower() and 'SLURM_JOB_CPUS_PER_NODE' in r.message + for r in caplog.records) def test_the_workers_are_started_only_after_the_scheduler_is_ready(self, monkeypatch, tmp_path): """Ordering is load-bearing: a worker started before the scheduler file @@ -1751,7 +1988,7 @@ def fake_sleep(seconds): monkeypatch.setattr(cluster, 'Popen', spy) monkeypatch.setattr(cluster.time, 'sleep', fake_sleep) - cluster.Cluster.setup_srun_cluster(str(sched_file), str(tmp_path), 1) + cluster.Cluster.setup_srun_cluster(str(sched_file), str(tmp_path), ['n1']) assert launched_at == [1, 1, 1] # only the scheduler was running while waiting assert len(spy.calls) == 2 @@ -1769,7 +2006,7 @@ def test_a_stale_scheduler_file_is_removed_before_the_scheduler_starts(self, mon lambda cmd, **k: seen.append(sched_file.exists()) or spy(cmd, **k)) monkeypatch.setattr(cluster.time, 'sleep', spy.sleep) - cluster.Cluster.setup_srun_cluster(str(sched_file), str(tmp_path), 1) + cluster.Cluster.setup_srun_cluster(str(sched_file), str(tmp_path), ['n1']) assert seen[0] is False # gone before the scheduler started assert json.loads(sched_file.read_text())['address'] == 'tcp://live:8786' @@ -1784,7 +2021,7 @@ def test_a_scheduler_that_dies_takes_no_srun_with_it(self, monkeypatch, tmp_path with pytest.raises(printing.PybnfError, match='scheduler exited'): cluster.Cluster.setup_srun_cluster( - str(tmp_path / 'dask_scheduler.json'), str(tmp_path), 2) + str(tmp_path / 'dask_scheduler.json'), str(tmp_path), ['n1', 'n2']) assert len(spy.calls) == 1 # srun was never launched assert spy.procs[0].terminated is True @@ -1795,7 +2032,7 @@ def test_a_missing_scheduler_file_directory_is_refused_up_front(self, tmp_path): it can be a configuration error naming the path.""" with pytest.raises(printing.PybnfError, match='scheduler file .* does not exist'): cluster.Cluster.setup_srun_cluster( - str(tmp_path / 'no_such_dir' / 's.json'), str(tmp_path), 1) + str(tmp_path / 'no_such_dir' / 's.json'), str(tmp_path), ['n1']) def test_a_missing_log_directory_is_refused_up_front(self, tmp_path): """Likewise for the directory the logs go in: opening that file is the first @@ -1803,7 +2040,7 @@ def test_a_missing_log_directory_is_refused_up_front(self, tmp_path): "an unknown error ... please report this bug".""" with pytest.raises(printing.PybnfError, match='logs .* does not exist'): cluster.Cluster.setup_srun_cluster( - str(tmp_path / 's.json'), str(tmp_path / 'no_such_dir'), 1) + str(tmp_path / 's.json'), str(tmp_path / 'no_such_dir'), ['n1']) class TestInitSrunDispatch: @@ -1823,15 +2060,15 @@ def test_srun_type_never_reaches_dask_ssh(self, monkeypatch): assert len(rec.srun_setup_calls) == 1 assert c.local is False - def test_brings_up_srun_with_the_node_count_and_connects_by_file(self, monkeypatch): - """The srun bring-up gets the scheduler-file path, the output directory, - the *number* of nodes (all srun needs -- it places the workers itself) and - parallel_count; the client then connects through that file.""" + def test_brings_up_srun_with_the_node_names_and_connects_by_file(self, monkeypatch): + """The srun bring-up gets the scheduler-file path, the output directory, the + *names* of the machines (needed to size each one, #617) and parallel_count; + the client then connects through that file.""" rec = _patch_init(monkeypatch, read_returns=('n1', 'n1 n2 n3')) c = _build(self._cfg_srun(output_dir='out', parallel_count=12)) expected_file = os.path.abspath(os.path.join('out', 'dask_scheduler.json')) - assert rec.srun_setup_calls == [(expected_file, 'out', 3, 12)] + assert rec.srun_setup_calls == [(expected_file, 'out', ['n1', 'n2', 'n3'], 12)] assert rec.client_calls == [((), {'scheduler_file': expected_file})] assert c._own_scheduler_file == expected_file @@ -1871,10 +2108,10 @@ def test_waits_for_the_workers_before_returning(self, monkeypatch): c = _build(self._cfg_srun(output_dir='out')) assert len(rec.srun_wait_calls) == 1 - client, srun_proc, worker_log = rec.srun_wait_calls[0] + client, worker_procs, expected, worker_logs = rec.srun_wait_calls[0] assert client is rec.last_client - assert srun_proc is c._dask_proc - assert worker_log == os.path.join('out', 'dask_workers.log') + assert worker_procs is c._srun_worker_procs + assert worker_logs == [os.path.join('out', 'dask_workers.log')] def test_a_failed_worker_wait_stops_what_it_started(self, monkeypatch): """A constructor that raises never becomes a Cluster, so no one else can @@ -1886,9 +2123,10 @@ def test_a_failed_worker_wait_stops_what_it_started(self, monkeypatch): faked_setup = cluster.Cluster.setup_srun_cluster # the recorder _patch_init installed def spy_setup(*args): - procs = faked_setup(*args) - started.extend(procs) - return procs + scheduler_proc, worker_procs, expected, logs = faked_setup(*args) + started.append(scheduler_proc) + started.extend(worker_procs) + return scheduler_proc, worker_procs, expected, logs monkeypatch.setattr(cluster.Cluster, 'setup_srun_cluster', staticmethod(spy_setup)) with pytest.raises(printing.PybnfError, match='no workers'):