如何可靠检测Python代码是否运行于Dask Worker实例中？

针对Dask distributed集群（如LocalCluster）环境下，尤其是Worker初始化阶段的代码，以下是几种不受contextvars更新影响的可靠检测方式：

1. 检查Dask Worker专属环境变量 #

Dask Worker启动时会自动注入专属环境变量，直接检查这些变量是否存在即可快速判断：

import os

def is_on_dask_worker():
    return "DASK_WORKER_NAME" in os.environ or "DASK_WORKER_ADDRESS" in os.environ

该方法轻量化且稳定，Worker进程启动时就会设置这些变量，无论初始化代码还是任务代码都能生效。

2. 尝试获取当前Worker实例 #

通过distributed.worker.get_worker函数尝试获取Worker对象，成功则说明处于Worker环境：

from distributed.worker import get_worker

def is_on_dask_worker():
    try:
        get_worker()
        return True
    except ValueError:
        return False

此方法直接关联Dask Worker实例，适合需要与Worker交互的场景（如初始化时配置资源），完全不受contextvars变更影响。

3. 验证进程名称（补充方案） #

Dask Worker进程名称通常包含dask-worker标识，可通过进程名称辅助判断：

import psutil

def is_on_dask_worker():
    current_proc = psutil.Process()
    return "dask-worker" in current_proc.name()

注：此方法依赖进程命名规则，可靠性略低于前两种，可作为辅助验证手段。

推荐场景 #

• Worker初始化阶段：优先使用环境变量检查，性能开销最小；
• 需要操作Worker实例时：使用获取Worker对象的方法，直接关联实例逻辑。

内容的提问来源于stack exchange，提问作者Alessio Arena