diff --git a/distributed/deploy/ssh.py b/distributed/deploy/ssh.py index 4acd072d07..6f85994d53 100644 --- a/distributed/deploy/ssh.py +++ b/distributed/deploy/ssh.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import copy import logging import sys @@ -28,6 +29,7 @@ class Process(ProcessInterface): def __init__(self, **kwargs): self.connection = None self.proc = None + self._logger_tasks: list[asyncio.Task] = [] super().__init__(**kwargs) async def start(self): @@ -35,9 +37,42 @@ async def start(self): weakref.finalize( self, self.proc.kill ) # https://github.com/ronf/asyncssh/issues/112 + self._start_forwarding_output() await super().start() + def _start_forwarding_output(self) -> None: + """Keep draining the remote process' stdout and stderr. + + Subclasses read stderr only until the remote process announces its address, + and never read stdout at all. Once that startup handshake is over nothing + consumes the SSH channel any more, so the receive window fills up and the + remote process blocks forever the next time it writes a log line. + """ + assert self.proc + self._logger_tasks = [ + asyncio.create_task(self._forward_output(stream)) + for stream in (self.proc.stdout, self.proc.stderr) + ] + + @staticmethod + async def _forward_output(stream: Any) -> None: + try: + while True: + line = await stream.readline() + if not line: # EOF; the remote process exited + break + logger.info(line.strip()) + except Exception: + # The connection dropping is an expected way for this to end, and it is + # already reported through other channels; never let it surface here. + logger.debug("Stopped forwarding output from %r", stream, exc_info=True) + async def close(self): + if self._logger_tasks: + for task in self._logger_tasks: + task.cancel() + await asyncio.gather(*self._logger_tasks, return_exceptions=True) + self._logger_tasks.clear() if self.proc: self.proc.kill() # https://github.com/ronf/asyncssh/issues/112 if self.connection: diff --git a/distributed/deploy/tests/test_ssh.py b/distributed/deploy/tests/test_ssh.py index 29fe3fd900..9ee8b18661 100644 --- a/distributed/deploy/tests/test_ssh.py +++ b/distributed/deploy/tests/test_ssh.py @@ -4,13 +4,14 @@ pytest.importorskip("asyncssh") +import asyncio import sys import dask from distributed import Client from distributed.compatibility import MACOS, WINDOWS -from distributed.deploy.ssh import SSHCluster +from distributed.deploy.ssh import Scheduler, SSHCluster from distributed.utils_test import gen_test pytestmark = [ @@ -19,6 +20,73 @@ ] +# asyncssh.create_server leaks 2 fds on its own, independently of anything under test +@pytest.mark.leaking("fds") +@gen_test() +async def test_remote_process_not_blocked_by_unread_output(): + """The remote process must keep running once its startup banner has been read. + + ``Scheduler.start`` reads stderr only until the remote announces its address and + never reads stdout at all. If nothing drains those streams afterwards, the SSH + receive window fills up and the remote process blocks forever the next time it + logs. See https://github.com/dask/distributed/issues/9033. + + This drives the real ``ssh.Scheduler`` against an in-process asyncssh server, so + it exercises genuine SSH flow control without needing a reachable sshd. + """ + import asyncssh + + # Enough to overflow asyncssh's 2 MiB default channel window several times over. + chunk = "distributed.core - INFO - " + "x" * 4000 + "\n" + n_chunks = 1500 + finished = asyncio.Event() + + async def handle_process(process): + if process.command == "uname": + process.stdout.write("Linux\n") + process.exit(0) + return + process.stderr.write( + "distributed.scheduler - INFO - Scheduler at: tcp://127.0.0.1:8786\n" + ) + await process.stderr.drain() + for _ in range(n_chunks): + process.stderr.write(chunk) + await process.stderr.drain() # blocks once the peer window is exhausted + finished.set() + process.exit(0) + + class _Server(asyncssh.SSHServer): + def begin_auth(self, username): + return False # the test server does not authenticate + + server = await asyncssh.create_server( + _Server, + "127.0.0.1", + 0, + server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")], + process_factory=handle_process, + ) + try: + port = next(iter(server.sockets)).getsockname()[1] + scheduler = Scheduler( + address="127.0.0.1", + connect_options={"port": port, "known_hosts": None, "username": "test"}, + kwargs={}, + ) + await scheduler.start() + connection = scheduler.connection + try: + assert scheduler.address == "tcp://127.0.0.1:8786" + await asyncio.wait_for(finished.wait(), timeout=30) + finally: + await scheduler.close() + await connection.wait_closed() + finally: + server.close() + await server.wait_closed() + + def test_ssh_hosts_None(): with pytest.raises(ValueError): SSHCluster(hosts=None)