`Scheduler`/`Worker` threads hang when calling `sys.exit()`

[jacobtomlinson]

Describe the issue:

When using the Scheduler or Worker class to start cluster components if the program is exited with sys.exit() (as is done in dask-mpi) the Python process hangs, likely due to a background thread holding the process open.

Minimal Complete Verifiable Example:

import sys
from distributed import Client, Scheduler
from distributed.utils import LoopRunner

async def main():
    async with Scheduler() as scheduler:
        async with Client(scheduler.address, asynchronous=True) as client:
            await client.shutdown()
    print("Done, exiting")
    sys.exit()  # Hangs at this line, comment this out and the program exits as expected

loop_runner = LoopRunner(loop=None, asynchronous=False)
loop_runner.run_sync(main)

I tried to strip things down as far as possible to still reproduce the issue, but I note this doesn't happen when using asyncio.run(main) instead of the LoopRunner that is commonly used in distributed.

Anything else we need to know?:

Environment:

• Dask version: 2024.4.2
• Python version: 3.11.9
• Operating System: Ubuntu
• Install method (conda, pip, source):

[jacobtomlinson]

I found a workaround for this behaviour using a signal to defer the shutdown. This example does not hang.

+import os
+import signal
 import sys
 from distributed import Client, Scheduler
 from distributed.utils import LoopRunner

 async def main():
     async with Scheduler() as scheduler:
         async with Client(scheduler.address, asynchronous=True) as client:
             await client.shutdown()
     print("Done, exiting")
-    sys.exit()  # Hangs at this line, comment this out and the program exits as expected
+    os.kill(os.getpid(), signal.SIGINT)  # Shutdown using a signal instead

+signal.signal(signal.SIGINT, lambda *_: sys.exit())  # Exit gracefully on signal
 loop_runner = LoopRunner(loop=None, asynchronous=False)
 loop_runner.run_sync(main)