Repository: dask/distributed

Issue:
SSHCluster expects conda environment to be at the same path on all systems
**What happened**:
When using `SSHCluster` on machines with different conda paths things fail.

**What you expected to happen**:
The correct conda environment should be activated.

It seems that the `SSHCluster` tries to call the python executable directly at the same path that the current Python is running at. It may be more robust to activate the conda environment with the same name as the current one and use the `python` provided with that. It would also be good to be able to specify a conda environment in the kwargs.

**Minimal Complete Verifiable Example**:

```shell
# On Host A
conda create -n test -p /tmp/condaA python ipython dask
conda activate /tmp/condaA/test

# On Host B
conda create -n test -p /tmp/condaB python ipython dask
conda activate /tmp/condaB/test
```

```python
# On Host A

from dask.distributed import SSHCluster
cluster = SSHCluster(["localhost", "HostB"]
```

```python-traceback
...
distributed.deploy.ssh - INFO - env: ‘/tmp/condaA/test/bin/python’: No such file or directory
...
```

<details>
<summary>Full Example Traceback</summary>

This traceback is from a real example to the paths don't quite match the simplified example above.

```python-traceback
distributed.deploy.ssh - INFO - distributed.scheduler - INFO - -----------------------------------------------
distributed.deploy.ssh - INFO - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
distributed.deploy.ssh - INFO - distributed.scheduler - INFO - -----------------------------------------------
distributed.deploy.ssh - INFO - distributed.scheduler - INFO - Clear task state
distributed.deploy.ssh - INFO - distributed.scheduler - INFO -   Scheduler at:   tcp://10.51.100.15:8786
distributed.deploy.ssh - INFO - env: ‘/Users/jtomlinson/miniconda3/envs/coiledstream/bin/python’: No such file or directory
Task exception was never retrieved
future: <Task finished name='Task-45' coro=<_wrap_awaitable() done, defined at /Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/asyncio/tasks.py:677> exception=Exception('Worker failed to start')>
Traceback (most recent call last):
  File "/Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/asyncio/tasks.py", line 684, in _wrap_awaitable
    return (yield from awaitable.__await__())
  File "/Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py", line 50, in _
    await self.start()
  File "/Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/ssh.py", line 129, in start
    raise Exception("Worker failed to start")
Exception: Worker failed to start
distributed.deploy.ssh - INFO - env: ‘/Users/jtomlinson/miniconda3/envs/coiledstream/bin/python’: No such file or directory
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-2-dbe5c7142de0> in <module>
----> 1 cluster = SSHCluster(["localhost", "10.51.0.32"], connect_options=[{}, {"username": "jacob"}])

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/ssh.py in SSHCluster(hosts, connect_options, worker_options, scheduler_options, worker_module, remote_python, **kwargs)
    352         for i, host in enumerate(hosts[1:])
    353     }
--> 354     return SpecCluster(workers, scheduler, name="SSHCluster", **kwargs)

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py in __init__(self, workers, scheduler, worker, asynchronous, loop, security, silence_logs, name)
    255             self._loop_runner.start()
    256             self.sync(self._start)
--> 257             self.sync(self._correct_state)
    258
    259     async def _start(self):

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/cluster.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    167             return future
    168         else:
--> 169             return sync(self.loop, func, *args, **kwargs)
    170
    171     async def _get_logs(self, scheduler=True, workers=True):

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
    337     if error[0]:
    338         typ, exc, tb = error[0]
--> 339         raise exc.with_traceback(tb)
    340     else:
    341         return result[0]

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/utils.py in f()
    321             if callback_timeout is not None:
    322                 future = asyncio.wait_for(future, callback_timeout)
--> 323             result[0] = yield future
    324         except Exception as exc:
    325             error[0] = sys.exc_info()

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/tornado/gen.py in run(self)
    733
    734                     try:
--> 735                         value = future.result()
    736                     except Exception:
    737                         exc_info = sys.exc_info()

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py in _correct_state_internal(self)
    333                 for w in workers:
    334                     w._cluster = weakref.ref(self)
--> 335                     await w  # for tornado gen.coroutine support
    336             self.workers.update(dict(zip(to_open, workers)))
    337

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py in _()
     48             async with self.lock:
     49                 if self.status == "created":
---> 50                     await self.start()
     51                     assert self.status == "running"
     52             return self

~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/ssh.py in start(self)
    127             line = await self.proc.stderr.readline()
    128             if not line.strip():
--> 129                 raise Exception("Worker failed to start")
    130             logger.info(line.strip())
    131             if "worker at" in line:

Exception: Worker failed to start
```

