From 0b0693a0a2d676124da59fae76ba337acfeb99bc Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 29 Jul 2026 13:18:11 +0800
Subject: [PATCH] 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

---
 astrai/tokenize/chat_template.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/astrai/tokenize/chat_template.py b/astrai/tokenize/chat_template.py
index ea04ca1..f63c360 100644
--- a/astrai/tokenize/chat_template.py
+++ b/astrai/tokenize/chat_template.py
@@ -38,12 +38,27 @@ class ChatTemplate:
         The compiled :class:`~jinja2.Template` holds a dynamically-generated
         ``root`` render function whose ``__module__`` is ``None``; under
         ``pickle`` it falls back to ``__main__`` and breaks ``spawn``-based
-        multiprocessing.  By deferring compilation to first access, the
-        default pickle protocol serialises only ``template_str``; each
-        worker rebuilds the cache on first render.
+        multiprocessing.  :meth:`__getstate__` drops the cached template so
+        that pickle serialises only ``template_str``; each worker rebuilds
+        the cache on first render.
         """
         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
     def from_string(
         cls,