Dask Worker与主线程导入不一致的通用解决方案咨询

问题描述 #

我有一个依赖子模块选项的Dask Delayed函数，第三个模块会在导入时修改这些选项。如果把导入操作放在__name__ == '__main__'之后（比如Notebook环境），分布式客户端运行函数时会直接忽略已修改的选项。有没有通用方法能让Worker在执行计算前，和主线程完成完全相同的导入操作？

最小复现示例 #

三个Python模块：

constants.py:

N = 1

add.py:

import dask
import constants as c

@dask.delayed
def add(da):
    return da + c.N

overhead.py:

import constants

constants.N = 4

测试场景：

1. 脚本模式（导入在顶部）：输出5，符合预期

from dask.distributed import Client
import dask
import add
import overhead

if __name__ == '__main__':
    c = Client(n_workers=2, threads_per_worker=2)
    print(dask.compute(add.add(1))[0])

2. Notebook/__main__内导入：输出2，不符合预期

from dask.distributed import Client
import dask

if __name__ == '__main__':
    import add
    import overhead
    c = Client(n_workers=2, threads_per_worker=2)
    print(dask.compute(add.add(1)))

3. 现有手动方案：有效但需指定模块，不够通用

from dask.distributed import Client
import dask

if __name__ == '__main__':
    import add
    import overhead
    c = Client(n_workers=2, threads_per_worker=2)
    c.run(lambda: overhead)
    print(dask.compute(add.add(1)))

通用解决方案 #

方法1：同步主线程已导入模块到Worker #

通过捕获主线程中已导入的自定义模块，让Worker重新导入这些模块，触发配置修改的代码执行。可以写一个通用函数来处理：

from dask.distributed import Client
import dask
import sys
import importlib

def sync_all_imports():
    # 筛选需要同步的模块（排除系统内置模块，可根据需求调整规则）
    target_modules = [
        mod_name for mod_name in sys.modules 
        if not mod_name.startswith('_') and '.' not in mod_name
    ]
    for name in target_modules:
        # 重新加载模块，执行其顶层代码
        importlib.reload(sys.modules[name])

if __name__ == '__main__':
    import add
    import overhead
    c = Client(n_workers=2, threads_per_worker=2)
    # 在所有Worker上执行同步逻辑
    c.run(sync_all_imports)
    print(dask.compute(add.add(1))[0])  # 输出5

方法2：用Worker启动钩子预加载配置 #

如果你的配置逻辑是固定的，可以在Worker启动时自动执行初始化代码，通过Client的worker_kwargs传递预加载函数：

from dask.distributed import Client
import dask

def worker_setup():
    # 在这里放Worker需要执行的导入和配置代码
    import add
    import overhead

if __name__ == '__main__':
    import add
    import overhead
    # 创建客户端时指定Worker初始化逻辑
    c = Client(
        n_workers=2, 
        threads_per_worker=2, 
        worker_kwargs={'preload': [worker_setup]}
    )
    print(dask.compute(add.add(1))[0])  # 输出5

方法3：避免依赖导入副作用，直接传递配置 #

最可靠的方式是摆脱对模块导入副作用的依赖，把配置作为参数直接传给Delayed函数：

修改add.py：

import dask

@dask.delayed
def add(da, N):
    return da + N

调用时传入修改后的配置：

from dask.distributed import Client
import dask

if __name__ == '__main__':
    import add
    import overhead
    import constants
    c = Client(n_workers=2, threads_per_worker=2)
    print(dask.compute(add.add(1, constants.N))[0])  # 输出5

原理说明 #

Dask Worker是独立的Python进程，启动时不会执行主线程__name__ == '__main__'块内的代码。所以main块里的导入和配置修改不会自动同步到Worker。通过Client.run或Worker启动钩子，能强制Worker执行相同的逻辑，确保运行环境和主线程一致。

内容的提问来源于stack exchange，提问作者Aule Mahal