Add alternative SSHCluster implementation

.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

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

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

distributed/deploy/local.py

@@ -198,6 +198,7 @@ def __init__(
             loop=loop,
             asynchronous=asynchronous,
             silence_logs=silence_logs,
+            security=security,
         )
 
     def __repr__(self):

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)

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 "<SSH %s: status=%s>" % (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)

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