taskiq-llm-queue 0.1.0a1__py3-none-any.whl

taskiq_llm_queue/broker/delivery.py

@@ -0,0 +1,50 @@
+"""One received SQS message and its local lifecycle: buffered, started, then acked or released.
+
+The broker, the heartbeat and `release()` all act on the same `Delivery`, so its `done`
+flag makes ack, release and heartbeat stop idempotent and race-free within one event loop.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import math
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from taskiq import TaskiqMessage
+
+    from taskiq_llm_queue.broker.queue import SQSQueue
+
+
+@dataclass(slots=True, eq=False)
+class Delivery:
+    """A received message: where it came from, its receipt handle and its local state.
+
+    `lost` is set when a heartbeat finds the receipt handle invalid (another consumer may
+    now own the message), so the running task can stop without writing a result (DS6).
+    `outcome` is "acked" or "released" once `done`; `after_ack` callbacks run only after
+    SQS confirmed the delete. `expires_at` is the `expiry` attribute (epoch seconds).
+    `message` is the screened, validated taskiq message until it is handed to the Receiver.
+    """
+
+    queue: SQSQueue
+    queue_url: str
+    message_id: str
+    receipt_handle: str
+    body: str
+    receive_count: int
+    received_at: float
+    visibility_secs: int
+    sent_at: float | None = None
+    started: bool = False
+    done: bool = False
+    extensions: int = 0
+    next_heartbeat_at: float = math.inf
+    lost: asyncio.Event = field(default_factory=asyncio.Event)
+    task_id: str | None = None
+    expires_at: float | None = None
+    outcome: str | None = None
+    after_ack: list[Callable[[], Awaitable[None]]] = field(default_factory=list)
+    message: TaskiqMessage | None = None

taskiq_llm_queue/broker/errors.py

@@ -0,0 +1,178 @@
+"""Broker exceptions and SQS error-family classification.
+
+ElasticMQ, moto and AWS spell the same failure differently (an oversized batch is
+`InvalidAttributeValue`, `InvalidParameterValue` or `BatchRequestTooLong`), so every broker
+decision uses `ErrorFamily`, never the raw code. The raw code is kept for logs only.
+"""
+
+from __future__ import annotations
+
+import random
+import re
+from collections.abc import Callable
+from enum import StrEnum
+
+from taskiq_llm_queue.errors import ConfigError, LLMQueueError, QueueUnavailable
+
+
+class BrokerConfigError(ConfigError):
+    """Invalid broker or queue settings (limits, priorities, names), raised at construction."""
+
+
+class InvalidLabelError(LLMQueueError, ValueError):
+    """A task label the broker reads (`delay`, `group_id`, `expiry`...) is out of range.
+
+    Raised by `kick`, i.e. by the caller's `.kiq()`, so a bad label never reaches SQS.
+    """
+
+
+class MessageTooLargeError(LLMQueueError, ValueError):
+    """One serialized message is over the SQS 1 MiB limit; offload the payload instead."""
+
+
+class SQSSendError(QueueUnavailable):
+    """SQS refused or failed a send; `family` says whether retrying can help."""
+
+    def __init__(self, message: str, *, family: ErrorFamily, code: str = "") -> None:
+        """Keep the error family for decisions and the provider's code for logs."""
+        super().__init__(message)
+        self.family = family
+        self.code = code
+
+
+class ErrorFamily(StrEnum):
+    """What kind of SQS failure this is, independent of which fake or region spelled it.
+
+    `TRANSIENT` (5xx, connection errors) and `THROTTLED` are both retried with backoff.
+    """
+
+    SIZE_LIMIT = "size_limit"
+    INVALID_HANDLE = "invalid_handle"
+    THROTTLED = "throttled"
+    TRANSIENT = "transient"
+    NOT_FOUND = "not_found"
+    OTHER = "other"
+
+    @property
+    def retryable(self) -> bool:
+        """True for families where the same call may succeed after a pause."""
+        return self in (ErrorFamily.THROTTLED, ErrorFamily.TRANSIENT)
+
+
+_SIZE_CODES = frozenset(
+    {"batchrequesttoolong", "messagetoolong", "requestentitytoolarge", "requesttoolarge"}
+)
+# Generic "bad parameter" codes that mean "too big" only when the message says so.
+_GENERIC_CODES = frozenset({"invalidattributevalue", "invalidparametervalue"})
+_SIZE_TEXT = re.compile(
+    r"too ?long|shorter than|too large|exceed\w* .*(size|bytes|limit)|maximum .*size", re.I
+)
+_HANDLE_CODES = frozenset(
+    {"receipthandleisinvalid", "invalidreceipthandle", "messagenotinflight", "invalididformat"}
+)
+_HANDLE_TEXT = re.compile(r"receipt ?handle|not ?in ?flight", re.I)
+_THROTTLE_CODES = frozenset(
+    {
+        "requestthrottled",
+        "throttlingexception",
+        "throttling",
+        "throttled",
+        "slowdown",
+        "overlimit",
+        "toomanyrequestsexception",
+        "requestlimitexceeded",
+        "kmsthrottlingexception",
+    }
+)
+_TRANSIENT_CODES = frozenset(
+    {
+        "internalerror",
+        "internalfailure",
+        "serviceunavailable",
+        "serviceunavailableexception",
+        "requesttimeout",
+        "requesttimeoutexception",
+    }
+)
+_NOT_FOUND_CODES = frozenset(
+    {
+        "queuedoesnotexist",
+        "nonexistentqueue",
+        "resourcenotfoundexception",
+        "queuedeletedrecently",
+    }
+)
+
+
+def classify_code(code: str, message: str = "", status: int | None = None) -> ErrorFamily:
+    """Map one SQS error code (and its message / HTTP status) to an `ErrorFamily`.
+
+    Accepts both legacy (`AWS.SimpleQueueService.NonExistentQueue`) and JSON-protocol
+    spellings; used for whole-call errors and for `Failed[]` batch entries alike.
+    """
+    key = code.rsplit(".", 1)[-1].lower()
+    if key in _SIZE_CODES or (key in _GENERIC_CODES and _SIZE_TEXT.search(message)):
+        return ErrorFamily.SIZE_LIMIT
+    if key in _HANDLE_CODES or (key in _GENERIC_CODES and _HANDLE_TEXT.search(message)):
+        return ErrorFamily.INVALID_HANDLE
+    if key in _THROTTLE_CODES or status == 429:
+        return ErrorFamily.THROTTLED
+    if key in _NOT_FOUND_CODES:
+        return ErrorFamily.NOT_FOUND
+    if key in _TRANSIENT_CODES or (status is not None and status >= 500):
+        return ErrorFamily.TRANSIENT
+    return ErrorFamily.OTHER
+
+
+def error_code(exc: BaseException) -> str:
+    """The provider's error code of a botocore `ClientError`, or the exception class name."""
+    response = getattr(exc, "response", None)
+    if isinstance(response, dict):
+        return str(response.get("Error", {}).get("Code", "")) or type(exc).__name__
+    return type(exc).__name__
+
+
+def classify_error(exc: BaseException) -> ErrorFamily:
+    """Classify any exception raised by an SQS call into an `ErrorFamily`.
+
+    `ClientError` is read through its response; connection and read-timeout errors are
+    `TRANSIENT`; anything else is `OTHER`.
+    """
+    response = getattr(exc, "response", None)
+    if isinstance(response, dict):
+        error = response.get("Error", {})
+        status = response.get("ResponseMetadata", {}).get("HTTPStatusCode")
+        return classify_code(
+            str(error.get("Code", "")),
+            str(error.get("Message", "")),
+            status if isinstance(status, int) else None,
+        )
+    # Imported here so this module stays importable (and testable) without botocore loaded.
+    from botocore.exceptions import (
+        ConnectionClosedError,
+        EndpointConnectionError,
+        ReadTimeoutError,
+    )
+
+    if isinstance(
+        exc,
+        EndpointConnectionError | ConnectionClosedError | ReadTimeoutError | ConnectionError,
+    ):
+        return ErrorFamily.TRANSIENT
+    if isinstance(exc, TimeoutError):
+        return ErrorFamily.TRANSIENT
+    return ErrorFamily.OTHER
+
+
+def backoff_delay(
+    attempt: int,
+    *,
+    base_secs: float,
+    cap_secs: float,
+    rand: Callable[[float, float], float] = random.uniform,
+) -> float:
+    """Full-jitter exponential backoff: uniform(0, min(cap, base * 2**attempt)).
+
+    Each poller and flusher keeps its own attempt counter, so pods spread their retries.
+    """
+    return rand(0.0, min(cap_secs, base_secs * (2 ** min(attempt, 30))))

