python-pq 0.7.1__tar.gz → 0.8.0__tar.gz

{python_pq-0.7.1 → python_pq-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-pq
-Version: 0.7.1
+Version: 0.8.0
 Summary: Postgres-backed job queue for Python
 Keywords: postgres,job-queue,task-queue,background-jobs,worker
 Author: ricwo

{python_pq-0.7.1 → python_pq-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "python-pq"
-version = "0.7.1"
+version = "0.8.0"
 description = "Postgres-backed job queue for Python"
 readme = "README.md"
 authors = [

{python_pq-0.7.1 → python_pq-0.8.0}/src/pq/__init__.py

@@ -5,7 +5,7 @@ from pq.models import Periodic, Task, TaskStatus
 from pq.priority import Priority
 from pq.worker import PostExecuteHook, PreExecuteHook, TaskTimeoutError
 
-__version__ = "0.7.1"
+__version__ = "0.8.0"
 
 __all__ = [
     "PQ",

{python_pq-0.7.1 → python_pq-0.8.0}/src/pq/client.py

@@ -1,6 +1,7 @@
 """PQ client - main interface for task queue."""
 
 import importlib.resources
+import math
 from collections.abc import Callable, Set
 from contextlib import contextmanager
 from datetime import UTC, datetime, timedelta
@@ -10,6 +11,7 @@ from typing import Any, Self
 from croniter import croniter
 from croniter.croniter import CroniterBadCronError
 from loguru import logger
+import sqlalchemy as sa
 from sqlalchemy import create_engine, delete, func, select, update
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.engine import Engine
@@ -21,6 +23,45 @@ from pq.registry import get_function_path
 from pq.serialization import serialize
 
 
+def _validate_max_runtime(max_runtime_seconds: float | None) -> None:
+    """Reject obviously-degenerate per-task ``max_runtime_seconds`` values.
+
+    ``None`` means "use the worker's configured default" and is the
+    common case — nothing to validate. ``0``, negative, NaN, and
+    infinity are all rejected because the downstream enforcement is
+    undefined for those:
+
+    - ``signal.alarm(int(max_runtime_seconds) + 1)`` would fire at 1 s for
+      ``0``, at 0 s (which disables the alarm entirely) for any
+      negative integer ≥ -1, and raise ``ValueError`` from inside
+      the worker for NaN / infinity.
+    - ``max_runtime_seconds * 2`` in the stale-reaper SQL would be ≤ 0 or
+      NaN, making the per-row threshold a no-op or NULL (the global
+      default would apply) — silent surprise vs. the call site's
+      intent.
+    - Periodic ``lock_duration`` already guards with a 3600 s
+      fallback for ``≤ 0`` (``worker.py``), but that fallback exists
+      for the worker-level default, not as a contract for callers.
+
+    None of these are useful behaviours to expose. Failing fast at
+    enqueue/upsert/schedule time keeps the bug at the call site,
+    where it's trivially fixable, instead of producing a confusing
+    runtime surprise hours later.
+    """
+    if max_runtime_seconds is None:
+        return
+    if math.isnan(max_runtime_seconds) or math.isinf(max_runtime_seconds):
+        raise ValueError(
+            f"max_runtime_seconds must be a finite positive number (got {max_runtime_seconds!r}); "
+            f"pass None to use the worker's configured default."
+        )
+    if max_runtime_seconds <= 0:
+        raise ValueError(
+            f"max_runtime_seconds must be > 0 (got {max_runtime_seconds!r}); pass None to "
+            f"use the worker's configured default."
+        )
+
+
 class PQ:
     """Postgres-backed task queue client."""
 
@@ -125,6 +166,7 @@ class PQ:
         run_at: datetime | None = None,
         priority: Priority = Priority.NORMAL,
         client_id: str | None = None,
+        max_runtime_seconds: float | None = None,
         **kwargs: Any,
     ) -> int:
         """Enqueue a one-off task.
@@ -135,15 +177,24 @@ class PQ:
             run_at: When to run the task. Defaults to now.
             priority: Task priority. Higher = higher priority. Defaults to NORMAL.
             client_id: Optional client-provided identifier. Must be unique if provided.
+            max_runtime_seconds: Per-task wall-clock cap in seconds. When ``None`` (the
+                common case), the worker that picks up this task uses its
+                configured default. When set, this specific task gets the
+                supplied ceiling AND the stale-reaper threshold scales with
+                it (``max_runtime_seconds * 2`` if larger than the reaper's global
+                default). Use this for occasionally-long tasks in a fleet
+                whose worker default is sized for typical short work.
             **kwargs: Keyword arguments to pass to the handler.
 
         Returns:
             Task ID.
 
         Raises:
-            ValueError: If task is a lambda, closure, or cannot be imported.
+            ValueError: If task is a lambda, closure, or cannot be imported,
+                or if ``max_runtime_seconds`` is not None and ``<= 0``.
             IntegrityError: If client_id already exists.
         """
+        _validate_max_runtime(max_runtime_seconds)
         name = get_function_path(task)
         payload = serialize(args, kwargs)
 
@@ -156,6 +207,7 @@ class PQ:
             run_at=run_at,
             priority=priority,
             client_id=client_id,
+            max_runtime_seconds=max_runtime_seconds,
         )
 
         with self.session() as session:
@@ -170,6 +222,7 @@ class PQ:
         run_at: datetime | None = None,
         priority: Priority = Priority.NORMAL,
         client_id: str,
+        max_runtime_seconds: float | None = None,
         **kwargs: Any,
     ) -> int:
         """Enqueue a task, updating if client_id already exists.
@@ -183,14 +236,19 @@ class PQ:
             run_at: When to run the task. Defaults to now.
             priority: Task priority. Higher = higher priority. Defaults to NORMAL.
             client_id: Client-provided identifier. Required for conflict resolution.
+            max_runtime_seconds: Per-task wall-clock cap in seconds. See ``enqueue`` for
+                details. ``None`` (the default) leaves the worker's own
+                configured ``max_runtime_seconds`` in effect.
             **kwargs: Keyword arguments to pass to the handler.
 
         Returns:
             Task ID.
 
         Raises:
-            ValueError: If task is a lambda, closure, or cannot be imported.
+            ValueError: If task is a lambda, closure, or cannot be imported,
+                or if ``max_runtime_seconds`` is not None and ``<= 0``.
         """
+        _validate_max_runtime(max_runtime_seconds)
         name = get_function_path(task)
         payload = serialize(args, kwargs)
 
@@ -206,6 +264,7 @@ class PQ:
                 priority=priority,
                 status=TaskStatus.PENDING,
                 run_at=run_at,
+                max_runtime_seconds=max_runtime_seconds,
             )
             .on_conflict_do_update(
                 index_elements=["client_id"],
@@ -215,6 +274,7 @@ class PQ:
                     "priority": priority,
                     "status": TaskStatus.PENDING,
                     "run_at": run_at,
+                    "max_runtime_seconds": max_runtime_seconds,
                     "attempts": 0,
                     "started_at": None,
                     "completed_at": None,
@@ -239,6 +299,7 @@ class PQ:
         max_concurrent: int | None = 1,
         key: str = "",
         active: bool = True,
+        max_runtime_seconds: float | None = None,
         **kwargs: Any,
     ) -> int:
         """Schedule a periodic task.
@@ -260,6 +321,14 @@ class PQ:
                 Defaults to "" (empty string).
             active: Whether the task is active. Inactive tasks are not executed.
                 Defaults to True.
+            max_runtime_seconds: Per-execution wall-clock cap in seconds. When ``None``
+                (the common case), each execution uses the worker's configured
+                default. When set, this schedule's executions get the
+                supplied ceiling, and the ``locked_until`` window (used when
+                ``max_concurrent`` is in effect) is extended to match so the
+                lock doesn't expire while the execution is legitimately still
+                running. Periodic tasks are not subject to the stale-task
+                reaper, so this knob is purely about the wall-clock cap.
             **kwargs: Keyword arguments to pass to the handler.
 
         Returns:
@@ -270,6 +339,7 @@ class PQ:
             ValueError: If cron expression is invalid.
             ValueError: If max_concurrent is greater than 1.
             ValueError: If task is a lambda, closure, or cannot be imported.
+            ValueError: If ``max_runtime_seconds`` is not None and ``<= 0``.
             IntegrityError: If client_id already exists.
         """
         if run_every is None and cron is None:
