Dask读取超内存HDF5文件修改后写入新文件报序列化错误问题

我正在处理一个大小超出内存容量的HDF5文件，因此选择使用Dask完成文件修改操作，目标为加载源文件、执行自定义修改操作（无需强制保留原有数据形状）后，将结果保存到新的HDF5文件中。

测试源文件创建代码 #

import h5py as h5
import numpy as np

source_file = "source.hdf5"
x = np.zeros((3, 3))  # 实际场景中x的大小会超出内存容量
with h5.File(source_file, "w") as f:
    f.create_dataset("/x", data=x, compression="gzip")

初始实现代码 #

from dask import array as da
import h5py as h5
from dask.distributed import Client


if __name__ == "__main__":
    dask_client = Client(n_workers=1)  # 不需要并行，仅用Dask解决内存不足问题

    source_file = "source.hdf5"
    temp_filename = "target.hdf5"

    # 加载数据
    f = h5.File(source_file, "r")
    x_da = da.from_array(f["/x"])

    # 执行自定义修改
    x_da = x_da * 2

    # 保存到目标文件
    x_da.to_hdf5(temp_filename, "/x", compression="gzip")

    # 关闭源文件
    f.close()

报错信息 #

运行上述代码时抛出如下错误：

TypeError: ('Could not serialize object of type Dataset.', '<HDF5 dataset "x": shape (3, 3), type "<f8">') distributed.comm.utils - ERROR - ('Could not serialize object of type Dataset.', '<HDF5 dataset "x": shape (3, 3), type "<f8">')

核心疑问 #

该问题是操作方式存在错误，还是Dask本身不支持该类操作场景？如果支持的话，是否存在可行的替代实现方案？

这个报错是操作方式错误导致的，Dask本身完全支持超内存HDF5文件的读写修改场景。
报错的核心原因是：在主进程中手动打开了h5py文件句柄、获取了Dataset对象，再基于这个对象创建Dask数组交给分布式Client调度时，Dask需要把这个Dataset对象序列化后发送给worker进程。但h5py的Dataset是绑定了底层C级文件句柄的对象，本身不支持跨进程序列化，哪怕只启动1个worker，分布式Client默认的多进程通信机制还是会触发序列化流程，最终抛出错误。

• 方案1：移除分布式Client初始化代码，使用Dask默认同步调度器。如果不需要分布式/多进程并行能力，只是需要Dask的分块读写能力避免内存占满，直接删除dask_client = Client(n_workers=1)这行代码即可。默认调度器在单进程内执行任务，不需要序列化h5py对象，不会触发该错误。
• 方案2：保留分布式Client的前提下，不要在主进程手动打开HDF5文件，直接调用Dask封装好的HDF5读取接口从文件路径创建数组，让worker进程自行管理文件句柄。将加载数据的代码替换为x_da = da.from_hdf5(source_file, "/x")即可，这也是官方推荐的分布式场景下HDF5文件读取方式。
• 方案3：如果需要自定义HDF5读取逻辑，将文件打开、数据读取的逻辑全部封装到延迟执行的Dask任务函数中，保证h5py文件句柄和Dataset对象是在worker进程内部创建的，不要在主进程创建这类对象后传递给Dask任务。

修正后可运行代码示例（对应方案2） #

from dask import array as da
from dask.distributed import Client


if __name__ == "__main__":
    dask_client = Client(n_workers=1)

    source_file = "source.hdf5"
    temp_filename = "target.hdf5"

    # 直接通过文件路径创建Dask数组，不在主进程持有h5py文件句柄
    x_da = da.from_hdf5(source_file, "/x")

    # 执行自定义修改操作
    x_da = x_da * 2

    # 保存结果到目标HDF5文件
    x_da.to_hdf5(temp_filename, "/x", compression="gzip")

    dask_client.close()

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