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:
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user