@@ -281,6 +351,7 @@ class PQ:
                 f"max_concurrent must be 1 or None, got {max_concurrent} "
                 "(values > 1 reserved for future use)"
             )
+        _validate_max_runtime(max_runtime_seconds)
 
         # Validate and normalize cron expression
         cron_expr: str | None = None
@@ -321,6 +392,7 @@ class PQ:
                     client_id=client_id,
                     max_concurrent=max_concurrent,
                     active=active,
+                    max_runtime_seconds=max_runtime_seconds,
                 )
                 .on_conflict_do_update(
                     index_elements=["name", "key"],
@@ -332,6 +404,7 @@ class PQ:
                         "next_run": next_run,
                         "max_concurrent": max_concurrent,
                         "active": active,
+                        "max_runtime_seconds": max_runtime_seconds,
                     },
                 )
                 .returning(Periodic.id)
@@ -516,39 +589,85 @@ class PQ:
         process remains to update their status. This method detects those
         orphaned rows and transitions them to FAILED.
 
+        Tasks enqueued with a per-task ``max_runtime_seconds`` override get a
+        proportionally larger reaper window: the effective threshold for
+        a row is ``max(threshold, max_runtime_seconds * 2)``. This avoids
+        reaping a legitimately long-running task whose declared budget
+        exceeds the global default. NULL ``max_runtime_seconds`` (the
+        common case) falls back to the supplied ``threshold`` unchanged.
+
         Args:
