Ferchau Gupta
04/03/2026, 9:42 AMMarvin
04/03/2026, 9:42 AMMarvin
04/03/2026, 9:45 AMfrom prefect import task
from sqlalchemy import create_engine, text
@task(timeout_seconds=30)
def run_query():
engine = create_engine(
"<postgresql+psycopg2://user:pass@host:5432/dbname>",
connect_args={"options": "-c statement_timeout=25s -c lock_timeout=5s"},
pool_pre_ping=True,
)
with engine.connect() as conn:
# Optional: reinforce per-session
conn.exec_driver_sql("SET statement_timeout = '25s'")
conn.exec_driver_sql("SET lock_timeout = '5s'")
rows = conn.execute(text("SELECT ...")).all()
return rows
psycopg2 directly:
from prefect import task
import psycopg2
@task(timeout_seconds=30)
def run_query():
conn = psycopg2.connect("<postgresql://user:pass@host:5432/dbname>")
try:
with conn.cursor() as cur:
cur.execute("SET statement_timeout = 25000;")
cur.execute("SET lock_timeout = 5000;")
cur.execute("SELECT ...;")
return cur.fetchall()
finally:
conn.close()
Notes:
- Set the DB statement_timeout lower than the Prefect task timeout.
- If you use pgbouncer in transaction pooling, prefer SET LOCAL statement_timeout = '25s' inside each transaction or use the options connect arg so it’s applied per-connection.
- Consider idle_in_transaction_session_timeout and lock_timeout too if lock waits are an issue.
2) Proactively cancel the query when Prefect cancels the task
- You can run the query in a helper thread and periodically join; when Prefect injects cancellation into your task, you’ll catch it between joins and call conn.cancel() to ask Postgres to stop the running query.
from prefect import task
from prefect.exceptions import CancelledError, TimeoutError as PrefectTimeout
import psycopg2, threading
@task(timeout_seconds=30)
def cancellable_query(sql):
conn = psycopg2.connect("<postgresql://user:pass@host:5432/dbname>")
result = {"rows": None, "err": None}
def worker():
try:
with conn.cursor() as cur:
cur.execute("SET statement_timeout = 25000;")
cur.execute(sql)
result["rows"] = cur.fetchall()
except Exception as e:
result["err"] = e
t = threading.Thread(target=worker, daemon=True)
t.start()
try:
while t.is_alive():
t.join(0.2) # allow Prefect to raise cancellation between joins
except (CancelledError, PrefectTimeout):
try:
conn.cancel() # ask server to cancel in-flight query
except Exception:
pass
raise
finally:
conn.close()
if result["err"]:
raise result["err"]
return result["rows"]
3) Last resort: isolate the query in a subprocess and kill the process on cancel
- Heavier-weight, but guarantees termination even for stubborn drivers. Useful if DB-side timeouts are not an option.
Why it looked different locallyMarvin
04/03/2026, 9:45 AMFerchau Gupta
04/03/2026, 9:51 AMMarvin
04/03/2026, 9:51 AMMarvin
04/03/2026, 9:52 AMfrom prefect import task
from sqlalchemy import create_engine, text
from sqlalchemy.exc import OperationalError
# 25s DB timeout < 30s Prefect task timeout
engine = create_engine(
"<postgresql+psycopg2://user:pass@host:5432/dbname>",
pool_pre_ping=True,
connect_args={"options": "-c statement_timeout=25s -c lock_timeout=5s"},
)
@task(timeout_seconds=30)
def run_query(sql: str):
# Use an explicit transaction so we can SET LOCAL if desired
with engine.begin() as conn:
try:
# Optional: reinforce per-transaction timeouts (helps with pgbouncer txn pooling)
conn.exec_driver_sql("SET LOCAL statement_timeout = '25s'")
conn.exec_driver_sql("SET LOCAL lock_timeout = '5s'")
rows = conn.execute(text(sql)).all()
return rows
except OperationalError as e:
# Statement canceled or timed out — invalidate this DBAPI connection
# Common messages: "canceling statement due to user request", "statement timeout"
try:
conn.invalidate()
finally:
raise
Notes
- connect_args={"options": "-c statement_timeout=25s -c lock_timeout=5s"} applies a server-side timeout immediately on connection (works well even with pools).
- If you’re behind pgbouncer in transaction pooling mode, per-connection SET may not persist; keep the SET LOCAL ... inside each transaction as shown.
- After a cancellation/timeout, calling conn.invalidate() ensures a clean pool state. The context manager will roll back automatically.
Optional: streaming/chunked fetch to add interruption points
- Large result sets can be pulled in chunks to give Prefect a chance to raise the cancellation exception between fetches (still need the DB timeout for long-running execution on the server).
SQLAlchemy Core streaming:
from sqlalchemy import text
@task(timeout_seconds=30)
def stream_query(sql: str):
with engine.begin() as conn:
conn.exec_driver_sql("SET LOCAL statement_timeout = '25s'")
result = conn.execution_options(stream_results=True).execute(text(sql))
rows = []
try:
while True:
chunk = result.fetchmany(10_000) # adjust batch size
if not chunk:
break
rows.extend(chunk)
return rows
except Exception:
conn.invalidate()
raise
If you need to proactively cancel on Prefect timeout
- In addition to DB-side timeouts, you can run the query in a helper thread and on cancel call conn.connection.cancel() (psycopg2), which asks Postgres to abort the in-flight statement. Example available if you want it, but the server-side statement_timeout is usually sufficient and simpler.
A couple of quick checks
- Are you using pgbouncer (and which pooling mode)? If transaction pooling, keep the SET LOCAL inside each transaction.
- Are you using ORM Session or Core Connection? I can translate the pattern to a Session if that’s your setup.
If you share those details, I’ll tailor the snippet precisely to your environment.Ferchau Gupta
04/03/2026, 9:54 AMMarvin
04/03/2026, 9:54 AMMarvin
04/03/2026, 9:56 AMfrom prefect import flow, task
from prefect.task_runners import ProcessPoolTaskRunner
@task(timeout_seconds=30)
def slow_db_call(...):
# psycopg2 query here; still recommend statement_timeout
...
@flow(task_runner=ProcessPoolTaskRunner())
def my_flow(...):
slow_db_call.submit(...)
Notes:
- Works best on Linux (e.g., in Kubernetes). Windows can’t use signal-based interruption.
- Even with processes, DB-level timeouts are still recommended.
If you’d like, I can sketch the minimal changes to move your current 2.20.6 flow to Prefect 3 with ProcessPoolTaskRunner, plus the exact SQLAlchemy 1.4 + psycopg2 engine settings for statement_timeout in your environment.Ferchau Gupta
04/03/2026, 9:58 AMsession.execute(text(f"SET LOCAL statement_timeout = '{FLOW_RUN_TIMEOUT_SECONDS}s'"))Marvin
04/03/2026, 9:58 AMMarvin
04/03/2026, 10:00 AMSET LOCAL statement_timeout = 'Xs' “does nothing”, it’s almost always because it didn’t apply to the exact transaction/connection where the long query ran.
Key gotchas with SET LOCAL
- It only applies inside an active transaction. If there isn’t one, it’s ignored.
- It only applies to statements in the same transaction on the same connection. If SQLAlchemy checks out a different connection for the query, it won’t inherit the setting.
- With PgBouncer in transaction pooling, you must set it for every transaction that needs it (SET LOCAL inside the same BEGIN/COMMIT block). Session- or connection-level settings won’t persist.
- After a cancellation/timeout, invalidate the connection so a potentially bad connection isn’t returned to the pool.
A robust pattern (SQLAlchemy 1.4 + psycopg2, ORM Session)
from sqlalchemy.orm import Session
from sqlalchemy import text
from sqlalchemy.exc import OperationalError
def run_with_timeout(engine, sql: str, timeout_ms: int = 25000):
with Session(bind=engine) as session:
# Pin to a single connection + single transaction
with session.begin():
# Verify we’re in a tx and set timeout for THIS tx only
session.execute(text("SET LOCAL statement_timeout = :ms").bindparams(ms=timeout_ms))
# Optional: avoid long lock waits
session.execute(text("SET LOCAL lock_timeout = :ms").bindparams(ms=5000))
try:
result = session.execute(text(sql))
rows = result.fetchall()
return rows
except OperationalError:
# e.g. "canceling statement due to statement timeout"
# Ensure the DBAPI connection is not reused after cancel
session.get_bind().invalidate()
raise
Core-style (no Session)
from sqlalchemy import text
def run_with_timeout_core(engine, sql: str, timeout_ms: int = 25000):
# engine.begin() starts a transaction and ensures single connection scope
with engine.begin() as conn:
conn.exec_driver_sql(f"SET LOCAL statement_timeout = {timeout_ms}")
conn.exec_driver_sql("SET LOCAL lock_timeout = 5000")
try:
return conn.execute(text(sql)).all()
except Exception:
conn.invalidate()
raise
Quick validation inside your k8s pod
- This proves whether the environment honors the timeout:
with engine.begin() as conn:
conn.exec_driver_sql("SET LOCAL statement_timeout = '2s'")
try:
conn.execute(text("SELECT pg_sleep(5)")).all()
print("Unexpected: did not timeout")
except Exception as e:
print("Timed out as expected:", e)
- Also confirm the effective value right before the query:
show = conn.exec_driver_sql("SHOW statement_timeout").scalar_one()
print("Current statement_timeout:", show)
If you’re using PgBouncer
- Transaction pooling: always use SET LOCAL ... inside the same transaction as the target query (as shown). Per-session SET or connect_args={"options": "-c statement_timeout=..."} won’t persist across transactions.
- Session pooling or direct connections: you can set a default at connection time as a safety net:
engine = create_engine(
"<postgresql+psycopg2://user:pass@host:5432/db>",
pool_pre_ping=True,
connect_args={"options": "-c statement_timeout=25s -c lock_timeout=5s"},
)
Still keep SET LOCAL before sensitive queries if anything in your stack might reset parameters.
If it still doesn’t work, please share:
- Whether you’re behind PgBouncer (and which pooling mode)
- Whether you’re using ORM Session or Core Connection
- A minimal snippet showing how you open the Session/Connection and run both the SET LOCAL and the query
- The output of SHOW statement_timeout just before the query
- Whether SELECT pg_sleep(5) with SET LOCAL '2s' times out in your pod
ReferencesMarvin
04/03/2026, 10:00 AMFerchau Gupta
04/03/2026, 10:12 AMMarvin
04/03/2026, 10:12 AMMarvin
04/03/2026, 10:14 AMfrom sqlalchemy import create_engine, text
engine = create_engine("<postgresql+psycopg2://user:pass@host:5432/db>", pool_pre_ping=True)
with engine.begin() as conn: # ensures a transaction AND pins a single connection
print("Before:", conn.exec_driver_sql("SHOW statement_timeout").scalar_one())
conn.exec_driver_sql("SET LOCAL statement_timeout = '2s'")
print("After SET LOCAL (same tx):", conn.exec_driver_sql("SHOW statement_timeout").scalar_one())
try:
conn.execute(text("SELECT pg_sleep(5)")).all()
print("Unexpected: did not timeout")
except Exception as e:
print("Timed out as expected:", e)
Interpretation:
- If the second SHOW still prints the default (not 2s), you’re not actually in a transaction or PgBouncer is in statement pooling and splitting statements.
- If SHOW shows 2s but the sleep doesn’t time out, you’re likely not on the same connection for the SELECT (pool/Session scoping issue).
Make the pattern robust (SQLAlchemy 1.4 + psycopg2)
- ORM Session:
from sqlalchemy.orm import Session
from sqlalchemy import text
from sqlalchemy.exc import OperationalError
def run_with_timeout(engine, sql: str, timeout_s: int = 25):
with Session(bind=engine) as session:
# Single connection + single transaction
with session.begin():
session.execute(text("SET LOCAL statement_timeout = :t").bindparams(t=f"{timeout_s}s"))
session.execute(text("SET LOCAL lock_timeout = '5s'"))
try:
rows = session.execute(text(sql)).all()
return rows
except OperationalError:
session.get_bind().invalidate() # don’t return a cancelled connection to pool
raise
- Core:
from sqlalchemy import text
def run_with_timeout_core(engine, sql: str, timeout_s: int = 25):
with engine.begin() as conn: # single connection + transaction
conn.exec_driver_sql(f"SET LOCAL statement_timeout = '{timeout_s}s'")
conn.exec_driver_sql("SET LOCAL lock_timeout = '5s'")
try:
return conn.execute(text(sql)).all()
except Exception:
conn.invalidate()
raise
If you’re behind PgBouncer
- Transaction pooling: Use the patterns above (SET LOCAL inside the same transaction as the query). This should work.
- Statement pooling: SET LOCAL will not carry over to the SELECT — it’s effectively useless. Options:
- Switch that service/workload to transaction (or session) pooling.
- Set a default statement_timeout at the server or role/database level:
- ALTER ROLE myuser SET statement_timeout = '25s';
- ALTER DATABASE mydb SET statement_timeout = '25s';
- Or configure PgBouncer to enforce server_settings like statement_timeout=25s on all server conns.
- If you tried connect_args={"options": "-c statement_timeout=25s"} and it didn’t work, check PgBouncer’s ignore_startup_parameters — remove options from that list so it’s forwarded.
Other pitfalls to checkMarvin
04/03/2026, 10:14 AMFerchau Gupta
04/03/2026, 10:34 AMMarvin
04/03/2026, 10:34 AMMarvin
04/03/2026, 10:37 AMprefect worker start -p <your-pool> --limit 4
- Or set an env var in your k8s manifest:
env:
- name: PREFECT_WORKER_CONCURRENCY
value: "4"
- Alternatively, run multiple worker pods (e.g., 4 replicas), each with a limit of 1 or more.
- Verify the worker is polling the correct work queue(s) and seeing eligible runs (check logs for “Found X ready flow runs”). A queue mismatch will also lead to Late runs.
- Make sure no tag-based/global concurrency limits are effectively capping you at 1.
If you share:
- How you start the worker in k8s (command, env, replicas)
- Whether you set a work pool concurrency limit vs worker limit
- The worker logs around polling
…I can pinpoint the exact reason and provide the concrete k8s manifest/helm changes.Ferchau Gupta
04/03/2026, 10:40 AMprefect work-pool set-concurrency-limit Optimize_pool 4
echo "Starting workers for each work pools..."
prefect worker start --pool Optimize_poolMarvin
04/03/2026, 10:40 AMMarvin
04/03/2026, 10:43 AMprefect worker start --pool Optimize_pool --limit 4
- Or set it via env in your k8s manifest:
env:
- name: PREFECT_WORKER_CONCURRENCY
value: "4"
- Alternatively, run multiple worker pods (e.g., 4 replicas), each with a limit of 1–2.
Also verify
- The worker is polling the correct pool/queues (add --work-queue <name> if you use named queues).
- There aren’t other concurrency caps (tag-based/global) effectively reducing capacity to 1.
- Check worker logs; when at capacity, you’ll see it skip submitting additional runs, which explains runs going Late.
If you share your k8s manifest (env, command, replicas) and whether you use work queues, I can suggest the exact changes.Ferchau Gupta
04/03/2026, 12:29 PMMarvin
04/03/2026, 12:30 PMMarvin
04/03/2026, 12:31 PM--limit, the worker has no local concurrency cap and can start multiple flow runs in parallel (subject to Python resources and your infrastructure). That’s why you’re seeing multiple runs in parallel even though you didn’t set a worker limit.
How the pieces interact
- Worker limit (local): Controlled only by --limit on prefect worker start. If omitted, there is no per-worker cap. Example to cap at 4:
prefect worker start --pool Optimize_pool --limit 4
- Work pool concurrency (global): Your prefect work-pool set-concurrency-limit Optimize_pool 4 caps the total number of Running/Pending flow runs across all workers in that pool. It does not raise a worker’s capacity.
- Effective concurrency = min(pool limit, sum of all workers’ limits). With no worker limit, the worker can consume up to the pool’s global cap on its own.
Why runs still go Late with a long runner
- If the pool limit is 4 but one long flow run monopolizes resources inside the worker environment (e.g., a single-process task runner, GIL-bound tasks, or external bottlenecks), the worker may be technically “starting” runs but not making progress fast enough, causing queued runs to age into Late.
- If you do set --limit 1 on the worker, then a single long run will indeed block others and they’ll go Late until another worker instance picks them up.
Recommendations
- Decide explicitly how much parallelism you want per worker and set --limit to that number to avoid accidental over- or under-utilization.
- If you want strict isolation, run multiple worker pods with smaller per-worker limits (e.g., 4 replicas, --limit 1–2 each).
- Confirm there aren’t other caps (tag/global concurrency limits) or queue mismatches contributing to Late.
If you share your worker’s exact start command in k8s (including --limit, replicas, and task runner you use in deployments), I can suggest the precise tuning to match your intent.Ferchau Gupta
04/03/2026, 12:37 PMMarvin
04/03/2026, 12:37 PMMarvin
04/03/2026, 12:37 PMFerchau Gupta
04/03/2026, 12:37 PMMarvin
04/03/2026, 12:37 PMMarvin
04/03/2026, 12:38 PM--limit is low (e.g., 1), one long run can occupy the only slot and other runs go Late.
- Kubernetes worker: the worker just submits a Kubernetes Job (new pod) per flow run. This gives you per-run pod isolation and lets the cluster schedule multiple runs independently. It won’t block the worker as long as the worker’s --limit and the work pool’s concurrency allow more submissions.
What does not change
- Prefect’s sync task timeout still cannot preempt a single blocking psycopg2 call. Running the flow in a separate pod does not change Python’s inability to interrupt that call. Keep using Postgres statement_timeout (and fix the scoping/pooling details we discussed), or consider redesigns (async driver, subprocess isolation for the query, or in Prefect 3 use ProcessPoolTaskRunner).
To improve throughput with Kubernetes jobs
- Create/use a Kubernetes work pool and set the pool limit:
prefect work-pool create Optimize_k8s --type kubernetes
prefect work-pool set-concurrency-limit Optimize_k8s 4
- Start a Kubernetes worker with an explicit per-worker limit (how many Jobs it submits concurrently):
prefect worker start --pool Optimize_k8s --limit 4
- Point your deployment at the Optimize_k8s pool.
Optional safety net for stuck pods
- Add activeDeadlineSeconds to the Job template to hard-kill an overlong flow run pod (nuclear option; kills the whole run, not just the task).
Bottom line
- For the “Late” runs / parallelism: yes, running as K8s jobs can help, provided you set the worker --limit and pool concurrency appropriately and your cluster has capacity.
- For the blocking DB task timing out: no, that won’t change—fix with DB-side timeouts (ensuring correct scoping with SET LOCAL in the same transaction/connection) or an architectural change.Ferchau Gupta
04/03/2026, 12:40 PMMarvin
04/03/2026, 12:40 PMMarvin
04/03/2026, 12:40 PM--limit too low (e.g., 1), worker not polling the right queue, pool/queue/tag/global concurrency caps reached, or worker lacks resources to start more runs.
What changes with K8s Jobs vs process
- Process jobs: each run is a subprocess inside the worker pod. If the worker has --limit 1 or the pod is resource-starved, one long run can effectively block others.
- K8s Jobs: the worker just submits a Job per run; each run gets its own pod. This improves isolation and reduces contention in the worker pod. However, the worker’s --limit still governs how many Jobs it will submit concurrently, and the work pool’s concurrency still caps total concurrent runs.
What to do to resolve the queue blockage
- Set a per-worker limit high enough to submit multiple runs:
prefect worker start --pool Optimize_k8s --limit 4
- Keep/set the work pool concurrency at or above your target (you already did 4):
prefect work-pool set-concurrency-limit Optimize_k8s 4
- Optionally run multiple worker replicas for more submission capacity.
- Ensure the worker is polling the correct work queue(s) and has permissions to create Jobs in the cluster.
- Watch worker logs for “Found X ready flow runs” and “Submitted run …” to confirm it’s picking up more than one.
Bottom line
- If your current blockage is due to the worker pod only running one process job (limit/resource), moving to K8s Jobs plus setting --limit 4 will allow the worker to submit multiple runs and the cluster to execute them in parallel, preventing Late runs.
- If --limit remains 1 (or queues/limits mismatch), switching to K8s Jobs alone won’t prevent Late runs.
Docs: https://docs.prefect.io/latest/concepts/work-pools/#managing-concurrencyFerchau Gupta
04/03/2026, 12:46 PMMarvin
04/03/2026, 12:47 PMMarvin
04/03/2026, 12:47 PMFerchau Gupta
04/03/2026, 12:47 PMMarvin
04/03/2026, 12:47 PMMarvin
04/03/2026, 12:55 PMprefect worker start --pool <pool> --limit 4
- Then confirm the worker logs show it submitting up to 4 per cycle.
- Consider moving to Kubernetes Jobs for clean per-run isolation if process-level contention persists.
If you can share:
- Worker logs around polling/submission while one long run is active
- Worker pod CPU/memory limits/requests
- The command/env used to start the worker
- Any errors in the worker pod when additional runs should start
…I can pinpoint whether this is capacity, proposal, or resource contention and recommend the exact fix.