Dask Distributed如何处理TaskStepMethWrapper类特殊对象的序列化？

背景 #

我正尝试用dask distributed结合旧项目遗留代码编写数据管道。get_data函数接收url:str和session:ClientSession作为参数，返回pandas DataFrame。

from dask.distributed import Client
from aiohttp import ClientSession
client = Client()
session: ClientSession = connector.session_factory()

futures = client.map(
    get_data, # 获取数据的函数（接收url和http会话）
    urls,
    [session for _ in range(len(urls))],  # 问题出在这里
    retries=5,
)
r = client.map(loader.job, futures)
_ = client.gather(r)

问题 #

我遇到以下错误：

File "/home/zar3bski/.cache/pypoetry/virtualenvs/poc-dask-iG-N0GH5-py3.10/lib/python3.10/site-packages/distributed/worker.py", line 2952, in warn_dumps
    b = dumps(obj)
  File "/home/zar3bski/.cache/pypoetry/virtualenvs/poc-dask-iG-N0GH5-py3.10/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 58, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
  File "/home/zar3bski/.cache/pypoetry/virtualenvs/poc-dask-iG-N0GH5-py3.10/lib/python3.10/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps
    cp.dump(obj)
  File "/home/zar3bski/.cache/pypoetry/virtualenvs/poc-dask-iG-N0GH5-py3.10/lib/python3.10/site-packages/cloudpickle/cloudpickle_fast.py", line 632, in dump
    return Pickler.dump(self, obj)
TypeError: cannot pickle 'TaskStepMethWrapper' object
Unclosed client session
client_session: <aiohttp.client.ClientSession object at 0x7f3042b2fa00>

我尝试为该特殊对象注册序列化器和反序列化器：

from distributed.protocol import dask_serialize, dask_deserialize
@dask_serialize.register(TaskStepMethWrapper)
def serialize(ctx: TaskStepMethWrapper) -> Tuple[Dict, List[bytes]]:
    header = {} #?
    frames = [] #?
    return header, frames


@dask_deserialize.register(TaskStepMethWrapper)
def deserialize(header: Dict, frames: List[bytes]) -> TaskStepMethWrapper:
    return TaskStepMethWrapper(frames) #?

但问题是我不知道从哪里导入TaskStepMethWrapper类。我知道它和asyncio相关：

grep -rnw './' -e '.*TaskStepMethWrapper.*'
grep: ./lib-dynload/_asyncio.cpython-310-x86_64-linux-gnu.so : fichiers binaires correspondent

我在site-packages/aiohttp里找不到它的定义。还试过用Client(asynchronous=True)，结果出现TypeError: cannot pickle '_contextvars.Context' object错误。

请问在Dask中该如何处理特殊对象的序列化？应该扩展Dask序列化器还是使用额外的序列化族？

client = Client('tcp://scheduler-address:8786',
                serializers=['dask', 'pickle'], # 该选哪一个？
                deserializers=['dask', 'msgpack']) # 该选哪一个？

解决方案 #

核心问题：不要跨进程传递ClientSession #

ClientSession是绑定事件循环的异步会话对象，本身不支持序列化，不能直接通过Dask传递给远程Worker。当前错误的本质是试图把本地会话对象序列化后发给Worker，而这类异步相关对象（包括TaskStepMethWrapper——它是asyncio底层C扩展实现的内部包装类，无法直接导入或序列化）根本无法被pickle处理。

正确做法：在Worker端创建会话 #

不要在主进程创建会话后传递给Worker，而是让每个Worker自行创建ClientSession，有两种实现方式：

1. 任务内部初始化会话
   修改get_data函数，在函数内部创建并使用会话：

def get_data_sync(url: str):
    import asyncio
    from aiohttp import ClientSession
    
    async def _fetch():
        async with ClientSession() as session:
            # 原有数据获取逻辑
            response = await session.get(url)
            data = await response.json()
            # 转换为pandas DataFrame
            ...
            return df
    
    return asyncio.run(_fetch())

2. Worker生命周期内复用会话
   如果需要在Worker运行期间复用会话，可注册Worker回调，在Worker启动时创建全局会话，任务中直接调用：

def setup_worker():
    global session
    from aiohttp import ClientSession
    session = ClientSession()

def cleanup_worker():
    import asyncio
    asyncio.run(session.close())

# 注册Worker启动/销毁时的回调
client.register_worker_callbacks(setup=setup_worker, teardown=cleanup_worker)

# 修改get_data使用全局会话
def get_data(url: str):
    global session
    import asyncio
    
    async def _fetch():
        response = await session.get(url)
        data = await response.json()
        ...
        return df
    
    return asyncio.run(_fetch())

序列化器选择建议 #

• 无需自定义TaskStepMethWrapper的序列化器，这类内部对象本就不应该被跨进程传递。
• 无特殊自定义对象时，默认的['dask', 'pickle']组合足够使用——Dask会优先用自身序列化方式处理支持的对象，无法处理时回退到pickle。
• 不要随意修改序列化器列表，除非明确需要处理特定格式的数据。

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