<@ULVA73B9P> if I have a flow with a ProcessPoolTa...
# ask-marvin
a
@Marvin if I have a flow with a ProcessPoolTaskRunner, is there a way of performing any hooks on startup of a process? to be specific, I'm wanting to run
resource.setrlimit
to set a hard memory limit on the process. Can I do this in a startup hook, or will it work if I just run it at the beginnig of a
@task
that I submit to that pool?
m
thought for 422.3 seconds
Short answer: - There isn’t a supported “process startup” hook in Prefect 3.x’s ProcessPoolTaskRunner. - Calling
resource.setrlimit
at the start of a task will run inside the worker process and works. - Be aware: limits you set will persist for the lifetime of that worker process and can affect subsequent tasks executed by the same process unless you reset them. Details and references: - ProcessPoolTaskRunner uses a ProcessPoolExecutor with the “spawn” start method and does not expose `initializer`/`initargs` like Python’s ProcessPoolExecutor does. Source: https://github.com/PrefectHQ/prefect/blob/main/src/prefect/task_runners.py#L864-L873 - Task runner concepts: https://docs.prefect.io/v3/concepts/task-runners/ - Python ProcessPoolExecutor (shows `initializer`/`initargs` that Prefect does not expose): https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor Recommended pattern (per-task, and reversible so you don’t affect later tasks on the same worker):
Copy code
import resource
from contextlib import contextmanager
from prefect import task, flow
from prefect.task_runners import ProcessPoolTaskRunner

@contextmanager
def limit_mem(soft_bytes=None, hard_bytes=None, rlimit=resource.RLIMIT_AS):
    old_soft, old_hard = resource.getrlimit(rlimit)
    try:
        new_soft = soft_bytes if soft_bytes is not None else old_soft
        new_hard = hard_bytes if hard_bytes is not None else old_hard
        resource.setrlimit(rlimit, (new_soft, new_hard))
        yield
    finally:
        # Restore original limits so the worker process doesn't carry them to the next task
        resource.setrlimit(rlimit, (old_soft, old_hard))

@task
def do_work():
    with limit_mem(hard_bytes=2_000_000_000):  # 2 GB hard limit (Linux: typically RLIMIT_AS)
        # your task body
        ...

@flow(task_runner=ProcessPoolTaskRunner())
def my_flow():
    do_work.submit()
Notes: - Which limit to use depends on your OS and what you want to cap: - On Linux,
RLIMIT_AS
(address space) is most commonly used for a memory cap;
RLIMIT_RSS
is often ignored. - On macOS, enforcement of some limits can be inconsistent. - If your production runs use Docker/Kubernetes workers, consider container-level memory limits (cgroups) for stronger isolation: - Docker worker job config supports memory limits (e.g.,
memory
,
memory_swap
) in the job template. Docs: https://docs.prefect.io/v3/deploy/infrastructure/docker/ If you really need a pool-wide initializer (run once per worker when it starts), Prefect doesn’t currently expose that on ProcessPoolTaskRunner. You could hack around it by subclassing and constructing the underlying ProcessPoolExecutor with an initializer, but that’s not supported and could break with Prefect updates. For most cases, the per-task context wrapper above is the safest path. If you can share your OS and which specific rlimit you plan to set (e.g., RLIMIT_AS vs RLIMIT_DATA), I can tailor the snippet accordingly.
a
@Marvin are worker processes long-lived and handle multiple tasks that get submitted, or do they start up a new one every time a task is submitted?
m
thought for 64.0 seconds
Short answer: they’re long-lived for the lifetime of the task runner within a single flow run. Prefect creates a single ProcessPoolExecutor when the ProcessPoolTaskRunner context opens and reuses its worker processes for all tasks you submit during that flow run; the pool is shut down when the runner exits. Details with source: - Pool creation (on enter): https://github.com/PrefectHQ/prefect/blob/main/src/prefect/task_runners.py#L864-L873 - Pool shutdown (on exit/cancel): https://github.com/PrefectHQ/prefect/blob/main/src/prefect/task_runners.py#L847-L884 Other notes: - Prefect does not set max_tasks_per_child or recycle workers; no built-in per-worker task limit is applied. - Start method is “spawn”, so each worker is a fresh interpreter, but it’s reused across tasks until the pool is torn down.
a
@Marvin I'm inside docker (ECS) but I want to prevent a runaway malloc inside a single task from taking down the entire container
m
thought for 95.2 seconds
Great question. In Docker/ECS, a single process that goes wild can push the container over its cgroup limit and get the whole task killed. To contain a single Prefect task’s memory use, you need to bound the memory of the process actually doing the work. Key points: - ProcessPoolTaskRunner reuses worker processes. If a worker process gets OOM-killed, the pool can become unusable (BrokenProcessPool), impacting the entire flow run. - Putting
resource.setrlimit
at the top of the task does apply to that worker process, but if it still gets killed, you risk breaking the pool. Two robust patterns to avoid taking down the whole container: 1) Run the risky code in a short-lived subprocess inside the task (recommended) - Spawn a child process from within your Prefect task. - In the child, set strict rlimits (e.g., RLIMIT_AS) before running the risky code. - If the child hits the limit, it fails or is killed, but your Prefect worker process stays alive and you can return a controlled error. Example:
Copy code
import os
import traceback
import resource
from multiprocessing import get_context
from prefect import task, flow
from prefect.task_runners import ProcessPoolTaskRunner

