diff --git a/distributed/__init__.py b/distributed/__init__.py index b624db8ad1..2725aa96e5 100644 --- a/distributed/__init__.py +++ b/distributed/__init__.py @@ -36,6 +36,7 @@ from distributed.deploy import ( Adaptive, LocalCluster, + LocalEnvCluster, SpecCluster, SSHCluster, SubprocessCluster, diff --git a/distributed/deploy/__init__.py b/distributed/deploy/__init__.py index ffbeb6a8c6..3e3dbbd620 100644 --- a/distributed/deploy/__init__.py +++ b/distributed/deploy/__init__.py @@ -3,6 +3,7 @@ from distributed.deploy.adaptive import Adaptive from distributed.deploy.cluster import Cluster from distributed.deploy.local import LocalCluster +from distributed.deploy.local_env import LocalEnvCluster from distributed.deploy.spec import ProcessInterface, SpecCluster from distributed.deploy.ssh import SSHCluster from distributed.deploy.subprocess import SubprocessCluster diff --git a/distributed/deploy/local_env.py b/distributed/deploy/local_env.py new file mode 100644 index 0000000000..2538b8a502 --- /dev/null +++ b/distributed/deploy/local_env.py @@ -0,0 +1,274 @@ +import asyncio +import logging +from typing import Any as AnyType + +import dask +import dask.config + +from distributed.compatibility import WINDOWS +from distributed.core import Status +from distributed.scheduler import Scheduler as _Scheduler +from distributed.utils import cli_keywords +from distributed.worker import Worker as _Worker + +from .spec import ProcessInterface, SpecCluster + +logger = logging.getLogger(__name__) + + +class LocalEnvProcess(ProcessInterface): + """A superclass for Workers and Nannies run by a specified Python executable + + See Also + -------- + Worker + Scheduler + """ + + def __init__(self, **kwargs): + if WINDOWS: + raise RuntimeError( + "LocalEnvProcess relies on subprocesses, which are not supported on Windows." + ) + self.proc = None + super().__init__(**kwargs) + + async def start(self): + await super().start() + + async def close(self): + try: + self.proc.kill() + except ProcessLookupError: + pass + await super().close() + + def __repr__(self): + return f"" + + async def _set_env_helper(self): + """Helper function to locate existing dask internal config for the remote + Scheduler and Workers to inherit when started. + + Returns + ------- + Dask config to inherit, if any. + """ + proc = await asyncio.create_subprocess_shell("uname", **self.connect_options) + await proc.communicate() + if proc.returncode == 0: + set_env = f'env DASK_INTERNAL_INHERIT_CONFIG="{dask.config.serialize(dask.config.global_config)}"' + else: # pragma: nocover + raise RuntimeError(f"{self.__class__.__name__} failed to set DASK_INTERNAL_INHERIT_CONFIG variable") + return set_env + + async def _get_address(self, search_str): + # We watch stderr in order to get the address, then we return + name = self.__class__.__name__ + while True: + line = await self.proc.stderr.readline() + if not line.decode("ascii").strip(): + raise RuntimeError(f"{name} failed to start") + else: + line = line.decode("ascii").strip() + logger.info(line) + if search_str in line: + self.address = line.split(f"{search_str}:")[1].strip() + if name == "Worker": + self.status = Status.running + break + logger.debug("%s", line) + + +class LocalEnvWorker(LocalEnvProcess): + """A Remote Dask Worker run by a specified Python executable + + Parameters + ---------- + scheduler: str + The address of the scheduler + python_executable: str + Full path to Python executable to run this worker + connect_options: dict + kwargs to be passed to asyncio subprocess connections + kwargs: dict + These will be passed through the dask-worker CLI to the + dask.distributed.Worker class + worker_module: str + The python module to run to start the worker + name: str + Optionally specify a name for this worker + """ + + def __init__( + self, + scheduler: str, + python_executable: str, + connect_options: dict[str, AnyType] | None = None, + worker_module: str = "distributed.cli.dask_worker", + name: None = None, + kwargs: dict[str, AnyType] | None = None, + ): + super().__init__() + + self.scheduler = scheduler + self.python_executable = python_executable + self.connect_options = connect_options if connect_options is not None else {} + self.kwargs = kwargs if kwargs is not None else {} + self.worker_module = worker_module + self.name = name + + async def start(self): + set_env = await self._set_env_helper() + + cmd = " ".join( + [ + set_env, + self.python_executable, + "-m", + self.worker_module, + self.scheduler, + "--name", + str(self.name), + ] + + cli_keywords(self.kwargs, cls=_Worker, cmd=self.worker_module) + ) + + self.proc = await asyncio.create_subprocess_shell( + cmd, stderr=asyncio.subprocess.PIPE, **self.connect_options + ) + + search_string = "worker at" + await self._get_address(search_string) + await super().start() + + +class LocalEnvScheduler(LocalEnvProcess): + """A Remote Dask Scheduler run by a specified Python executable + + Parameters + ---------- + python_executable: str + Full path to Python executable to run this scheduler + connect_options: dict + kwargs to be passed to asyncio subprocess connections + kwargs: dict + These will be passed through the dask-scheduler CLI to the + dask.distributed.Scheduler class + """ + + def __init__( + self, + python_executable: str, + connect_options: dict, + kwargs: dict[str, AnyType] | None = None, + ): + super().__init__() + + self.python_executable = python_executable + self.kwargs = kwargs if kwargs is not None else {} + self.connect_options = connect_options if connect_options is not None else {} + + async def start(self): + logger.debug("Created Scheduler") + + set_env = await self._set_env_helper() + + cmd = " ".join( + [set_env, self.python_executable, "-m", "distributed.cli.dask_scheduler"] + + cli_keywords(self.kwargs, cls=_Scheduler) + ) + + self.proc = await asyncio.create_subprocess_shell( + cmd, stderr=asyncio.subprocess.PIPE, **self.connect_options + ) + + search_string = "Scheduler at" + await self._get_address(search_string) + await super().start() + + +def LocalEnvCluster( + python_executable: str, + n_workers: int = 1, + connect_options: list[dict] | dict | None = None, + worker_options: dict | None = None, + scheduler_options: dict | None = None, + worker_module: str = "distributed.cli.dask_worker", + **kwargs: AnyType, +) -> SpecCluster: + """Deploy a Dask cluster that utilises a different Python executable + + The LocalEnvCluster function deploys a Dask Scheduler and Workers for you on + your local machine, running in the Python environment specified by + `python_executable`. This allows you to run a Dask cluster in a different + Python environment to the host Python environment; particularly useful when + combined with `dask-labextension` as you can run a Dask Scheduler and + Workers in a different Python env to the env hosting JupyterLab. + + You may configure the scheduler and workers by passing + ``scheduler_options`` and ``worker_options`` dictionary keywords. See the + ``dask.distributed.Scheduler`` and ``dask.distributed.Worker`` classes for + details on the available options, but the defaults should work in most + situations. + + You may configure how to connect to the local Scheduler and Workers using + the ``connect_options`` keyword, which passes values to the + ``asyncio.create_subprocess_shell`` function. For more information on this + see the documentation on the + [asyncio library](https://docs.python.org/3/library/asyncio-subprocess.html#asyncio-subprocess). + + Parameters + ---------- + python_executable : str + Full path to a Python executable, for example from another + conda env or Python venv. + n_workers : int + Number of workers to start + connect_options : dict, optional + Keywords to pass through to :func:`asyncio.create_subprocess_shell`. + See docs for :func:`asyncio.create_subprocess_shell` for full information. + worker_options : dict, optional + Keywords to pass on to workers. + scheduler_options : dict, optional + Keywords to pass on to scheduler. + worker_module : str, optional + Python module to call to start the worker. + + Example + ------- + >>> from dask.distributed import Client, LocalEnvCluster + >>> cluster = LocalEnvCluster( + ... "/Users/user/miniconda3/envs/myenv/bin/python", + ... worker_options={"nthreads": 1}, + ... ) + >>> client = Client(cluster) + + See Also + -------- + dask.distributed.Scheduler + dask.distributed.Worker + asyncio.create_subprocess_shell + """ + + scheduler = { + "cls": LocalEnvScheduler, + "options": { + "python_executable": python_executable, + "connect_options": connect_options, + "kwargs": scheduler_options, + }, + } + workers = { + i: { + "cls": LocalEnvWorker, + "options": { + "python_executable": python_executable, + "connect_options": connect_options, + "kwargs": worker_options, + "worker_module": worker_module, + }, + } + for i in range(n_workers) + } + return SpecCluster(workers, scheduler, name="LocalEnvCluster", **kwargs) diff --git a/distributed/deploy/tests/test_local_env.py b/distributed/deploy/tests/test_local_env.py new file mode 100644 index 0000000000..53f015f7f4 --- /dev/null +++ b/distributed/deploy/tests/test_local_env.py @@ -0,0 +1,93 @@ +import sys + +import pytest + +import dask + +from distributed import Client +from distributed.compatibility import WINDOWS +from distributed.core import Status +from distributed.deploy.local_env import LocalEnvCluster +from distributed.utils_test import gen_test + + +@pytest.mark.skipif(WINDOWS, reason="distributed#7434") +@gen_test() +async def test_basic(): + async with LocalEnvCluster( + sys.executable, + asynchronous=True, + scheduler_options={"idle_timeout": "5s", "port": 0, "dashboard_address": ":0"}, + worker_options={"death_timeout": "5s"}, + ) as cluster: + assert len(cluster.workers) == 1 + assert cluster.scheduler.python_executable == sys.executable + assert cluster.workers[0].python_executable == sys.executable + assert "LocalEnv" in repr(cluster) + assert cluster.status == Status.closed + +@pytest.mark.skipif(WINDOWS, reason="distributed#7434") +@gen_test() +async def test_job_submission(): + async with LocalEnvCluster( + sys.executable, + asynchronous=True, + scheduler_options={"idle_timeout": "5s", "port": 0, "dashboard_address": ":0"}, + worker_options={"death_timeout": "5s"}, + ) as cluster: + async with Client(cluster, asynchronous=True) as client: + result = await client.submit(lambda x: x + 1, 10) + assert result == 11 + +@pytest.mark.skipif(WINDOWS, reason="distributed#7434") +@gen_test() +async def test_multiple_workers(): + n_workers = 2 + async with LocalEnvCluster( + sys.executable, + n_workers=n_workers, + asynchronous=True, + scheduler_options={"idle_timeout": "5s", "port": 0, "dashboard_address": ":0"}, + worker_options={"death_timeout": "5s"}, + ) as cluster: + assert len(cluster.workers) == n_workers + +@pytest.mark.skipif(WINDOWS, reason="distributed#7434") +@gen_test() +async def test_bad_executable(): + with pytest.raises(RuntimeError): + async with LocalEnvCluster( + "/foo/bar/baz/python", + asynchronous=True, + scheduler_options={ + "idle_timeout": "5s", + "port": 0, + "dashboard_address": ":0", + }, + worker_options={"death_timeout": "5s"}, + ) as cluster: + assert cluster + cluster.close() + +@pytest.mark.skipif(WINDOWS, reason="distributed#7434") +@gen_test() +async def test_set_env(): + value = 100 + + def f(): + return dask.config.get("foo") + + with dask.config.set(foo=value): + async with LocalEnvCluster( + sys.executable, + asynchronous=True, + scheduler_options={ + "idle_timeout": "5s", + "port": 0, + "dashboard_address": ":0", + }, + worker_options={"death_timeout": "5s"}, + ) as cluster: + async with Client(cluster, asynchronous=True) as client: + result = await client.submit(f) + assert result == value diff --git a/pyproject.toml b/pyproject.toml index 8876a833ce..c23649e243 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,9 +138,13 @@ filterwarnings = [ '''ignore:unclosed event loop <_(Unix|Windows)SelectorEventLoop.*:ResourceWarning''', '''ignore:unclosed file <_io.BufferedWriter.*:ResourceWarning''', '''ignore:unclosed file <_io.TextIOWrapper.*:ResourceWarning''', + '''ignore:unclosed file <_io.FileIO.*:ResourceWarning''', '''ignore:unclosed transport <_SelectorSocketTransport.*:ResourceWarning''', '''ignore:unclosed transport .*sem_unlink.*FileNotFoundError:pytest.PytestUnraisableExceptionWarning''', '''ignore:(?s)Exception ignored in.