PyBNF cluster setup (pybnf.cluster)

Functions for managing dask cluster setup and teardown on distributed computing systems

Two launchers bring up a multi-machine cluster, chosen by cluster_type:

• the SSH launcher (cluster_type = slurm, or scheduler_node / worker_nodes set by hand), which runs dask ssh and therefore logs in to every node; and

• the srun launcher (cluster_type = slurm-srun, #614), which never logs in anywhere. It starts the scheduler here, has SLURM place one dask worker per node inside the allocation the scheduler already granted, and connects through the scheduler file the scheduler writes.

The srun launcher exists because the SSH one cannot work at all on a cluster whose nodes authenticate to each other by host-based or Kerberos (GSSAPI) SSH: dask ssh logs in with paramiko, which offers only public-key and password authentication – it has no host-based support and dask never enables its GSSAPI support – so on such a cluster the login fails no matter what the user configures, and no amount of ssh-keygen helps. See docs/adr/0122 for the full argument. When that login is what fails, the SSH launcher says so, quotes what dask ssh said, and names both ways of running that need no login (#618) – the failure ends the run, so the message is the whole of what the user gets.

Both launchers size their default worker pool from what the job was granted rather than from how big the machine is, and record which number they used and where it came from (#616); Cluster.cpus_per_node is the one place that decides it.

class pybnf.cluster.Cluster(config, log_prefix, debug, log_level_name)

Class handling the setup and teardown of the dask Client used to submit simulation jobs The client is accessible

static captured_text(handle)

Read back everything a launched process wrote to its capture file, as plain text.

Colour escapes are removed: dask colours its own failure lines, and those bytes would otherwise reach a log file and an error message as literal characters.

Parameters:

handle – The open binary file the process was given as its output

Returns:

the captured text, or ‘’ if nothing was captured or it cannot be read

Return type:

str

static cpus_per_node()

The number of CPUs the running job was granted on a node, and where that came from.

Both launchers size their default worker pool with this, because the number that decides how many processes to start has to describe what the job holds, not what the machine has (#616). multiprocessing.cpu_count() answers the second question: it reports every processor on the machine whatever the scheduler granted, so a job given 4 CPUs of a 128-processor node is told 128, and one worker process per processor oversubscribes it 32-fold – 32 times the memory, and workers competing for time rather than a fit that runs faster. The two numbers agree only when whole nodes were allocated, which is why the defect stayed hidden.

Four sources are consulted, best first:

• $SLURM_CPUS_ON_NODE – what the allocation granted on a node. Preferred because it describes the allocation rather than the process doing the asking, so it is still the right number for a worker the SSH launcher starts on some other machine. (When nodes differ in size it describes this node; per-node counts are issue #617.)

• $SLURM_JOB_CPUS_PER_NODE – the same question answered for the job as a whole, as a per-node list, reduced here to its smallest entry. SLURM sets the variable above only inside a step running on an allocated node, so it is empty when the launching process is not on one: salloc opens its shell on the login node on many clusters, and the launcher then runs there while the allocation is held elsewhere (#642). This variable is set correctly in that shell, and the two machine-level numbers below are not – they describe a login node that is not in the allocation at all, and is typically far larger than what the job holds. The smallest entry is taken because this single number both sizes a pool started on every machine and is what an srun step asks for on every machine in it: asking for less than a machine holds costs speed, while asking for more than the smallest machine holds is refused outright.

• dask.system.CPU_COUNT – what the operating system will let this process run on: the machine’s processors narrowed by CPU affinity and by any cgroup CPU quota. This is the number a local run already sizes itself by, and it is the right one whenever the job is confined on the machine PyBNF is running on but no scheduler published a count.

• multiprocessing.cpu_count() – the whole machine, correct only when nothing is limiting the job at all, and reached only if no number above is usable.

The srun launcher does not merely count workers with this: it also asks SLURM for that many CPUs per task, and a request larger than the allocation is refused outright.

Returns:

CPUs granted per node, and a phrase naming where that number came from

Return type:

tuple

static dask_scheduler_command(scheduler_file)

Build the dask scheduler invocation that starts the scheduler on this node.

A one-line command, named anyway so that it sits beside srun_worker_command() and can be read – and checked against the dask that is actually installed (#619) – without starting a cluster to see it.

Parameters:

scheduler_file (str) – Path the scheduler should write its connection information to

Returns:

the dask scheduler argument list

Return type:

list

static fold_traceback_frames(text, lines=40)

The same output with Python traceback frames folded away, for an error message.

What dask ssh writes when a login fails is mostly traceback: one per node per retry, three retries each, every one of them a dozen frames of dask’s and paramiko’s own source. Quoted whole, the sentences that say what happened – dask’s “SSH reported this exception: …”, and the exception line each traceback ends with – are buried in code the user did not write and cannot act on, and are the first thing a length limit throws away. Folding the frames left 32 lines of a measured 137, and lost none of the sentences.

Each traceback is recognized by its header and ends at the first line that is not indented, which is the exception itself; that line is kept, the frames between are dropped. The full text still goes to the log.

Parameters:

• text (str) – The captured output

• lines (int) – Number of trailing lines to keep after folding

Returns:

the output without traceback frames, at most lines lines

Return type:

str

static local_cluster_kwargs(parallel_count)

Build the LocalCluster keyword arguments for a local (non-cluster) run.

threads_per_worker is 1 unconditionally (#526). PyBNF’s simulation backends hold process-wide state that is not advertised as thread-safe – a C++ engine plus code generation with module-level caches – so two worker threads in one process can race (issue #525 caught exactly that: concurrent emissions through bngsim’s cached sympy->C printer intermittently reported ordinary quotients as non-differentiable, which killed a trf fit). Every other client PyBNF builds is already single-threaded per worker: both dask ssh branches pass --nthreads 1, and the manual-setup documentation recommends the same. Only the local default used to let dask pick, so a user who set nothing got the less safe configuration.

n_workers is left to dask when parallel_count is None: given one thread per worker, dask sizes the pool at one worker per available core (dask.system.CPU_COUNT, which honors CPU affinity and cgroup quotas), matching the dask ssh default of --nworkers <cores> --nthreads 1. Total concurrency is therefore unchanged from the old default – the same number of jobs run at once, each in its own process.

Parameters:

parallel_count (int or None) – Number of parallel jobs requested, or None for one per core

Returns:

kwargs for distributed.LocalCluster

Return type:

dict

static log_tail(log_path, lines=20)

The last few lines of a log file, for inclusion in an error message.

Parameters:

• log_path (str) – File to read

• lines (int) – Number of trailing lines to keep

Returns:

the trailing text, or ‘’ if the file cannot be read

Return type:

str

static 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 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 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.

Parameters:

node_names (list) – The machines in the allocation, in the order SLURM lists them

Returns:

a CPU count for each machine, and a phrase naming where the counts came from

Return type:

tuple

static popen_logged(cmd, log_path)

Launch cmd with no shell, sending its output to log_path.

Both srun-launcher processes outlive this call and neither has a terminal, so their output has to go somewhere a user can read it afterwards – and somewhere that cannot fill up and deadlock the writer, as an undrained pipe would. stderr is merged into stdout so one file reads in order.

Parameters:

• cmd (list) – The argument list to run

• log_path (str) – File to write the process output to

Returns:

subprocess.Popen

static read_node_names(config)

Reads the available node names, if running on a cluster. If not running on a cluster, returns None for both.

Parameters:

config (pybnf.config.Configuration) – PyBNF configuration

Returns:

scheduler node, string composed of all available nodes

static require_slurm_allocation()

Refuse the srun launcher when this process is not inside a SLURM allocation.

Outside an allocation srun does not place a task, it submits a job and then waits for the scheduler to grant it – so the failure this catches would otherwise look like PyBNF hanging with no output, possibly for hours. The allocation is what makes the launcher credential-free, so its absence is a configuration error.

Raises:

PybnfError – if no SLURM allocation is visible in the environment

static setup_cluster(node_string, out_dir, parallel_count=None)

Sets up a Dask cluster using the dask ssh command

Parameters:

• node_string – A string composed of a list of compute nodes

• out_dir – A directory for cluster logging output

• parallel_count – Total number of single-threaded worker processes over all nodes, divided evenly among them. If None, one worker per CPU the job was granted on a node (cpus_per_node)

Returns:

a tuple of the dask ssh process, the number of worker processes it should bring up (one per node times the per-node count), and the open file its output was captured to. The caller waits on that worker count and quotes the captured file if the bring-up fails (wait_for_ssh_workers(), #398)

Return type:

tuple

static 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 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).

Parameters:

• scheduler_file (str) – Path the scheduler should write its connection information to

• out_dir (str) – Directory for the scheduler and worker logs

• node_names (list) – The machines in the allocation, in the order SLURM lists them

• parallel_count (int or None) – Total number of worker processes over all nodes, or None

Returns:

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

Return type:

tuple

static srun_scheduler_file(config)

The scheduler file the srun launcher writes, as an absolute path.

Under this launcher the scheduler file is an output: PyBNF starts the scheduler that writes it. scheduler_file therefore chooses where it goes (useful when the output directory is not the filesystem you want the workers to read connection information from), and defaults to the output directory, which a cluster run already requires to be shared.

Parameters:

config (pybnf.config.Configuration) – PyBNF configuration

Returns:

absolute path of the scheduler file

Return type:

str

static srun_worker_command(scheduler_file, node_count, parallel_count=None, granted=None, source=None)

Build the srun invocation that starts one dask worker process group per node.

granted is how many CPUs the job holds on a node, which both sizes the default worker pool and is what this step asks SLURM for. The caller passes it when it already knows – srun_worker_layout() reads the allocation’s own per-node list before it decides which command to build, and that list is right even when the launching process is not on an allocated node, which is where re-deriving the number here went wrong (#642). It is derived from cpus_per_node() only when no caller supplied it.

Parameters:

• scheduler_file (str) – Path of the scheduler file the workers should read

• node_count (int) – Number of nodes in the allocation

• parallel_count (int or None) – Total number of worker processes over all nodes, or None for one per granted CPU

• granted (int or None) – CPUs the job holds on a node, or None to work it out here

• source (str or None) – Phrase naming where granted came from, for the log

Returns:

the srun argument list

Return type:

list

static 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 (#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 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 srun_worker_command() caps its own, so an oversubscribed count still starts every worker rather than being refused by SLURM.

Parameters:

• scheduler_file (str) – Path of the scheduler file the workers should read

• nodes (list) – The machines in this group, all granted the same CPU count

• cpus (int) – CPUs granted on each machine in the group

• n_workers (int or None) – Workers to start on each machine, or None for one per granted CPU

Returns:

the srun argument list for this group

Return type:

list

static 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, #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.

Parameters:

• scheduler_file (str) – Path of the scheduler file the workers should read

• out_dir (str) – Directory the worker logs are written in

• node_names (list) – The machines in the allocation, in the order SLURM lists them

• parallel_count (int or None) – Total number of worker processes over all nodes, or None

Returns:

the srun command list(s), the matching log path(s), and the total worker count

Return type:

tuple

static ssh_bringup_hints(output)

What to suggest to a user whose SSH bring-up failed (#618).

Two things a bare exit code does not tell them. First, whether the login is the problem: on the cluster this was reported from it was, and no part of the message said so. A login failure is worth naming outright because the obvious remedy – creating SSH keys – fixes only one of its causes, and because ssh succeeding from the same shell makes the failure look impossible (dask ssh does not run ssh; it logs in with paramiko, which offers a public key or a password and nothing else).

Second, that a failed login is not the end of the run: two of the ways PyBNF can use several machines never log in anywhere, and both remain open. They are named whatever the cause, since anything that stops dask ssh leaves them as the ways forward.

Parameters:

output (str) – What dask ssh wrote before exiting

Returns:

suggested remedies, most specific first

Return type:

list

stop_own_processes()

Terminate the cluster processes PyBNF started, and remove the scheduler file it wrote

Separate from teardown() because bring-up can fail before there is a client to close, and the processes started up to that point still have to be stopped.

static stop_process(proc, description, timeout=30.0)

Ask a process PyBNF started to stop, and wait until it actually has (#398).

This is what replaces the fixed 10 s sleep that used to follow every cluster teardown. The process is asked to terminate and then waited on, so teardown returns as soon as it is really gone rather than after a fixed guess. A process that ignores the request is killed after timeout seconds, so a stuck one cannot hold the run open forever.

Parameters:

• proc – The process to stop

• description (str) – What to call it in the log

• timeout (float) – Seconds to wait for it to exit before killing it

teardown()

Terminates the processes PyBNF started for this run, after the fitting run completes

The worker launcher goes first and the scheduler second: terminating srun signals the workers, which are the processes that would otherwise be left talking to a scheduler that is already gone. Nothing here touches a cluster PyBNF did not start – scheduler_file and scheduler_node runs have no processes of their own.

static wait_for_scheduler_file(scheduler_file, scheduler_proc, scheduler_log, timeout=60.0, poll=0.25)

Wait until the scheduler has written a usable scheduler file, and return its address.

The file is written without being renamed into place, so a reader can catch it half-written; requiring it to parse as JSON carrying an address is what makes its appearance a readiness signal rather than a race. The scheduler process is checked on every pass, so a scheduler that dies (an occupied port, a bad interpreter) is reported immediately instead of after the timeout.

Parameters:

• scheduler_file (str) – Path the scheduler was told to write

• scheduler_proc – The running scheduler process

• scheduler_log (str) – Path of the scheduler’s log, quoted if it failed

• timeout (float) – Seconds to wait before giving up

• poll (float) – Seconds between checks

Returns:

the scheduler address recorded in the file

Return type:

str

Raises:

PybnfError – if the scheduler exits, or the file does not appear in time

static wait_for_srun_workers(client, worker_procs, expected, worker_logs, timeout=120.0, poll=0.25)

Wait until all of the srun-launched workers have registered with the scheduler.

The waiting itself is _poll_for_workers(), the loop both launchers share; this 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.

Parameters:

• client – The connected dask Client

• worker_procs (list) – The running srun process(es), one per machine size

• expected (int) – Total number of workers that should register across all the steps

• worker_logs (list) – The srun output log(s), quoted if the workers never all arrive

• timeout (float) – Seconds to wait before giving up

• poll (float) – Seconds between checks

Returns:

the number of workers that had registered

Return type:

int

Raises:

PybnfError – if an srun step exits, or the workers do not all register in time

static wait_for_ssh_workers(client, dask_proc, expected, output_file, out_dir, timeout=120.0, poll=0.25)

Wait until the workers dask ssh is bringing up have registered with the scheduler (#398).

This is what replaces the fixed 10 s sleep the SSH launcher used to take after starting dask ssh. The waiting itself is _poll_for_workers(), the loop the srun launcher also uses, so the behaviour this adds is exercised on a real cluster through the srun path even though the SSH login cannot be (#398). What this adds is specific to dask ssh. If dask ssh exits, the login or launch failed, and its captured output is quoted the way an immediate failure was before – naming the login as the likely cause when it looks like one, and naming the ways of running that need no login (#618).

The full expected count is required rather than just one worker. A run that connected with fewer workers than were asked for is the silent-degradation problem of #200: it does not fail, it just quietly uses less than was reserved. Waiting for all of them turns that into a clear, bounded error instead.

Parameters:

• client – The connected dask Client

• dask_proc – The running dask ssh process

• expected (int) – Number of worker processes dask ssh should bring up

• output_file – Open file dask ssh’s output was captured to, quoted on failure

• out_dir (str) – Directory dask ssh logged each node to, named when it said nothing here

• timeout (float) – Seconds to wait before giving up

• poll (float) – Seconds between checks

Returns:

the number of workers that had registered

Return type:

int

Raises:

PybnfError – if dask ssh exits, or fewer than expected workers register in time

pybnf.cluster.check_dask_subcommand(subcommand)

Confirm that dask <subcommand> can be run by this interpreter, before running it.

A missing command surfaces as FileNotFoundError or an unrecognized-command exit from a subprocess PyBNF launched – either an “unknown error … please report this bug” traceback or, worse, a bring-up that fails ten seconds later for a reason the user has to go digging for (#615). This is the same question dask’s own CLI asks when it assembles its command group: subcommands come from the dask_cli entry point group, so anything this check cannot see, dask cannot run either.

Parameters:

subcommand (str) – The dask subcommand PyBNF is about to run, e.g. ‘ssh’

Raises:

PybnfError – if this interpreter has no dask CLI, or no such subcommand

pybnf.cluster.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.

Parameters:

spec (str or None) – The value of $SLURM_JOB_CPUS_PER_NODE, or None

Returns:

one CPU count per node, or None if the text is empty or does not parse

Return type:

list or None

pybnf.cluster.uses_srun(cluster_type)

Whether this cluster_type selects the srun launcher (#614).

Parameters:

cluster_type (str or None) – The configured cluster_type, or None for a local run

Returns:

True if workers should be started with srun rather than over SSH

Return type:

bool