def _child_with_limits(target_fn_path, args, kwargs, hard_bytes, queue):
    # Set per-process virtual memory cap (Linux). Use RLIMIT_AS for a hard cap.
    resource.setrlimit(resource.RLIMIT_AS, (hard_bytes, hard_bytes))
    try:
        # Dynamically import the target (module:function)
        mod_name, fn_name = target_fn_path.rsplit(":", 1)
        mod = __import__(mod_name, fromlist=[fn_name])
        fn = getattr(mod, fn_name)
        result = fn(*args, **kwargs)
        queue.put(("ok", result))
    except Exception as e:
        queue.put(("err", (type(e).__name__, str(e), traceback.format_exc())))

# Example risky function you control (must be importable at module level)
def mypkg_risky_func(n):
    # ... code that might allocate a lot ...
    return n * 2

@task
def safe_wrapper_for_risky(n, mem_hard_bytes=2_000_000_000, timeout_s=300):
    ctx = get_context("spawn")
    q = ctx.Queue()
    p = ctx.Process(
        target=_child_with_limits,
        args=("your_module_path:mypkg_risky_func", (n,), {}, mem_hard_bytes, q),
    )
    p.start()
    p.join(timeout_s)

    if p.exitcode is None:
        p.terminate()
        raise TimeoutError(f"Child process timed out after {timeout_s}s")

    # If the child was killed by the kernel due to limits, exitcode may be negative (signal)
    if p.exitcode != 0 and q.empty():
        raise RuntimeError(f"Child exited abnormally with code {p.exitcode}")

    status, payload = q.get_nowait()
    if status == "ok":
        return payload
    else:
        exc_type, msg, tb = payload
        raise RuntimeError(f"Child error: {exc_type}: {msg}\n{tb}")

@flow(task_runner=ProcessPoolTaskRunner())
def my_flow():
    f = safe_wrapper_for_risky.submit(123, mem_hard_bytes=1_500_000_000)
    return f.result()
