Add SchemaSerializer.__reduce__ method to enable pickle serialization

[edoakes]

Change Summary

Currently pydantic model definitions cannot be serialized using cloudpickle. One of the underlying issues is that the SchemaSerializer object is not serializable (even using vanilla pickle).

This PR enables serialization for the SchemaSerializer by:

1. Holding a reference to the Python schema and config objects required to initialize the SchemaSerializer, and
2. defining a custom __reduce__ method that instructs cloudpickle to reconstruct the SchemaSerializer from these objects.

I've also updated the existing SchemaValidator codepath to match this one, including fixing a bug where the config object was being lost on serialization.

Related issue number

pydantic/pydantic#6763

Checklist

• Unit tests for the changes exist
• Documentation reflects the changes where applicable
• Pydantic tests pass with this pydantic-core (except for expected changes)
• My PR is ready to review, please add a comment including the phrase "please review" to assign reviewers

[edoakes]

@davidhewitt I am unable to assign reviewers, but PTAL.

• I'm not very familiar with PyO3 or Rust in general, so please spot check that I'm not doing something nonsensical.
• Let me know what testing you'd like me to add.

[codspeed-hq]

CodSpeed Performance Report

Merging #1006 will improve performances by 42.09%

_{Comparing edoakes:cloudpickle-support (722bbdf) with main (8e66bd9)}

Summary

⚡ 1 improvements
✅ 139 untouched benchmarks

Benchmarks breakdown

	Benchmark	main	edoakes:cloudpickle-support	Change
⚡	test_decimal_from_string_core	74.8 µs	52.6 µs	+42.09%

[codecov]

Codecov Report

Merging #1006 (722bbdf) into main (8e66bd9) will increase coverage by 0.01%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1006      +/-   ##
==========================================
+ Coverage   93.05%   93.06%   +0.01%     
==========================================
  Files         106      106              
  Lines       15714    15739      +25     
  Branches       35       35              
==========================================
+ Hits        14622    14648      +26     
+ Misses       1085     1084       -1     
  Partials        7        7              

Files	Coverage Δ
src/serializers/mod.rs	98.58% <100.00%> (+0.10%)	⬆️
src/validators/mod.rs	94.84% <100.00%> (+0.12%)	⬆️

... and 1 file with indirect coverage changes

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8e66bd9...722bbdf. Read the comment docs.

[adriangb]

What you're suggesting generally makes sense. I'd take a look at how SchemaValidator does it:

pydantic-core/src/validators/mod.rs

Lines 144 to 149 in d93482e

	pub fn __reduce__(slf: &PyCell<Self>) -> PyResult<PyObject> {
	let py = slf.py();
	let args = (slf.try_borrow()?.schema.to_object(py),);
	let cls = slf.getattr("__class__")?;
	Ok((cls, args).into_py(py))
	}

[adriangb]

Curious why this was necessary?

[edoakes]

Wasn't necessary, I just wanted to ensure that there was coverage of the garbage collection path when a config was passed given that I'm holding a reference to it now (existing tests didn't take a config).

[adriangb]

You're pretty close! Please ignore the benchmarks and pydantic-integration test.

[edoakes]

@adriangb oh wow, thanks for the pointer, it would've saved me a lot of trouble if I'd seen SchemaValidator.__reduce__ 😓

Looks like I came to essentially the same solution. The differences are:

1. SchemaValidator.__reduce__ is using different syntax to get the reference to the class and the Python schema argument. I don't understand the differences between these, let me know which I should opt for. To an uninformed eye, SchemaSerializer::type_object(py).to_object(py) looks less hacky than let cls = slf.getattr("__class__")?;. I haven't a clue what slf.try_borrow()?.schema.to_object(py) means.
2. I'm also saving and reconstructing the "config" object which SchemaValidator.__reduce__ ignores. I'm not sure the impact of this.

Let me know if/how you'd like me to update these to match more closely.

[davidhewitt]

Generally looks great, thanks! Just posted my suggestion of how to implement __reduce__, and if you're happy with it then it'd be great to update SchemaValidator to work the same way. I also think it's probably a bug if we're not pickling the config for SchemaValidator.

[davidhewitt]

Here's my take:

Suggested change

	pub fn __reduce__(&self, py: Python) -> PyResult<(PyObject, (PyObject, PyObject))> {
	// Enables support for `cloudpickle` serialization.
	Ok((
	SchemaSerializer::type_object(py).to_object(py),
	(self.py_schema.to_object(py), self.py_config.to_object(py)),
	))
	}
	pub fn __reduce__(slf: &PyCell<Self>) -> PyResult<(&PyType, (PyObject, PyObject))> {
	let py = slf.py();
	let args = (slf.get().schema.to_object(py), slf.get().py_config.to_object(py));
	let cls = slf.get_type();
	Ok((cls, args))
	}

You'll need to update the #[pyclass] definition above for SchemaSerializer to include #[pyclass(frozen)].

Justifications:

1. slf: &PyCell<Self> is equivalent to self Python object
2. slf.get() gives read access to the Rust data as long as it's frozen
3. I think we want to look up type(self) rather than hard-code SchemaSerializer type in case we allow subclassing later. slf.get_type() is a cleaner way to do it than trying to use getattr.

[edoakes]

Makes sense, thanks for the pointers, I'll update!

[edoakes]

Ok @davidhewitt I've updated the PR to include the suggested changes.

• Now using slf.get_type() to get the reference to the object type in __reduce__.
• Updated SchemaValidator.__reduce__ to match and to fix the bug with the config not being serialized (added a regression test there as well).
• I also updated some naming to make them consistent, please let me know if I've taken too much liberty.
• Added test cases as test_pickling.py for both the serializer and validator.

[edoakes]

This is moved to test_pickling.py

[edoakes]

I verified that this newly-added test case fails without my changes to SchemaValidator.__reduce__ (config object is lost).

[davidhewitt]

Looks great to me, many thanks 👍