Repository: dask/distributed

Issue:
`Event loop stopped before Future completed` on Jupyter shutdown 
If I run JupyterLab on the scheduler with `dask scheduler --jupyter`, then navigate to the the JupyterLab session (same as Bokeh dashboard link, but with `jupyter/lab` instead of `status`), and then shutdown the Jupyter session (File > Shut Down), I get the following output 

```
[I 2023-09-27 13:34:08.341 ServerApp] Shutting down on /api/shutdown request.
[I 2023-09-27 13:34:08.341 ServerApp] Shutting down 4 extensions
2023-09-27 13:34:08,344 - distributed.scheduler - INFO - End scheduler
Traceback (most recent call last):
  File "/Users/james/mambaforge/envs/distributed/bin/dask-scheduler", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/Users/james/mambaforge/envs/distributed/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/james/mambaforge/envs/distributed/lib/python3.11/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/Users/james/mambaforge/envs/distributed/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/james/mambaforge/envs/distributed/lib/python3.11/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/james/projects/dask/distributed/distributed/cli/dask_scheduler.py", line 251, in main
    asyncio_run(run(), loop_factory=get_loop_factory())
  File "/Users/james/projects/dask/distributed/distributed/compatibility.py", line 204, in asyncio_run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/Users/james/mambaforge/envs/distributed/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/james/mambaforge/envs/distributed/lib/python3.11/asyncio/base_events.py", line 651, in run_until_complete
    raise RuntimeError('Event loop stopped before Future completed.')
RuntimeError: Event loop stopped before Future completed.
```

@graingert any idea what might be going on here with the event loop? 


Discussion:
this is because jupyter_server's shutdown handler abruptly stops the event loop: https://github.com/jupyter-server/jupyter_server/blob/cf8df8a6b26455977f1eb946dc3dead069981ea4/jupyter_server/services/shutdown.py#L24

we could override the shutdown handler to do something else (eg call `await scheduler.close(...)`)

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -35,6 +35,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, cast, overload
 
 import psutil
+import tornado.web
 from sortedcontainers import SortedDict, SortedSet
 from tlz import (
     concat,
@@ -3589,6 +3590,7 @@ def __init__(
             distributed.dashboard.scheduler.connect(
                 self.http_application, self.http_server, self, prefix=http_prefix
             )
+        scheduler = self
         if self.jupyter:
             try:
                 from jupyter_server.serverapp import ServerApp
@@ -3599,6 +3601,30 @@ def __init__(
                 )
             from traitlets.config import Config
 
+            """HTTP handler to shut down the Jupyter server.
+            """
+            try:
+                from jupyter_server.auth import authorized
+            except ImportError:
+
+                def authorized(c):
+                    return c
+
+            from jupyter_server.base.handlers import JupyterHandler
+
+            class ShutdownHandler(JupyterHandler):
+                """A shutdown API handler."""
+
+                auth_resource = "server"
+
+                @tornado.web.authenticated
+                @authorized
+                async def post(self):
+                    """Shut down the server."""
+                    self.log.info("Shutting down on /api/shutdown request.")
+
+                    await scheduler.close(reason="shutdown requested via Jupyter")
+
             j = ServerApp.instance(
                 config=Config(
                     {
@@ -3620,6 +3646,11 @@ def __init__(
                 argv=[],
             )
             self._jupyter_server_application = j
+            shutdown_app = tornado.web.Application(
+                [(r"/jupyter/api/shutdown", ShutdownHandler)]
+            )
+            shutdown_app.settings = j.web_app.settings
+            self.http_application.add_application(shutdown_app)
             self.http_application.add_application(j.web_app)
 
         # Communication state


Test patch:
diff --git a/distributed/tests/test_jupyter.py b/distributed/tests/test_jupyter.py
--- a/distributed/tests/test_jupyter.py
+++ b/distributed/tests/test_jupyter.py
@@ -1,27 +1,24 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
+import subprocess
 from datetime import datetime, timedelta, timezone
 from time import perf_counter
 
 import pytest
+from tornado.httpclient import AsyncHTTPClient
 
+from distributed import Client, Scheduler
 from distributed.core import Status
+from distributed.utils import open_port
+from distributed.utils_test import gen_test, popen
 
 pytest.importorskip("requests")
-
 import requests
 
-from distributed import Client
-
 pytest.importorskip("jupyter_server")
 
-from tornado.httpclient import AsyncHTTPClient
-
-from distributed import Scheduler
-from distributed.utils import open_port
-from distributed.utils_test import gen_test, popen
-
 pytestmark = pytest.mark.filterwarnings("ignore:Jupyter is migrating its paths")
 
 
@@ -100,3 +97,48 @@ async def test_jupyter_idle_timeout_returned():
         assert s.check_idle() is None
         # ^ NOTE: this probably should be `== next_idle`;
         # see discussion in https://github.com/dask/distributed/pull/7687#discussion_r1145095196
+
+
+@pytest.mark.slow
+def test_shutsdown_cleanly(loop):
+    port = open_port()
+    with concurrent.futures.ThreadPoolExecutor() as tpe:
+        subprocess_fut = tpe.submit(
+            subprocess.run,
+            [
+                "dask",
+                "scheduler",
+                "--jupyter",
+                "--no-dashboard",
+                "--host",
+                f"127.0.0.1:{port}",
+            ],
+            check=True,
+            capture_output=True,
+            timeout=10,
+            encoding="utf8",
+        )
+
+        # wait until scheduler is running
+        with Client(f"127.0.0.1:{port}", loop=loop):
+            pass
+
+        with requests.Session() as session:
+            session.get("http://127.0.0.1:8787/jupyter/lab").raise_for_status()
+            session.post(
+                "http://127.0.0.1:8787/jupyter/api/shutdown",
+                headers={"X-XSRFToken": session.cookies["_xsrf"]},
+            ).raise_for_status()
+
+        stderr = subprocess_fut.result().stderr
+        assert "Traceback" not in stderr
+        assert (
+            "distributed.scheduler - INFO - Scheduler closing due to shutdown "
+            "requested via Jupyter...\n" in stderr
+        )
+        assert "Shutting down on /api/shutdown request.\n" in stderr
+        assert (
+            "distributed.scheduler - INFO - Stopped scheduler at "
+            f"'tcp://127.0.0.1:{port}'\n" in stderr
+        )
+        assert stderr.endswith("distributed.scheduler - INFO - End scheduler\n")