Notes: - Use RLIMIT_AS for Linux to cap virtual memory; adjust as needed. - The risky code runs in a one-off child. If it dies, the Prefect worker process (and pool) remain healthy. - Make sure the function you call is importable by name in the child (module:function). 2) Use container-level isolation by running the risky step as its own flow/deployment - Move the risky work into a separate flow and call it via
run_deployment
. - With a Docker or Kubernetes worker, that subflow runs in its own container with its own memory limit, so any OOM only affects that container. - Docs: Docker infra/worker config https://docs.prefect.io/v3/deploy/infrastructure/docker/ What about just calling resource.setrlimit at the start of a task? - It can work, but if the worker process dies from the limit, the entire process pool may become broken for the rest of the flow run. If you do this, wrap it in a context manager and restore the limits after the risky section to avoid impacting later tasks in the same worker process.
If you share: - Linux distro inside ECS, and - Whether you can split the risky step into a subflow, I can tailor the snippet (RLIMIT choice, error handling) and/or show a
run_deployment
pattern for clean container-level isolation.
a
@Marvin inside the container it's just the python-slim container, so I guess debian. At the moment we have only one flow for the entire process, but I'm open to splitting this bad step into a separate subflow. crucially, this is also part of a django app, so I need to do django setup at the beginning of the flow. I don't know if this is necessary in each subtask which runs in a new process, or if part of the process setup will perform this setup process
m
thought for 187.5 seconds
Thanks — that context helps. Key point for Django + ProcessPoolTaskRunner: - Prefect’s process pool uses the “spawn” start method. Each worker is a fresh Python interpreter. Any global state (including Django app registry) is not inherited from the parent flow process. - Therefore, you need to initialize Django once per process that will touch Django (each task process and any short-lived subprocess you create). A safe, idempotent helper you can call at the top of tasks/subprocesses:
Copy code
# myproject/django_setup.py
import django
from django.apps import apps

def ensure_django():
    # Requires DJANGO_SETTINGS_MODULE in env (ECS task def / Prefect deployment env)
    if not apps.ready:
        django.setup()
Option A: Keep one flow, isolate the risky step in a short-lived subprocess with a hard memory cap - Use RLIMIT_AS on Debian (python-slim). RLIMIT_RSS is typically ignored on Linux. - Initialize Django inside the child before running your code.
Copy code
import resource
import traceback
from multiprocessing import get_context
from queue import Empty
from myproject.django_setup import ensure_django

def _child_entry(target_mod_fn, args, kwargs, hard_bytes, queue):
    # Per-process hard memory cap
    resource.setrlimit(resource.RLIMIT_AS, (hard_bytes, hard_bytes))
    ensure_django()
    try:
        mod_name, fn_name = target_mod_fn.rsplit(":", 1)
        mod = __import__(mod_name, fromlist=[fn_name])
        fn = getattr(mod, fn_name)
        res = fn(*args, **kwargs)
        queue.put(("ok", res))
    except Exception as e:
        queue.put(("err", (type(e).__name__, str(e), traceback.format_exc())))

def run_with_mem_cap(target_mod_fn, *args, hard_bytes, timeout_s=300, **kwargs):
    ctx = get_context("spawn")
    q = ctx.Queue()
    p = ctx.Process(target=_child_entry, args=(target_mod_fn, args, kwargs, hard_bytes, q))
    p.start()
    p.join(timeout_s)

    if p.exitcode is None:
        p.terminate()
        raise TimeoutError(f"Timed out after {timeout_s}s")
    if p.exitcode != 0 and q.empty():
        # Likely killed by signal (e.g., out of memory)
        raise RuntimeError(f"Child exited with code {p.exitcode}")

    status, payload = q.get_nowait()
    if status == "ok":
        return payload
    else:
        etype, msg, tb = payload
        raise RuntimeError(f"Child error: {etype}: {msg}\n{tb}")
Use it in your Prefect task (Django used in parent task too):
Copy code
from prefect import task

@task
def risky_step(n: int):
    # Parent process: ensure Django if you also touch it here
    from myproject.django_setup import ensure_django
    ensure_django()

    # Run the real work in a capped child
    return run_with_mem_cap("myproject.risky:do_risky_work", n, hard_bytes=2_000_000_000)
Why this helps: - If the risky code allocates too much, only the child dies — the Prefect worker process (and pool) stays healthy, so the flow can handle the failure gracefully. - Because workers are long-lived, avoid setting process-wide limits directly in the task without restoring them; otherwise later tasks on the same worker inherit the cap. Option B: Split the risky step into a separate flow/deployment so it runs in its own container - This gives the strongest isolation: set container memory limits and let cgroups enforce them. - In the risky subflow’s module, call
ensure_django()
at the top or at the first task that uses Django. - From the main flow, call the risky flow’s deployment as a subflow:
Copy code
from prefect import flow
from prefect.deployments import run_deployment