taskiq_llm_queue/broker/events.py

@@ -0,0 +1,51 @@
+"""Broker event hooks (A7): counts and timings for metrics, without a metrics dependency.
+
+Subclass `BrokerHooks` and override what you need; the broker calls every hook through
+`safe_call`, so a failing hook is logged and never breaks receiving, acking or sending.
+Queue names are the `SQSQueue.name` routing keys, safe to use as a `queue` metric label.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from taskiq_llm_queue._logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class BrokerHooks:
+    """No-op callbacks for broker events; override the ones a metrics layer cares about."""
+
+    def on_receive(self, queue: str, count: int) -> None:
+        """`count` messages arrived from one `ReceiveMessage` call (0 for an empty poll)."""
+
+    def on_queue_wait(self, queue: str, wait_secs: float) -> None:
+        """One message waited `wait_secs` in SQS (now - `SentTimestamp`)."""
+
+    def on_ack(self, queue: str, count: int) -> None:
+        """`count` messages were deleted (acks after the task, and expired-message deletes)."""
+
+    def on_heartbeat(self, queue: str, extended: int, lost: int) -> None:
+        """A heartbeat batch extended `extended` messages; `lost` had invalid handles."""
+
+    def on_release(self, queue: str, count: int, reason: str) -> None:
+        """Messages made visible again; reason is shutdown, drain, unknown_task, manual..."""
+
+    def on_expired(self, queue: str, count: int) -> None:
+        """Messages past their `expiry` label were deleted without running (P11)."""
+
+    def on_error(self, queue: str, operation: str, family: str) -> None:
+        """An SQS call failed; `family` is an `ErrorFamily` value."""
+
+    def on_send(self, queue: str, count: int) -> None:
+        """`count` messages were sent (one call, batched or not)."""
+
+
+def safe_call(hook: Callable[..., Any], *args: Any) -> None:
+    """Run one hook, logging (not raising) anything it throws."""
+    try:
+        hook(*args)
+    except Exception:
+        logger.exception("broker_hook_failed", hook=getattr(hook, "__name__", repr(hook)))

taskiq_llm_queue/broker/filters.py

@@ -0,0 +1,32 @@
+"""Pre-receive checks that decide whether a message may reach taskiq's Receiver (DS9).
+
+taskiq's Receiver acks (deletes) messages whose task it does not know, which would lose
+work during a rolling deploy. The broker runs these checks first and releases rejected
+messages with a short visibility instead, so a pod that knows the task picks them up.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from taskiq_llm_queue.envelope import ENVELOPE_VERSION
+
+if TYPE_CHECKING:
+    from taskiq.message import TaskiqMessage
+
+MessageFilter = Callable[["TaskiqMessage"], bool]
+
+
+def envelope_supported(message: TaskiqMessage) -> bool:
+    """False when any argument is an envelope (`v` + `request_id`) newer than this worker.
+
+    Older-but-supported and malformed envelopes pass; the task itself rejects malformed ones.
+    """
+    values: list[Any] = [*message.args, *message.kwargs.values()]
+    for value in values:
+        if isinstance(value, dict) and "request_id" in value:
+            version = value.get("v")
+            if isinstance(version, int) and version > ENVELOPE_VERSION:
+                return False
+    return True

