<@ULVA73B9P> within a task, is there a way to use ...
# ask-marvin
s
@Marvin within a task, is there a way to use multiprocessing to increase parallelism and use multi-core, without creation sub-tasks ? I tried to use a standard ProcessPoolExecutor, but the function didn't start as expected. In addition, how can I propagate the task and flow context to capture logs.
m
thought for 221.9 seconds
Short answer: yes, you can use multiprocessing inside a task, but you’ll need a few guardrails. Prefect’s flow/task context does not automatically exist in child processes, so you must serialize and re-hydrate the context to get proper logging/metadata. Also ensure you’re using a “spawn” start method and only submit top-level, picklable functions. Prefer Prefect-native parallelism if you can: - If you’re okay with creating multiple task runs, use mapped/background tasks (your current .delay() + asyncio.gather pattern) — this gives you observability, retries, and clean logs “for free”. - For multi-core parallelism at the flow level, use a process-based task runner (e.g., ProcessPoolTaskRunner, DaskTaskRunner, RayTaskRunner): Task runners If you must parallelize within a single task run (no sub-tasks), here’s a working pattern for context and logs. Key rules - Use spawn start method for portability and stability. - Worker functions must be defined at module top-level (no closures/lambdas). - Don’t pass Prefect objects (logger/client) directly to subprocesses; pass a serialized context dict. Synchronous example with ProcessPoolExecutor
Copy code
import os
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed

from prefect import flow, task, get_run_logger
from prefect.context import serialize_context, hydrated_context

def cpu_bound_worker(x: int, serialized_ctx: dict) -> int:
    # Rehydrate Prefect context in the child process
    with hydrated_context(serialized_ctx):
        logger = get_run_logger()
        <http://logger.info|logger.info>(f"Child PID {os.getpid()} processing {x}")
        # do CPU work
        return x * x

@task
def compute_squares(n: int) -> int:
    logger = get_run_logger()
    <http://logger.info|logger.info>(f"Parent PID {os.getpid()} starting ProcessPoolExecutor")

    # Serialize the current flow/task context so logs/metadata are attributed correctly
    serialized = serialize_context()

    ctx = mp.get_context("spawn")
    results = []
    with ProcessPoolExecutor(max_workers=os.cpu_count(), mp_context=ctx) as ex:
        futures = [ex.submit(cpu_bound_worker, i, serialized) for i in range(n)]
        for f in as_completed(futures):
            # calling result() surfaces child exceptions immediately if any
            results.append(f.result())

    <http://logger.info|logger.info>(f"Got {len(results)} results")
    return sum(results)

@flow
def run():
    total = compute_squares(10)
    print(total)

if __name__ == "__main__":
    run()
Async task version (using the same worker)
Copy code
import asyncio
import functools
import os
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor

from prefect import task, get_run_logger
from prefect.context import serialize_context

def cpu_bound_worker(x: int, serialized_ctx: dict) -> int:
    from prefect import get_run_logger  # safe import here too
    from prefect.context import hydrated_context
    with hydrated_context(serialized_ctx):
        logger = get_run_logger()
        <http://logger.info|logger.info>(f"Child PID {os.getpid()} processing {x}")
        return x * x

@task
async def compute_squares_async(n: int) -> int:
    logger = get_run_logger()
    <http://logger.info|logger.info>(f"Parent PID {os.getpid()} starting ProcessPoolExecutor")

    serialized = serialize_context()
    loop = asyncio.get_running_loop()
    ctx = mp.get_context("spawn")

    with ProcessPoolExecutor(max_workers=os.cpu_count(), mp_context=ctx) as ex:
        tasks = [
            loop.run_in_executor(ex, functools.partial(cpu_bound_worker, i, serialized))
            for i in range(n)
        ]
        results = await asyncio.gather(*tasks)

    <http://logger.info|logger.info>(f"Got {len(results)} results")
    return sum(results)
