Support async preload click commands (#4170)

jrbourbeau · web-flow · commit 3e5b506ed94d · 2020-10-20T12:04:35.000-05:00
diff --git a/distributed/preloading.py b/distributed/preloading.py
@@ -172,7 +172,12 @@ async def start(self):
                 context = dask_setup.make_context(
                     "dask_setup", list(self.argv), allow_extra_args=False
                 )
-                dask_setup.callback(self.dask_server, *context.args, **context.params)
+                result = dask_setup.callback(
+                    self.dask_server, *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)
                 if inspect.isawaitable(future):
diff --git a/distributed/tests/test_preload.py b/distributed/tests/test_preload.py
@@ -106,6 +106,36 @@ def check_worker():
         shutil.rmtree(tmpdir)
 
 
+@pytest.mark.asyncio
+async def test_worker_preload_click(cleanup, tmpdir):
+    CLICK_PRELOAD_TEXT = """
+import click
+
+@click.command()
+def dask_setup(worker):
+    worker.foo = 'setup'
+"""
+    async with Scheduler(port=0) as s:
+        async with Worker(s.address, preload=CLICK_PRELOAD_TEXT) as w:
+            assert w.foo == "setup"
+
+
+@pytest.mark.asyncio
+async def test_worker_preload_click_async(cleanup, tmpdir):
+    # Ensure we allow for click commands wrapping coroutines
+    # https://github.com/dask/distributed/issues/4169
+    CLICK_PRELOAD_TEXT = """
+import click
+
+@click.command()
+async def dask_setup(worker):
+    worker.foo = 'setup'
+"""
+    async with Scheduler(port=0) as s:
+        async with Worker(s.address, preload=CLICK_PRELOAD_TEXT) as w:
+            assert w.foo == "setup"
+
+
 @pytest.mark.asyncio
 async def test_preload_import_time(cleanup):
     text = """
