diff --git a/docs/cluster.rst b/docs/cluster.rst index baca6c11..3bebf260 100644 --- a/docs/cluster.rst +++ b/docs/cluster.rst @@ -70,7 +70,9 @@ If you use one of the first two, the login must succeed without a password promp * It raises ``AuthenticationException`` while plain ``ssh`` to the same node succeeds: your cluster authenticates its nodes to each other by a method paramiko cannot use -- most often host-based or Kerberos (GSSAPI) SSH. **Creating SSH keys cannot fix this**, because the cluster is not asking for a key. Use `Starting workers without SSH`_. * It asks for a password, or fails in the same way plain ``ssh`` does: this is the case SSH keys do fix. Create a key pair with ``ssh-keygen`` (documented in many places, such as `here `__) and append the public half to ``~/.ssh/authorized_keys``. Where the nodes share your home directory, that covers all of them at once. Then run the check above again. -**What a failed login looks like.** PyBNF stops about ten seconds into the run rather than carrying on with fewer machines than you asked for, and quotes what ``dask ssh`` said -- including dask's own account of the failure, which names the node it was connecting to and the exception paramiko raised. When that output reads as a refused credential, the message says the login is the likely cause, says what PyBNF logs in with, and names the two ways of running that need no login at all. The same output is in the log file, tracebacks and all. +**What a failed login looks like.** PyBNF stops as soon as ``dask ssh`` gives up on the login, rather than carrying on with fewer machines than you asked for, and quotes what ``dask ssh`` said -- including dask's own account of the failure, which names the node it was connecting to and the exception paramiko raised. When that output reads as a refused credential, the message says the login is the likely cause, says what PyBNF logs in with, and names the two ways of running that need no login at all. The same output is in the log file, tracebacks and all. + +**When the login succeeds** but the workers are slow to start, PyBNF waits for all of the workers you asked for to register before the run begins, up to a time limit (``SSH_WORKER_TIMEOUT`` in ``pybnf/cluster.py``, two minutes by default). If fewer than that arrive in time, it stops and says how many of how many registered rather than running on a smaller cluster than you reserved. A slow or busy cluster may need a longer limit. If SSH cannot be made to work for some other reason, `Starting workers without SSH`_ and `Manual configuration with Dask`_ both avoid it entirely. diff --git a/pybnf/cluster.py b/pybnf/cluster.py index a73bc9c0..962342e2 100644 --- a/pybnf/cluster.py +++ b/pybnf/cluster.py @@ -67,16 +67,29 @@ SRUN_SCHEDULER_LOG = 'dask_scheduler.log' SRUN_WORKER_LOG = 'dask_workers.log' -# Readiness limits for the srun launcher. Both waits are polls on a real signal -- the -# scheduler file appearing, and a worker registering with the scheduler -- rather than a -# fixed sleep, and both also watch the launched process, so a bring-up that fails outright -# is reported in the time it takes to fail rather than after the whole timeout (#398 asks -# for the same treatment of the SSH path's two fixed 10 s sleeps, which this does not -# touch). +# Readiness limits for both launchers. Every wait is a poll on a real signal -- the +# scheduler file appearing, a worker registering with the scheduler -- rather than a fixed +# sleep, and every one also watches the launched process, so a bring-up that fails outright +# is reported in the time it takes to fail rather than after the whole timeout. +# +# SSH_WORKER_TIMEOUT is the SSH launcher's share of this (#398). It replaces a fixed 10 s +# sleep in setup_cluster that assumed the cluster was up after ten seconds whatever the +# truth: on a real SLURM cluster the workers took 26-59 s to register (measured, see #398), +# so ten seconds connected before they were ready; on a fast cluster it was longer than +# needed. The wait now polls for the workers dask ssh brings up and reports a login that +# fails as soon as dask ssh exits. Sized like SRUN_WORKER_TIMEOUT, with headroom over the +# slowest bring-up observed. SCHEDULER_FILE_TIMEOUT = 60. SRUN_WORKER_TIMEOUT = 120. +SSH_WORKER_TIMEOUT = 120. READINESS_POLL_INTERVAL = 0.25 +# How long teardown waits for a process PyBNF started (the worker launcher, the scheduler) to +# exit after being asked to stop, before killing it outright (#398). This replaces a fixed +# 10 s sleep that followed every cluster teardown: teardown now returns as soon as the +# processes are really gone, and still bounds the wait for one that will not stop. +TEARDOWN_TIMEOUT = 30. + # How PyBNF invokes dask's command line interface (#615). Two things are decided here. # @@ -196,6 +209,13 @@ def __init__(self, config, log_prefix, debug, log_level_name): self._scheduler_proc = None self._own_scheduler_file = None self._srun_worker_log = None + # 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 + # it was told to log each node to. + self._ssh_output_file = None + self._ssh_expected_workers = None + self._ssh_out_dir = None # Where the client should look for the scheduler. This is the user's # ``scheduler_file`` when they are attaching to a cluster of their own, and the @@ -236,7 +256,11 @@ def __init__(self, config, log_prefix, debug, log_level_name): scheduler_node, node_string = self.read_node_names(config) if node_string: - self._dask_proc = self.setup_cluster(node_string, os.getcwd(), config.config['parallel_count']) + ssh_out_dir = os.getcwd() + self._ssh_out_dir = ssh_out_dir + (self._dask_proc, self._ssh_expected_workers, + self._ssh_output_file) = self.setup_cluster(node_string, ssh_out_dir, + config.config['parallel_count']) logger.info(f'Initializing dask Client with dask v{daskv}, distributed v{distributedv}') @@ -264,6 +288,20 @@ def __init__(self, config, log_prefix, debug, log_level_name): logger.info(f'Creating a client by connecting to the scheduler node {scheduler_node}:8786') self.client = Client(f'{scheduler_node}:8786') self.local = False + if self._dask_proc is not None: + # PyBNF started this cluster with dask ssh, so it is the one that can tell + # whether the workers arrived. Connecting to the scheduler says nothing about + # that -- a scheduler with no workers connects fine -- so without this wait a + # failed or slow bring-up would surface later as a fit that submits jobs and + # never gets one back (#398, and the silent-degradation risk of #200). Anything + # that fails here leaves the dask ssh process running, so stop it first. + try: + self.wait_for_ssh_workers(self.client, self._dask_proc, + self._ssh_expected_workers, self._ssh_output_file, + self._ssh_out_dir) + except Exception: + self.stop_own_processes() + raise else: # One local branch, one thread policy (#526). `parallel_count` chooses how many # worker *processes* there are; it never decides how many threads run inside one. @@ -375,10 +413,11 @@ def setup_cluster(node_string, out_dir, parallel_count=None): :param 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``) - :return: subprocess.Popen - :raises PybnfError: if ``dask ssh`` has already exited when the bring-up wait is - over -- quoting everything it said, naming the login as the likely cause when - it looks like one, and naming the ways of running that need no login (#618) + :return: 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 (:meth:`wait_for_ssh_workers`, #398) + :rtype: tuple """ # Ask before launching, so a dask that cannot do this reads as a configuration # error rather than as a FileNotFoundError traceback from Popen (#615). @@ -437,28 +476,12 @@ def setup_cluster(node_string, out_dir, parallel_count=None): dask_ssh_out = TemporaryFile() dask_ssh_proc = Popen(dask_ssh_cmd, stdout=dask_ssh_out, stderr=STDOUT, env=dict(os.environ, PYTHONUNBUFFERED='1')) - time.sleep(10) - # If dask ssh has already exited, the cluster never came up. Surface the - # failure here instead of letting it resurface later as an opaque dask - # Client connection error. - returncode = dask_ssh_proc.poll() - if returncode is not None: - output = Cluster.captured_text(dask_ssh_out) - dask_ssh_out.close() - # The log keeps every line, including the traceback frames; the message keeps - # what a user can act on. - logger.error('dask ssh exited with code %s during cluster bring-up. Output:\n%s' - % (returncode, output if output else '(it produced none)')) - raise PybnfError( - 'Could not start the workers on the other machines: dask ssh exited with code ' - '%s during cluster bring-up.\n%s' - % (returncode, - ('This is what it said:\n%s' % Cluster.fold_traceback_frames(output)) - if output else - ('It produced no output at all. Anything the nodes themselves wrote is ' - 'in %s on each of them.' % out_dir)), - hint=Cluster.ssh_bringup_hints(output)) - return dask_ssh_proc + # No fixed wait here any more (#398). The caller connects a Client and then calls + # wait_for_ssh_workers, which polls for these workers to register and watches this + # process, so a failed login is reported as soon as dask ssh exits and a healthy + # bring-up proceeds the moment the workers are up rather than after a fixed guess. + expected_workers = n_per_node * len(nodes) + return dask_ssh_proc, expected_workers, dask_ssh_out @staticmethod def captured_text(handle): @@ -868,12 +891,61 @@ def wait_for_scheduler_file(scheduler_file, scheduler_proc, scheduler_log, raise PybnfError('The dask scheduler did not write its connection file %s within %s s. ' 'See %s for details.' % (scheduler_file, timeout, scheduler_log)) + @staticmethod + def _poll_for_workers(client, worker_proc, expected, timeout, poll): + """ + Poll the scheduler until enough workers have registered, watching the process (#398). + + This is the readiness mechanism both launchers share. Whether PyBNF started the workers + with srun or with dask ssh, the question is the same: have the workers registered with + the scheduler yet, and is the process that was meant to bring them up still running? + Both are checked on every pass, so a bring-up that dies is caught in the time it takes + to die rather than after the whole timeout. What differs between the launchers -- how + many workers to wait for, and what to say when it goes wrong -- is left to the caller, + which is why this returns an outcome rather than raising: the srun path quotes its log, + the SSH path quotes what dask ssh said and names the login, and neither vocabulary + belongs here. + + The srun launcher is the one path that runs on a cluster whose nodes need no login, so + this shared loop is what a real multi-machine run actually exercises (#398's own reason + 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 expected: Number of registered workers that counts as ready + :type expected: int + :param timeout: Seconds to wait before giving up + :type timeout: float + :param poll: Seconds between checks + :type poll: float + :return: a tuple of the outcome (``'ready'``, ``'exited'`` or ``'timeout'``), the number + of workers that had registered, and the process return code (only set on ``'exited'``) + :rtype: tuple + """ + 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 + try: + n_workers = len(client.scheduler_info()['workers']) + except Exception: + n_workers = 0 + if n_workers >= expected: + return 'ready', n_workers, None + time.sleep(poll) + return 'timeout', n_workers, None + @staticmethod def wait_for_srun_workers(client, srun_proc, worker_log, timeout=SRUN_WORKER_TIMEOUT, poll=READINESS_POLL_INTERVAL): """ Wait until at least one srun-launched worker has 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. + :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 @@ -886,23 +958,18 @@ def wait_for_srun_workers(client, srun_proc, worker_log, :rtype: int :raises PybnfError: if srun exits, or no worker registers in time """ - for _ in range(max(1, int(timeout / poll))): - returncode = srun_proc.poll() - if returncode is not None: - details = Cluster.log_tail(worker_log) - logger.error('srun exited with code %s before any worker started. Log:\n%s' - % (returncode, details)) - raise PybnfError('srun exited with code %s before any dask worker started. %s' - % (returncode, ('Details:\n%s' % details) if details - else 'See %s for details.' % worker_log)) - try: - n_workers = len(client.scheduler_info()['workers']) - except Exception: - n_workers = 0 - if n_workers: - logger.info('%i dask worker process(es) registered with the scheduler' % n_workers) - return n_workers - time.sleep(poll) + outcome, n_workers, returncode = Cluster._poll_for_workers( + client, srun_proc, expected=1, timeout=timeout, poll=poll) + if outcome == 'ready': + logger.info('%i dask worker process(es) registered with the scheduler' % n_workers) + return n_workers + if outcome == 'exited': + details = Cluster.log_tail(worker_log) + logger.error('srun exited with code %s before any worker started. Log:\n%s' + % (returncode, details)) + raise PybnfError('srun exited with code %s before any dask worker 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 ' @@ -911,6 +978,77 @@ def wait_for_srun_workers(client, srun_proc, worker_log, 'A message about job step creation means another step already holds ' 'the allocation; run PyBNF as the only job step.']) + @staticmethod + def wait_for_ssh_workers(client, dask_proc, expected, output_file, out_dir, + timeout=SSH_WORKER_TIMEOUT, poll=READINESS_POLL_INTERVAL): + """ + 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 :meth:`_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. + + :param client: The connected dask Client + :param dask_proc: The running dask ssh process + :param expected: Number of worker processes dask ssh should bring up + :type expected: int + :param output_file: Open file dask ssh's output was captured to, quoted on failure + :param out_dir: Directory dask ssh logged each node to, named when it said nothing here + :type out_dir: str + :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 dask ssh exits, or fewer than ``expected`` workers register in time + """ + outcome, n_workers, returncode = Cluster._poll_for_workers( + client, dask_proc, expected=expected, timeout=timeout, poll=poll) + if outcome == 'ready': + logger.info('%i of %i dask worker process(es) registered with the scheduler' + % (n_workers, expected)) + return n_workers + if outcome == 'exited': + # dask ssh exited during bring-up, so the cluster never came up -- almost always + # a refused login (see ssh_bringup_hints). The log keeps every line, including + # the traceback frames; the message keeps what a user can act on. + output = Cluster.captured_text(output_file) + logger.error('dask ssh exited with code %s during cluster bring-up. Output:\n%s' + % (returncode, output if output else '(it produced none)')) + raise PybnfError( + 'Could not start the workers on the other machines: dask ssh exited with code ' + '%s during cluster bring-up.\n%s' + % (returncode, + ('This is what it said:\n%s' % Cluster.fold_traceback_frames(output)) + if output else + ('It produced no output at all. Anything the nodes themselves wrote is ' + 'in %s on each of them.' % out_dir)), + hint=Cluster.ssh_bringup_hints(output)) + output = Cluster.captured_text(output_file) + logger.error('Only %i of the %i expected dask worker process(es) registered within %s s. ' + 'dask ssh output:\n%s' % (n_workers, expected, timeout, + output if output else '(it produced none)')) + raise PybnfError( + 'Only %i of the %i expected worker process(es) started over SSH registered with the ' + 'scheduler within %s s.%s' + % (n_workers, expected, timeout, + ('\nThis is what dask ssh said:\n%s' % Cluster.fold_traceback_frames(output)) + if output else ''), + hint=['Anything the workers themselves wrote is in the log directory %s on each ' + 'node.' % out_dir, + 'A slow or busy cluster may need longer; the wait is SSH_WORKER_TIMEOUT in ' + 'pybnf/cluster.py.']) + def teardown(self): """ Terminates the processes PyBNF started for this run, after the fitting run completes @@ -932,13 +1070,19 @@ def stop_own_processes(self): close, and the processes started up to that point still have to be stopped. """ if self._dask_proc: - logger.info('Closing the worker launcher subprocess') - self._dask_proc.terminate() + self.stop_process(self._dask_proc, 'worker launcher subprocess') self._dask_proc = None if self._scheduler_proc: - logger.info('Closing the dask scheduler subprocess') - self._scheduler_proc.terminate() + self.stop_process(self._scheduler_proc, 'dask scheduler subprocess') self._scheduler_proc = None + if self._ssh_output_file is not None: + # The captured dask ssh output is only needed while waiting for the workers; once + # the run is over, close it so the temporary file is released. + try: + self._ssh_output_file.close() + except OSError: + pass + self._ssh_output_file = None if self._own_scheduler_file and os.path.exists(self._own_scheduler_file): # The file names a scheduler that is being shut down, so leaving it in place # would leave a live-looking connection file behind for the next run to find. @@ -947,3 +1091,32 @@ def stop_own_processes(self): except OSError: logger.debug('Could not remove the scheduler file %s' % self._own_scheduler_file) self._own_scheduler_file = None + + @staticmethod + def stop_process(proc, description, timeout=TEARDOWN_TIMEOUT): + """ + 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. + + :param proc: The process to stop + :param description: What to call it in the log + :type description: str + :param timeout: Seconds to wait for it to exit before killing it + :type timeout: float + """ + logger.info('Closing the %s' % description) + proc.terminate() + try: + proc.wait(timeout=timeout) + except TimeoutExpired: + logger.warning('The %s did not exit within %s s of being asked to; killing it.' + % (description, timeout)) + proc.kill() + try: + proc.wait(timeout=timeout) + except TimeoutExpired: + logger.error('The %s did not exit even after being killed.' % description) diff --git a/pybnf/pybnf.py b/pybnf/pybnf.py index 8e836542..48eed5a6 100644 --- a/pybnf/pybnf.py +++ b/pybnf/pybnf.py @@ -601,12 +601,12 @@ def _reap_running_sims(): def _teardown_cluster(cluster): """Tear down the dask cluster after a run, logging (not raising on) any failure.""" - # Stop the cluster's worker launcher (dask ssh, or srun) regardless of success + # Stop the cluster's worker launcher (dask ssh, or srun) regardless of success. The wait + # for those processes to actually exit lives in Cluster.teardown now, which waits on each + # one rather than sleeping a fixed ten seconds and hoping (#398). if cluster: try: cluster.teardown() - if not cluster.local: - time.sleep(10) # wait for teardown before continuing except Exception: logging.exception('Failed to tear down cluster') else: diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 541a43ff..5d686013 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -44,6 +44,7 @@ ``reinit_logging`` workaround is asserted *to be called*, not pinned), so they remain a valid safety net across the dask-unpinning upgrade. """ +import io import json import os import re @@ -79,17 +80,27 @@ def __init__(self, stdout): class _FakeDaskProc: - """Stand-in for a launched Popen object: the bring-up paths only poll it, and - terminate it when they have to abandon a partly-built cluster.""" + """Stand-in for a launched Popen object: the bring-up paths poll it, and + terminate then wait on it when they have to abandon a partly-built cluster.""" def __init__(self, returncode=None): self._returncode = returncode self.terminated = False + self.killed = False def poll(self): return self._returncode def terminate(self): self.terminated = True + if self._returncode is None: + self._returncode = -15 + + def wait(self, timeout=None): + return self._returncode + + def kill(self): + self.killed = True + self._returncode = -9 class TestReadNodeNames: @@ -424,8 +435,8 @@ def test_every_option_pybnf_passes_is_still_declared_by_dask(self, subcommand, m """The half of #619 the subcommand check cannot reach: a command that still exists but no longer takes the option PyBNF hands it. ``--nworkers`` is itself a survivor of that -- distributed renamed it from ``--nprocs`` -- - and the next rename would otherwise reach a user as a bring-up that fails - ten seconds in, with dask's complaint in a log nobody is reading.""" + and the next rename would otherwise reach a user as a bring-up that fails, + with dask's complaint in a log nobody is reading.""" _, dask_argv = _split_at_dask_cli(DASK_COMMAND_BUILDERS[subcommand](monkeypatch)) declared = _declared_options(_help_text((*cluster.DASK_CLI, subcommand))) missing = [opt for opt in _options_in(dask_argv) if opt not in declared] @@ -529,6 +540,21 @@ def fake_popen(*args, **kwargs): monkeypatch.setattr(cluster, 'cpu_count', lambda: cpu) return popen_calls + @staticmethod + def _bringup_then_wait(node_string='node9', out_dir='/log', parallel_count=1): + """Launch dask ssh and run the readiness wait that watches it (#398). + + setup_cluster no longer waits or reports a failure by itself; it launches dask ssh + and hands back the process, the worker count, and the captured-output file. The + readiness wait is what then watches the process and reports a bring-up that has + already died, reading the output setup_cluster captured. These tests drive the two + together, with a client reporting no workers so the only thing that can happen is the + dead-process branch.""" + proc, expected, out_file = cluster.Cluster.setup_cluster(node_string, out_dir, + parallel_count=parallel_count) + return cluster.Cluster.wait_for_ssh_workers(_ClientStub(workers=()), proc, expected, + out_file, out_dir) + def test_default_worker_count_is_what_the_job_was_granted(self, monkeypatch): """parallel_count=None ⇒ one single-threaded worker per CPU **the job holds on a node**: ``--nthreads 1 --nworkers {cpus_per_node()}`` (note this @@ -537,9 +563,10 @@ def test_default_worker_count_is_what_the_job_was_granted(self, monkeypatch): entry) with SLURM having granted 7, over an affinity count of 4 and a machine of 64.""" popen_calls = self._patch(monkeypatch, granted=7, affinity=4, cpu=64) - proc = cluster.Cluster.setup_cluster('n1 n2', '/out', parallel_count=None) + proc, expected, out_file = cluster.Cluster.setup_cluster('n1 n2', '/out', parallel_count=None) assert proc.poll() is None + assert expected == 14 # 7 workers on each of 2 nodes (args, kwargs), = popen_calls assert args[0] == [*DASK_SSH, 'n1', 'n2', '--log-directory', '/out', '--nthreads', '1', '--nworkers', '7'] @@ -624,12 +651,16 @@ def test_a_missing_subcommand_is_refused_before_launching(self, monkeypatch): cluster.Cluster.setup_cluster('n1', '/log', parallel_count=1) assert popen_calls == [] - def test_running_proc_is_returned_without_raising(self, monkeypatch): - """The happy path: dask ssh is still running after the startup wait - (poll() is None), so setup_cluster returns the proc rather than raising.""" + def test_setup_returns_the_proc_expected_count_and_output_file(self, monkeypatch): + """setup_cluster launches dask ssh and hands back what the readiness wait needs + (#398): the running process, the number of workers it should bring up (one per node + times the per-node count), and the open file its output was captured to. It no longer + waits or decides success itself.""" self._patch(monkeypatch) - proc = cluster.Cluster.setup_cluster('n1', '/log', parallel_count=1) + proc, expected, out_file = cluster.Cluster.setup_cluster('n1 n2', '/log', parallel_count=6) assert proc.poll() is None + assert expected == 6 # ceil(6/2)=3 per node, over 2 nodes + assert hasattr(out_file, 'read') def test_failed_bringup_raises_with_the_captured_output(self, monkeypatch): """If dask ssh has already exited after the startup wait, the cluster @@ -639,7 +670,7 @@ def test_failed_bringup_raises_with_the_captured_output(self, monkeypatch): self._patch(monkeypatch, returncode=1, output_bytes=b'ssh: connect to host node9 port 22: Connection refused') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') msg = str(exc.value) assert 'code 1' in msg assert 'Connection refused' in msg @@ -656,7 +687,7 @@ def test_bringup_captures_the_stream_dask_explains_itself_on(self, monkeypatch): b'node9:22\n SSH reported this exception: ' b'Authentication failed.\n') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') assert 'SSH reported this exception: Authentication failed.' in str(exc.value) def test_dask_is_run_unbuffered_so_its_own_account_survives(self, monkeypatch): @@ -691,7 +722,7 @@ def test_traceback_frames_are_folded_out_of_the_message(self, monkeypatch): b'paramiko.ssh_exception.AuthenticationException: Authentication failed.\n' b' Retrying... (attempt 1/3)\n')) with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') message = exc.value.message assert 'SSH connection error when connecting to node9:22' in message assert 'AuthenticationException: Authentication failed.' in message @@ -709,7 +740,7 @@ def test_the_log_keeps_the_frames_the_message_folds(self, monkeypatch, caplog): b'paramiko.ssh_exception.AuthenticationException: Authentication failed.\n')) with caplog.at_level('ERROR', logger='pybnf.cluster'): with pytest.raises(printing.PybnfError): - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') line, = [r.message for r in caplog.records if 'dask ssh exited' in r.message] assert 'old_ssh.py' in line @@ -724,7 +755,7 @@ def test_a_login_failure_is_named_as_one(self, monkeypatch): output_bytes=b'paramiko.ssh_exception.AuthenticationException: ' b'Authentication failed.') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') message = exc.value.message assert 'login' in message assert 'paramiko' in message @@ -740,7 +771,7 @@ def test_a_network_failure_is_not_blamed_on_the_login(self, monkeypatch): output_bytes=b'paramiko.ssh_exception.NoValidConnectionsError: ' b'[Errno None] Unable to connect to port 22 on 10.0.0.9') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') message = exc.value.message assert 'failed login' not in message # no diagnosis is offered assert 'public key' not in message @@ -755,7 +786,7 @@ def test_failure_names_both_ways_of_running_without_a_login(self, monkeypatch): no-login-vocabulary case.""" self._patch(monkeypatch, returncode=1, output_bytes=b'exit status 127') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') message = exc.value.message assert 'slurm-srun' in message assert 'scheduler_file' in message @@ -768,7 +799,7 @@ def test_output_is_reported_even_when_there_is_none(self, monkeypatch): the directory the nodes write to is named.""" self._patch(monkeypatch, returncode=1, output_bytes=b'') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/logdir', parallel_count=1) + self._bringup_then_wait('node9', '/logdir') message = exc.value.message assert 'no output' in message assert '/logdir' in message @@ -779,7 +810,7 @@ def test_captured_output_is_logged_as_well_as_raised(self, monkeypatch, caplog): self._patch(monkeypatch, returncode=1, output_bytes=b'Authentication failed.') with caplog.at_level('ERROR', logger='pybnf.cluster'): with pytest.raises(printing.PybnfError): - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') line, = [r.message for r in caplog.records if 'dask ssh exited' in r.message] assert 'Authentication failed.' in line @@ -789,7 +820,7 @@ def test_colour_codes_are_stripped_from_the_quoted_output(self, monkeypatch): self._patch(monkeypatch, returncode=1, output_bytes=b'\x1b[91mSSH connection failed after 3 retries.\x1b[0m') with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') assert 'SSH connection failed after 3 retries.' in str(exc.value) assert '\x1b' not in exc.value.message @@ -799,7 +830,7 @@ def test_only_the_tail_of_a_long_output_is_quoted(self, monkeypatch): self._patch(monkeypatch, returncode=1, output_bytes=b'\n'.join(b'line %i' % i for i in range(200))) with pytest.raises(printing.PybnfError) as exc: - cluster.Cluster.setup_cluster('node9', '/log', parallel_count=1) + self._bringup_then_wait('node9', '/log') quoted = str(exc.value) assert 'line 199' in quoted assert 'line 0\n' not in quoted @@ -844,23 +875,15 @@ def test_node_names_passed_as_literal_argv_no_shell(self, monkeypatch): assert args[0][:len(DASK_SSH) + 2] == [*DASK_SSH, 'n1$(whoami)', 'n2'] # literal, unexpanded assert kwargs.get('shell', False) is False - def test_sleeps_ten_seconds_for_startup(self, monkeypatch): - """After launching dask ssh, setup_cluster waits 10s for workers to come - up before returning the proc. Oracle: time.sleep called once with 10.""" - self._patch(monkeypatch) - sleeps = [] - monkeypatch.setattr(cluster.time, 'sleep', lambda s: sleeps.append(s)) - cluster.Cluster.setup_cluster('n1', '/log', parallel_count=1) - assert sleeps == [10] - - # --------------------------------------------------------------------------- # # __init__ — node-detection dispatch + Client-construction dispatch # --------------------------------------------------------------------------- # class _ProcStub: - """Stand-in for a process PyBNF started and later terminates.""" + """Stand-in for a process PyBNF started and later terminates. Teardown asks it to + stop and then waits on it, so it records both and reports itself as exited.""" def __init__(self, returncode=None): self.terminated = False + self.killed = False self._returncode = returncode def poll(self): @@ -868,6 +891,15 @@ def poll(self): def terminate(self): self.terminated = True + if self._returncode is None: + self._returncode = -15 + + def wait(self, timeout=None): + return self._returncode + + def kill(self): + self.killed = True + self._returncode = -9 class _ClientStub: @@ -899,6 +931,8 @@ def __init__(self): self.reinit_logging_calls = [] self.read_calls = [] 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) @@ -913,7 +947,7 @@ def LocalCluster(self, *args, **kwargs): return self.last_lc -def _patch_init(monkeypatch, read_returns=(None, None), srun_raises=None): +def _patch_init(monkeypatch, read_returns=(None, None), srun_raises=None, ssh_raises=None): rec = _Recorder() monkeypatch.setattr(cluster, 'Client', rec.Client) monkeypatch.setattr(cluster, 'LocalCluster', rec.LocalCluster) @@ -928,12 +962,22 @@ def fake_read(config): def fake_setup(node_string, out_dir, parallel_count): rec.setup_calls.append((node_string, out_dir, parallel_count)) - return 'PROC' + # setup_cluster returns the process, the worker count it should bring up, and the + # file its output was captured to (#398). The count and file are placeholders here; + # wait_for_ssh_workers, which consumes them, is faked below. + 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_ssh_wait(client, dask_proc, expected, output_file, out_dir, **kwargs): + rec.ssh_wait_calls.append((client, dask_proc, expected, output_file, out_dir)) + if ssh_raises: + raise ssh_raises + return expected + def fake_wait(client, srun_proc, worker_log, **kwargs): rec.srun_wait_calls.append((client, srun_proc, worker_log)) if srun_raises: @@ -942,6 +986,7 @@ def fake_wait(client, srun_proc, worker_log, **kwargs): monkeypatch.setattr(cluster.Cluster, 'read_node_names', staticmethod(fake_read)) monkeypatch.setattr(cluster.Cluster, 'setup_cluster', staticmethod(fake_setup)) + monkeypatch.setattr(cluster.Cluster, 'wait_for_ssh_workers', staticmethod(fake_ssh_wait)) monkeypatch.setattr(cluster.Cluster, 'setup_srun_cluster', staticmethod(fake_setup_srun)) monkeypatch.setattr(cluster.Cluster, 'wait_for_srun_workers', staticmethod(fake_wait)) monkeypatch.setattr(cluster.Cluster, 'require_slurm_allocation', staticmethod(lambda: None)) @@ -978,9 +1023,12 @@ def test_scheduler_node_plus_worker_nodes_joins_worker_list(self, monkeypatch): assert rec.read_calls == [] assert rec.setup_calls == [('w1 w2 w3', os.getcwd(), 12)] - assert c._dask_proc == 'PROC' + assert c._dask_proc is rec.setup_proc assert rec.client_calls == [(('head:8786',), {})] assert c.local is False + # Having started dask ssh itself, PyBNF waits for those workers to register (#398). + assert len(rec.ssh_wait_calls) == 1 + assert rec.ssh_wait_calls[0][1] is rec.setup_proc def test_scheduler_node_alone_detects_workers_via_read_node_names(self, monkeypatch): """scheduler_node set but no worker_nodes ⇒ the worker list comes from @@ -993,6 +1041,7 @@ def test_scheduler_node_alone_detects_workers_via_read_node_names(self, monkeypa assert len(rec.read_calls) == 1 assert rec.setup_calls == [('d1 d2', os.getcwd(), 8)] + assert c._dask_proc is rec.setup_proc assert rec.client_calls == [(('head:8786',), {})] assert c.local is False @@ -1005,9 +1054,22 @@ def test_detected_cluster_uses_both_outputs_of_read_node_names(self, monkeypatch assert len(rec.read_calls) == 1 assert rec.setup_calls == [('sched9 c1 c2', os.getcwd(), 4)] + assert c._dask_proc is rec.setup_proc assert rec.client_calls == [(('sched9:8786',), {})] assert c.local is False + def test_a_failed_worker_wait_stops_the_dask_ssh_it_started(self, monkeypatch): + """If the workers never register, the readiness wait raises and the constructor + never becomes a Cluster, so no one else can tear it down. The dask ssh process it + started is stopped on the way out rather than left running (#398).""" + rec = _patch_init(monkeypatch, read_returns=('sched9', 'sched9 c1 c2'), + ssh_raises=printing.PybnfError('workers never came up')) + with pytest.raises(printing.PybnfError, match='workers never came up'): + _build(_cfg(parallel_count=4)) + + assert len(rec.ssh_wait_calls) == 1 + assert rec.setup_proc.terminated is True + class TestInitClientDispatch: @@ -1135,6 +1197,7 @@ def _torn_down(client=None, dask_proc=None, scheduler_proc=None, scheduler_file= c._dask_proc = dask_proc c._scheduler_proc = scheduler_proc c._own_scheduler_file = scheduler_file + c._ssh_output_file = None return c @@ -1565,6 +1628,85 @@ def test_no_worker_in_time_names_the_log_and_the_step_hazard(self, monkeypatch, assert 'job step' in exc.value.message.lower() +class TestPollForWorkers: + """The readiness loop both launchers share (#398). The srun and SSH waits are thin wrappers + that add their own worker count and their own failure vocabulary on top of this; the loop + itself is what a real srun run exercises, so pin its three outcomes here. It reports the + outcome rather than raising, leaving each launcher to phrase the error its own way.""" + + def test_ready_when_the_expected_count_is_reached(self, monkeypatch): + """Enough workers registered, process still running: 'ready', with the count.""" + monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + outcome, n, rc = cluster.Cluster._poll_for_workers( + _ClientStub(workers=('tcp://n1:1', 'tcp://n2:1')), _FakeDaskProc(), + expected=2, timeout=5., poll=0.25) + assert (outcome, n, rc) == ('ready', 2, None) + + def test_exited_when_the_process_is_gone(self, monkeypatch): + """The bring-up process exited before the workers arrived: 'exited', with its code, + so the caller can quote whatever that launcher logged.""" + monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + outcome, n, rc = cluster.Cluster._poll_for_workers( + _ClientStub(workers=()), _FakeDaskProc(returncode=1), + expected=2, timeout=5., poll=0.25) + assert (outcome, rc) == ('exited', 1) + + def test_timeout_when_too_few_arrive_in_time(self, monkeypatch): + """Process still running but short of the count when time runs out: 'timeout', with + how many did register.""" + monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + outcome, n, rc = cluster.Cluster._poll_for_workers( + _ClientStub(workers=('tcp://n1:1',)), _FakeDaskProc(), + expected=3, timeout=1., poll=0.25) + assert (outcome, n, rc) == ('timeout', 1, None) + + +class TestWaitForSSHWorkers: + """The SSH launcher's startup readiness check (#398), driven directly. The dead-process + branch is covered through setup_cluster in TestSetupCluster; these pin the worker-count + branch: dask ssh still running, the scheduler filling up over time.""" + + def test_returns_once_the_full_expected_count_registers(self, monkeypatch): + """The readiness signal is all of the expected workers registering, not dask ssh + having been launched. Two asked for, two connected, so the wait returns two.""" + monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + n = cluster.Cluster.wait_for_ssh_workers( + _ClientStub(workers=('tcp://n1:1', 'tcp://n2:1')), _FakeDaskProc(), + expected=2, output_file=io.BytesIO(), out_dir='/log') + assert n == 2 + + def test_a_partial_cluster_is_not_ready_yet(self, monkeypatch): + """Fewer workers than were asked for is the silent-degradation case of #200: the poll + keeps waiting rather than returning a cluster smaller than reserved. 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) == 2: + client._workers = dict.fromkeys(('tcp://n1:1', 'tcp://n2:1'), {}) + + monkeypatch.setattr(cluster.time, 'sleep', fake_sleep) + n = cluster.Cluster.wait_for_ssh_workers( + client, _FakeDaskProc(), expected=2, output_file=io.BytesIO(), out_dir='/log') + assert n == 2 + assert len(polls) == 2 + + def test_too_few_in_time_names_the_expected_and_connected_counts(self, monkeypatch): + """dask ssh still running but short of the count means some workers never came up -- + a fit would quietly use less than was reserved. The message names how many of how + many arrived and where to read what the missing ones wrote.""" + monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + with pytest.raises(printing.PybnfError) as exc: + cluster.Cluster.wait_for_ssh_workers( + _ClientStub(workers=('tcp://n1:1',)), _FakeDaskProc(), + expected=3, output_file=io.BytesIO(), out_dir='/logdir', timeout=1.) + message = exc.value.message + assert '1 of the 3' in message + assert '/logdir' in ' '.join(exc.value.hints) + + class TestSetupSrunCluster: def test_starts_the_scheduler_here_then_the_workers_with_srun(self, monkeypatch, tmp_path):