From bc4b1f26abea289cad010d4d00debeb3b6fc3b2f Mon Sep 17 00:00:00 2001 From: bradyyie Date: Tue, 7 Jul 2026 01:37:05 -0400 Subject: [PATCH] Make @worker_task workers spawn-safe and default all platforms to spawn Fixes the macOS failure mode where every @worker_task worker subprocess died with SIGSEGV (exitcode=-11) under the forced 'fork' start method and was silently restarted by the TaskHandler monitor every ~5s, so no task ever completed (issues #264, #271). Worker pickling (required by spawn/forkserver, which pickle Process args): - Worker.api_client is now a lazy property: no eager httpx.Client/TLS/auth in the parent process, and no '_thread.lock' in pickled state - Worker.__getstate__/__setstate__ exclude process-local runtime state (pending-tasks lock, background event loop, pending async futures) and rebuild it in the child - @worker_task-decorated functions pickle as importable references (module, qualname, __wrapped__ depth) resolved in the child, fixing "PicklingError: it's not the same object as module.name" caused by the decorator rebinding the module-level name to its wrapper Start method: - Default is now 'spawn' on all platforms (was fork everywhere except Windows). fork is unsafe with native state: Apple frameworks SIGSEGV on macOS (CPython's own default is spawn there since 3.8, bpo-33725) and fork-with-held-lock deadlocks on POSIX. CONDUCTOR_MP_START_METHOD=fork remains as an escape hatch Diagnostics and robustness: - TaskHandler.start_processes() cleans up already-started subprocesses and raises an actionable error when worker state cannot be pickled (was: the non-daemon logger process kept the interpreter alive forever -> hang) - Worker processes killed by a signal now log the signal number and PYTHONFAULTHANDLER=1 guidance instead of restarting silently Tests: - tests/unit/worker/test_worker_spawn_safety.py: pickle round-trip, identity preservation, async detection, lazy api_client, runtime-state isolation, fail-fast for nested functions, and execution of a Worker in a real spawn child process - tests/unit/worker/test_worker_config_integration.py no longer replaces sys.modules['...task_handler'] with a Mock at import time; that leaked into the whole pytest session and broke pickle-by-reference for every later test that started a real process - examples/spawn_safety_probe.py: standalone before/after probe script --- CHANGELOG.md | 7 + docs/WORKER.md | 38 +++- examples/spawn_safety_probe.py | 125 +++++++++++ .../client/automator/task_handler.py | 65 +++++- src/conductor/client/worker/worker.py | 121 ++++++++++- tests/unit/resources/spawn_worker_helpers.py | 46 ++++ .../worker/test_worker_config_integration.py | 13 +- tests/unit/worker/test_worker_spawn_safety.py | 201 ++++++++++++++++++ 8 files changed, 592 insertions(+), 24 deletions(-) create mode 100644 examples/spawn_safety_probe.py create mode 100644 tests/unit/resources/spawn_worker_helpers.py create mode 100644 tests/unit/worker/test_worker_spawn_safety.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b37cdf97..32c88208 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,5 +15,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- **Multiprocessing start method now defaults to `spawn` on all platforms** (was `fork` everywhere except Windows). `fork` caused silently-restarting `SIGSEGV` worker subprocesses on macOS (exitcode `-11`) and fork-with-held-lock deadlocks on POSIX. Opt back in with `CONDUCTOR_MP_START_METHOD=fork`. Entrypoint scripts must use the standard `if __name__ == "__main__":` guard, and workers must be defined at module level - Legacy metrics emit unchanged by default; no env var required - `metrics_collector.py` is now a compatibility shim; `from conductor.client.telemetry.metrics_collector import MetricsCollector` continues to work + +### Fixed + +- `@worker_task` workers are now picklable, making the decorator path work with the `spawn`/`forkserver` start methods (fixes `TypeError: cannot pickle '_thread.lock' object` and `PicklingError: it's not the same object as ...`; issues #264, #271): `Worker.api_client` is created lazily in the worker process, runtime state (locks, pending async tasks, background loop) is excluded from pickling and rebuilt in the child, and decorated functions are pickled as importable references resolved in the child +- `TaskHandler.start_processes()` no longer hangs the interpreter when a worker fails to start (e.g., unpicklable state under `spawn`); it now cleans up already-started subprocesses and raises with actionable guidance +- Worker processes killed by a signal now log a diagnostic hint (signal number, `PYTHONFAULTHANDLER=1` guidance) instead of restarting silently diff --git a/docs/WORKER.md b/docs/WORKER.md index 848e9ccf..7aca6cf6 100644 --- a/docs/WORKER.md +++ b/docs/WORKER.md @@ -229,11 +229,6 @@ from conductor.client.configuration.configuration import Configuration from conductor.client.automator.task_handler import TaskHandler from conductor.client.worker.worker import Worker -#### Add these lines if running on a mac#### -from multiprocessing import set_start_method -set_start_method('fork') -############################################ - SERVER_API_URL = 'http://localhost:8080/api' KEY_ID = '' KEY_SECRET = '' @@ -261,10 +256,39 @@ workers = [ # TaskHandler scans for @worker_task decorated workers by default. # Set scan_for_annotated_workers=False if you want to disable auto-discovery. -with TaskHandler(workers, configuration, scan_for_annotated_workers=True) as task_handler: - task_handler.start_processes() +# The __main__ guard is required: worker processes use the 'spawn' +# multiprocessing start method by default, which re-imports this script. +if __name__ == '__main__': + with TaskHandler(workers, configuration, scan_for_annotated_workers=True) as task_handler: + task_handler.start_processes() ``` +### Multiprocessing start method + +`TaskHandler` runs each worker in its own process using Python's `multiprocessing` +with the **`spawn`** start method by default on all platforms. `spawn` starts each +worker as a fresh interpreter, which avoids the crashes and deadlocks `fork` +causes in processes holding native state — most visibly on macOS, where forked +workers can die with `SIGSEGV` (exitcode `-11`) inside Apple system frameworks +and get silently restarted by the supervisor. + +Requirements under `spawn` (standard Python multiprocessing rules): + +- Guard your entrypoint with `if __name__ == '__main__':` +- Define workers at module level (not nested inside functions or lambdas) +- `configuration`, `metrics_settings`, and `event_listeners` passed to + `TaskHandler` must be picklable + +To opt back into the previous behavior (e.g., a Linux deployment relying on +fork's copy-on-write memory sharing): + +```shell +export CONDUCTOR_MP_START_METHOD=fork # spawn (default) | fork | forkserver +``` + +If a worker process dies from a signal repeatedly at startup, set +`PYTHONFAULTHANDLER=1` to capture the crashing stack trace from the child. + ### Resilience: auto-restart and health checks If you run workers as a long-lived service (e.g., alongside FastAPI/Uvicorn), you can optionally enable process diff --git a/examples/spawn_safety_probe.py b/examples/spawn_safety_probe.py new file mode 100644 index 00000000..69ffab4f --- /dev/null +++ b/examples/spawn_safety_probe.py @@ -0,0 +1,125 @@ +""" +spawn_safety_probe.py - @worker_task spawn-safety probe (GitHub issues #264/#271). + +Run from the repo root: + PYTHONPATH=src python3 examples/spawn_safety_probe.py + +Forces the 'spawn' multiprocessing start method (the SDK default; also the +only safe method on macOS) and then checks, in order: + + 1. Which Worker attributes survive pickling (attribute-level diagnostic) + 2. Whole-Worker pickle round-trip + task execution on the restored worker + 3. Worker transferred into a REAL 'spawn' child process and executed there + (exactly what TaskHandler.start_processes() does) + +Broken SDK output: FAIL lines for api_client / _pending_tasks_lock / + _execute_function; round-trip and spawn-child checks fail. +Fixed SDK output: everything PASS. +""" +import os +import sys + +# Force spawn BEFORE importing any conductor module: task_handler pins the +# start method at import time from this env var. +os.environ["CONDUCTOR_MP_START_METHOD"] = "spawn" + +import multiprocessing +import pickle + +from conductor.client.worker.worker_task import worker_task +from conductor.client.http.models.task import Task + + +@worker_task(task_definition_name="pickle_probe_task") +def pickle_probe_task(x: int) -> dict: + return {"doubled": x * 2} + + +def _make_task() -> Task: + return Task( + task_id="probe-task-id", + workflow_instance_id="probe-wf-id", + task_def_name="pickle_probe_task", + input_data={"x": 21}, + status="IN_PROGRESS", + ) + + +def run_in_child(worker, result_queue) -> None: + """Spawn-child entry point: unpickling `worker` here IS the test.""" + result = worker.execute(_make_task()) + result_queue.put((str(result.status), result.output_data)) + + +def get_probe_worker(): + from conductor.client.automator.task_handler import get_registered_workers + return next(w for w in get_registered_workers() + if w.task_definition_name == "pickle_probe_task") + + +def main() -> int: + failures = 0 + print(f"python={sys.version.split()[0]} platform={sys.platform}") + + worker = get_probe_worker() # importing task_handler pins the start method + print(f"multiprocessing start method: {multiprocessing.get_start_method(allow_none=True)}\n") + + # -- 1. attribute-level diagnostic -------------------------------------- + # Pickle what pickling actually serializes: the __getstate__ output. + # (Raw __dict__ values legitimately contain process-local objects - a + # threading.Lock, the original decorated function - which __getstate__ + # excludes or substitutes. If a future change adds unpicklable state and + # forgets to handle it in __getstate__, this section catches it.) + print("[1] Worker pickled-state (__getstate__) picklability:") + state = worker.__getstate__() if hasattr(worker, "__getstate__") else dict(worker.__dict__) + for k, v in state.items(): + try: + pickle.dumps(v) + print(f" PASS {k}") + except Exception as e: + failures += 1 + print(f" FAIL {k}: {type(e).__name__}: {str(e)[:70]}") + + # -- 2. whole-worker round-trip + execute -------------------------------- + print("\n[2] Whole-Worker pickle round-trip:") + try: + restored = pickle.loads(pickle.dumps(worker)) + same_fn = restored.execute_function is worker.execute_function + result = restored.execute(_make_task()) + if str(result.status) == "COMPLETED" and result.output_data == {"doubled": 42}: + print(f" PASS round-trip (same function object: {same_fn}, " + f"result: {result.output_data})") + else: + failures += 1 + print(f" FAIL unexpected result: {result.status} {result.output_data}") + except Exception as e: + failures += 1 + print(f" FAIL {type(e).__name__}: {str(e)[:100]}") + + # -- 3. real spawn child -------------------------------------------------- + print("\n[3] Execute in a real 'spawn' child process:") + ctx = multiprocessing.get_context("spawn") + q = ctx.Queue() + p = ctx.Process(target=run_in_child, args=(get_probe_worker(), q)) + try: + p.start() + status, output = q.get(timeout=60) + p.join(timeout=10) + if status == "COMPLETED" and output == {"doubled": 42}: + print(f" PASS child exitcode={p.exitcode}, result={output}") + else: + failures += 1 + print(f" FAIL status={status} output={output}") + except Exception as e: + failures += 1 + print(f" FAIL {type(e).__name__}: {str(e)[:100]}") + if p.is_alive(): + p.terminate() + p.join(timeout=5) + + print(f"\n{'ALL CHECKS PASSED' if failures == 0 else f'{failures} CHECK(S) FAILED'}") + return 0 if failures else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/conductor/client/automator/task_handler.py b/src/conductor/client/automator/task_handler.py index 27a0e887..6d24b516 100644 --- a/src/conductor/client/automator/task_handler.py +++ b/src/conductor/client/automator/task_handler.py @@ -4,10 +4,11 @@ import inspect import logging import os +import pickle import signal import threading import time -from multiprocessing import Process, freeze_support, Queue, set_start_method +from multiprocessing import Process, freeze_support, Queue, set_start_method, get_start_method from sys import platform from typing import List, Optional, Any, Dict @@ -37,12 +38,25 @@ _mp_fork_set = False if not _mp_fork_set: try: - # The prometheus_client library holds a module-level threading lock; - # forking while that lock is held causes a deadlock in child processes. - # Set CONDUCTOR_MP_START_METHOD=spawn to avoid this if you hit the - # deadlock. Default is fork for backward compatibility (spawn requires - # all Process arguments to be picklable). - _default_method = "spawn" if platform == "win32" else "fork" + # Default start method: "spawn" on every platform. + # + # fork() is fundamentally unsafe in a process that holds native + # framework state or non-main threads: + # - macOS: Apple frameworks (CFNetwork, Security, libdispatch) may + # SIGSEGV in forked children, producing silently-restarting worker + # processes with exitcode=-11. CPython switched its own macOS + # default to spawn in 3.8 for the same reason (bpo-33725). + # - all POSIX: forking while another thread holds a lock (e.g. the + # prometheus_client module-level lock, or the TaskHandler monitor + # thread restarting a worker) leaves that lock permanently held in + # the child -> silent deadlock. + # Workers (including the @worker_task decorator path) are pickle-safe, + # so spawn works everywhere. Deployments that rely on fork's + # copy-on-write inheritance can opt back in with + # CONDUCTOR_MP_START_METHOD=fork (re-exposing the hazards above). + # NOTE: spawn requires the standard `if __name__ == "__main__":` guard + # in the entrypoint script. + _default_method = "spawn" _method = os.environ.get("CONDUCTOR_MP_START_METHOD", "").strip().lower() or _default_method if _method not in _VALID_MP_START_METHODS: logger.warning( @@ -341,8 +355,26 @@ def start_processes(self) -> None: logger.info("Starting worker processes...") freeze_support() self._monitor_stop_event.clear() - self.__start_task_runner_processes() - self.__start_metrics_provider_process() + try: + self.__start_task_runner_processes() + self.__start_metrics_provider_process() + except (TypeError, pickle.PicklingError, AttributeError) as e: + # Under the 'spawn'/'forkserver' start methods, Process arguments + # are pickled at start(); unpicklable state fails here. Clean up + # everything already started - otherwise the non-daemon logger + # subprocess keeps the interpreter alive forever after the + # exception (observed as a hang after "cannot pickle + # '_thread.lock' object"). + logger.error( + "Failed to start worker processes with start method %r: %s. " + "Workers must be defined at module level (not nested inside " + "functions), and configuration, metrics_settings and " + "event_listeners must be picklable. " + "See https://github.com/conductor-oss/conductor-python/issues/264", + get_start_method(allow_none=True), e, + ) + self.stop_processes() + raise self.__start_monitor_thread() logger.info("Started all processes") @@ -445,6 +477,21 @@ def __check_and_restart_processes(self) -> None: worker = self.workers[i] if i < len(self.workers) else None worker_name = worker.get_task_definition_name() if worker is not None else f"worker[{i}]" logger.warning("Worker process exited (worker=%s, pid=%s, exitcode=%s)", worker_name, process.pid, exitcode) + if exitcode < 0: + # Negative exitcode == killed by signal (e.g. -11 is + # SIGSEGV). Under the 'fork' start method this is the + # classic symptom of fork-unsafe native libraries in the + # parent (Apple frameworks on macOS, numpy/Accelerate, + # grpc, ...). + logger.warning( + "Worker process %s was killed by signal %s. If this " + "repeats on every restart, the process is likely " + "crashing at startup: use the 'spawn' start method " + "(CONDUCTOR_MP_START_METHOD=spawn, the default) and " + "set PYTHONFAULTHANDLER=1 to capture the crashing " + "stack trace.", + worker_name, -exitcode, + ) if not self.restart_on_failure: continue self.__restart_worker_process(i) diff --git a/src/conductor/client/worker/worker.py b/src/conductor/client/worker/worker.py index a5855a2e..c26378ad 100644 --- a/src/conductor/client/worker/worker.py +++ b/src/conductor/client/worker/worker.py @@ -2,8 +2,10 @@ import asyncio import atexit import dataclasses +import importlib import inspect import logging +import sys import threading import time import traceback @@ -275,6 +277,79 @@ def _cleanup(self): logger.debug("Background event loop stopped") +class _ExecuteFunctionReference: + """Pickle surrogate for functions whose module-level name was rebound. + + The ``@worker_task`` decorator registers the *original* function but rebinds + the module-level name to its wrapper. Standard pickle-by-reference then + fails with ``PicklingError: it's not the same object as module.name`` when + the ``spawn``/``forkserver`` start methods pickle Process arguments + (GitHub issues #264/#271). + + Instead of the raw function we pickle ``(module, qualname, unwrap_depth)`` + and re-resolve in the child process: import the module, walk the qualname, + then follow ``__wrapped__`` (set by ``functools.wraps``) ``unwrap_depth`` + times to get back to the registered function. + """ + + __slots__ = ("module_name", "qualname", "unwrap_depth") + + def __init__(self, module_name: str, qualname: str, unwrap_depth: int): + self.module_name = module_name + self.qualname = qualname + self.unwrap_depth = unwrap_depth + + def resolve(self) -> Callable: + module = importlib.import_module(self.module_name) + obj: Any = module + for part in self.qualname.split("."): + obj = getattr(obj, part) + for _ in range(self.unwrap_depth): + obj = obj.__wrapped__ + return obj + + +_MAX_UNWRAP_DEPTH = 32 + + +def _importable_function_reference(func: Any) -> Optional[_ExecuteFunctionReference]: + """Return a pickle surrogate for *func* if it needs one, else None. + + None is returned when normal pickle-by-reference works (module-level + function whose name still points at itself) or when the function is not + resolvable by import at all (lambdas, closures, bound methods); in the + latter case pickling fails with the standard, descriptive pickle error. + """ + if not inspect.isfunction(func): + return None + module_name = getattr(func, "__module__", None) + qualname = getattr(func, "__qualname__", None) + if not module_name or not qualname or "" in qualname or "" in qualname: + return None + module = sys.modules.get(module_name) + if module is None: + return None + obj: Any = module + try: + for part in qualname.split("."): + obj = getattr(obj, part) + except AttributeError: + return None + if obj is func: + # Plain module-level function: default pickling works, no surrogate. + return None + # The name was rebound (typically by @worker_task). Walk the wrapper chain + # to find the registered function and record how many hops it takes. + depth = 0 + current = obj + while hasattr(current, "__wrapped__") and depth < _MAX_UNWRAP_DEPTH: + current = current.__wrapped__ + depth += 1 + if current is func: + return _ExecuteFunctionReference(module_name, qualname, depth) + return None + + def is_callable_input_parameter_a_task(callable: ExecuteTaskFunction, object_type: Any) -> bool: parameters = inspect.signature(callable).parameters if len(parameters) != 1: @@ -306,7 +381,11 @@ def __init__(self, register_schema: bool = False ) -> Self: super().__init__(task_definition_name) - self.api_client = ApiClient() + # Lazily constructed (see the api_client property): an eager ApiClient + # holds an httpx.Client and threading locks, which made Worker + # unpicklable under the 'spawn' start method and initialized TLS state + # in the parent process before fork() (unsafe on macOS). + self._api_client = None if poll_interval is None: self.poll_interval = DEFAULT_POLLING_INTERVAL else: @@ -335,6 +414,46 @@ def __init__(self, # Add thread lock for safe concurrent access to _pending_async_tasks self._pending_tasks_lock = threading.Lock() + @property + def api_client(self) -> ApiClient: + """ApiClient used for output serialization, created on first use. + + Lazy so that Worker instances remain picklable (required by the + 'spawn'/'forkserver' multiprocessing start methods) and so the parent + process never initializes HTTP/TLS state on behalf of workers. + """ + if self._api_client is None: + self._api_client = ApiClient() + return self._api_client + + @api_client.setter + def api_client(self, value: ApiClient) -> None: + self._api_client = value + + def __getstate__(self): + state = self.__dict__.copy() + # Process-local runtime state must not cross process boundaries: the + # 'spawn'/'forkserver' start methods pickle Process args (issues + # #264/#271: "cannot pickle '_thread.lock' object"). Everything below + # is rebuilt lazily in the child process. + state["_api_client"] = None + state["_background_loop"] = None + state["_pending_async_tasks"] = {} + state["_pending_tasks_lock"] = None + ref = _importable_function_reference(state.get("_execute_function")) + if ref is not None: + state["_execute_function"] = ref + return state + + def __setstate__(self, state): + self.__dict__.update(state) + if self._pending_tasks_lock is None: + self._pending_tasks_lock = threading.Lock() + if isinstance(self._execute_function, _ExecuteFunctionReference): + # Assign through the property setter so the signature-derived + # flags are recomputed against the resolved function. + self.execute_function = self._execute_function.resolve() + def execute(self, task: Task) -> TaskResult: task_input = {} task_output = None diff --git a/tests/unit/resources/spawn_worker_helpers.py b/tests/unit/resources/spawn_worker_helpers.py new file mode 100644 index 00000000..93e7ed51 --- /dev/null +++ b/tests/unit/resources/spawn_worker_helpers.py @@ -0,0 +1,46 @@ +""" +Module-level helpers for spawn start-method regression tests. + +These MUST live in an importable module (not the test file's local scope): +the 'spawn' start method pickles Process arguments by reference, and the +child process re-imports this module to resolve them. + +Regression tests for GitHub issues #264 / #271: +@worker_task workers were unpicklable ("cannot pickle '_thread.lock' object" / +"it's not the same object as module.name"), so TaskHandler could not start +worker subprocesses with CONDUCTOR_MP_START_METHOD=spawn (the only safe start +method on macOS). +""" +from conductor.client.http.models.task import Task +from conductor.client.worker.worker_task import worker_task + + +@worker_task(task_definition_name="spawn_pickle_task", thread_count=2, domain="spawn-test") +def spawn_pickle_task(x: int) -> dict: + return {"doubled": x * 2} + + +@worker_task(task_definition_name="spawn_pickle_async_task") +async def spawn_pickle_async_task(x: int) -> dict: + return {"echo": x} + + +def plain_function_worker(x: int) -> dict: + """Module-level function NOT wrapped by @worker_task.""" + return {"plain": x} + + +def run_worker_in_child(worker, result_queue) -> None: + """Spawn-child entry point: execute one task on the unpickled worker. + + Receiving `worker` here already exercises unpickling inside the child. + """ + task = Task( + task_id="spawn-child-task-id", + workflow_instance_id="spawn-child-wf-id", + task_def_name=worker.get_task_definition_name(), + input_data={"x": 21}, + status="IN_PROGRESS", + ) + result = worker.execute(task) + result_queue.put((str(result.status), result.output_data)) diff --git a/tests/unit/worker/test_worker_config_integration.py b/tests/unit/worker/test_worker_config_integration.py index d3c315cc..88ed5b79 100644 --- a/tests/unit/worker/test_worker_config_integration.py +++ b/tests/unit/worker/test_worker_config_integration.py @@ -3,15 +3,14 @@ """ import os -import sys import unittest -import asyncio -from unittest.mock import Mock, patch -# Prevent actual task handler initialization -sys.modules['conductor.client.automator.task_handler'] = Mock() - -from conductor.client.worker.worker_task import worker_task +# NOTE: do NOT replace sys.modules['conductor.client.automator.task_handler'] +# with a Mock here. pytest imports every test module at collection time, so a +# module-level sys.modules substitution leaks into the whole test session and +# breaks pickling of task_handler functions by reference ("it's not the same +# object as ...") - which the 'spawn' multiprocessing start method (the SDK +# default) relies on to start worker/logger processes. from conductor.client.worker.worker_config import resolve_worker_config diff --git a/tests/unit/worker/test_worker_spawn_safety.py b/tests/unit/worker/test_worker_spawn_safety.py new file mode 100644 index 00000000..62b7ecb9 --- /dev/null +++ b/tests/unit/worker/test_worker_spawn_safety.py @@ -0,0 +1,201 @@ +""" +Spawn start-method safety tests for Worker / @worker_task (issues #264, #271). + +Before the fix, Worker instances built from @worker_task functions could not +be pickled, so TaskHandler could not start worker subprocesses under the +'spawn' multiprocessing start method (the only safe method on macOS, and the +SDK default): + + TypeError: cannot pickle '_thread.lock' object (Worker/ApiClient locks) + PicklingError: it's not the same object as module.name (decorator rebinding) + +...and on failure the parent process hung forever instead of exiting. +""" +import inspect +import multiprocessing +import pickle +import unittest + +import conductor.client.automator.task_handler as task_handler +from conductor.client.http.models.task import Task +from conductor.client.http.models.task_result_status import TaskResultStatus +from conductor.client.worker.worker import ( + Worker, + _ExecuteFunctionReference, + _importable_function_reference, +) +from conductor.client.worker.worker_task import worker_task + +from tests.unit.resources import spawn_worker_helpers as helpers + +# Importing the helpers module registered its @worker_task functions in the +# global registry. Scrub them immediately so tests that scan for annotated +# workers (scan_for_annotated_workers=True) are not polluted. +for _key in [k for k in list(task_handler._decorated_functions) if k[0].startswith("spawn_pickle_")]: + task_handler._decorated_functions.pop(_key, None) + + +def _make_decorated_worker() -> Worker: + """Build a Worker exactly as TaskHandler does for a @worker_task function: + from the ORIGINAL (pre-wrapper) function object.""" + return Worker( + task_definition_name="spawn_pickle_task", + execute_function=inspect.unwrap(helpers.spawn_pickle_task), + domain="spawn-test", + thread_count=2, + poll_interval=250, + ) + + +def _make_task(x: int = 21) -> Task: + return Task( + task_id="task-id-1", + workflow_instance_id="wf-id-1", + task_def_name="spawn_pickle_task", + input_data={"x": x}, + status="IN_PROGRESS", + ) + + +class TestWorkerPickleRoundTrip(unittest.TestCase): + def test_decorated_worker_pickle_roundtrip_restores_original_function(self): + worker = _make_decorated_worker() + restored = pickle.loads(pickle.dumps(worker)) + + # Same function object: resolved via module reference + __wrapped__. + self.assertIs(restored.execute_function, worker.execute_function) + self.assertEqual(restored.get_task_definition_name(), "spawn_pickle_task") + self.assertEqual(restored.domain, "spawn-test") + self.assertEqual(restored.thread_count, 2) + self.assertEqual(restored.poll_interval, 250) + + def test_unpickled_worker_executes_task(self): + restored = pickle.loads(pickle.dumps(_make_decorated_worker())) + result = restored.execute(_make_task(x=21)) + self.assertEqual(result.status, TaskResultStatus.COMPLETED) + self.assertEqual(result.output_data, {"doubled": 42}) + + def test_async_decorated_worker_preserves_coroutine_detection(self): + worker = Worker( + task_definition_name="spawn_pickle_async_task", + execute_function=inspect.unwrap(helpers.spawn_pickle_async_task), + ) + restored = pickle.loads(pickle.dumps(worker)) + # TaskHandler routes to AsyncTaskRunner based on this check. + self.assertTrue(inspect.iscoroutinefunction(restored.execute_function)) + + def test_plain_function_worker_pickles_by_reference(self): + worker = Worker( + task_definition_name="plain_task", + execute_function=helpers.plain_function_worker, + ) + restored = pickle.loads(pickle.dumps(worker)) + self.assertIs(restored.execute_function, helpers.plain_function_worker) + + def test_nested_decorated_function_fails_fast_in_parent(self): + """A @worker_task function defined inside another function can never be + re-imported by a spawn child; pickling must fail in the PARENT (clear, + immediate) rather than crash-looping the child at unpickle time.""" + + @worker_task(task_definition_name="nested_probe") + def nested_probe(x: int) -> int: + return x + + # Scrub the registration side effect of the local decorator above. + task_handler._decorated_functions.pop(("nested_probe", None), None) + + worker = Worker( + task_definition_name="nested_probe", + execute_function=inspect.unwrap(nested_probe), + ) + # PicklingError or AttributeError depending on Python version; both + # are caught by TaskHandler.start_processes' actionable error handler. + with self.assertRaises((pickle.PicklingError, AttributeError)): + pickle.dumps(worker) + + +class TestWorkerRuntimeStateIsolation(unittest.TestCase): + def test_api_client_is_lazy_and_never_pickled(self): + worker = _make_decorated_worker() + self.assertIsNone(worker._api_client) # not built in the parent + + restored = pickle.loads(pickle.dumps(worker)) + self.assertIsNone(restored._api_client) # not built by unpickling + + # Property builds it on first use; setter allows injection. + sentinel = object() + restored.api_client = sentinel + self.assertIs(restored.api_client, sentinel) + + def test_lock_and_async_state_rebuilt_fresh_in_child(self): + worker = _make_decorated_worker() + worker._pending_async_tasks["stale-task-id"] = object() + + restored = pickle.loads(pickle.dumps(worker)) + + self.assertIsNotNone(restored._pending_tasks_lock) + self.assertIsNot(restored._pending_tasks_lock, worker._pending_tasks_lock) + self.assertEqual(restored._pending_async_tasks, {}) + self.assertIsNone(restored._background_loop) + # Lock must be functional. + with restored._pending_tasks_lock: + pass + + def test_original_worker_unaffected_by_pickling(self): + worker = _make_decorated_worker() + pickle.dumps(worker) + result = worker.execute(_make_task(x=5)) + self.assertEqual(result.output_data, {"doubled": 10}) + + +class TestExecuteFunctionReference(unittest.TestCase): + def test_reference_created_for_decorator_rebound_function(self): + original = inspect.unwrap(helpers.spawn_pickle_task) + ref = _importable_function_reference(original) + self.assertIsInstance(ref, _ExecuteFunctionReference) + self.assertIs(ref.resolve(), original) + + def test_no_reference_for_plain_module_function(self): + self.assertIsNone(_importable_function_reference(helpers.plain_function_worker)) + + def test_no_reference_for_lambda_nested_or_non_function(self): + self.assertIsNone(_importable_function_reference(lambda x: x)) + self.assertIsNone(_importable_function_reference("not-a-function")) + self.assertIsNone(_importable_function_reference(None)) + + def local_fn(x): + return x + + self.assertIsNone(_importable_function_reference(local_fn)) + + +class TestWorkerInRealSpawnChild(unittest.TestCase): + """The definitive regression test: transfer a @worker_task Worker into a + real 'spawn' child process (fresh interpreter, re-imports modules) and + execute a task there. This is exactly what TaskHandler does when + CONDUCTOR_MP_START_METHOD=spawn (default), and exactly what failed in + issues #264/#271.""" + + def test_worker_executes_in_spawn_child_process(self): + ctx = multiprocessing.get_context("spawn") + result_queue = ctx.Queue() + process = ctx.Process( + target=helpers.run_worker_in_child, + args=(_make_decorated_worker(), result_queue), + ) + process.start() + try: + status, output = result_queue.get(timeout=60) + finally: + process.join(timeout=10) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + self.assertEqual(process.exitcode, 0) + self.assertEqual(status, "COMPLETED") + self.assertEqual(output, {"doubled": 42}) + + +if __name__ == "__main__": + unittest.main()