@flow
def main_flow(...):
    # do normal work...
    # then run risky part in its own container with its own memory limit
    fr = run_deployment(
        name="risky-flow/risky-deployment",
        parameters={"n": 123},
        as_subflow=True,
        timeout=900,
    )
    # optional: poll fr.state or just rely on exceptions bubbling up
- For Docker-based runs, set memory in the risky deployment’s job_variables (so its container can’t exceed, say, 2GiB). Docs: https://docs.prefect.io/v3/deploy/infrastructure/docker/ - Ensure DJANGO_SETTINGS_MODULE is set for that deployment environment too. What to set up where: - Environment: Set
DJANGO_SETTINGS_MODULE
(and any secrets) in ECS task definition or Prefect deployment environment. Prefect will pass env vars to subprocesses, but each process must still call
django.setup()
(via
ensure_django()
). - In-process tasks (ProcessPoolTaskRunner): Call
ensure_django()
at the top of any task that uses Django. - Subprocess child (Option A): Call
ensure_django()
inside the child entry before running the risky function. - Separate subflow (Option B): Call
ensure_django()
in that flow’s code as well. If you tell me: - Which memory cap you want for the risky step (e.g., 1.5–2 GiB), - Whether you’d prefer Option A (subprocess) or Option B (separate container via subflow), I’ll tailor a drop-in snippet for your repo structure and show exactly where to place
ensure_django()
and env settings.
a
@Marvin if the risky step doesn't actually require django then perhaps I could use a task task_runner, and then just make a quick subprocess (with a memory limit) for that single step, is that a reasonable idea?
m
thought for 192.8 seconds
That’s a reasonable idea, with one caveat. - In Prefect 3.x, the task runner is flow-scoped, not per-task. You cannot assign a different task runner to just one task. Docs: https://docs.prefect.io/v3/concepts/task-runners/ Viable pattern: - Keep your flow’s task runner as-is (e.g., ThreadPoolTaskRunner for most tasks or even Sequential for simplicity). - For the single risky step that doesn’t need Django, spawn a short-lived subprocess inside that task and set a strict memory limit (RLIMIT_AS on Debian). If it OOMs, only that child dies; the Prefect worker process (and thus the pool and the rest of the flow) stays healthy. Skeleton:
Copy code
# Inside your risky Prefect task:
from multiprocessing import get_context
import resource, traceback
from queue import Empty

def _child_entry(mod_fn, args, kwargs, hard_bytes, q):
    # Cap per-process virtual memory (Linux/Debian)
    resource.setrlimit(resource.RLIMIT_AS, (hard_bytes, hard_bytes))
    try:
        mod, fn = mod_fn.rsplit(":", 1)
        m = __import__(mod, fromlist=[fn])
        res = getattr(m, fn)(*args, **kwargs)
        q.put(("ok", res))
    except Exception as e:
        q.put(("err", (type(e).__name__, str(e), traceback.format_exc())))

def run_with_mem_cap(mod_fn, *args, hard_bytes, timeout_s=300, **kwargs):
    ctx = get_context("spawn")
    q = ctx.Queue()
    p = ctx.Process(target=_child_entry, args=(mod_fn, args, kwargs, hard_bytes, q))
    p.start(); p.join(timeout_s)
    if p.exitcode is None:
        p.terminate()
        raise TimeoutError(f"Timed out after {timeout_s}s")
    if p.exitcode != 0 and q.empty():
        raise RuntimeError(f"Child exited with code {p.exitcode}")
    status, payload = q.get_nowait()
    if status == "ok":
        return payload
    etype, msg, tb = payload
    raise RuntimeError(f"Child error: {etype}: {msg}\n{tb}")
Use it:
Copy code
from prefect import task

@task
def risky_without_django(n: int):
    return run_with_mem_cap("yourpkg.risky:do_risky_work", n, hard_bytes=2_000_000_000)
Why this works well here: - No Django init required in the child since the risky step doesn’t need it. - Failure is contained to the one-off child process. - You keep your existing flow structure and runner choice. If later you want even stronger isolation or resource control, consider splitting the risky step into its own subflow/deployment so it runs in a separate container with its own ECS memory limit, and call it via
run_deployment
from the main flow.