Syméon del Marmol
10/23/2025, 6:10 AMMarvin
10/23/2025, 6:10 AMMarvin
10/23/2025, 6:14 AMimport 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)
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.Marvin
10/23/2025, 6:14 AMmy_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.Syméon del Marmol
10/23/2025, 6:36 AMFinished 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.Marvin
10/23/2025, 6:37 AMMarvin
10/23/2025, 6:41 AMfrom 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
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.Marvin
10/23/2025, 6:41 AM