Rust extension module breaks multiprocessing no Pickle support

I have added a rust extension to a Python project using pyo3 and maturin.
I have constructed a class as below, essentially creating a calendar object.

#[pyclass]
#[derive(Clone, Default, Debug, PartialEq)]
pub struct Cal {
    holidays: IndexSet<NaiveDateTime>,
    week_mask: HashSet<Weekday>,
}

(With the implemented methods not shown) Almost everything works well.
However, I have one failing test, which uses the Python mulitprocessing library. The code boils down to something like:

def _func(obj_with_cal, *args, **kwargs):
    return obj_with_cal.value(*args, **kwargs)

from multiprocessing import Pool
from functools import partial

func = partial(_func, *args, **kwargs))
p = Pool(defaults.pool)
results = p.map(func, list_of_objs_with_cal)
p.close()

The Error I receive is:

TypeError: cannot pickle 'builtins.Cal' object

I have never had use for pickle in Python, nor serde in Rust, so am unfamiliar with these in general. The issue: Pickle Support · Issue #100 · PyO3/pyo3 · GitHub seems to discuss this but at a technical level that I don't currently understand.

I have tried the derive the Serialize and Deserialize traits for Cal but their container types dont support it.

Can anyone give a holistic, helpful description of what I am facing here, and possibly suggest the most straightforward way ti implment a fix, if indeed one exists?

Are you using indexmap::IndexSet? In that case, you can enable the serde feature of indexmap to add (de-)serialization support. That should enable you to derive Serialize and Deserialize. AFAIU the discussion in the issue you linked, the next step would be to follow the gist that was shared in this reply, possibly enhancing it as described here.

Ok, thanks for the pointers. This is what I did:

  1. Added the necessary features:
// Cargo.toml
[dependencies]
chrono = { version = "0.4.38", features = ["serde"] }
serde = {version = "1.0", features = ["derive"] }
indexmap = {version = "1.9.3", features = ["serde"]
  1. Added derivations:
use serde::{Serialize, Deserialize};

#[pyclass]
#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
pub struct Cal {
    holidays: IndexSet<NaiveDateTime>,
    week_mask: HashSet<Weekday>,
}
  1. Implemented the __setstate__ and __getstate__ methods on the struct using 0.21 new Bound terminology.
    pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
        *self = deserialize(state.as_bytes()).unwrap();
        Ok(())
    }
    pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        Ok(PyBytes::new_bound(py, &serialize(&self).unwrap()))
    }

All of this compiled without errors or warnings.

The new error in my Python test suite is:

_pickle.PicklingError: Can't pickle <class 'builtins.Cal'>: attribute lookup Cal on builtins failed

?

OK I managed to find Can't pickle <class 'builtins.RustClass'>: attribute lookup RustClass on builtins failed · Issue #1517 · PyO3/pyo3 · GitHub, which indicated that the namespace was the problem.

Since I am building a Python project with rust extensions my rust extension is called "py-project.rust-extension" When I added this to the struct it could find it and pickle it:

#[pyclass(module = "py-project.rust-extension")]
#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
pub struct Cal {
    holidays: IndexSet<NaiveDateTime>,
    week_mask: HashSet<Weekday>,
}

The next problem I face is unpickling! I wrote a test that fails only a successful pickle.

    def test_pickle(self, cal):
        import pickle
        pickled_cal = pickle.dumps(cal)
>       unpickled_cal = pickle.loads(pickled_cal)
E       TypeError: Cal.__new__() missing 2 required positional arguments: 'holidays' and 'week_mask'

And finally made it:

pub fn __getnewargs__(&self) -> PyResult<(Vec<NaiveDateTime>, Vec<u8>)> {
        Ok((
            self.clone().holidays.into_iter().collect(),
            self.clone().week_mask.into_iter().map(|x| x.num_days_from_monday() as u8).collect()
        ))
    }

Thanks for the initial pointer @jofas , put me on the right path.

1 Like

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.