From c15c48385076b220043adb7cb69a4f9d803f3012 Mon Sep 17 00:00:00 2001 From: m0g3r <87276771+m0g3r@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:10:18 +0200 Subject: [PATCH] Keep draining remote stdout and stderr in SSHCluster Worker.start and Scheduler.start read the remote process' stderr only until it 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. Forward both streams to the logger for the lifetime of the process instead, and cancel the forwarding tasks when the process is closed. Closes #9033 Co-Authored-By: Claude --- distributed/deploy/ssh.py | 35 ++++++++++++++ distributed/deploy/tests/test_ssh.py | 70 +++++++++++++++++++++++++++- 2 files changed, 104 insertions(+), 1 deletion(-) 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)