From 18eead3dc5cca0bc0e055a0d22b0063d0156bbb8 Mon Sep 17 00:00:00 2001 From: Kowser Date: Mon, 6 Jul 2026 19:56:48 -0700 Subject: [PATCH] Make @worker_task workers spawn-safe (macOS fork/spawn fix) --- src/conductor/client/worker/worker.py | 61 +++++++++++++++++++ .../automator/test_worker_spawn_safety.py | 57 +++++++++++++++++ .../worker/test_worker_config_integration.py | 13 +++- 3 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 tests/unit/automator/test_worker_spawn_safety.py diff --git a/src/conductor/client/worker/worker.py b/src/conductor/client/worker/worker.py index a5855a2e..0f3b59cf 100644 --- a/src/conductor/client/worker/worker.py +++ b/src/conductor/client/worker/worker.py @@ -335,6 +335,67 @@ def __init__(self, # Add thread lock for safe concurrent access to _pending_async_tasks self._pending_tasks_lock = threading.Lock() + def _worker_task_registry_key(self): + # If this worker's execute_function is a @worker_task-registered function, + # return its (name, domain) key in _decorated_functions; else None. + # Matched by function identity so it is robust to env-based domain overrides. + from conductor.client.automator.task_handler import _decorated_functions + fn = self._execute_function + for key, record in _decorated_functions.items(): + if record.get("func") is fn: + return key + return None + + def __getstate__(self): + # Make the Worker picklable so multiprocessing 'spawn' can send it to a + # worker subprocess. Two kinds of state block stdlib pickle: + # + # 1. Transient, process-local objects holding unpicklable _thread.lock + # objects (ApiClient connection pool, the locks, the event loop). + # These are dropped here and rebuilt fresh in the child (__setstate__). + # 2. The execute_function: @worker_task returns a functools.wraps wrapper + # that shadows the raw registered function at module scope, so the raw + # function cannot be pickled by reference. For decorated workers we drop + # the function and store its registry key so __setstate__ can re-resolve + # the raw function in the child (which preserves async-worker detection). + # Non-decorated workers keep execute_function and pickle it by reference. + state = self.__dict__.copy() + state.pop("api_client", None) + state.pop("_pending_tasks_lock", None) + state.pop("_background_loop", None) + state.pop("_pending_async_tasks", None) + + key = self._worker_task_registry_key() + if key is not None: + state["_wt_registry_key"] = key + state.pop("_execute_function", None) + return state + + def __setstate__(self, state): + key = state.pop("_wt_registry_key", None) + self.__dict__.update(state) + # Rebuild the transient members exactly as __init__ does, in the child. + self.api_client = ApiClient() + self._background_loop = None + self._pending_async_tasks = {} + self._pending_tasks_lock = threading.Lock() + + if key is not None: + # Re-resolve the raw @worker_task function from this process's registry. + # Requires the worker's module to be imported in the child, which the + # 'spawn' start method already requires for annotated workers. + from conductor.client.automator.task_handler import _decorated_functions + record = _decorated_functions.get(key) + if record is None: + raise RuntimeError( + f"Cannot restore @worker_task worker {key!r} in this process: its " + f"module was not imported before the worker started. Import the " + f"worker module (e.g. via TaskHandler(import_modules=[...]) or an " + f"explicit import) so it is registered in this process." + ) + # Setting via the property also restores the input/return-type flags. + self.execute_function = record["func"] + def execute(self, task: Task) -> TaskResult: task_input = {} task_output = None diff --git a/tests/unit/automator/test_worker_spawn_safety.py b/tests/unit/automator/test_worker_spawn_safety.py new file mode 100644 index 00000000..d7f88c8d --- /dev/null +++ b/tests/unit/automator/test_worker_spawn_safety.py @@ -0,0 +1,57 @@ +import pickle +import unittest + +import conductor.client.automator.task_handler as task_handler +from conductor.client.worker.worker import Worker +from conductor.client.worker.worker_task import worker_task + +_TASK_NAME = "spawn_safety_probe" + + +class TestWorkerSpawnSafety(unittest.TestCase): + """Reproduces the macOS fork/spawn bug for @worker_task workers. + + On macOS the multiprocessing 'fork' start method is unsafe, and the + documented escape hatch (CONDUCTOR_MP_START_METHOD=spawn) requires every + Process argument to be picklable. The Worker built by @worker_task held a + live ApiClient (whose connection pool owns a _thread.lock) and a + threading.Lock, so pickling it failed with: + + TypeError: cannot pickle '_thread.lock' object + + and its execute_function (a wrapper-shadowed function) could not be pickled + by reference either. This test fails before the spawn-safety fix and passes + after it. + """ + + def setUp(self) -> None: + # Hermetic registry: some other tests replace _decorated_functions with a + # Mock and don't restore it. Swap in a fresh dict and restore in tearDown. + self._saved_registry = task_handler._decorated_functions + task_handler._decorated_functions = {} + + def tearDown(self) -> None: + task_handler._decorated_functions = self._saved_registry + + def test_worker_task_worker_is_picklable_for_spawn(self): + # Register a decorated worker fresh (other tests clear the registry). + @worker_task(task_definition_name=_TASK_NAME) + def spawn_safety_probe(x: int) -> int: + return x + 1 + + # Build the Worker exactly as the SDK does for a decorated function, + # straight from the registry record (avoids helpers other tests mock). + record = task_handler._decorated_functions[(_TASK_NAME, None)] + worker = Worker( + task_definition_name=_TASK_NAME, + execute_function=record["func"], + ) + + restored = pickle.loads(pickle.dumps(worker)) + + self.assertEqual(restored.get_task_definition_name(), _TASK_NAME) + self.assertEqual(restored.execute_function(1), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/worker/test_worker_config_integration.py b/tests/unit/worker/test_worker_config_integration.py index d3c315cc..7f27b760 100644 --- a/tests/unit/worker/test_worker_config_integration.py +++ b/tests/unit/worker/test_worker_config_integration.py @@ -8,12 +8,21 @@ import asyncio from unittest.mock import Mock, patch -# Prevent actual task handler initialization -sys.modules['conductor.client.automator.task_handler'] = Mock() +# Prevent actual task handler initialization during the imports below, but do +# NOT leave the module mocked in sys.modules afterwards — that poisons every +# other test module that imports task_handler later in the same process. +_TASK_HANDLER = 'conductor.client.automator.task_handler' +_original_task_handler = sys.modules.get(_TASK_HANDLER) +sys.modules[_TASK_HANDLER] = Mock() from conductor.client.worker.worker_task import worker_task from conductor.client.worker.worker_config import resolve_worker_config +if _original_task_handler is not None: + sys.modules[_TASK_HANDLER] = _original_task_handler +else: + del sys.modules[_TASK_HANDLER] + class TestWorkerConfigWithDecorator(unittest.TestCase): """Test worker configuration resolution with @worker_task decorator"""