使用pybind11函数与Dask时出现cannot pickle 'PyCapsule' object错误

问题原因 #

使用dask的processes调度器时，任务函数需要被序列化（pickle）后传递给子进程。pybind11导出的C++函数对象底层为PyCapsule类型，默认不支持pickle序列化，因此触发报错。

解决方案 #

方案1：子进程内重新导入模块（推荐） #

将pybind11模块的导入从全局作用域移到任务函数内部，避免序列化C++函数对象。子进程会自行导入模块并获取函数，同时修复原函数缺少返回值的问题：

修改example.py中的add_column_values函数：

def add_column_values(row: pd.Series) -> pd.Series:
    # 将导入移到函数内部，避免全局序列化C++函数对象
    from python_example import add
    row['sum'] = add(int(row['i']), int(row['j']))
    return row  # 必须返回修改后的Series，否则apply会丢失结果

同时建议修复原Python版函数的返回值问题：

def add_column_values_python(row: pd.Series) -> pd.Series:
    row['sum'] = python_add(row['i'], row['j'])
    return row

方案2：切换到线程调度器 #

如果你的C++函数能释放GIL（pybind11默认导出的函数会自动释放GIL），可将dask的调度器改为threads，绕过跨进程序列化问题：

dataframe = d_dataframe.map_partitions(
    lambda df: df.apply(add_column_values, axis=1), meta=pd.Series(dtype='float64')
).compute(scheduler='threads')

注：线程调度受Python GIL限制，但C++函数执行时会释放GIL，因此仍能实现并行加速。

方案3：为pybind11函数添加pickle支持（进阶） #

若必须传递C++函数对象，可通过pybind11自定义pickle逻辑，序列化时保存模块和函数名，反序列化时重新导入：

修改src/main.cpp中的模块定义部分：

#include <pybind11/pybind11.h>
#include <pybind11/pickle.h>

#define STRINGIFY(x) #x
#define MACRO_STRINGIFY(x) STRINGIFY(x)

int add(int i, int j) {
    return i + j;
}

namespace py = pybind11;

// 自定义pickle序列化逻辑
py::object add_pickle(py::pickle_protocol) {
    return py::make_tuple("python_example", "add");
}

PYBIND11_MODULE(python_example, m) {
    m.doc() = R"pbdoc(
        Pybind11 example plugin
        -----------------------

        .. currentmodule:: python_example

        .. autosummary::
        :toctree: _generate

        add
        subtract
    )pbdoc";

    py::object add_obj = m.def("add", &add, R"pbdoc(
        Add two numbers

        Some other explanation about the add function.
    )pbdoc");
    // 注册pickle方法
    add_obj.attr("__reduce__") = py::cpp_function(add_pickle, py::is_method(add_obj));

    m.def("subtract", [](int i, int j) { return i - j; }, R"pbdoc(
        Subtract two numbers

        Some other explanation about the subtract function.
    )pbdoc");

#ifdef VERSION_INFO
    m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
#else
    m.attr("__version__") = "dev";
#endif
}

该方案需修改C++代码，实现复杂度较高，适合必须传递函数对象的场景。

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