-            threshold: A task is considered stale when
-                ``started_at + threshold < now()``. Should be at least
-                2x ``max_runtime`` to avoid reaping legitimately running
-                tasks, e.g. ``timedelta(seconds=max_runtime * 2)``.
+            threshold: Default per-task stale window for tasks that don't
+                have a ``max_runtime_seconds`` override (or whose override
+                is smaller than this default). A task is stale when
+                ``started_at + effective_threshold < now()``. Should be
+                at least 2x the worker's own ``max_runtime_seconds``, e.g.
+                ``timedelta(seconds=max_runtime_seconds * 2)``.
 
         Returns:
             Number of tasks reaped.
         """
         now = datetime.now(UTC)
-        cutoff = now - threshold
+        threshold_seconds = threshold.total_seconds()
+        # Per-row effective stale window: the larger of the supplied
+        # default and 2x the row's per-task override (NULL → just the
+        # default). Computed in SQL with ``GREATEST`` + ``COALESCE``
+        # so the reaper stays a single round-trip even when the row
+        # set is heterogeneous in declared budget.
+        effective_threshold_seconds = sa.func.greatest(
+            threshold_seconds,
+            sa.func.coalesce(Task.max_runtime_seconds * 2, threshold_seconds),
+        )
+        stale_cutoff = sa.func.now() - sa.func.make_interval(
+            sa.literal(0),  # years
+            sa.literal(0),  # months
+            sa.literal(0),  # weeks
+            sa.literal(0),  # days
+            sa.literal(0),  # hours
+            sa.literal(0),  # mins
+            effective_threshold_seconds,  # secs
+        )
+        # Per-row ``error`` message includes the values that actually
+        # decided this row's fate: the row's own
+        # ``max_runtime_seconds`` (or ``NULL`` if unset) and the
+        # effective threshold the reaper applied. This avoids the
+        # operator needing a separate query to figure out "why did MY
+        # long-budget task get reaped?" — the answer is in the row.
+        # Built with Postgres ``format()`` so the message is computed
+        # per row in the same UPDATE round-trip.
+        per_row_error = sa.func.format(
+            "Reaped: task still RUNNING past its stale window "
+            "(default threshold=%s s, per-task max_runtime_seconds=%s, "
+            "effective threshold=%s s). Worker likely died.",
+            threshold_seconds,
+            sa.func.coalesce(sa.cast(Task.max_runtime_seconds, sa.String), "NULL"),
+            effective_threshold_seconds,
+        )
         with self.session() as session:
             stmt = (
                 update(Task)
                 .where(
                     Task.status == TaskStatus.RUNNING,
-                    Task.started_at < cutoff,
+                    Task.started_at < stale_cutoff,
                 )
                 .values(
                     status=TaskStatus.FAILED,
                     completed_at=now,
-                    error=(
-                        f"Reaped: task still RUNNING after {threshold} "
-                        f"(cutoff={cutoff.isoformat()}). Worker likely died."
-                    ),
+                    error=per_row_error,
+                )
+                .returning(
+                    Task.id,
+                    Task.name,
+                    Task.started_at,
+                    Task.max_runtime_seconds,
                 )
-                .returning(Task.id, Task.name, Task.started_at)
             )
             reaped = list(session.execute(stmt).all())
-            for task_id, name, started_at in reaped:
+            for task_id, name, started_at, max_runtime_seconds in reaped:
                 logger.warning(
                     f"Reaped stale task '{name}' (id={task_id},"
-                    f" started_at={started_at})"
+                    f" started_at={started_at},"
+                    f" max_runtime_seconds={max_runtime_seconds})"
                 )
             return len(reaped)
 
@@ -559,11 +678,17 @@ class PQ:
         poll_interval: float = 1.0,
         max_runtime: float = 30 * 60,
         priorities: Set[Priority] | None = None,
+        drain_timeout: float = 20.0,
     ) -> None:
         """Run the worker loop (blocking).
 
         Each task executes in a forked child process for memory isolation.
 
+        On SIGTERM or SIGINT the worker stops claiming new tasks, waits up
+        to ``drain_timeout`` seconds for in-flight tasks to finish, then
+        kills any still-running task and marks its row FAILED with an
+        explicit shutdown error. Interrupted tasks are not re-queued.
+
         Args:
             concurrency: Maximum number of tasks to process simultaneously.
                 Default: 1 (sequential processing).
@@ -571,6 +696,10 @@ class PQ:
             max_runtime: Maximum execution time per task in seconds. Default: 30 min.
             priorities: If set, only process tasks with these priority levels.
                 Use this to dedicate workers to specific priority tiers.
