Python多进程中numpy.busdaycalendar无法Pickle的问题问询

在多进程中执行包含numpy.busdaycalendar对象参数的函数时，触发如下错误：

TypeError: cannot pickle 'numpy.busdaycalendar' object

复现问题的简化代码：

from concurrent.futures import ProcessPoolExecutor
import numpy as np


def shift_one_day(x, constant):
    return np.busday_offset(x, 1, busdaycal=constant)


def main():
    constant = np.busdaycalendar()

    variables = [np.datetime64('2023-03-21'), np.datetime64('2023-03-22'), np.datetime64('2023-03-23')]

    with ProcessPoolExecutor() as executor:
        futures_list = []
        for v in variables:
            futures_list.append(executor.submit(shift_one_day, v, constant))
        result_list = [future.result() for future in futures_list]
    print(result_list)


if __name__ == '__main__':
    main()

环境信息：Windows 10系统（无多进程fork选项），Python=3.12.2、numpy=1.26.4，包来自conda-forge。

用户疑问 #

1. 是否存在方法让busdaycalendar对象可Pickle且保留原有功能？
2. 若不可行，在所有shift_one_day函数调用中constant对象均相同的情况下，有何可行的替代方案？

问题1：让busdaycalendar可Pickle #

可以通过自定义Pickle逻辑实现。numpy.busdaycalendar的核心配置是weekmask和holidays两个属性，我们可以基于这两个属性实现序列化/反序列化：

import pickle
import numpy as np

def pickle_busdaycalendar(cal):
    return (np.busdaycalendar, (cal.weekmask, cal.holidays))

def unpickle_busdaycalendar(args):
    cls, (weekmask, holidays) = args
    return cls(weekmask=weekmask, holidays=holidays)

# 注册Pickle处理逻辑
pickle.register(np.busdaycalendar, pickle_busdaycalendar, unpickle_busdaycalendar)

将这段代码放在主程序开头，之后busdaycalendar对象即可正常被Pickle，且保留原有功能。原理是：Pickle时保存对象的构造参数，反序列化时重新调用构造函数生成功能完全一致的新对象。

问题2：替代方案（无需修改Pickle逻辑） #

既然所有函数调用都使用同一个constant，可以在子进程中初始化该对象，避免跨进程传递：

方案1：利用进程池初始化参数传递配置 #

使用ProcessPoolExecutor的initializer和initargs参数，在每个子进程启动时初始化全局的busdaycalendar对象：

from concurrent.futures import ProcessPoolExecutor
import numpy as np

global_cal = None

def init_worker(weekmask, holidays):
    global global_cal
    global_cal = np.busdaycalendar(weekmask=weekmask, holidays=holidays)

def shift_one_day(x):
    return np.busday_offset(x, 1, busdaycal=global_cal)

def main():
    constant = np.busdaycalendar()
    variables = [np.datetime64('2023-03-21'), np.datetime64('2023-03-22'), np.datetime64('2023-03-23')]

    with ProcessPoolExecutor(initializer=init_worker, initargs=(constant.weekmask, constant.holidays)) as executor:
        result_list = list(executor.map(shift_one_day, variables))
    print(result_list)

if __name__ == '__main__':
    main()

方案2：直接在子进程函数内构造对象 #

如果busdaycalendar的配置固定（比如默认工作日历），可以直接在shift_one_day里构造对象，省去跨进程传递的步骤：

from concurrent.futures import ProcessPoolExecutor
import numpy as np

def shift_one_day(x):
    cal = np.busdaycalendar()
    return np.busday_offset(x, 1, busdaycal=cal)

def main():
    variables = [np.datetime64('2023-03-21'), np.datetime64('2023-03-22'), np.datetime64('2023-03-23')]

    with ProcessPoolExecutor() as executor:
        result_list = list(executor.map(shift_one_day, variables))
    print(result_list)

if __name__ == '__main__':
    main()

该方案适合工作日历配置固定的场景，无需传递额外参数，实现简单直接。

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