Repository: dask/distributed

Issue:
add graceful signal handling to `dask spec`
It appears that this is missing from `dask spec` by oversight

There is a pattern one can use, see [dask/distributed@de3f755/distributed/cli/dask_worker.py#L418-L455](https://github.com/dask/distributed/blob/de3f755243ecd4ebbe436af7a9ba9ea6098d6b3e/distributed/cli/dask_worker.py#L418-L455)

that needs to be applied to dask-spec here [dask/distributed@de3f755/distributed/cli/dask_spec.py#L37-L44](https://github.com/dask/distributed/blob/de3f755243ecd4ebbe436af7a9ba9ea6098d6b3e/distributed/cli/dask_spec.py#L37-L44)
add graceful signal handling to `dask spec`
It appears that this is missing from `dask spec` by oversight

There is a pattern one can use, see [dask/distributed@de3f755/distributed/cli/dask_worker.py#L418-L455](https://github.com/dask/distributed/blob/de3f755243ecd4ebbe436af7a9ba9ea6098d6b3e/distributed/cli/dask_worker.py#L418-L455)

that needs to be applied to dask-spec here [dask/distributed@de3f755/distributed/cli/dask_spec.py#L37-L44](https://github.com/dask/distributed/blob/de3f755243ecd4ebbe436af7a9ba9ea6098d6b3e/distributed/cli/dask_spec.py#L37-L44)


Discussion:

Resolving patch:
diff --git a/distributed/cli/dask_spec.py b/distributed/cli/dask_spec.py
--- a/distributed/cli/dask_spec.py
+++ b/distributed/cli/dask_spec.py
@@ -7,9 +7,10 @@
 import click
 import yaml
 
+from distributed._signals import wait_for_signals
 from distributed.compatibility import asyncio_run
 from distributed.config import get_loop_factory
-from distributed.deploy.spec import run_spec
+from distributed.deploy.spec import init_spec
 
 
 @click.command(name="spec", context_settings=dict(ignore_unknown_options=True))
@@ -34,13 +35,31 @@ def main(args: list, spec: str, spec_file: str) -> None:
     if "cls" in _spec:  # single worker spec
         _spec = {_spec["opts"].get("name", 0): _spec}
 
-    async def run():
-        servers = await run_spec(_spec, *args)
-        try:
+    async def run() -> None:
+        servers = init_spec(_spec, *args)
+
+        async def wait_for_servers_to_finish() -> None:
+            await asyncio.gather(*servers.values())
             await asyncio.gather(*(w.finished() for w in servers.values()))
-        except KeyboardInterrupt:
+
+        async def wait_for_signals_and_close() -> None:
+            await wait_for_signals()
             await asyncio.gather(*(w.close() for w in servers.values()))
 
+        wait_for_signals_and_close_task = asyncio.create_task(
+            wait_for_signals_and_close()
+        )
+        wait_for_servers_to_finish_task = asyncio.create_task(
+            wait_for_servers_to_finish()
+        )
+
+        done, _ = await asyncio.wait(
+            [wait_for_signals_and_close_task, wait_for_servers_to_finish_task],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        # Re-raise exceptions from done tasks
+        [task.result() for task in done]
+
     asyncio_run(run(), loop_factory=get_loop_factory())
 
 
diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
--- a/distributed/deploy/spec.py
+++ b/distributed/deploy/spec.py
@@ -674,14 +674,18 @@ def from_name(cls, name: str) -> ProcessInterface:
         raise NotImplementedError()
 
 
-async def run_spec(spec: dict[str, Any], *args: Any) -> dict[str, Worker | Nanny]:
+def init_spec(spec: dict[str, Any], *args: Any) -> dict[str, Worker | Nanny]:
     workers = {}
     for k, d in spec.items():
         cls = d["cls"]
         if isinstance(cls, str):
             cls = import_term(cls)
         workers[k] = cls(*args, **d.get("opts", {}))
+    return workers
 
+
+async def run_spec(spec: dict[str, Any], *args: Any) -> dict[str, Worker | Nanny]:
+    workers = init_spec(spec, *args)
     if workers:
         await asyncio.gather(*workers.values())
     return workers


Test patch:
diff --git a/distributed/cli/tests/test_dask_spec.py b/distributed/cli/tests/test_dask_spec.py
--- a/distributed/cli/tests/test_dask_spec.py
+++ b/distributed/cli/tests/test_dask_spec.py
@@ -1,8 +1,15 @@
 from __future__ import annotations
 
+import asyncio
+import json
+import signal
+import subprocess
+
+import pytest
 import yaml
 
 from distributed import Client
+from distributed.compatibility import WINDOWS
 from distributed.utils import open_port
 from distributed.utils_test import gen_cluster, gen_test, popen
 
@@ -82,3 +89,73 @@ def test_errors():
         line = proc.stdout.readline().decode()
         assert "exactly one" in line
         assert "--spec" in line and "--spec-file" in line
+
+
+@pytest.mark.skipif(WINDOWS, reason="POSIX only")
+@pytest.mark.parametrize("worker_type", ["Nanny", "Worker"])
+@pytest.mark.parametrize("sig", [signal.SIGINT, signal.SIGTERM])
+@gen_cluster(client=True, nthreads=[])
+async def test_signal_handling_worker(c, s, worker_type, sig):
+    worker = await asyncio.create_subprocess_exec(
+        "dask",
+        "spec",
+        "--spec",
+        json.dumps(
+            {
+                "cls": f"dask.distributed.{worker_type}",
+                "opts": {"scheduler_ip": s.address, "nthreads": 3},
+            },
+        ),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    try:
+        await c.wait_for_workers(1)
+
+        worker.send_signal(sig)
+        stdout, stderr = await worker.communicate()
+        assert stdout == b""
+        logs = stderr.decode().lower()
+        assert sig.name.lower() in logs
+        assert worker.returncode == 0
+        if worker_type == "Nanny":
+            assert "closing nanny" in logs
+            assert "stopping worker" in logs
+        else:
+            assert "nanny" not in logs
+        assert "timed out" not in logs
+        assert "error" not in logs
+        assert "exception" not in logs
+    finally:
+        await worker.wait()
+
+
+@pytest.mark.skipif(WINDOWS, reason="POSIX only")
+@pytest.mark.parametrize("sig", [signal.SIGINT, signal.SIGTERM])
+@gen_test()
+async def test_signal_handling_scheduler(sig):
+    port = open_port()
+    scheduler = await asyncio.create_subprocess_exec(
+        "dask",
+        "spec",
+        "--spec",
+        json.dumps({"cls": "dask.distributed.Scheduler", "opts": {"port": port}}),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    try:
+        async with Client(f"tcp://localhost:{port}", asynchronous=True) as client:
+            pass
+
+        scheduler.send_signal(sig)
+        stdout, stderr = await scheduler.communicate()
+        assert stdout == b""
+        logs = stderr.decode().lower()
+        assert sig.name.lower() in logs
+        assert scheduler.returncode == 0
+        assert "scheduler closing all comms" in logs
+        assert "timed out" not in logs
+        assert "error" not in logs
+        assert "exception" not in logs
+    finally:
+        await scheduler.wait()