Fix celery#405

GregoirePelegrin · Jun 26, 2024 · c7ebbc0 · c7ebbc0
1 parent 81cc942
commit c7ebbc0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/billiard/common.py b/billiard/common.py
@@ -6,6 +6,7 @@
 import signal
 import sys
 
+import dill
 import pickle
 
 from .exceptions import RestartFreqExceeded
@@ -14,6 +15,9 @@
 pickle_load = pickle.load
 pickle_loads = pickle.loads
 
+dill_load = dill.load
+dill_loads = dill.loads
+
 # cPickle.loads does not support buffer() objects,
 # but we can just create a StringIO and use load.
 from io import BytesIO

diff --git a/billiard/pool.py b/billiard/pool.py
@@ -27,7 +27,7 @@
 from . import cpu_count, get_context
 from . import util
 from .common import (
-    TERM_SIGNAL, human_status, pickle_loads, reset_signals, restart_state,
+    TERM_SIGNAL, human_status, dill_loads, reset_signals, restart_state,
 )
 from .compat import get_errno, mem_rss, send_offset
 from .einfo import ExceptionInfo
@@ -441,7 +441,7 @@ def _make_recv_method(self, conn):
             if hasattr(conn, 'get_payload') and conn.get_payload:
                 get_payload = conn.get_payload
 
-                def _recv(timeout, loads=pickle_loads):
+                def _recv(timeout, loads=dill_loads):
                     return True, loads(get_payload())
             else:
                 def _recv(timeout):  # noqa
@@ -456,7 +456,7 @@ def _recv(timeout):  # noqa
                     return False, None
         return _recv
 
-    def _make_child_methods(self, loads=pickle_loads):
+    def _make_child_methods(self, loads=dill_loads):
         self.wait_for_job = self._make_protected_receive(self.inq)
         self.wait_for_syn = (self._make_protected_receive(self.synq)
                              if self.synq else None)