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/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/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 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..f94f909cb59 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] @@ -214,11 +225,10 @@ 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 {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 +250,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..0f8823cdab8 --- /dev/null +++ b/distributed/deploy/ssh2.py @@ -0,0 +1,171 @@ +import asyncio +import logging +import sys +import warnings +import weakref + +import asyncssh + +from .spec import SpecCluster + +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 + + 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 + connect_kwargs: dict + kwargs to be passed to asyncssh connections + kwargs: + TODO + """ + + 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 + + super().__init__() + + async def start(self): + self.connection = await asyncssh.connect(self.address, **self.connect_kwargs) + 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 + connect_kwargs: dict + kwargs to be passed to asyncssh connections + kwargs: + TODO + """ + + def __init__(self, address: str, connect_kwargs: dict, **kwargs): + self.address = address + self.kwargs = kwargs + self.connect_kwargs = connect_kwargs + + super().__init__() + + async def start(self): + logger.debug("Created Scheduler Connection") + + 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"]) + ) + + # 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, connect_kwargs, **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 + 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], "connect_kwargs": connect_kwargs}, + } + workers = { + 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 new file mode 100644 index 00000000000..beb1c6ef91e --- /dev/null +++ b/distributed/deploy/tests/test_ssh2.py @@ -0,0 +1,17 @@ +import pytest + +pytest.importorskip("asyncssh") + +from dask.distributed import Client +from distributed.deploy.ssh2 import SSHCluster + + +@pytest.mark.asyncio +async def test_basic(): + async with SSHCluster( + ["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: + result = await client.submit(lambda x: x + 1, 10) + assert result == 11