+            drain_timeout: Seconds to wait for in-flight tasks on
+                SIGTERM/SIGINT. Default: 20. Set below the orchestrator's
+                termination grace period. ``0`` skips the wait: in-flight
+                tasks are killed as soon as the shutdown is noticed.
         """
         from pq.worker import run_worker
 
@@ -580,6 +709,7 @@ class PQ:
             poll_interval=poll_interval,
             max_runtime=max_runtime,
             priorities=priorities,
+            drain_timeout=drain_timeout,
         )
 
     def run_worker_once(

python_pq-0.8.0/src/pq/migrations/versions/20260527T170136Z_aee3e8e7e647_add_task_max_runtime_seconds.py

@@ -0,0 +1,50 @@
+"""add max_runtime_seconds to pq_tasks and pq_periodic
+
+Revision ID: aee3e8e7e647
+Revises: c3d4e5f6a7b8
+Create Date: 2026-05-27 17:01:36 Z
+
+Per-task override for the worker's wall-clock ``max_runtime`` cap, on
+both one-off tasks (``pq_tasks``) and periodic schedules
+(``pq_periodic``). Nullable on both so every existing row (and every
+existing enqueue / schedule call site that doesn't pass the new
+parameter) keeps the previous behaviour exactly — a NULL value tells
+the worker to use its own default.
+
+Schema impact on populated tables is intentionally minimal: PostgreSQL
+11+ treats ``ADD COLUMN ... DEFAULT NULL`` as a catalog-only operation
+(no table rewrite, brief ACCESS EXCLUSIVE lock on the system catalog
+row, typically < 100 ms even for tables with millions of rows). No
+backfill — NULL is the intended default.
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "aee3e8e7e647"
+down_revision: Union[str, Sequence[str], None] = "c3d4e5f6a7b8"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Add the nullable ``max_runtime_seconds`` column to both task tables."""
+    op.add_column(
+        "pq_tasks",
+        sa.Column("max_runtime_seconds", sa.Float(), nullable=True),
+    )
+    op.add_column(
+        "pq_periodic",
+        sa.Column("max_runtime_seconds", sa.Float(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    """Remove the ``max_runtime_seconds`` column from both task tables."""
+    op.drop_column("pq_periodic", "max_runtime_seconds")
+    op.drop_column("pq_tasks", "max_runtime_seconds")

{python_pq-0.7.1 → python_pq-0.8.0}/src/pq/models.py

@@ -9,6 +9,7 @@ from sqlalchemy import (
     Boolean,
     DateTime,
     Enum,
+    Float,
     Identity,
     Index,
     Integer,
@@ -70,6 +71,14 @@ class Task(Base):
     )
     error: Mapped[str | None] = mapped_column(Text, nullable=True)
     attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+    # Per-task override of the worker-level ``max_runtime``. When NULL
+    # (the common case), the worker uses its configured default. When
+    # set, it caps this specific task's wall-clock at the given number
+    # of seconds AND extends the stale-reaper threshold proportionally
+    # (the reaper picks the larger of its global default and
+    # ``max_runtime_seconds * 2``). Useful for occasionally-long tasks
+    # in a fleet whose default is sized for typical short work.
+    max_runtime_seconds: Mapped[float | None] = mapped_column(Float, nullable=True)
 
 
 class Periodic(Base):
@@ -94,6 +103,17 @@ class Periodic(Base):
     next_run: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
     max_concurrent: Mapped[int | None] = mapped_column(SmallInteger, nullable=True)
     active: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default="true")
+    # Per-task override of the worker-level ``max_runtime``. Same semantics
+    # as ``Task.max_runtime_seconds`` — NULL means "use the worker's
+    # configured default", a set value caps this specific schedule's
+    # wall-clock at the given number of seconds and also extends the
+    # ``locked_until`` window (when ``max_concurrent`` is in effect) so
+    # the lock doesn't expire while the task is legitimately still
+    # running. Periodic tasks are not subject to the stale-task reaper
+    # (they're guarded by ``locked_until`` and the natural re-fire of
+    # the schedule), so this knob is purely about the per-execution
+    # wall-clock cap.
+    max_runtime_seconds: Mapped[float | None] = mapped_column(Float, nullable=True)
     last_run: Mapped[datetime | None] = mapped_column(
         DateTime(timezone=True), nullable=True
     )

{python_pq-0.7.1 → python_pq-0.8.0}/src/pq/worker.py

@@ -8,14 +8,18 @@ from __future__ import annotations
 
 import asyncio
 import inspect
+import math
 import os
 import select as select_module
 import signal
 import sys
+import threading
 import time
 import traceback
+from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
+from types import FrameType
 from typing import TYPE_CHECKING, Any, NamedTuple, Protocol
 
 from croniter import croniter
@@ -27,11 +31,16 @@ from pq.registry import resolve_function_path
 from pq.serialization import deserialize
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Set
+    from collections.abc import Set
 
     from pq.client import PQ
     from pq.priority import Priority
 
+# What ``signal.getsignal`` returns / ``signal.signal`` accepts.
+_SignalHandler = (
+    Callable[[int, FrameType | None], object] | int | signal.Handlers | None
+)
+
 
 class PreExecuteHook(Protocol):
     """Protocol for pre-execution hooks called before task handler runs.
@@ -83,6 +92,13 @@ DEFAULT_STALE_TASK_TIMEOUT: timedelta = timedelta(hours=1)
 # How often the worker checks for stale RUNNING tasks.
 DEFAULT_REAPER_INTERVAL: float = 300
 
+# Default drain timeout: 20 seconds
+# On SIGTERM/SIGINT the worker stops claiming and waits up to this long for
+# in-flight tasks to finish before SIGKILLing them and marking them FAILED.
+# Sits under the common 30 s orchestrator termination grace period, leaving
+# headroom for the final status writes.
+DEFAULT_DRAIN_TIMEOUT: float = 20.0
+
 # Exit codes for child process
 EXIT_SUCCESS = 0
 EXIT_FAILURE = 1
@@ -134,6 +150,103 @@ class TaskKilledError(WorkerError):
     pass
 
 
