diff --git a/zetta_utils/mazepa/worker.py b/zetta_utils/mazepa/worker.py index ba271b2fc..f83ea601a 100644 --- 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 index 19403e163..8e639952c 100644 --- 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