</details>

**Environment**:

- Dask version: 2.23.0
- Python version: 3.8.5
- Operating System: macOS 10.14 and Ubuntu 10.04
- Install method (conda, pip, source): conda



Discussion:
We might also consider splitting worker kwargs by worker, perhaps with an
optional more complex input.  We do something similar today with a hostfile

On Fri, Aug 21, 2020 at 8:18 AM Jacob Tomlinson <notifications@github.com>
wrote:

> *What happened*:
> When using SSHCluster on machines with different conda paths things fail.
>
> *What you expected to happen*:
> The correct conda environment should be activated.
>
> It seems that the SSHCluster tries to call the python executable directly
> at the same path that the current Python is running at. It may be more
> robust to activate the conda environment with the same name as the current
> one and use the python provided with that. It would also be good to be
> able to specify a conda environment in the kwargs.
>
> *Minimal Complete Verifiable Example*:
>
> # On Host A
>
> conda create -n test -p /tmp/condaA python ipython dask
>
> conda activate /tmp/condaA/test
>
>
> # On Host B
>
> conda create -n test -p /tmp/condaB python ipython dask
>
> conda activate /tmp/condaB/test
>
> # On Host A
>
>
> from dask.distributed import SSHCluster
> cluster = SSHCluster(["localhost", "HostB"]
>
> ...
>
> distributed.deploy.ssh - INFO - env: ‘/tmp/condaA/test/bin/python’: No such file or directory
>
> ...
>
> Full Example Traceback
>
> This traceback is from a real example to the paths don't quite match the
> simplified example above.
>
> distributed.deploy.ssh - INFO - distributed.scheduler - INFO - -----------------------------------------------
>
> distributed.deploy.ssh - INFO - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
>
> distributed.deploy.ssh - INFO - distributed.scheduler - INFO - -----------------------------------------------
>
> distributed.deploy.ssh - INFO - distributed.scheduler - INFO - Clear task state
>
> distributed.deploy.ssh - INFO - distributed.scheduler - INFO -   Scheduler at:   tcp://10.51.100.15:8786
>
> distributed.deploy.ssh - INFO - env: ‘/Users/jtomlinson/miniconda3/envs/coiledstream/bin/python’: No such file or directory
>
> Task exception was never retrieved
> future: <Task finished name='Task-45' coro=<_wrap_awaitable() done, defined at /Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/asyncio/tasks.py:677> exception=Exception('Worker failed to start')>
>
> Traceback (most recent call last):
>
>   File "/Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/asyncio/tasks.py", line 684, in _wrap_awaitable
>
>     return (yield from awaitable.__await__())
>
>   File "/Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py", line 50, in _
>
>     await self.start()
>
>   File "/Users/jtomlinson/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/ssh.py", line 129, in start
>
>     raise Exception("Worker failed to start")
> Exception: Worker failed to start
>
> distributed.deploy.ssh - INFO - env: ‘/Users/jtomlinson/miniconda3/envs/coiledstream/bin/python’: No such file or directory
>
> ---------------------------------------------------------------------------
>
> Exception                                 Traceback (most recent call last)
>
> <ipython-input-2-dbe5c7142de0> in <module>
>
> ----> 1 cluster = SSHCluster(["localhost", "10.51.0.32"], connect_options=[{}, {"username": "jacob"}])
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/ssh.py in SSHCluster(hosts, connect_options, worker_options, scheduler_options, worker_module, remote_python, **kwargs)
>
>     352         for i, host in enumerate(hosts[1:])
>
>     353     }
>
> --> 354     return SpecCluster(workers, scheduler, name="SSHCluster", **kwargs)
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py in __init__(self, workers, scheduler, worker, asynchronous, loop, security, silence_logs, name)
>
>     255             self._loop_runner.start()
>
>     256             self.sync(self._start)
>
> --> 257             self.sync(self._correct_state)
>
>     258
>
>     259     async def _start(self):
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/cluster.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
>
>     167             return future
>
>     168         else:
>
> --> 169             return sync(self.loop, func, *args, **kwargs)
>
>     170
>
>     171     async def _get_logs(self, scheduler=True, workers=True):
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
>
>     337     if error[0]:
>
>     338         typ, exc, tb = error[0]
>
> --> 339         raise exc.with_traceback(tb)
>
>     340     else:
>
>     341         return result[0]
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/utils.py in f()
>
>     321             if callback_timeout is not None:
>
>     322                 future = asyncio.wait_for(future, callback_timeout)
>
> --> 323             result[0] = yield future
>
>     324         except Exception as exc:
>
>     325             error[0] = sys.exc_info()
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/tornado/gen.py in run(self)
>
>     733
>
>     734                     try:
>
> --> 735                         value = future.result()
>
>     736                     except Exception:
>
>     737                         exc_info = sys.exc_info()
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py in _correct_state_internal(self)
>
>     333                 for w in workers:
>
>     334                     w._cluster = weakref.ref(self)
>
> --> 335                     await w  # for tornado gen.coroutine support
>
>     336             self.workers.update(dict(zip(to_open, workers)))
>
>     337
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/spec.py in _()
>
>      48             async with self.lock:
>
>      49                 if self.status == "created":
>
> ---> 50                     await self.start()
>
>      51                     assert self.status == "running"
>
>      52             return self
>
>
>
> ~/miniconda3/envs/coiledstream/lib/python3.8/site-packages/distributed/deploy/ssh.py in start(self)
>
>     127             line = await self.proc.stderr.readline()
>
>     128             if not line.strip():
>
> --> 129                 raise Exception("Worker failed to start")
>
>     130             logger.info(line.strip())
>
>     131             if "worker at" in line:
>
>
> Exception: Worker failed to start
>
> *Environment*:
>
>    - Dask version: 2.23.0
>    - Python version: 3.8.5
>    - Operating System: macOS 10.14 and Ubuntu 10.04
>    - Install method (conda, pip, source): conda
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/4063>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTG7A6G2LJLTS3EXLQ3SB2F35ANCNFSM4QHMP35A>
> .
>

