Repository: dask/distributed

Issue:
Client preload script
Today schedulers, workers, and nannies all have the option to have a preload script run automatically during their initialization process. This allows users to run custom setup code when various nodes in their cluster are spinning up. We should consider adding the same functionality for clients too.


Discussion:

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -52,6 +52,7 @@
 from tornado import gen
 from tornado.ioloop import PeriodicCallback
 
+from . import preloading
 from . import versions as version_module  # type: ignore
 from .batched import BatchedSend
 from .cfexecutor import ClientExecutor
@@ -751,6 +752,8 @@ class Client(SyncMethodMixin):
 
     _default_event_handlers = {"print": _handle_print, "warn": _handle_warn}
 
+    preloads: list[preloading.Preload]
+
     def __init__(
         self,
         address=None,
@@ -917,6 +920,10 @@ def __init__(
         for ext in extensions:
             ext(self)
 
+        preload = dask.config.get("distributed.client.preload")
+        preload_argv = dask.config.get("distributed.client.preload-argv")
+        self.preloads = preloading.process_preloads(self, preload, preload_argv)
+
         self.start(timeout=timeout)
         Client._instances.add(self)
 
@@ -1178,6 +1185,9 @@ async def _start(self, timeout=no_default, **kwargs):
         for topic, handler in Client._default_event_handlers.items():
             self.subscribe_topic(topic, handler)
 
+        for preload in self.preloads:
+            await preload.start()
+
         self._handle_report_task = asyncio.create_task(self._handle_report())
 
         return self
@@ -1459,6 +1469,9 @@ async def _close(self, fast=False):
 
         self.status = "closing"
 
+        for preload in self.preloads:
+            await preload.teardown()
+
         with suppress(AttributeError):
             for pc in self._periodic_callbacks.values():
                 pc.stop()
diff --git a/distributed/preloading.py b/distributed/preloading.py
--- a/distributed/preloading.py
+++ b/distributed/preloading.py
@@ -10,7 +10,7 @@
 from collections.abc import Iterable
 from importlib import import_module
 from types import ModuleType
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import click
 
@@ -19,6 +19,10 @@
 from .core import Server
 from .utils import import_file
 
+if TYPE_CHECKING:
+    # This has to be inside this guard to avoid a circular import
+    from .client import Client
+
 logger = logging.getLogger(__name__)
 
 
@@ -72,10 +76,10 @@ def validate_preload_argv(ctx, param, value):
 
 
 def is_webaddress(s: str) -> bool:
-    return any(s.startswith(prefix) for prefix in ("http://", "https://"))
+    return s.startswith(("http://", "https://"))
 
 
-def _import_module(name, file_dir=None) -> ModuleType:
+def _import_module(name: str, file_dir: str | None = None) -> ModuleType:
     """Imports module and extract preload interface functions.
 
     Import modules specified by name and extract 'dask_setup'
@@ -153,16 +157,20 @@ class Preload:
         Path of a directory where files should be copied
     """
 
-    dask_server: Server
+    dask_object: Server | Client
     name: str
     argv: list[str]
     file_dir: str | None
     module: ModuleType
 
     def __init__(
-        self, dask_server: Server, name: str, argv: Iterable[str], file_dir: str | None
+        self,
+        dask_object: Server | Client,
+        name: str,
+        argv: Iterable[str],
+        file_dir: str | None,
     ):
-        self.dask_server = dask_server
+        self.dask_object = dask_object
         self.name = name
         self.argv = list(argv)
         self.file_dir = file_dir
@@ -182,13 +190,13 @@ async def start(self):
                     "dask_setup", self.argv, allow_extra_args=False
                 )
                 result = dask_setup.callback(
-                    self.dask_server, *context.args, **context.params
+                    self.dask_object, *context.args, **context.params
                 )
                 if inspect.isawaitable(result):
                     await result
                 logger.info("Run preload setup click command: %s", self.name)
             else:
-                future = dask_setup(self.dask_server)
+                future = dask_setup(self.dask_object)
                 if inspect.isawaitable(future):
                     await future
                 logger.info("Run preload setup function: %s", self.name)
@@ -197,13 +205,13 @@ async def teardown(self):
         """Run when the server starts its close method"""
         dask_teardown = getattr(self.module, "dask_teardown", None)
         if dask_teardown:
-            future = dask_teardown(self.dask_server)
+            future = dask_teardown(self.dask_object)
             if inspect.isawaitable(future):
                 await future
 
 
 def process_preloads(
-    dask_server,
+    dask_server: Server | Client,
     preload: str | list[str],
     preload_argv: list[str] | list[list[str]],
     *,
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -27,6 +27,7 @@
 from hashlib import md5
 from importlib.util import cache_from_source
 from time import sleep
+from types import ModuleType
 from typing import Any as AnyType
 from typing import ClassVar
 
@@ -942,12 +943,12 @@ def open_port(host=""):
     return port
 
 
-def import_file(path):
+def import_file(path: str):
     """Loads modules for a file (.py, .zip, .egg)"""
     directory, filename = os.path.split(path)
     name, ext = os.path.splitext(filename)
-    names_to_import = []
-    tmp_python_path = None
+    names_to_import: list[str] = []
+    tmp_python_path: str | None = None
 
     if ext in (".py",):  # , '.pyc'):
         if directory not in sys.path:
@@ -963,7 +964,7 @@ def import_file(path):
         names = (mod_info.name for mod_info in pkgutil.iter_modules([path]))
         names_to_import.extend(names)
 
-    loaded = []
+    loaded: list[ModuleType] = []
     if not names_to_import:
         logger.warning("Found nothing to import from %s", filename)
     else:


Test patch:
diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py
--- a/distributed/tests/test_preload.py
+++ b/distributed/tests/test_preload.py
@@ -5,6 +5,7 @@
 import tempfile
 import urllib.error
 import urllib.request
+from textwrap import dedent
 from time import sleep
 
 import pytest
@@ -288,3 +289,88 @@ async def test_web_preload_worker(cleanup, worker_preload):
     async with Scheduler(port=8786, host="localhost") as s:
         async with Nanny(preload_nanny=["http://127.0.0.1:12346/preload"]) as nanny:
             assert nanny.scheduler_addr == s.address
+
+
+# This test is blocked on https://github.com/dask/distributed/issues/5819
+@pytest.mark.xfail(
+    reason="The preload argument to the client isn't supported yet", strict=True
+)
+@gen_cluster(nthreads=[])
+async def test_client_preload_text(s: Scheduler):
+    text = dedent(
+        """\
+        def dask_setup(client):
+            client.foo = "setup"
+
+
+        def dask_teardown(client):
+            client.foo = "teardown"
+        """
+    )
+    async with Client(address=s.address, asynchronous=True, preload=text) as c:
+        assert c.foo == "setup"
+    assert c.foo == "teardown"
+
+
+@gen_cluster(nthreads=[])
+async def test_client_preload_config(s):
+    text = dedent(
+        """\
+        def dask_setup(client):
+            client.foo = "setup"
+
+
+        def dask_teardown(client):
+            client.foo = "teardown"
+        """
+    )
+    with dask.config.set({"distributed.client.preload": [text]}):
+        async with Client(address=s.address, asynchronous=True) as c:
+            assert c.foo == "setup"
+        assert c.foo == "teardown"
+
+
+# This test is blocked on https://github.com/dask/distributed/issues/5819
+@pytest.mark.xfail(
+    reason="The preload argument to the client isn't supported yet", strict=True
+)
+@gen_cluster(nthreads=[])
+async def test_client_preload_click(s):
+    text = dedent(
+        """\
+        import click
+
+        @click.command()
+        @click.argument("value")
+        def dask_setup(client, value):
+            client.foo = value
+        """
+    )
+    value = "setup"
+    async with Client(
+        address=s.address, asynchronous=True, preload=text, preload_argv=[[value]]
+    ) as c:
+        assert c.foo == value
+
+
+@gen_cluster(nthreads=[])
+async def test_client_preload_config_click(s):
+    text = dedent(
+        """\
+        import click
+
+        @click.command()
+        @click.argument("value")
+        def dask_setup(client, value):
+            client.foo = value
+        """
+    )
+    value = "setup"
+    with dask.config.set(
+        {
+            "distributed.client.preload": [text],
+            "distributed.client.preload-argv": [[value]],
+        }
+    ):
+        async with Client(address=s.address, asynchronous=True) as c:
+            assert c.foo == value
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -1587,7 +1587,7 @@ async def test_custom_metrics(c, s, a, b):
 
 @gen_cluster(client=True)
 async def test_register_worker_callbacks(c, s, a, b):
-    # preload function to run
+    # plugin function to run
     def mystartup(dask_worker):
         dask_worker.init_variable = 1
 
@@ -1597,7 +1597,7 @@ def mystartup2():
         os.environ["MY_ENV_VALUE"] = "WORKER_ENV_VALUE"
         return "Env set."
 
-    # Check that preload function has been run
+    # Check that plugin function has been run
     def test_import(dask_worker):
         return hasattr(dask_worker, "init_variable")
         #       and dask_worker.init_variable == 1
@@ -1619,7 +1619,7 @@ def test_startup2():
     assert list(result.values()) == [False]
     await worker.close()
 
-    # Add a preload function
+    # Add a plugin function
     response = await c.register_worker_callbacks(setup=mystartup)
     assert len(response) == 2
 
@@ -1633,7 +1633,7 @@ def test_startup2():
     assert list(result.values()) == [True]
     await worker.close()
 
-    # Register another preload function
+    # Register another plugin function
     response = await c.register_worker_callbacks(setup=mystartup2)
     assert len(response) == 2