Add with_context() for logging from subprocesses (#21304)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: desertaxle

.github/workflows/python-tests.yaml

@@ -131,6 +131,7 @@ jobs:
               tests/test_schedules.py
               tests/test_serializers.py
               tests/test_states.py
+              tests/test_subprocess_logging.py
               tests/test_task_engine.py
               tests/test_task_runs.py
               tests/test_task_worker.py

docs/v3/how-to-guides/workflows/add-logging.mdx

@@ -1,7 +1,7 @@
 ---
 title: How to add logging to a workflow
 sidebarTitle: Add logging
-keywords: ["logging", "log", "print", "debug", "log_prints", "get_run_logger"]
+keywords: ["logging", "log", "print", "debug", "log_prints", "get_run_logger", "subprocess", "multiprocessing", "with_context"]
 ---
 
 ### Emit custom logs
@@ -88,6 +88,43 @@ You can configure the default `log_prints` setting for all Prefect flow and task
 prefect config set PREFECT_LOGGING_LOG_PRINTS=True
 ```
 
+## Log from subprocesses
+
+When you spawn subprocesses inside a flow or task — for example, with `multiprocessing.Pool`
+or `concurrent.futures.ProcessPoolExecutor` — the Prefect run context is not automatically
+available in the child process. This means `get_run_logger()` raises a `MissingContextError`.
+
+Use `with_context` from `prefect.context` to propagate the current run context into
+subprocess workers. Logs emitted with `get_run_logger()` in the child process are
+associated with the parent flow run and task run and appear in the Prefect UI.
+
+{/* pmd-metadata: notest */}
+```python
+import multiprocessing
+from prefect import flow, task
+from prefect.context import with_context
+from prefect.logging import get_run_logger
+
+
+def process_item(item):
+    logger = get_run_logger()
+    logger.info(f"Processing {item}")
+    return item * 2
+
+
+@task
+def parallel_task(items):
+    with multiprocessing.Pool() as pool:
+        return pool.map(with_context(process_item), items)
+
+
+@flow
+def my_flow():
+    results = parallel_task([1, 2, 3, 4])
+```
+
+`with_context` also works with `concurrent.futures.ProcessPoolExecutor` and `multiprocessing.Process`.
+
 ## Access logs from the command line
 
 You can retrieve logs for a specific flow run ID using Prefect's CLI:

src/prefect/context.py

@@ -60,6 +60,53 @@
     from prefect.tasks import Task
 
 
+class _ContextWrappedCallable:
+    """Picklable callable that hydrates Prefect context before calling the
+    wrapped function.  The serialized context is stored as cloudpickle
+    bytes so that standard pickle (used by `multiprocessing`) can handle it."""
+
+    def __init__(
+        self, fn: Callable[..., Any], serialized_context: dict[str, Any]
+    ) -> None:
+        import cloudpickle
+
+        self.fn = fn
+        self._ctx_bytes = cloudpickle.dumps(serialized_context)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        import cloudpickle
+
+        ctx = cloudpickle.loads(self._ctx_bytes)
+        with hydrated_context(ctx):
+            return self.fn(*args, **kwargs)
+
+
+def with_context(fn: Callable[..., Any]) -> _ContextWrappedCallable:
+    """Wrap a function so it runs with the current Prefect context when
+    called in a subprocess.
+
+    Use this to enable `get_run_logger()` and other context-dependent
+    APIs in functions executed via `multiprocessing.Pool`,
+    `ProcessPoolExecutor`, or `multiprocessing.Process`.
+
+    Example:
+        ```python
+        from prefect.context import with_context
+
+        def worker(item):
+            logger = get_run_logger()
+            logger.info(f"Processing {item}")
+
+        @task
+        def my_task():
+            with multiprocessing.Pool() as pool:
+                pool.map(with_context(worker), items)
+        ```
+    """
+    ctx = serialize_context()
+    return _ContextWrappedCallable(fn, ctx)
+
+
 def serialize_context(
     asset_ctx_kwargs: Union[dict[str, Any], None] = None,
 ) -> dict[str, Any]:

tests/test_subprocess_logging.py

@@ -0,0 +1,179 @@
+"""Tests for `with_context` — propagating Prefect context to subprocesses."""
+
+import multiprocessing
+from concurrent.futures import ProcessPoolExecutor
+
+import pytest
+
+from prefect import flow, task
+from prefect.context import with_context
+from prefect.exceptions import MissingContextError
+from prefect.logging import get_run_logger
+
+# ---------------------------------------------------------------------------
+# Helper functions executed in child processes
+# ---------------------------------------------------------------------------
+
+
+def _log_and_return_ids(item: int) -> dict:
+    """Worker that uses get_run_logger() and returns context IDs."""
+    from prefect.context import FlowRunContext, TaskRunContext
+
+    logger = get_run_logger()
+    logger.info(f"Processing {item}")
+
+    flow_ctx = FlowRunContext.get()
+    task_ctx = TaskRunContext.get()
+    return {
+        "item": item,
+        "flow_run_id": str(flow_ctx.flow_run.id)
+        if flow_ctx and flow_ctx.flow_run
+        else None,
+        "task_run_id": str(task_ctx.task_run.id) if task_ctx else None,
+    }
+
+
+def _just_get_logger(_: object = None) -> str:
+    """Worker that calls get_run_logger() and returns its name."""
+    logger = get_run_logger()
+    return logger.name
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestWithContextPoolMap:
+    """multiprocessing.Pool.map with with_context."""
+
+    def test_get_run_logger_works_in_pool_worker(self):
+        @task
+        def my_task():
+            wrapped = with_context(_just_get_logger)
+            with multiprocessing.get_context("spawn").Pool(1) as pool:
+                results = pool.map(wrapped, [1])
+            return results
+
+        @flow
+        def my_flow():
+            return my_task()
+
+        results = my_flow()
+        assert len(results) == 1
+        assert isinstance(results[0], str)
+
+    def test_correct_run_ids_in_pool_worker(self):
+        @task
+        def my_task():
+            from prefect.context import FlowRunContext, TaskRunContext
+
+            flow_ctx = FlowRunContext.get()
+            task_ctx = TaskRunContext.get()
+            parent_flow_run_id = (
+                str(flow_ctx.flow_run.id) if flow_ctx and flow_ctx.flow_run else None
+            )
+            parent_task_run_id = str(task_ctx.task_run.id) if task_ctx else None
+
+            wrapped = with_context(_log_and_return_ids)
+            with multiprocessing.get_context("spawn").Pool(1) as pool:
+                results = pool.map(wrapped, [1, 2])
+            return results, parent_flow_run_id, parent_task_run_id
+
+        @flow
+        def my_flow():
+            return my_task()
+
+        results, parent_flow_id, parent_task_id = my_flow()
+        assert len(results) == 2
+        for r in results:
+            assert r["flow_run_id"] == parent_flow_id
+            assert r["task_run_id"] == parent_task_id
+
+
+class TestWithContextProcess:
+    """multiprocessing.Process with with_context."""
+
+    def test_get_run_logger_works_in_process(self):
+        @task
+        def my_task():
+            ctx = multiprocessing.get_context("spawn")
+            q: multiprocessing.Queue = ctx.Queue()  # type: ignore[type-arg]
+
+            wrapped = with_context(_target_with_queue)
+            p = ctx.Process(target=wrapped, args=(q,))
+            p.start()
+            p.join(timeout=30)
+            assert q.get_nowait() == "ok"
+
+        @flow
+        def my_flow():
+            return my_task()
+
+        my_flow()
+
+
+def _target_with_queue(queue: multiprocessing.Queue) -> None:  # type: ignore[type-arg]
+    """Process target that puts a value on a queue after using get_run_logger."""
+    logger = get_run_logger()
+    logger.info("hello from subprocess")
+    queue.put("ok")
+
+
+class TestWithContextProcessPoolExecutor:
+    """concurrent.futures.ProcessPoolExecutor with with_context."""
+
+    def test_get_run_logger_works_in_executor(self):
+        @task
+        def my_task():
+            wrapped = with_context(_just_get_logger)
+            ctx = multiprocessing.get_context("spawn")
+            with ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
+                future = executor.submit(wrapped)
+                return future.result(timeout=30)
+
+        @flow
+        def my_flow():
+            return my_task()
+
+        result = my_flow()
+        assert isinstance(result, str)
+
+
+class TestWithContextOutsideRun:
+    """with_context called outside a flow/task run should raise."""
+
+    def test_raises_missing_context_error(self):
+        """with_context still works outside a run — it serializes whatever
+        context exists (which may be empty). The error only happens when
+        the subprocess tries to call get_run_logger() without a run context."""
+        wrapped = with_context(_just_get_logger)
+        # The wrapper itself is created fine, but calling it in a subprocess
+        # without run context will fail when get_run_logger is called.
+        with pytest.raises(MissingContextError):
+            wrapped()
+
+
+class TestWithContextFlowOnly:
+    """with_context inside a flow (no task) propagates flow context."""
+
+    def test_flow_context_propagated(self):
+        @flow
+        def my_flow():
+            from prefect.context import FlowRunContext
+
+            flow_ctx = FlowRunContext.get()
+            parent_flow_run_id = (
+                str(flow_ctx.flow_run.id) if flow_ctx and flow_ctx.flow_run else None
+            )
+
+            wrapped = with_context(_log_and_return_ids)
+            ctx = multiprocessing.get_context("spawn")
+            with ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
+                future = executor.submit(wrapped, 42)
+                result = future.result(timeout=30)
+            return result, parent_flow_run_id
+
+        result, parent_flow_id = my_flow()
+        assert result["flow_run_id"] == parent_flow_id
+        assert result["task_run_id"] is None