From f250c81ee74f9c4bb63d681a6b0bd224530c9749 Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Sat, 6 Jul 2019 16:55:59 +0100 Subject: [PATCH 1/8] Add alternative SSHCluster implementation This is a proof of concept here for two reasons: 1. It opens up a possible alternative for SSH deployment (which was surprisingly popular in the user survey) 2. It is the first non-local application of `SpecCluster` and so serves as a proof of concept for other future deployments that are mostly defined by creating a remote Worker/Scheduler object This forced some changes in `SpecCluster`, notably we now have an `rpc` object that does remote calls rather than accessing the scheduler directly. Also, we're going to have to figure out how to handle all of the keyword arguments. In this case we need to pass them from Python down to the CLI, and presumably we'll also want a `dask-ssh` CLI command which has to translate the other way. --- distributed/deploy/local.py | 1 + distributed/deploy/spec.py | 28 +++-- distributed/deploy/ssh2.py | 144 ++++++++++++++++++++++++++ distributed/deploy/tests/test_ssh2.py | 13 +++ 4 files changed, 179 insertions(+), 7 deletions(-) create mode 100644 distributed/deploy/ssh2.py create mode 100644 distributed/deploy/tests/test_ssh2.py diff --git a/distributed/deploy/local.py b/distributed/deploy/local.py index 554459e43ac..dfe4d963c5e 100644 --- a/distributed/deploy/local.py +++ b/distributed/deploy/local.py @@ -198,6 +198,7 @@ def __init__( loop=loop, asynchronous=asynchronous, silence_logs=silence_logs, + security=security, ) def __repr__(self): diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py index bb46f81db88..98cb9fc4271 100644 --- a/distributed/deploy/spec.py +++ b/distributed/deploy/spec.py @@ -5,8 +5,10 @@ from tornado import gen from .cluster import Cluster +from ..core import rpc from ..utils import LoopRunner, silence_logging, ignoring from ..scheduler import Scheduler +from ..security import Security class SpecCluster(Cluster): @@ -107,6 +109,7 @@ def __init__( worker=None, asynchronous=False, loop=None, + security=None, silence_logs=False, ): self._created = weakref.WeakSet() @@ -125,6 +128,8 @@ def __init__( self.workers = {} self._i = 0 self._asynchronous = asynchronous + self.security = security or Security() + self.scheduler_comm = None if silence_logs: self._old_logging_level = silence_logging(level=silence_logs) @@ -156,6 +161,10 @@ async def _start(self): self._lock = asyncio.Lock() self.status = "starting" self.scheduler = await self.scheduler + self.scheduler_comm = rpc( + self.scheduler.address, + connection_args=self.security.get_connection_args("client"), + ) self.status = "running" def _correct_state(self): @@ -174,11 +183,13 @@ async def _correct_state_internal(self): pre = list(set(self.workers)) to_close = set(self.workers) - set(self.worker_spec) if to_close: - await self.scheduler.retire_workers(workers=list(to_close)) + if self.scheduler.status == "running": + await self.scheduler_comm.retire_workers(workers=list(to_close)) tasks = [self.workers[w].close() for w in to_close] await asyncio.wait(tasks) for task in tasks: # for tornado gen.coroutine support - await task + with ignoring(RuntimeError): + await task for name in to_close: del self.workers[name] @@ -216,9 +227,10 @@ async def _(): async def _wait_for_workers(self): # TODO: this function needs to query scheduler and worker state # remotely without assuming that they are local - while {d["name"] for d in self.scheduler.identity()["workers"].values()} != set( - self.workers - ): + while { + str(d["name"]) + for d in (await self.scheduler_comm.identity())["workers"].values() + } != set(map(str, self.workers)): if ( any(w.status == "closed" for w in self.workers.values()) and self.scheduler.status == "running" @@ -240,12 +252,14 @@ async def _close(self): return self.status = "closing" - async with self._lock: - await self.scheduler.close(close_workers=True) self.scale(0) await self._correct_state() + async with self._lock: + await self.scheduler_comm.close(close_workers=True) + await self.scheduler.close() for w in self._created: assert w.status == "closed" + self.scheduler_comm.close_rpc() if hasattr(self, "_old_logging_level"): silence_logging(self._old_logging_level) diff --git a/distributed/deploy/ssh2.py b/distributed/deploy/ssh2.py new file mode 100644 index 00000000000..d9dd26ade45 --- /dev/null +++ b/distributed/deploy/ssh2.py @@ -0,0 +1,144 @@ +import asyncio +import logging +import sys +import weakref + +import asyncssh + +from .spec import SpecCluster + +logger = logging.getLogger(__name__) + + +class Process: + """ A superclass for SSH Workers and Nannies + + See Also + -------- + Worker + Scheduler + """ + + def __init__(self): + self.lock = asyncio.Lock() + self.connection = None + self.proc = None + self.status = "created" + + def __await__(self): + async def _(): + async with self.lock: + if not self.connection: + await self.start() + assert self.connection + weakref.finalize(self, self.proc.terminate) + return self + + return _().__await__() + + async def close(self): + self.proc.terminate() + self.connection.close() + self.status = "closed" + + def __repr__(self): + return "" % (type(self).__name__, self.status) + + +class Worker(Process): + """ A Remote Dask Worker controled by SSH + + Parameters + ---------- + scheduler: str + The address of the scheduler + address: str + The hostname where we should run this worker + kwargs: + TODO + """ + + def __init__(self, scheduler: str, address: str, **kwargs): + self.address = address + self.scheduler = scheduler + self.kwargs = kwargs + + super().__init__() + + async def start(self): + self.connection = await asyncssh.connect(self.address) + self.proc = await self.connection.create_process( + " ".join( + [ + sys.executable, + "-m", + "distributed.cli.dask_worker", + self.scheduler, + "--name", # we need to have name for SpecCluster + str(self.kwargs["name"]), + ] + ) + ) + + # We watch stderr in order to get the address, then we return + while True: + line = await self.proc.stderr.readline() + if "worker at" in line: + self.address = line.split("worker at:")[1].strip() + self.status = "running" + break + logger.debug("%s", line) + + +class Scheduler(Process): + """ A Remote Dask Scheduler controled by SSH + + Parameters + ---------- + address: str + The hostname where we should run this worker + kwargs: + TODO + """ + + def __init__(self, address: str, **kwargs): + self.address = address + self.kwargs = kwargs + + super().__init__() + + async def start(self): + logger.debug("Created Scheduler Connection") + self.connection = await asyncssh.connect(self.address) + self.proc = await self.connection.create_process( + " ".join([sys.executable, "-m", "distributed.cli.dask_scheduler"]) + ) + + # We watch stderr in order to get the address, then we return + while True: + line = await self.proc.stderr.readline() + if "Scheduler at" in line: + self.address = line.split("Scheduler at:")[1].strip() + break + logger.debug("%s", line) + + +def SSHCluster(hosts, **kwargs): + """ Deploy a Dask cluster using SSH + + Parameters + ---------- + hosts: List[str] + List of hostnames or addresses on which to launch our cluster + The first will be used for the scheduler and the rest for workers + + TODO + ---- + This doesn't handle any keyword arguments yet. It is a proof of concept + """ + scheduler = {"cls": Scheduler, "options": {"address": hosts[0]}} + workers = { + i: {"cls": Worker, "options": {"address": host}} + for i, host in enumerate(hosts[1:]) + } + return SpecCluster(workers, scheduler, **kwargs) diff --git a/distributed/deploy/tests/test_ssh2.py b/distributed/deploy/tests/test_ssh2.py new file mode 100644 index 00000000000..0adbb08261c --- /dev/null +++ b/distributed/deploy/tests/test_ssh2.py @@ -0,0 +1,13 @@ +from distributed.deploy.ssh2 import SSHCluster +from dask.distributed import Client + +import pytest + + +@pytest.mark.asyncio +async def test_basic(): + async with SSHCluster(["localhost"] * 3, asynchronous=True) as cluster: + assert len(cluster.workers) == 2 + async with Client(cluster, asynchronous=True) as client: + result = await client.submit(lambda x: x + 1, 10) + assert result == 11 From a9cc08aafdaeb645ea120cd59b59fa3b6aab3eca Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Sat, 6 Jul 2019 16:14:24 -0500 Subject: [PATCH 2/8] add asyncssh to CI --- continuous_integration/travis/install.sh | 1 + distributed/deploy/tests/test_ssh2.py | 8 +++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/continuous_integration/travis/install.sh b/continuous_integration/travis/install.sh index 2ab9724db25..b2fab6afb52 100644 --- a/continuous_integration/travis/install.sh +++ b/continuous_integration/travis/install.sh @@ -67,6 +67,7 @@ pip install -q git+https://github.com/dask/s3fs.git --upgrade --no-deps pip install -q git+https://github.com/dask/zict.git --upgrade --no-deps pip install -q sortedcollections msgpack --no-deps pip install -q keras --upgrade --no-deps +pip install -q asyncssh if [[ $CRICK == true ]]; then conda install -q cython diff --git a/distributed/deploy/tests/test_ssh2.py b/distributed/deploy/tests/test_ssh2.py index 0adbb08261c..a9957430185 100644 --- a/distributed/deploy/tests/test_ssh2.py +++ b/distributed/deploy/tests/test_ssh2.py @@ -1,8 +1,10 @@ -from distributed.deploy.ssh2 import SSHCluster -from dask.distributed import Client - import pytest +pytest.importorskip("asyncssh") + +from dask.distributed import Client +from distributed.deploy.ssh2 import SSHCluster + @pytest.mark.asyncio async def test_basic(): From df6eef1a2c5bff68d0074c1c45cf90c9b87c40bf Mon Sep 17 00:00:00 2001 From: Benjamin Zaitlen Date: Tue, 9 Jul 2019 10:17:54 -0400 Subject: [PATCH 3/8] validate host set to None for disabling host key validation --- distributed/deploy/ssh2.py | 31 ++++++++++++++++++++++++++----- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/distributed/deploy/ssh2.py b/distributed/deploy/ssh2.py index d9dd26ade45..b541d29f4df 100644 --- a/distributed/deploy/ssh2.py +++ b/distributed/deploy/ssh2.py @@ -44,6 +44,19 @@ async def close(self): def __repr__(self): return "" % (type(self).__name__, self.status) +async def _create_connection(address, validate_host): + if not validate_host: + connection = await asyncssh.connect(address, known_hosts=None) + else: + try: + connection = await asyncssh.connect(address) + except asyncssh.misc.HostKeyNotVerifiable: + # improve debug message when there is more support for host validation + logger.debug("Explicitly validating host") + raise asyncssh.misc.HostKeyNotVerifiable + + return connection + class Worker(Process): """ A Remote Dask Worker controled by SSH @@ -55,18 +68,21 @@ class Worker(Process): address: str The hostname where we should run this worker kwargs: - TODO + validate_host: bool + Validate host key is trusted (default is False) """ def __init__(self, scheduler: str, address: str, **kwargs): self.address = address self.scheduler = scheduler self.kwargs = kwargs + self.validate_host = self.kwargs.get('validate_host', False) + super().__init__() async def start(self): - self.connection = await asyncssh.connect(self.address) + self.connection = await _create_connection(self.address, self.validate_host) self.proc = await self.connection.create_process( " ".join( [ @@ -99,17 +115,22 @@ class Scheduler(Process): The hostname where we should run this worker kwargs: TODO + validate_host: bool + Validate host key is trusted (default is False) """ def __init__(self, address: str, **kwargs): self.address = address self.kwargs = kwargs + self.validate_host = self.kwargs.get('validate_host', False) super().__init__() async def start(self): logger.debug("Created Scheduler Connection") - self.connection = await asyncssh.connect(self.address) + + self.connection = await _create_connection(self.address, self.validate_host) + self.proc = await self.connection.create_process( " ".join([sys.executable, "-m", "distributed.cli.dask_scheduler"]) ) @@ -136,9 +157,9 @@ def SSHCluster(hosts, **kwargs): ---- This doesn't handle any keyword arguments yet. It is a proof of concept """ - scheduler = {"cls": Scheduler, "options": {"address": hosts[0]}} + scheduler = {"cls": Scheduler, "options": {"address": hosts[0], "validate_host": False}} workers = { - i: {"cls": Worker, "options": {"address": host}} + i: {"cls": Worker, "options": {"address": host, "validate_host": False}} for i, host in enumerate(hosts[1:]) } return SpecCluster(workers, scheduler, **kwargs) From d6bca1d4083e2d99ffd7feac153120ac920829e4 Mon Sep 17 00:00:00 2001 From: Benjamin Zaitlen Date: Tue, 9 Jul 2019 11:03:34 -0400 Subject: [PATCH 4/8] use connect_kwargs instead of explicity args and pass through to asyncssh connection --- distributed/deploy/ssh2.py | 58 +++++++++++++-------------- distributed/deploy/tests/test_ssh2.py | 4 +- 2 files changed, 32 insertions(+), 30 deletions(-) diff --git a/distributed/deploy/ssh2.py b/distributed/deploy/ssh2.py index b541d29f4df..841441bcd90 100644 --- a/distributed/deploy/ssh2.py +++ b/distributed/deploy/ssh2.py @@ -44,19 +44,6 @@ async def close(self): def __repr__(self): return "" % (type(self).__name__, self.status) -async def _create_connection(address, validate_host): - if not validate_host: - connection = await asyncssh.connect(address, known_hosts=None) - else: - try: - connection = await asyncssh.connect(address) - except asyncssh.misc.HostKeyNotVerifiable: - # improve debug message when there is more support for host validation - logger.debug("Explicitly validating host") - raise asyncssh.misc.HostKeyNotVerifiable - - return connection - class Worker(Process): """ A Remote Dask Worker controled by SSH @@ -67,22 +54,22 @@ class Worker(Process): The address of the scheduler address: str The hostname where we should run this worker + connect_kwargs: dict + kwargs to be passed to asyncssh connections kwargs: - validate_host: bool - Validate host key is trusted (default is False) + TODO """ - def __init__(self, scheduler: str, address: str, **kwargs): + def __init__(self, scheduler: str, address: str, connect_kwargs: dict, **kwargs): self.address = address self.scheduler = scheduler + self.connect_kwargs = connect_kwargs self.kwargs = kwargs - self.validate_host = self.kwargs.get('validate_host', False) - super().__init__() async def start(self): - self.connection = await _create_connection(self.address, self.validate_host) + self.connection = await asyncssh.connect(self.address, **self.connect_kwargs) self.proc = await self.connection.create_process( " ".join( [ @@ -113,23 +100,23 @@ class Scheduler(Process): ---------- address: str The hostname where we should run this worker + connect_kwargs: dict + kwargs to be passed to asyncssh connections kwargs: TODO - validate_host: bool - Validate host key is trusted (default is False) """ - def __init__(self, address: str, **kwargs): + def __init__(self, address: str, connect_kwargs: dict, **kwargs): self.address = address self.kwargs = kwargs - self.validate_host = self.kwargs.get('validate_host', False) + self.connect_kwargs = connect_kwargs super().__init__() async def start(self): logger.debug("Created Scheduler Connection") - self.connection = await _create_connection(self.address, self.validate_host) + self.connection = await asyncssh.connect(self.address, **self.connect_kwargs) self.proc = await self.connection.create_process( " ".join([sys.executable, "-m", "distributed.cli.dask_scheduler"]) @@ -144,7 +131,7 @@ async def start(self): logger.debug("%s", line) -def SSHCluster(hosts, **kwargs): +def SSHCluster(hosts, connect_kwargs, **kwargs): """ Deploy a Dask cluster using SSH Parameters @@ -152,14 +139,27 @@ def SSHCluster(hosts, **kwargs): hosts: List[str] List of hostnames or addresses on which to launch our cluster The first will be used for the scheduler and the rest for workers - - TODO + connect_kwargs: + known_hosts: List[str] or None + The list of keys which will be used to validate the server host + key presented during the SSH handshake. If this is not specified, + the keys will be looked up in the file .ssh/known_hosts. If this + is explicitly set to None, server host key validation will be disabled. + TODO + kwargs: + TODO ---- This doesn't handle any keyword arguments yet. It is a proof of concept """ - scheduler = {"cls": Scheduler, "options": {"address": hosts[0], "validate_host": False}} + scheduler = { + "cls": Scheduler, + "options": {"address": hosts[0], "connect_kwargs": connect_kwargs}, + } workers = { - i: {"cls": Worker, "options": {"address": host, "validate_host": False}} + i: { + "cls": Worker, + "options": {"address": host, "connect_kwargs": connect_kwargs}, + } for i, host in enumerate(hosts[1:]) } return SpecCluster(workers, scheduler, **kwargs) diff --git a/distributed/deploy/tests/test_ssh2.py b/distributed/deploy/tests/test_ssh2.py index a9957430185..56f6236f4c2 100644 --- a/distributed/deploy/tests/test_ssh2.py +++ b/distributed/deploy/tests/test_ssh2.py @@ -8,7 +8,9 @@ @pytest.mark.asyncio async def test_basic(): - async with SSHCluster(["localhost"] * 3, asynchronous=True) as cluster: + async with SSHCluster( + ["localhost"] * 3, connect_kwargs=dict(known_hosts=None), asynchronous=True + ) as cluster: assert len(cluster.workers) == 2 async with Client(cluster, asynchronous=True) as client: result = await client.submit(lambda x: x + 1, 10) From d2f32f33e07c3b7be48536de55f691c461b1e04c Mon Sep 17 00:00:00 2001 From: Benjamin Zaitlen Date: Tue, 9 Jul 2019 11:52:26 -0400 Subject: [PATCH 5/8] use 127.0.0.1 instead of localhost --- distributed/deploy/tests/test_ssh2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distributed/deploy/tests/test_ssh2.py b/distributed/deploy/tests/test_ssh2.py index 56f6236f4c2..beb1c6ef91e 100644 --- a/distributed/deploy/tests/test_ssh2.py +++ b/distributed/deploy/tests/test_ssh2.py @@ -9,7 +9,7 @@ @pytest.mark.asyncio async def test_basic(): async with SSHCluster( - ["localhost"] * 3, connect_kwargs=dict(known_hosts=None), asynchronous=True + ["127.0.0.1"] * 3, connect_kwargs=dict(known_hosts=None), asynchronous=True ) as cluster: assert len(cluster.workers) == 2 async with Client(cluster, asynchronous=True) as client: From dbdd250002b4c516522e90114d72cea1ff97a2d1 Mon Sep 17 00:00:00 2001 From: Benjamin Zaitlen Date: Tue, 9 Jul 2019 13:44:32 -0400 Subject: [PATCH 6/8] add ssh-key for travis ci --- .travis.yml | 1 + continuous_integration/travis/setup-ssh.sh | 2 ++ 2 files changed, 3 insertions(+) create mode 100644 continuous_integration/travis/setup-ssh.sh diff --git a/.travis.yml b/.travis.yml index bcc09351eff..35f4383748e 100644 --- a/.travis.yml +++ b/.travis.yml @@ -19,6 +19,7 @@ matrix: install: - if [[ $TESTS == true ]]; then source continuous_integration/travis/install.sh ; fi + - if [[ $TESTS == true ]]; then source continuous_integration/travis/setup-ssh.sh ; fi script: - if [[ $TESTS == true ]]; then source continuous_integration/travis/run_tests.sh ; fi diff --git a/continuous_integration/travis/setup-ssh.sh b/continuous_integration/travis/setup-ssh.sh new file mode 100644 index 00000000000..f102612bc96 --- /dev/null +++ b/continuous_integration/travis/setup-ssh.sh @@ -0,0 +1,2 @@ +ssh-keygen -t rsa -f ~/.ssh/id_rsa -N "" -q +cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys \ No newline at end of file From c13cffd7fe963d267bea7ac2d7136a1bf098150b Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Wed, 17 Jul 2019 15:09:35 -0500 Subject: [PATCH 7/8] warn on ssh2 import --- distributed/deploy/ssh2.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/distributed/deploy/ssh2.py b/distributed/deploy/ssh2.py index 841441bcd90..0f8823cdab8 100644 --- a/distributed/deploy/ssh2.py +++ b/distributed/deploy/ssh2.py @@ -1,6 +1,7 @@ import asyncio import logging import sys +import warnings import weakref import asyncssh @@ -9,6 +10,11 @@ logger = logging.getLogger(__name__) +warnings.warn( + "the distributed.deploy.ssh2 module is experimental " + "and will move/change in the future without notice" +) + class Process: """ A superclass for SSH Workers and Nannies From 91356dc1889522f7c377f259aba0493a3ab837ef Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Wed, 17 Jul 2019 15:09:59 -0500 Subject: [PATCH 8/8] remove todo in comments --- distributed/deploy/spec.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py index 98cb9fc4271..f94f909cb59 100644 --- a/distributed/deploy/spec.py +++ b/distributed/deploy/spec.py @@ -225,8 +225,6 @@ async def _(): return _().__await__() async def _wait_for_workers(self): - # TODO: this function needs to query scheduler and worker state - # remotely without assuming that they are local while { str(d["name"]) for d in (await self.scheduler_comm.identity())["workers"].values()