SSHCluster() giving out JSONDecodeError

[morikplay]

Hello,
cross-posting from stackoverflow (https://stackoverflow.com/questions/71290821/dask-scheduler-sshcluster-not-establishing-connection) to see whether authors would be so kind as to shed some light on behavior i'm observing.

Use-case:: SSH'ing into self (localhost) to launch Scheduler + 1 worker
Error: JSONDecodeError at the time of scheduler launch
Failed Command generated by code: 'set xyz && && C:\Users\maulik\miniconda3\envs\mls\python.exe -m distributed.cli.dask_spec --spec '{"cls": "distributed.Scheduler", "opts": {"port": 8786}}'`
env:

• windows server 2019
• python 3.10.1
• dask 2022.2.1
• asyncssh 2.9.0

miniconda3 env: conda install geopandas dask=2022.2.1 sqlalchemy psycopg2 pykml fpdf openpyxl geoalchemy2 asyncssh paramiko

Python paths:

c:\Users\maulik\code\tests (<-- main() fn to test sshcluster is here)
C:\Users\maulik\miniconda3\envs\mls\python310.zip
C:\Users\maulik\miniconda3\envs\mls\DLLs
C:\Users\maulik\miniconda3\envs\mls\lib
C:\Users\maulik\miniconda3\envs\mls
C:\Users\maulik\miniconda3\envs\mls\lib\site-packages
C:\Users\maulik\miniconda3\envs\mls\python.exe

(mls) C:\Users\maulik\code> c: && cd c:\Users\maulik\code && cmd /C "C:\Users\maulik\miniconda3\envs\mls\python.exe c:\Users\maulik\.vscode\extensions\ms-python.python-2022.0.1814523869\pythonFiles\lib\python\debugpy\launcher 65392 -- c:\Users\maulik\code\tests\SSHClusterSetup.py "
c:\Users\maulik\code\tests
C:\Users\maulik\miniconda3\envs\mls\python310.zip
C:\Users\maulik\miniconda3\envs\mls\DLLs
C:\Users\maulik\miniconda3\envs\mls\lib
C:\Users\maulik\miniconda3\envs\mls
C:\Users\maulik\miniconda3\envs\mls\lib\site-packages
C:\Users\maulik\miniconda3\envs\mls\python.exe

code to invoke SSHCluster

import os, sys, pathlib, time, functools, logging, subprocess
from random import randint
import dask, dask.array as da
import dask.dataframe as ddf
from dask.delayed import delayed
from dask.distributed import Client, SSHCluster
import numpy as np

py_path = os.path.join("C:" + os.sep, "Users", "maulik", "miniconda3", "envs", "mls", "python.exe")
print(py_path)
co={"preferred_auth": "publickey"}
so={"port": 8786}
wo={
    "n_workers": 1, # only works in newer version of dask e.g. 2022.2.1
    # "nprocs": 1,
    "nthreads": 1,
    "memory_limit": "1GB"
    }

deploy_ssh = SSHCluster(["localhost", "localhost"],connect_options=co, scheduler_options=so,worker_options=wo,remote_python=py_path)

client = Client()
z = client.get_versions(check=True)
print("Version Info = %s", z)

client.run_on_scheduler(lambda: sys.executable)
client.run(lambda: sys.executable)

[bryanwweber]

@morikplay Can you add the appropriate import lines to your example to make it complete copy-pasteable? Thank you 😄

[morikplay]

Done 🙏🏽

[morikplay]

Upon closer examination of the dask code, it seems to me that perhaps either ssh.py or dask_spec.py may need changes?
In async def start() of ssh.py we see:

        cmd = " ".join(
            [
                set_env,
                self.remote_python,
                "-m",
                "distributed.cli.dask_spec",
                "--spec",
                "'%s'" % dumps({"cls": "distributed.Scheduler", "opts": self.kwargs}),
            ]
        )

and in dask_spec.py, we see:

    if spec and spec_file or not spec and not spec_file:
        print("Must specify exactly one of --spec and --spec-file")
        sys.exit(1)
    _spec = {}
    if spec_file:
        with open(spec_file) as f:
            _spec.update(yaml.safe_load(f))

    if spec:
        _spec.update(json.loads(spec))

    if "cls" in _spec:  # single worker spec
        _spec = {_spec["opts"].get("name", 0): _spec}

following the logic above:

1. spec is not provided in yaml so first if is skipped
2. because ssh.py generates --spec json loads() will occur as part of 2nd if condition. This has cls with value distributed.cli.scheduler. Per error message, this is failing.
3. even if 2nd step was not failing, 3rd if condition will also evaluate to true even though cls is for scheduler not worker. There is no name attribute for scheduler class provided as an input option (ln#3679 def __init__() of scheduler.py`)

which then calls run_spec in spec.py. Since the cls only contained scheduler not worker specs, ln#637 of worker.py will evaluate to nothing. The fn will then return empty worker, which eventually causes ln#184 in ssh.py.

A few ways to solve the issue, imho, would be:
a. change ssh.py to provide provide full spec for both scheduler + worker as input to dask_spec; launch workers as part of scheduler OR
b. split creation of workers and schedulers (like it was in 2021.9.1 - i think?)

Either that or perhaps my syntax for worker and scheduler options is wrong .. hehehe

to allay the concern of incorrect input, if i add

logger.info("Worker = %s", str(workers))

just before return SpecCluster(blah blah), i see the following output:

distributed.deploy.ssh - INFO - Worker = {0: {'cls': <class 'distributed.deploy.ssh.Worker'>, 'options': {'address': 'localhost', 'connect_options': {'preferred_auth': 'publickey'}, 'kwargs': {'n_workers': 1, 'nthreads': 1, 'memory_limit': '1GB', 'name': 'WhateverIWant'}, 'worker_class': 'distributed.Nanny', 'remote_python': 'C:\\Users\\maulik\\miniconda3\\envs\\mls\\python.exe'}}}

(edit) @bryanwweber @jrbourbeau
the issue is definitely in ssh.py in the way the command is constructed using .join(). A modified list of arguments, when issued from command line, create the scheduler as so:

(mls) C:\Users\maulik\miniconda3\envs\mls\Lib\site-packages\distributed\cli>python.exe -m distributed.cli.dask_spec --spec "{\"cls\": \"distributed.Scheduler\", \"opts\": {\"port\": 8786}}"
C:\Users\maulik\miniconda3\envs\mls\lib\site-packages\distributed\cli\dask_spec.py:39: DeprecationWarning: There is no current event loop
  asyncio.get_event_loop().run_until_complete(run())
distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
distributed.scheduler - INFO - Clear task state
distributed.scheduler - INFO -   Scheduler at:    tcp://<bleh<:8786
distributed.scheduler - INFO -   dashboard at:                     :8787

Aborted!

Notice, the construction of string argument. Anything within ' ' is not liked by default JSONDecoder. It has to be " " within which all parameters + values have to be escaped :-)

Fix here's the temporary fix for Scheduler in ssh.py. Essentially a dumps of dumps.

        maulik_obj = {"cls": "distributed.Scheduler", "opts": self.kwargs}
        logger.info("Scheduler result       = %s", maulik_obj)
        required_obj = dumps(maulik_obj)
        logger.info("Scheduler dumps result = "+ required_obj)
        dumps_new_obj = dumps(required_obj)
        logger.info("Scheduler dumps result = "+ dumps_new_obj)
        cmd = " ".join(
            [
                set_env,
                self.remote_python,
                "-m",
                "distributed.cli.dask_spec",
                "--spec",
                "%s" % dumps_new_obj,
            ]
        )

Result

distributed.deploy.ssh - INFO - Scheduler result       = {'cls': 'distributed.Scheduler', 'opts': {'port': 8786}}
distributed.deploy.ssh - INFO - Scheduler dumps result = {"cls": "distributed.Scheduler", "opts": {"port": 8786}}
distributed.deploy.ssh - INFO - Scheduler dumps result = "{\"cls\": \"distributed.Scheduler\", \"opts\": {\"port\": 8786}}"
...
...
&& C:\Users\maulik\miniconda3\envs\mls\python.exe -m distributed.cli.dask_spec --spec "{\"cls\": \"distributed.Scheduler\", \"opts\": {\"port\": 8786}}"
distributed.deploy.ssh - INFO - C:\Users\maulik\miniconda3\envs\mls\lib\site-packages\distributed\cli\dask_spec.py:39: DeprecationWarning: There is no current event loop
distributed.deploy.ssh - INFO - asyncio.get_event_loop().run_until_complete(run())
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 - Clear task state
distributed.deploy.ssh - INFO - distributed.scheduler - INFO -   Scheduler at:    tcp://<ip>:8786

Same change in scheduler, and atleast the cluster will launc (1). There will however be a side-effect of the chosen design. As the number of workers increase, so will the size of the command. At some value (e.g. 100 workers) this is definitely bound to fail. To circumvent it, I modified ssh.py to output --spec-file instead of --spec (2). Modified files attached for distributed authors to glean over. It covers both cases (1) and (2) above.

mod_base_dask_code.zip