Alastair
02/20/2026, 3:38 PMresource.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?Marvin
02/20/2026, 3:39 PMMarvin
02/20/2026, 3:46 PMresource.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):
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.Alastair
02/20/2026, 3:46 PMMarvin
02/20/2026, 3:46 PMMarvin
02/20/2026, 3:47 PMAlastair
02/20/2026, 3:48 PMMarvin
02/20/2026, 3:48 PMMarvin
02/20/2026, 3:49 PMresource.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:
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.Marvin
02/20/2026, 3:49 PMrun_deployment pattern for clean container-level isolation.Alastair
02/20/2026, 3:52 PMMarvin
02/20/2026, 3:52 PMMarvin
02/20/2026, 3:55 PM# 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.
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):
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:
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 upMarvin
02/20/2026, 3:55 PMDJANGO_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.Alastair
02/20/2026, 4:13 PMMarvin
02/20/2026, 4:13 PMMarvin
02/20/2026, 4:16 PM# 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:
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.