fix: make ChatTemplate picklable for spawn multiprocessing

- Add __getstate__/__setstate__ to drop cached _compiled Jinja2 template
- Jinja2 Template.root_render_func is a dynamic closure unpicklable by reference
- cached_property rebuilds the template lazily on first render after unpickle
This commit is contained in:
2026-07-29 13:24:13 +08:00
parent 115192c67c
commit 0b0693a0a2
+18 -3
View File
@@ -38,12 +38,27 @@ class ChatTemplate:
The compiled :class:`~jinja2.Template` holds a dynamically-generated The compiled :class:`~jinja2.Template` holds a dynamically-generated
``root`` render function whose ``__module__`` is ``None``; under ``root`` render function whose ``__module__`` is ``None``; under
``pickle`` it falls back to ``__main__`` and breaks ``spawn``-based ``pickle`` it falls back to ``__main__`` and breaks ``spawn``-based
multiprocessing. By deferring compilation to first access, the multiprocessing. :meth:`__getstate__` drops the cached template so
default pickle protocol serialises only ``template_str``; each that pickle serialises only ``template_str``; each worker rebuilds
worker rebuilds the cache on first render. the cache on first render.
""" """
return Template(self.template_str) return Template(self.template_str)
def __getstate__(self) -> Dict[str, Any]:
"""Exclude the cached Jinja2 template from pickling.
``Template.root_render_func`` is a dynamically generated closure
that cannot be pickled by reference. Dropping ``_compiled`` here
lets :class:`cached_property` rebuild it on first access after
unpickle.
"""
state = self.__dict__.copy()
state.pop("_compiled", None)
return state
def __setstate__(self, state: Dict[str, Any]) -> None:
self.__dict__.update(state)
@classmethod @classmethod
def from_string( def from_string(
cls, cls,