+class TaskInterruptedError(WorkerError):
+    """Raised when a task is killed because the worker is shutting down."""
+
+    pass
+
+
+@dataclass
+class _ShutdownState:
+    """Graceful-shutdown coordination for the worker loops.
+
+    Written by the signal handler (which runs in the main thread between
+    bytecodes, so no locking is needed) and read by the worker loops and
+    ``_wait_for_child``. Module-level because a worker owns its process;
+    ``run_worker`` resets it on entry.
+    """
+
+    requested: bool = False
+    deadline: float | None = None  # time.monotonic() deadline, set on signal
+    drain_timeout: float = DEFAULT_DRAIN_TIMEOUT
+
+    def reset(self, drain_timeout: float) -> None:
+        self.requested = False
+        self.deadline = None
+        self.drain_timeout = drain_timeout
+
+    def deadline_passed(self) -> bool:
+        """True once shutdown was requested and the drain budget is spent."""
+        if not self.requested:
+            return False
+        return self.deadline is None or time.monotonic() >= self.deadline
+
+
+_shutdown = _ShutdownState()
+
+
+def _handle_shutdown_signal(signum: int, frame: FrameType | None) -> None:
+    """SIGTERM/SIGINT handler: request a drain, never raise.
+
+    Raising from here would land at arbitrary bytecode and could orphan a
+    task (e.g. between claim-commit and child registration). Repeat signals
+    during the drain are ignored so they can't abandon remaining children.
+    """
+    if _shutdown.requested:
+        return
+    _shutdown.deadline = time.monotonic() + _shutdown.drain_timeout
+    _shutdown.requested = True
+    logger.info(
+        f"Received signal {signum}, draining"
+        f" (drain_timeout={_shutdown.drain_timeout}s)..."
+    )
+
+
+def _install_shutdown_handlers() -> dict[int, _SignalHandler] | None:
+    """Register the drain handler for SIGTERM/SIGINT.
+
+    Returns the previous handlers (for restoration), or ``None`` when not
+    running in the main thread — ``signal.signal`` is only allowed there, so
+    in that case the worker keeps the legacy KeyboardInterrupt-only behavior.
+    """
+    if threading.current_thread() is not threading.main_thread():
+        logger.warning(
+            "Worker is not running in the main thread; SIGTERM graceful"
+            " shutdown is disabled (only KeyboardInterrupt is handled)."
+        )
+        return None
+    previous: dict[int, _SignalHandler] = {}
+    for signum in (signal.SIGTERM, signal.SIGINT):
+        previous[signum] = signal.getsignal(signum)
+        signal.signal(signum, _handle_shutdown_signal)
+    return previous
+
+
+def _restore_shutdown_handlers(previous: dict[int, _SignalHandler] | None) -> None:
+    """Restore the signal handlers saved by ``_install_shutdown_handlers``."""
+    if previous is None:
+        return
+    for signum, handler in previous.items():
+        signal.signal(signum, handler)
+
+
+def _kill_child_group(child_pid: int) -> None:
+    """SIGKILL a task child's process group (best effort).
+
+    The child calls ``os.setpgrp()`` first thing, so its pgid equals its
+    pid. If the signal lands before ``setpgrp()`` the process group doesn't
+    exist yet — fall back to killing the pid directly. ESRCH everywhere
+    means the child already exited; the caller reaps it normally.
+    """
+    try:
+        os.killpg(child_pid, signal.SIGKILL)
+    except ProcessLookupError:
+        try:
+            os.kill(child_pid, signal.SIGKILL)
+        except ProcessLookupError:
+            pass
+
+
 def _child_timeout_handler(signum: int, frame: Any) -> None:
     """Signal handler for timeout in child process."""
     os._exit(EXIT_TIMEOUT)
