创建LocalCluster遇中断时如何自动关闭Dask集群与客户端？

问题背景 #

我正在编写一个用Dask处理数据的Python函数，想用with ... as:语句确保任何错误或异常发生时自动关闭Dask集群和客户端，函数结构如下：

def func(input:str, # path to input
         output:str, # path to output
        ):
    with LocalCluster() as cluster, Client(cluster) as client:
        # 加载输入数据
        # 构建计算图
        da.compute([...])

在数据处理阶段（执行da.compute()）触发KeyboardInterrupt时，集群和客户端能正常关闭，但在创建集群阶段（调用LocalCluster()）触发中断时，集群无法关闭，Worker会残留，下次运行函数时出现端口占用警告：

/users/kangl/miniconda3/envs/work/lib/python3.10/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.
Perhaps you already have a cluster running?
Hosting the HTTP server on port 37963 instead
  warnings.warn(

需要解决的问题是：如何在创建LocalCluster时发生中断时自动关闭集群？

解决方案 #

拆分with语句并添加异常捕获 #

把集群和客户端的创建拆分成独立的with块，同时在最外层包裹try/except捕获KeyboardInterrupt，确保即使在集群初始化阶段中断，也能手动触发关闭：

def func(input: str, output: str):
    try:
        with LocalCluster() as cluster:
            with Client(cluster) as client:
                # 数据处理逻辑
                da.compute([...])
    except KeyboardInterrupt:
        # 手动关闭集群（如果集群已初始化）
        if 'cluster' in locals():
            cluster.close()
        raise  # 重新抛出中断，保持原有行为

这种方式覆盖了集群创建过程中触发的中断场景，try块包裹整个集群初始化和处理流程，触发KeyboardInterrupt后会进入except块，检查集群是否已创建，若已创建则执行关闭操作。

使用signal模块注册中断处理函数 #

提前注册信号处理函数，当收到SIGINT（键盘中断）时，检查是否存在已创建的集群实例并关闭：

import signal

def func(input: str, output: str):
    cluster = None
    def handle_interrupt(signum, frame):
        nonlocal cluster
        if cluster is not None:
            cluster.close()
        raise KeyboardInterrupt

    # 注册信号处理
    original_handler = signal.signal(signal.SIGINT, handle_interrupt)
    try:
        with LocalCluster() as cluster:
            with Client(cluster) as client:
                da.compute([...])
    finally:
        # 恢复原信号处理函数
        signal.signal(signal.SIGINT, original_handler)

这种方法能更早捕获中断信号，即使在LocalCluster初始化的底层过程中触发中断，也能通过信号处理函数尝试关闭已部分创建的集群。

清理所有Dask残留进程 #

如果端口占用问题频繁出现，可以使用dask.distributed的cleanup工具彻底清理残留进程：

from distributed import cleanup

def func(input: str, output: str):
    try:
        with LocalCluster() as cluster, Client(cluster) as client:
            da.compute([...])
    except KeyboardInterrupt:
        cleanup()  # 清理当前用户下所有Dask相关进程
        raise

cleanup()函数会扫描并终止当前用户下的Dask Worker、Scheduler等相关进程，彻底解决残留问题，但注意这会关闭当前用户所有的Dask集群，适合单任务场景。

内容的提问来源于stack exchange，提问作者Kang Liang