Why your ProcessPoolExecutor may not have started - Function not picklable: it was nested or a lambda. Move it to module top-level.
- Using fork on Linux: can lead to deadlocks or silent hangs with async/logging. Use spawn via mp.get_context("spawn") and pass mp_context=ctx to ProcessPoolExecutor. - Passing non-picklable args (e.g., Prefect logger/client) to the subprocess. - Exceptions in children are only raised when you call future.result(); if you never awaited/collected results, it can look like “didn’t start”. - Running from an interactive environment (e.g., notebooks) without proper module import semantics; prefer running from a Python module/script. About logging and context in subprocesses - Use serialize_context() before spawning and hydrated_context(...) inside the worker. Then get_run_logger() will work and logs will be attributed to the parent task/flow run. - Standard logging in subprocesses won’t be captured by Prefect unless the context is hydrated (or you manually configure logging similarly); hydrated_context is the simplest path. Trade-offs - Multiprocessing inside a task keeps one task run, but you lose Prefect-native retries/caching/visibility for each item. - If you want each unit of work to have its own name, retries, and logs, keep using your background tasks pattern:
my_task.with_options(name=...).delay(...)
+
await asyncio.gather(*(t.result_async() ...))
. - For multi-core on a single machine without writing multiprocessing code, consider a process-based task runner at the flow level: Task runners. For large-scale CPU parallelism, Dask/Ray runners handle context/logging out of the box. Links - Task runners overview: Prefect task runners - Running work concurrently: How to run work concurrently - Context API (serialize_context, hydrated_context): prefect.context API - Logging API: prefect.logging.loggers Questions to tailor the advice - Is your workload strictly CPU-bound, and is it okay for all logs to roll up under a single task run? - Which OS are you running workers on (Windows/macOS/Linux)? This affects the start method and any extra guardrails needed.
s
@Marvin during the serialization of the context, I get the following error:
Finished in state Failed("Flow run encountered an exception: PicklingError: Can't pickle <function dummy_flow at 0xffff99915080>: it's not the same object as dummy.dummy_flow")
dummy_flow
is the name of the flow currently running, which is in the
dummy.py
file.
m
thought for 262.8 seconds
Thanks for the error detail — that PicklingError is a classic multiprocessing gotcha when the parent process sees your flow function as main.dummy_flow but the child process imports it as dummy.dummy_flow. With Prefect context serialization, the serialized dict includes a reference to the Flow (which holds the flow function), so when multiprocessing pickles that dict it tries to pickle the function and hits the mismatch. You’ve got a few solid ways to fix this: Option A — Make the serialized context “pickle-safe” by dropping callable-heavy fields - Remove the flow object from the serialized flow_run_context and the task object from task_run_context before passing to subprocesses. - hydrated_context() will still hydrate the run contexts sufficiently for logging and metadata. Example
Copy code
from prefect.context import serialize_context, hydrated_context
from prefect import task, flow, get_run_logger
from concurrent.futures import ProcessPoolExecutor
import multiprocessing as mp
import os

def worker(x, serialized_ctx):
    from prefect import get_run_logger
    from prefect.context import hydrated_context
    with hydrated_context(serialized_ctx):
        logger = get_run_logger()
        <http://logger.info|logger.info>(f"Child PID {os.getpid()} processing {x}")
        return x * x

@task
def do_work(n: int) -> int:
    logger = get_run_logger()
    serialized = serialize_context()

    # Strip objects that cause pickling issues
    frc = serialized.get("flow_run_context") or {}
    frc.pop("flow", None)
    trc = serialized.get("task_run_context") or {}
    trc.pop("task", None)

    ctx = mp.get_context("spawn")  # safer across OSes
    with ProcessPoolExecutor(max_workers=os.cpu_count(), mp_context=ctx) as ex:
        futures = [ex.submit(worker, i, serialized) for i in range(n)]
        return sum(f.result() for f in futures)

@flow
def dummy_flow():
    return do_work(10)

if __name__ == "__main__":
    dummy_flow()
Option B — Pre-serialize the whole context with cloudpickle and send bytes - This avoids the standard pickle trying to re-import your function via module name. - You send bytes (which are pickleable), then cloudpickle.loads() in the child. Example
Copy code
import cloudpickle, pickle
from prefect.context import serialize_context, hydrated_context

def worker(x, ctx_blob: bytes):
    serialized = cloudpickle.loads(ctx_blob)
    with hydrated_context(serialized):
        from prefect import get_run_logger
        logger = get_run_logger()
        <http://logger.info|logger.info>(f"processing {x}")
        return x * x

@task
def do_work(n: int) -> int:
    ctx_blob = cloudpickle.dumps(serialize_context(), protocol=pickle.HIGHEST_PROTOCOL)
    ctx = mp.get_context("spawn")
    with ProcessPoolExecutor(max_workers=os.cpu_count(), mp_context=ctx) as ex:
        futures = [ex.submit(worker, i, ctx_blob) for i in range(n)]
        return sum(f.result() for f in futures)
Option C — Fix the module import mismatch - Ensure your flow function lives in an importable module (not just main). - Run your script in a way that preserves its module name (e.g., python -m yourpkg.dummy), or install your package and import the flow from it. - Always guard your entrypoint with if name == "__main__": to avoid unintended import-time side effects when using spawn. Other guardrails - Use spawn: ctx = multiprocessing.get_context("spawn") and pass mp_context=ctx to ProcessPoolExecutor. - Only pass picklable arguments to subprocesses; don’t pass Prefect loggers/clients directly. - Define worker functions at module top level (no lambdas/closures). If your goal is simply multi-core with clean Prefect observability, also consider: - Using a process-based task runner at the flow level (e.g., ProcessPoolTaskRunner, or Dask/Ray runners) so you don’t need to hand-roll multiprocessing: Task runners - Or keep your current background task fan-out (.delay() + .result_async()) to get per-unit retries, names, and logs.
Relevant docs - Context API (serialize_context, hydrated_context): prefect.context API - Running work concurrently: How to run work concurrently If you can share how you invoke dummy.py (python dummy.py vs python -m pkg.dummy) and where the flow is defined, I can point to the most targeted fix for your setup.