@@ -156,6 +269,12 @@ def _run_in_child(
     # Create new process group so we don't get parent's signals
     os.setpgrp()
 
+    # Restore default signal disposition: the fork inherits the parent's
+    # drain handler, which in the child would silently swallow a direct
+    # SIGTERM/SIGINT (e.g. from systemd's cgroup-wide kill).
+    signal.signal(signal.SIGTERM, signal.SIG_DFL)
+    signal.signal(signal.SIGINT, signal.SIG_DFL)
+
     # Set up timeout
     signal.signal(signal.SIGALRM, _child_timeout_handler)
     signal.alarm(int(max_runtime) + 1)  # +1 buffer for async timeout
@@ -241,6 +360,15 @@ def _fork_child(
 def _wait_for_child(child_pid: int, read_fd: int) -> _ChildResult:
     """Wait for a forked child process and return the result.
 
+    The wait is drain-aware: once shutdown has been requested and the drain
+    deadline has passed, the child's process group is SIGKILLed and the
+    result reports the ``"shutdown"`` exit kind. A blocking ``os.wait4``
+    would never notice the shutdown flag (PEP 475 retries the syscall after
+    the non-raising signal handler), so this polls with ``WNOHANG`` and
+    sleeps in ``select`` on the child's error pipe — the pipe closes when
+    the child exits, so child completion wakes the wait immediately, while
+    the bounded timeout keeps flag/deadline checks prompt.
+
     Args:
         child_pid: PID of the child process.
         read_fd: Read end of the error pipe.
@@ -248,10 +376,30 @@ def _wait_for_child(child_pid: int, read_fd: int) -> _ChildResult:
     Returns:
         _ChildResult with task status, error message, and exit kind.
     """
-    _, raw_status, rusage = os.wait4(child_pid, 0)
-
-    # Read error message from pipe
+    killed_by_drain = False
     error_bytes = b""
+    while True:
+        pid, raw_status, rusage = os.wait4(child_pid, os.WNOHANG)
+        if pid != 0:
+            break
+        if not killed_by_drain and _shutdown.deadline_passed():
+            _kill_child_group(child_pid)
+            killed_by_drain = True
+        timeout = 1.0
+        if _shutdown.requested and _shutdown.deadline is not None:
+            timeout = min(timeout, max(0.0, _shutdown.deadline - time.monotonic()))
+            timeout = max(timeout, 0.01)  # killed child needs a moment to exit
+        ready, _, _ = select_module.select([read_fd], [], [], timeout)
+        if ready:
+            # Drain incrementally: an error message larger than the pipe
+            # buffer would otherwise block the child in write() while we
+            # spin on an always-readable fd.
+            try:
+                error_bytes += os.read(read_fd, 65536)
+            except OSError:
+                pass
+
+    # Read any remaining error output from the pipe
     try:
         while True:
             chunk = os.read(read_fd, 4096)
@@ -267,6 +415,13 @@ def _wait_for_child(child_pid: int, read_fd: int) -> _ChildResult:
 
     if os.WIFSIGNALED(raw_status):
         signal_num = os.WTERMSIG(raw_status)
+        if killed_by_drain and signal_num == signal.SIGKILL:
+            return _ChildResult(
+                TaskStatus.FAILED,
+                "Task interrupted by worker shutdown"
+                f" (exceeded drain_timeout={_shutdown.drain_timeout}s)",
+                "shutdown",
+            )
         if signal_num == signal.SIGKILL:
             max_rss_kb = rusage.ru_maxrss
             if sys.platform == "darwin":
@@ -351,9 +506,28 @@ def _execute_in_fork(
         raise TaskOOMError(result.error_msg or "Task killed (likely OOM)")
     if result.exit_kind == "killed":
         raise TaskKilledError(result.error_msg or "Task killed by signal")
+    if result.exit_kind == "shutdown":
+        raise TaskInterruptedError(
+            result.error_msg or "Task interrupted by worker shutdown"
+        )
     raise Exception(result.error_msg or "Task failed")
 
 
+def _interruptible_sleep(seconds: float) -> None:
+    """Sleep up to ``seconds``, returning early once shutdown is requested.
+
+    A plain ``time.sleep`` would run to completion despite SIGTERM (PEP 475
+    retries it after the non-raising handler), delaying an idle worker's
+    shutdown by up to ``poll_interval``.
+    """
+    deadline = time.monotonic() + seconds
+    while not _shutdown.requested:
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return
+        time.sleep(min(0.1, remaining))
+
+
 def _maybe_run_cleanup(
     pq: PQ,
     retention_days: int,
@@ -424,6 +598,7 @@ def run_worker(
     retention_days: int = DEFAULT_RETENTION_DAYS,
     cleanup_interval: float = DEFAULT_CLEANUP_INTERVAL,
     stale_task_timeout: timedelta | None = DEFAULT_STALE_TASK_TIMEOUT,
+    drain_timeout: float = DEFAULT_DRAIN_TIMEOUT,
     pre_execute: PreExecuteHook | None = None,
     post_execute: PostExecuteHook | None = None,
 ) -> None:
@@ -431,6 +606,14 @@ def run_worker(
 
     Each task executes in a forked child process for memory isolation.
 
+    On SIGTERM or SIGINT the worker shuts down gracefully: it stops
+    claiming new tasks, waits up to ``drain_timeout`` seconds for in-flight
+    tasks to finish (writing their statuses as usual), then SIGKILLs any
+    still-running task children and marks their rows FAILED with an explicit
+    shutdown error. Interrupted tasks are NOT re-queued — pq's at-most-once
+    semantics are preserved; applications that need redelivery must
+    re-enqueue such tasks themselves.
+
     Args:
         pq: PQ client instance.
         concurrency: Maximum number of tasks to process simultaneously.
@@ -445,11 +628,24 @@ def run_worker(
         stale_task_timeout: Mark RUNNING tasks older than this as FAILED.
             Catches orphaned tasks whose worker died mid-execution.
             Default: 1 hour. Set to ``None`` to disable.
+        drain_timeout: Seconds to wait for in-flight tasks when shutting
+            down on SIGTERM/SIGINT. Default: 20. Must be set below the
+            orchestrator's termination grace period, with headroom for the
+            final status writes. ``0`` skips the wait entirely: in-flight
+            tasks are killed as soon as the shutdown is noticed (within
+            about a second).
         pre_execute: Called in forked child BEFORE task execution.
             Use for initializing fork-unsafe resources (OTel, DB connections).
         post_execute: Called in forked child AFTER task execution (success or failure).
             Use for cleanup/flushing (OTel traces, etc.).
+
+    Raises:
+        ValueError: If ``drain_timeout`` is negative or not finite.
     """
+    if math.isnan(drain_timeout) or math.isinf(drain_timeout) or drain_timeout < 0:
+        raise ValueError(
+            f"drain_timeout must be a finite number >= 0 (got {drain_timeout!r})"
+        )
     if priorities:
         priority_names = ", ".join(p.name for p in sorted(priorities, reverse=True))
         logger.info(
@@ -462,40 +658,55 @@ def run_worker(
             " fork isolation enabled)..."
         )
 
-    if concurrency > 1:
-        _run_concurrent(
-            pq,
-            concurrency=concurrency,
-            poll_interval=poll_interval,
-            max_runtime=max_runtime,
-            priorities=priorities,
-            pre_execute=pre_execute,
-            post_execute=post_execute,
-            retention_days=retention_days,
-            cleanup_interval=cleanup_interval,
-            stale_task_timeout=stale_task_timeout,
-        )
-        return
-
-    last_cleanup: list[float] = [0.0]
-    last_reap: list[float] = [0.0]
-
+    _shutdown.reset(drain_timeout)
+    previous_handlers = _install_shutdown_handlers()
     try:
-        while True:
-            if not run_worker_once(
+        if concurrency > 1:
+            _run_concurrent(
                 pq,
+                concurrency=concurrency,
+                poll_interval=poll_interval,
                 max_runtime=max_runtime,
                 priorities=priorities,
                 pre_execute=pre_execute,
                 post_execute=post_execute,
-            ):
-                _maybe_run_cleanup(pq, retention_days, cleanup_interval, last_cleanup)
-                _maybe_reap_stale(
-                    pq, stale_task_timeout, DEFAULT_REAPER_INTERVAL, last_reap
-                )
-                time.sleep(poll_interval)
-    except KeyboardInterrupt:
+                retention_days=retention_days,
+                cleanup_interval=cleanup_interval,
+                stale_task_timeout=stale_task_timeout,
+            )
+            return
+
+        last_cleanup: list[float] = [0.0]
+        last_reap: list[float] = [0.0]
+
+        try:
+            while not _shutdown.requested:
+                if not run_worker_once(
+                    pq,
+                    max_runtime=max_runtime,
+                    priorities=priorities,
+                    pre_execute=pre_execute,
+                    post_execute=post_execute,
+                ):
+                    if _shutdown.requested:
+                        break
+                    _maybe_run_cleanup(
+                        pq, retention_days, cleanup_interval, last_cleanup
+                    )
+                    _maybe_reap_stale(
+                        pq, stale_task_timeout, DEFAULT_REAPER_INTERVAL, last_reap
+                    )
+                    _interruptible_sleep(poll_interval)
+        except KeyboardInterrupt:
+            # Fallback for non-main-thread workers (no handlers installed).
+            pass
         logger.info("Worker stopped.")
+    finally:
+        _restore_shutdown_handlers(previous_handlers)
+        # Clear the drain state: a lingering past-deadline flag would make
+        # later run_worker_once calls in this process kill their children
+        # on the first _wait_for_child iteration.
+        _shutdown.reset(drain_timeout)
 
 
 def run_worker_once(
@@ -596,6 +807,10 @@ def _process_one_off_task(
             name = task.name
             payload = task.payload
             task_id = task.id
+            # Per-task ``max_runtime_seconds`` (NULL → use worker default).
+            # Read while still in the session because the row is expunged
+            # below for the forked-child handoff.
+            task_max_runtime = task.max_runtime_seconds
 
             # Flush to commit status changes, then expunge for forked process
             session.flush()
@@ -605,6 +820,14 @@ def _process_one_off_task(
         logger.error(f"Error claiming task: {e}")
         return False
 
+    # Effective per-task wall-clock ceiling: per-task value when set,
+    # otherwise the worker's configured default. The reaper applies the
+    # same rule on its own (see ``Client.reap_stale_tasks``), so the
+    # two stay in sync without the worker having to coordinate.
+    effective_max_runtime = (
+        task_max_runtime if task_max_runtime is not None else max_runtime
+    )
+
     # Phase 2: Execute handler in forked process
     start = time.perf_counter()
     status = TaskStatus.COMPLETED
@@ -617,7 +840,7 @@ def _process_one_off_task(
             handler,
             args,
             kwargs,
-            max_runtime=max_runtime,
+            max_runtime=effective_max_runtime,
             task=task,
             pre_execute=pre_execute,
             post_execute=post_execute,
@@ -635,6 +858,10 @@ def _process_one_off_task(
         status = TaskStatus.FAILED
         error_msg = str(e)
 
+    except TaskInterruptedError as e:
+        status = TaskStatus.FAILED
+        error_msg = str(e)
+
     except Exception as e:
         status = TaskStatus.FAILED
         error_msg = str(e)
@@ -739,10 +966,27 @@ def _process_periodic_task(
             payload = periodic.payload
             periodic_id = periodic.id
             periodic_max_concurrent = periodic.max_concurrent
+            # Per-schedule ``max_runtime_seconds`` (NULL → use worker default).
+            # Read while still in the session because the row is expunged
+            # below for the forked-child handoff.
+            periodic_max_runtime = periodic.max_runtime_seconds
+
+            # Effective per-execution wall-clock ceiling: per-schedule value
+            # when set, otherwise the worker's configured default. Drives
+            # both the in-child timeout enforcement AND the
+            # ``locked_until`` window so the lock doesn't expire while a
+            # legitimately-long execution is still running.
+            effective_max_runtime = (
+                periodic_max_runtime
+                if periodic_max_runtime is not None
+                else max_runtime
+            )
 
             # Set lock before execution if concurrency is limited
             if periodic.max_concurrent is not None:
-                lock_duration = max_runtime if max_runtime > 0 else 3600
+                lock_duration = (
+                    effective_max_runtime if effective_max_runtime > 0 else 3600
+                )
                 periodic.locked_until = func.now() + timedelta(seconds=lock_duration)
 
             # Advance schedule BEFORE execution
@@ -769,7 +1013,7 @@ def _process_periodic_task(
             handler,
             args,
             kwargs,
-            max_runtime=max_runtime,
+            max_runtime=effective_max_runtime,
             task=periodic,
             pre_execute=pre_execute,
             post_execute=post_execute,
@@ -789,6 +1033,10 @@ def _process_periodic_task(
         elapsed = time.perf_counter() - start
         logger.error(f"Periodic task '{name}' killed after {elapsed:.3f} s: {e}")
 
+    except TaskInterruptedError as e:
+        elapsed = time.perf_counter() - start
+        logger.warning(f"Periodic task '{name}' interrupted after {elapsed:.3f} s: {e}")
+
     except Exception as e:
         elapsed = time.perf_counter() - start
         logger.error(f"Periodic task '{name}' failed after {elapsed:.3f} s: {e}")
@@ -851,6 +1099,9 @@ def _claim_and_fork_one_off(
             name = task.name
             payload = task.payload
             task_id = task.id
+            # Per-task ``max_runtime_seconds`` (NULL → use worker default).
+            # Same handling as the sequential ``_process_one_off_task`` path.
+            task_max_runtime = task.max_runtime_seconds
 
             session.flush()
             session.expunge(task)
@@ -859,6 +1110,10 @@ def _claim_and_fork_one_off(
         logger.error(f"Error claiming task: {e}")
         return None
 
+    effective_max_runtime = (
+        task_max_runtime if task_max_runtime is not None else max_runtime
+    )
+
     try:
         handler = resolve_function_path(name)
         args, kwargs = deserialize(payload)
@@ -866,7 +1121,7 @@ def _claim_and_fork_one_off(
             handler,
             args,
             kwargs,
-            max_runtime=max_runtime,
+            max_runtime=effective_max_runtime,
             task=task,
             pre_execute=pre_execute,
             post_execute=post_execute,
@@ -934,9 +1189,20 @@ def _claim_and_fork_periodic(
             payload = periodic.payload
             periodic_id = periodic.id
             periodic_max_concurrent = periodic.max_concurrent
+            # Per-schedule ``max_runtime_seconds`` (NULL → use worker default).
+            # Same handling as the sequential ``_process_periodic_task`` path.
+            periodic_max_runtime = periodic.max_runtime_seconds
+
+            effective_max_runtime = (
+                periodic_max_runtime
+                if periodic_max_runtime is not None
+                else max_runtime
+            )
 
             if periodic.max_concurrent is not None:
-                lock_duration = max_runtime if max_runtime > 0 else 3600
+                lock_duration = (
+                    effective_max_runtime if effective_max_runtime > 0 else 3600
+                )
                 periodic.locked_until = func.now() + timedelta(seconds=lock_duration)
 
             periodic.last_run = func.now()
@@ -959,7 +1225,7 @@ def _claim_and_fork_periodic(
             handler,
             args,
             kwargs,
-            max_runtime=max_runtime,
+            max_runtime=effective_max_runtime,
             task=periodic,
             pre_execute=pre_execute,
             post_execute=post_execute,
@@ -1048,6 +1314,11 @@ def _reap_and_update(pq: PQ, slot: _ChildSlot) -> None:
                 f"Periodic task '{slot.name}' killed after {elapsed:.3f} s:"
                 f" {result.error_msg}"
             )
+        elif result.exit_kind == "shutdown":
+            logger.warning(
+                f"Periodic task '{slot.name}' interrupted after {elapsed:.3f} s:"
+                f" {result.error_msg}"
+            )
         else:
             logger.error(
                 f"Periodic task '{slot.name}' failed after {elapsed:.3f} s:"
@@ -1139,9 +1410,9 @@ def _run_concurrent(
     last_reap: list[float] = [0.0]
 
     try:
-        while True:
+        while not _shutdown.requested:
             # Step 1: Fill empty slots with new tasks
-            while len(children) < concurrency:
+            while len(children) < concurrency and not _shutdown.requested:
                 slot = _try_claim_and_fork(
                     pq,
                     max_runtime=max_runtime,
@@ -1163,7 +1434,10 @@ def _run_concurrent(
                     slot = children.pop(pid)
                     _reap_and_update(pq, slot)
             else:
-                time.sleep(poll_interval)
+                _interruptible_sleep(poll_interval)
+
+            if _shutdown.requested:
+                break
 
             # Maintenance (each rate-limited independently)
             _maybe_run_cleanup(pq, retention_days, cleanup_interval, last_cleanup)
@@ -1172,10 +1446,16 @@ def _run_concurrent(
             )
 
     except KeyboardInterrupt:
-        if children:
-            logger.info(
-                f"Shutting down, waiting for {len(children)} task(s) to finish..."
-            )
-            for slot in list(children.values()):
-                _reap_and_update(pq, slot)
-        logger.info("Worker stopped.")
+        # Fallback for non-main-thread workers (no handlers installed); the
+        # drain below still runs, just without a deadline.
+        pass
+
+    if children:
+        logger.info(f"Shutting down, draining {len(children)} in-flight task(s)...")
+        # _wait_for_child inside _reap_and_update enforces the drain
+        # deadline (absolute, shared by all children): each child either
+        # finishes in time and gets its real status, or is SIGKILLed and
+        # marked FAILED with the shutdown error.
+        for slot in list(children.values()):
+            _reap_and_update(pq, slot)
+    logger.info("Worker stopped.")