taskiq_llm_queue/broker/heartbeat.py

@@ -0,0 +1,157 @@
+"""Keeps received messages invisible while they wait or run, with batched extensions (DS6).
+
+One loop serves every tracked delivery: each is extended by its queue's visibility timeout
+every `interval` (≤ visibility / 3), so two missed beats still leave margin. Tracking stops
+on ack or release, after `max_extensions`, once the message's `expiry` (deadline) plus a grace
+has passed, or when SQS says the receipt handle is invalid; the last case sets
+`Delivery.lost` and calls `on_lost` so the task can be cancelled.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import functools
+import math
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.broker.acker import VisibilityResult
+from taskiq_llm_queue.broker.events import BrokerHooks, safe_call
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.broker.acker import Acker
+    from taskiq_llm_queue.broker.delivery import Delivery
+    from taskiq_llm_queue.broker.queue import Heartbeat
+
+logger = get_logger(__name__)
+
+
+class HeartbeatManager:
+    """Schedules visibility extensions for tracked deliveries through the `Acker`."""
+
+    def __init__(
+        self,
+        acker: Acker,
+        config: Heartbeat,
+        *,
+        hooks: BrokerHooks,
+        on_lost: Callable[[Delivery], None] | None = None,
+        wall_clock: Callable[[], float] = time.time,
+    ) -> None:
+        """Nothing runs until `start()`; `track` before `start` is allowed.
+
+        `wall_clock` (epoch seconds) is compared with `Delivery.expires_at`; injectable for tests.
+        """
+        self._acker = acker
+        self._wall_clock = wall_clock
+        self._config = config
+        self._hooks = hooks
+        self._on_lost = on_lost
+        self._tracked: dict[int, Delivery] = {}
+        self._wake = asyncio.Event()
+        self._task: asyncio.Task[None] | None = None
+
+    @property
+    def tracked(self) -> int:
+        """Deliveries currently being extended."""
+        return len(self._tracked)
+
+    def track(self, delivery: Delivery) -> None:
+        """Start extending a delivery; the first beat is one interval after it was received."""
+        interval = self._config.interval_for(delivery.visibility_secs)
+        delivery.next_heartbeat_at = delivery.received_at + interval
+        self._tracked[id(delivery)] = delivery
+        self._wake.set()
+
+    def untrack(self, delivery: Delivery) -> None:
+        """Stop extending a delivery (ack, release, or extensions exhausted)."""
+        self._tracked.pop(id(delivery), None)
+        delivery.next_heartbeat_at = math.inf
+
+    def start(self) -> None:
+        """Start the background loop; idempotent."""
+        if self._task is None or self._task.done():
+            self._task = asyncio.create_task(self._run(), name="llmq-heartbeat")
+
+    async def stop(self) -> None:
+        """Stop the loop; tracked deliveries are left to their visibility timeout."""
+        if self._task is not None:
+            self._task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+
+    def due(self, now: float) -> list[Delivery]:
+        """Deliveries whose next beat is at or before `now` (oldest first)."""
+        ready = [d for d in self._tracked.values() if d.next_heartbeat_at <= now]
+        return sorted(ready, key=lambda d: d.next_heartbeat_at)
+
+    def beat(self, now: float) -> int:
+        """Submit one extension for every due delivery; returns how many were submitted."""
+        submitted = 0
+        for delivery in self.due(now):
+            if delivery.done:
+                self.untrack(delivery)
+                continue
+            limit = self._config.max_extensions
+            if limit is not None and delivery.extensions >= limit:
+                logger.warning(
+                    "sqs_heartbeat_max_extensions",
+                    queue=delivery.queue.name,
+                    max_extensions=limit,
+                )
+                self.untrack(delivery)
+                continue
+            if (
+                delivery.expires_at is not None
+                and self._wall_clock() > delivery.expires_at + self._config.deadline_grace_secs
+            ):
+                # DS6: a delivery stuck past its deadline must not stay invisible forever; once
+                # it reappears the broker deletes it as expired (P11).
+                logger.warning(
+                    "sqs_heartbeat_deadline_passed",
+                    queue=delivery.queue.name,
+                    message_id=delivery.message_id,
+                    extensions=delivery.extensions,
+                )
+                self.untrack(delivery)
+                continue
+            delivery.extensions += 1
+            delivery.next_heartbeat_at = now + self._config.interval_for(delivery.visibility_secs)
+            future = self._acker.change_visibility(
+                delivery.queue.name,
+                delivery.queue_url,
+                delivery.receipt_handle,
+                delivery.visibility_secs,
+            )
+            future.add_done_callback(functools.partial(self._settle, delivery))
+            submitted += 1
+        return submitted
+
+    def _settle(self, delivery: Delivery, future: asyncio.Future[VisibilityResult]) -> None:
+        if future.cancelled() or future.exception() is not None:
+            return
+        result = future.result()
+        if result is VisibilityResult.OK:
+            safe_call(self._hooks.on_heartbeat, delivery.queue.name, 1, 0)
+            return
+        if result is VisibilityResult.LOST and not delivery.done:
+            safe_call(self._hooks.on_heartbeat, delivery.queue.name, 0, 1)
+            self.untrack(delivery)
+            delivery.lost.set()
+            if self._on_lost is not None:
+                safe_call(self._on_lost, delivery)
+
+    async def _run(self) -> None:
+        loop = asyncio.get_running_loop()
+        while True:
+            now = loop.time()
+            self.beat(now)
+            nxt = min((d.next_heartbeat_at for d in self._tracked.values()), default=math.inf)
+            self._wake.clear()
+            timeout = None if math.isinf(nxt) else max(0.0, nxt - loop.time())
+            with contextlib.suppress(TimeoutError):
+                await asyncio.wait_for(self._wake.wait(), timeout)

taskiq_llm_queue/broker/queue.py

@@ -0,0 +1,141 @@
+"""Queue, send-batching and heartbeat settings for `HybridSQSBroker`, validated on creation.
+
+Every SQS hard limit (10 messages per call, 20 s long poll, 900 s delay, 1 MiB payload,
+128-character FIFO ids) is checked here, so a bad value fails at boot instead of on the
+first send or receive (P3, P16).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from taskiq_llm_queue.broker.errors import BrokerConfigError
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.protocols import ReceiveGate
+
+MAX_BATCH_ENTRIES = 10
+MAX_WAIT_TIME_SECS = 20
+MAX_DELAY_SECS = 900
+MAX_VISIBILITY_SECS = 43_200
+MAX_MESSAGE_BYTES = 1_048_576
+MAX_BATCH_BYTES = 1_048_576
+# DS15: flush a send batch before 900 KB so attribute overhead never crosses the 1 MiB cap.
+DEFAULT_BATCH_BYTES = 900_000
+MAX_FIFO_ID_CHARS = 128
+
+
+@dataclass(frozen=True, slots=True)
+class Batching:
+    """Send batching for one queue (P8): up to `size` messages or `max_bytes` per call.
+
+    A partial batch is sent after `timeout_secs`; `kick` awaits the send so callers still
+    see failures.
+    """
+
+    size: int = MAX_BATCH_ENTRIES
+    timeout_secs: float = 0.05
+    max_bytes: int = DEFAULT_BATCH_BYTES
+
+    def __post_init__(self) -> None:
+        """Reject values SQS would refuse."""
+        errors = []
+        if not 1 <= self.size <= MAX_BATCH_ENTRIES:
+            errors.append(f"batching.size must be 1..{MAX_BATCH_ENTRIES}, got {self.size}")
+        if not 0 <= self.timeout_secs <= 10:
+            errors.append(f"batching.timeout_secs must be 0..10, got {self.timeout_secs}")
+        if not 1024 <= self.max_bytes <= MAX_BATCH_BYTES:
+            errors.append(f"batching.max_bytes must be 1024..{MAX_BATCH_BYTES}")
+        if errors:
+            raise BrokerConfigError("; ".join(errors))
+
+
+@dataclass(frozen=True, slots=True)
+class Heartbeat:
+    """Visibility extension for received messages (P13, DS6), sent in batches.
+
+    `interval_secs` defaults to (and is capped at) a third of the queue's visibility timeout;
+    `max_extensions=None` extends until the message is acked or released, but never past the
+    message's `expiry` (its `deadline_at`) plus `deadline_grace_secs` (DS6).
+    """
+
+    interval_secs: float | None = None
+    max_extensions: int | None = None
+    flush_secs: float = 0.5
+    deadline_grace_secs: float = 60.0
+
+    def __post_init__(self) -> None:
+        """Reject non-positive intervals and negative extension counts."""
+        if self.interval_secs is not None and self.interval_secs <= 0:
+            raise BrokerConfigError("heartbeat.interval_secs must be > 0")
+        if self.max_extensions is not None and self.max_extensions < 0:
+            raise BrokerConfigError("heartbeat.max_extensions must be >= 0")
+        if not 0 < self.flush_secs <= 1:
+            raise BrokerConfigError("heartbeat.flush_secs must be in (0, 1] (DS6)")
+        if self.deadline_grace_secs < 0:
+            raise BrokerConfigError("heartbeat.deadline_grace_secs must be >= 0")
+
+    def interval_for(self, visibility_secs: int) -> float:
+        """Seconds between extensions for a queue with this visibility timeout (≤ visibility/3)."""
+        limit = max(visibility_secs / 3, 0.1)
+        return limit if self.interval_secs is None else min(self.interval_secs, limit)
+
+
+@dataclass(frozen=True, slots=True)
+class SQSQueue:
+    """One SQS queue the broker sends to and receives from.
+
+    `name` is the routing key (`queue_name` label) and the SQS queue name unless `url` is
+    given. `priority` 0 is highest. `visibility_timeout_secs=None` reads the queue's own
+    setting at startup. `gate` pauses receiving while closed (A4).
+    """
+
+    name: str
+    priority: int = 0
+    url: str | None = None
+    wait_time_secs: int = MAX_WAIT_TIME_SECS
+    max_messages: int = MAX_BATCH_ENTRIES
+    visibility_timeout_secs: int | None = None
+    is_fifo: bool | None = None
+    is_fair: bool = False
+    default_delay_secs: int = 0
+    batching: Batching | None = None
+    is_declare: bool = False
+    attributes: Mapping[str, str] = field(default_factory=dict)
+    gate: ReceiveGate | None = None
+
+    def __post_init__(self) -> None:
+        """Validate every limit at once and derive `is_fifo` from the `.fifo` suffix."""
+        sqs_name = (self.url or self.name).rstrip("/").rsplit("/", 1)[-1]
+        if self.is_fifo is None:
+            object.__setattr__(self, "is_fifo", sqs_name.endswith(".fifo"))
+        errors = []
+        if not self.name:
+            errors.append("name is empty")
+        if self.priority < 0:
+            errors.append(f"priority must be >= 0, got {self.priority}")
+        if not 0 <= self.wait_time_secs <= MAX_WAIT_TIME_SECS:
+            errors.append(f"wait_time_secs must be 0..{MAX_WAIT_TIME_SECS}")
+        if not 1 <= self.max_messages <= MAX_BATCH_ENTRIES:
+            errors.append(f"max_messages must be 1..{MAX_BATCH_ENTRIES}")
+        if self.visibility_timeout_secs is not None and not (
+            0 <= self.visibility_timeout_secs <= MAX_VISIBILITY_SECS
+        ):
+            errors.append(f"visibility_timeout_secs must be 0..{MAX_VISIBILITY_SECS}")
+        if not 0 <= self.default_delay_secs <= MAX_DELAY_SECS:
+            errors.append(f"default_delay_secs must be 0..{MAX_DELAY_SECS}")
+        if self.is_fifo and not sqs_name.endswith(".fifo"):
+            errors.append("a FIFO queue name must end with '.fifo'")
+        if self.is_fifo and self.is_fair:
+            errors.append("is_fair applies to standard queues only")
+        if self.is_fifo and self.default_delay_secs:
+            errors.append("per-message delay is not supported on FIFO queues (DS19)")
+        if errors:
+            raise BrokerConfigError(f"queue {self.name!r}: " + "; ".join(errors))
+
+    @property
+    def fifo(self) -> bool:
+        """`is_fifo` as a plain bool (resolved in `__post_init__`)."""
+        return bool(self.is_fifo)