PR #3701 added `remote_python` keyword to `SSHCluster`. How about extending this keyword to accept a dict as {node: path to python on node}?
@abduhbm that would be a great start.

We do also need to consider the rest of the environment. The Python path needs to point to the correct place as well as executing the correct binary.
@jacobtomlinson I can work on a PR for this if it is fine with you. Thanks, 
Yes please that would be great!

Resolving patch:
diff --git a/distributed/deploy/ssh.py b/distributed/deploy/ssh.py
--- a/distributed/deploy/ssh.py
+++ b/distributed/deploy/ssh.py
@@ -108,10 +108,13 @@ async def start(self):
                     "Worker failed to set DASK_INTERNAL_INHERIT_CONFIG variable "
                 )
 
+        if not self.remote_python:
+            self.remote_python = sys.executable
+
         cmd = " ".join(
             [
                 set_env,
-                self.remote_python or sys.executable,
+                self.remote_python,
                 "-m",
                 self.worker_module,
                 self.scheduler,
@@ -186,10 +189,13 @@ async def start(self):
                     "Scheduler failed to set DASK_INTERNAL_INHERIT_CONFIG variable "
                 )
 
+        if not self.remote_python:
+            self.remote_python = sys.executable
+
         cmd = " ".join(
             [
                 set_env,
-                self.remote_python or sys.executable,
+                self.remote_python,
                 "-m",
                 "distributed.cli.dask_scheduler",
             ]
@@ -235,7 +241,7 @@ def SSHCluster(
     worker_options: dict = {},
     scheduler_options: dict = {},
     worker_module: str = "distributed.cli.dask_worker",
-    remote_python: str = None,
+    remote_python: Union[str, List[str]] = None,
     **kwargs,
 ):
     """Deploy a Dask cluster using SSH
@@ -274,7 +280,7 @@ def SSHCluster(
         Keywords to pass on to scheduler.
     worker_module: str, optional
         Python module to call to start the worker.
-    remote_python: str, optional
+    remote_python: str or list of str, optional
         Path to Python on remote nodes.
 
     Examples
@@ -326,6 +332,12 @@ def SSHCluster(
             "dictionary for each address."
         )
 
+    if isinstance(remote_python, list) and len(remote_python) != len(hosts):
+        raise RuntimeError(
+            "When specifying a list of remote_python you must provide a "
+            "path for each address."
+        )
+
     scheduler = {
         "cls": Scheduler,
         "options": {
@@ -334,7 +346,9 @@ def SSHCluster(
             if isinstance(connect_options, dict)
             else connect_options[0],
             "kwargs": scheduler_options,
-            "remote_python": remote_python,
+            "remote_python": remote_python[0]
+            if isinstance(remote_python, list)
+            else remote_python,
         },
     }
     workers = {
@@ -347,7 +361,9 @@ def SSHCluster(
                 else connect_options[i + 1],
                 "kwargs": worker_options,
                 "worker_module": worker_module,
-                "remote_python": remote_python,
+                "remote_python": remote_python[i + 1]
+                if isinstance(remote_python, list)
+                else remote_python,
             },
         }
         for i, host in enumerate(hosts[1:])


Test patch:
diff --git a/distributed/deploy/tests/test_ssh.py b/distributed/deploy/tests/test_ssh.py
--- a/distributed/deploy/tests/test_ssh.py
+++ b/distributed/deploy/tests/test_ssh.py
@@ -2,6 +2,7 @@
 
 pytest.importorskip("asyncssh")
 
+import sys
 import dask
 from dask.distributed import Client
 from distributed.deploy.ssh import SSHCluster
@@ -156,3 +157,43 @@ async def test_list_of_connect_options_raises():
             worker_options={"death_timeout": "5s"},
         ) as _:
             pass
+
+
+@pytest.mark.asyncio
+async def test_remote_python():
+    async with SSHCluster(
+        ["127.0.0.1"] * 3,
+        connect_options=[dict(known_hosts=None)] * 3,
+        asynchronous=True,
+        scheduler_options={"port": 0, "idle_timeout": "5s"},
+        worker_options={"death_timeout": "5s"},
+        remote_python=sys.executable,
+    ) as cluster:
+        assert cluster.workers[0].remote_python == sys.executable
+
+
+@pytest.mark.asyncio
+async def test_remote_python_as_dict():
+    async with SSHCluster(
+        ["127.0.0.1"] * 3,
+        connect_options=[dict(known_hosts=None)] * 3,
+        asynchronous=True,
+        scheduler_options={"port": 0, "idle_timeout": "5s"},
+        worker_options={"death_timeout": "5s"},
+        remote_python=[sys.executable] * 3,
+    ) as cluster:
+        assert cluster.workers[0].remote_python == sys.executable
+
+
+@pytest.mark.asyncio
+async def test_list_of_remote_python_raises():
+    with pytest.raises(RuntimeError):
+        async with SSHCluster(
+            ["127.0.0.1"] * 3,
+            connect_options=[dict(known_hosts=None)] * 3,
+            asynchronous=True,
+            scheduler_options={"port": 0, "idle_timeout": "5s"},
+            worker_options={"death_timeout": "5s"},
+            remote_python=[sys.executable] * 4,  # Mismatch in length 4 != 3
+        ) as _:
+            pass