feat: mazepa convenience fixes

ZettaAI · Feb 20, 2025 · 4bf3074 · 4bf3074
1 parent 4b46c35
commit 4bf3074
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/zetta_utils/mazepa/worker.py b/zetta_utils/mazepa/worker.py
@@ -8,7 +8,7 @@
 
 import tenacity
 
-from zetta_utils import log
+from zetta_utils import builder, log
 from zetta_utils.common import RepeatTimer
 from zetta_utils.mazepa import constants, exceptions
 from zetta_utils.mazepa.exceptions import MazepaCancel, MazepaTimeoutError
@@ -30,6 +30,7 @@ def __call__(self, task: Task):
 logger = log.get_logger("mazepa")
 
 
+@builder.register("run_worker")
 def run_worker(
     task_queue: MessageQueue[Task],
     outcome_queue: MessageQueue[OutcomeReport],

diff --git a/zetta_utils/message_queues/serialization.py b/zetta_utils/message_queues/serialization.py
@@ -10,7 +10,7 @@
 def serialize(obj):  # pragma: no cover
     try:
         result = _serialize(obj, pickle)
-    except pickle.PicklingError:
+    except (pickle.PicklingError, TypeError):
         result = _serialize(obj, dill)
 
     return result