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

taskiq_llm_queue/envelope.py

@@ -0,0 +1,136 @@
+"""Wire contract: the task envelope sent through SQS and the result returned through Redis.
+
+Both sides of a rolling deploy read these, so the format is versioned (`v`). A worker accepts
+its own version and the previous one; anything newer raises `UnsupportedEnvelopeError` so the
+worker releases the message instead of dropping it.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from typing import Any, Literal
+
+import orjson
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from taskiq_llm_queue.config import DEFAULT_ID_HASH_SALT, Lane
+from taskiq_llm_queue.errors import UnsupportedEnvelopeError
+
+ENVELOPE_VERSION = 1
+TASK_CALL = "taskiq_llm_queue.call"
+TASK_EMBED = "taskiq_llm_queue.embed"
+TASK_JOB = "taskiq_llm_queue.job"
+
+Op = Literal["converse", "invoke_model", "chat", "embeddings"]
+FallbackPolicy = Literal["same_model", "chain"]
+
+
+class Envelope(BaseModel):
+    """One LLM call as a task argument; `params` are provider-native and passed through as-is.
+
+    Large params travel by reference (`params_ref`, a Redis key) instead of inline.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    v: int = ENVELOPE_VERSION
+    request_id: str
+    parent_request_id: str | None = None
+    session_id: str | None = None
+    user_id: str | None = None
+    reply_to: str | None = None
+    lane: Lane = Lane.BG
+    op: Op
+    fallback: FallbackPolicy = "same_model"
+    model: str
+    params: dict[str, Any] | None = None
+    params_ref: str | None = None
+    est_tokens: int = Field(default=0, ge=0)
+    enqueued_at: float = Field(default_factory=time.time)
+    deadline_at: float
+    trace: dict[str, str] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def _params_xor_ref(self) -> Envelope:
+        if (self.params is None) == (self.params_ref is None):
+            raise ValueError("exactly one of params / params_ref must be set")
+        if self.deadline_at < self.enqueued_at:
+            raise ValueError("deadline_at is before enqueued_at")
+        return self
+
+
+class CallMeta(BaseModel):
+    """How a call was actually served; surfaced as `response_metadata` / `ConverseResult.meta`."""
+
+    model_config = ConfigDict(frozen=True)
+
+    served_model: str = ""
+    served_target: str = ""
+    fallback_used: bool = False
+    queue_wait_ms: int = 0
+    provider_ms: int = 0
+    attempts: int = 0
+    worker: str = ""
+
+
+class LLMReply(BaseModel):
+    """Result stored by the worker: the provider-native response, or a classified error."""
+
+    model_config = ConfigDict(frozen=True)
+
+    request_id: str
+    ok: bool
+    response: dict[str, Any] | None = None
+    meta: CallMeta = Field(default_factory=CallMeta)
+    error: dict[str, Any] | None = None
+    cancelled: bool = False
+
+
+def parse_envelope(data: bytes | str | dict[str, Any]) -> Envelope:
+    """Decode an envelope, enforcing the version rule before field validation.
+
+    Raises `UnsupportedEnvelopeError` for a newer version and `ValueError` for malformed input.
+    """
+    raw = orjson.loads(data) if isinstance(data, bytes | str) else data
+    if not isinstance(raw, dict):
+        raise ValueError("envelope is not a JSON object")
+    version = raw.get("v")
+    if not isinstance(version, int):
+        raise ValueError("envelope has no integer 'v'")
+    if version > ENVELOPE_VERSION:
+        raise UnsupportedEnvelopeError(f"envelope v{version} is newer than v{ENVELOPE_VERSION}")
+    if version < ENVELOPE_VERSION - 1:
+        raise ValueError(f"envelope v{version} is no longer supported")
+    return Envelope.model_validate(raw)
+
+
+def keyed_digest(value: str, salt: str | bytes = DEFAULT_ID_HASH_SALT) -> bytes:
+    """8-byte keyed BLAKE2b of `value`; equal on every pod that shares `salt` (max 64 bytes).
+
+    Keyed, so ids cannot be recovered by hashing guesses without the salt.
+    """
+    key = salt.encode("utf-8") if isinstance(salt, str) else salt
+    return hashlib.blake2b(value.encode("utf-8"), digest_size=8, key=key).digest()
+
+
+def hashed_id(value: str | None, salt: str | bytes = DEFAULT_ID_HASH_SALT) -> str | None:
+    """Stable, non-reversible id for labels, logs and metrics (raw ids stay in the body).
+
+    `salt` is `QueueConfig.id_hash_salt` (`LLMQ_ID_HASH_SALT`); set your own in production.
+    """
+    return None if value is None else keyed_digest(value, salt).hex()
+
+
+def task_labels(
+    envelope: Envelope, *, hash_ids: bool = True, salt: str | bytes = DEFAULT_ID_HASH_SALT
+) -> dict[str, str]:
+    """Taskiq labels for one envelope: routing and context fields, ids hashed by default."""
+    labels = {"lane": envelope.lane.value, "model": envelope.model}
+    for key, value in (("session_id", envelope.session_id), ("user_id", envelope.user_id)):
+        shown = hashed_id(value, salt) if hash_ids else value
+        if shown is not None:
+            labels[key] = shown
+    if envelope.reply_to:
+        labels["reply_to"] = envelope.reply_to
+    return labels

taskiq_llm_queue/errors.py

@@ -0,0 +1,71 @@
+"""Exception tree and provider-error classification shared by every layer.
+
+Callers catch one type, `LLMQueueError`; workers turn provider exceptions into `ErrorInfo` so
+the limiter, retry and fallback logic never depend on SDK-specific exception classes.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+
+
+class LLMQueueError(Exception):
+    """Base class for every error this package raises to callers."""
+
+
+class LLMQueueTimeout(LLMQueueError):  # public name is part of the frozen API
+    """The caller waited longer than its lane timeout or deadline for a result."""
+
+
+class LLMQueueCancelled(LLMQueueError):
+    """The call was cancelled (`LLMQueue.cancel`, or its caller went away) before it finished."""
+
+
+class QueueUnavailable(LLMQueueError):  # public name is part of the frozen API
+    """SQS or Redis/Valkey could not be reached and no fallback path was allowed."""
+
+
+class ConfigError(LLMQueueError):
+    """Invalid configuration, raised at startup (never on a live request path)."""
+
+
+class UnsupportedEnvelopeError(LLMQueueError):
+    """A message carries an envelope version newer than this worker understands.
+
+    Workers release such messages for another pod instead of acking them, so a rolling deploy
+    never drops work.
+    """
+
+
+class ProviderError(LLMQueueError):
+    """The provider call failed after the worker's retries and fallbacks were exhausted."""
+
+    def __init__(
+        self, message: str, *, retryable: bool, retry_after_secs: float | None = None
+    ) -> None:
+        """Keep whether the caller may retry later and any provider-advised wait."""
+        super().__init__(message)
+        self.retryable = retryable
+        self.retry_after_secs = retry_after_secs
+
+
+class ErrorKind(StrEnum):
+    """How the worker reacts to a provider error: shrink limits, retry, or give up."""
+
+    THROTTLE = "throttle"
+    RETRYABLE = "retryable"
+    FATAL = "fatal"
+
+
+@dataclass(frozen=True, slots=True)
+class ErrorInfo:
+    """Provider-neutral classification of one failed call, produced by `Provider.classify`.
+
+    `code` keeps the provider's own spelling for logs only; decisions use `kind`, because
+    fakes and real services spell the same failure differently.
+    """
+
+    kind: ErrorKind
+    code: str = ""
+    retry_after_secs: float | None = None

taskiq_llm_queue/failover.py

@@ -0,0 +1,286 @@
+"""SQS and Valkey failover state for the sender side (D73, D74, DS17, DS18).
+
+A failed SQS send puts that model on `direct` for `sqs_failover_secs`; a Valkey failure puts the
+whole pod on `direct` with a local per-pod cap until a health probe (or a backoff window) says
+Valkey is back. `LLMQueue` calls the `FailoverPolicy` methods; nothing here blocks or raises.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import threading
+import time
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.config import env_value
+from taskiq_llm_queue.errors import ConfigError
+from taskiq_llm_queue.metrics_names import FAILOVER_ACTIVE
+from taskiq_llm_queue.reliability.limiter import local_cap
+from taskiq_llm_queue.reliability.local_limiter import LocalLimiter
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.metrics import MetricsSink
+
+logger = get_logger(__name__)
+Probe = Callable[[], Awaitable[object]]
+
+
+@dataclass(frozen=True, slots=True)
+class FailoverSettings:
+    """Failover timings and the local cap (`LLMQ_SQS_FAILOVER_SECS`, `LLMQ_LOCAL_LIMIT_CAP`).
+
+    `live_pods` < 1 means unknown: the local cap is then `local_cap_fallback`, never
+    `last_known_limit` per pod, so an unknown pod count cannot overshoot the provider limit.
+    A running `LLMQueue` replaces it with the pod-registry count (`FailoverState.set_live_pods`).
+    """
+
+    sqs_failover_secs: float = 300.0
+    redis_retry_min_secs: float = 1.0
+    redis_retry_max_secs: float = 30.0
+    local_cap_fallback: int = 20
+    live_pods: int = 0
+
+    def __post_init__(self) -> None:
+        """Reject impossible values at construction (boot time)."""
+        if self.sqs_failover_secs <= 0 or self.local_cap_fallback < 1:
+            raise ConfigError("sqs_failover_secs must be > 0 and local_cap_fallback >= 1")
+        if not 0 < self.redis_retry_min_secs <= self.redis_retry_max_secs:
+            raise ConfigError("need 0 < redis_retry_min_secs <= redis_retry_max_secs")
+
+    @classmethod
+    def from_env(
+        cls, environ: Mapping[str, str] | None = None, *, live_pods: int = 0
+    ) -> FailoverSettings:
+        """Read `LLMQ_SQS_FAILOVER_SECS` and `LLMQ_LOCAL_LIMIT_CAP`; bad values raise `ConfigError`.
+
+        The deprecated `LLMQ_SQS_FAILOVER_S` is still read (warned once).
+        """
+        env = os.environ if environ is None else environ
+        values: dict[str, str] = {}
+        errors: list[str] = []
+        for name, conv in (("SQS_FAILOVER_SECS", float), ("LOCAL_LIMIT_CAP", int)):
+            raw = env_value(env, name)
+            if raw is None:
+                continue
+            try:
+                conv(raw)
+            except ValueError as exc:
+                errors.append(f"LLMQ_{name}: {exc}")
+            values[name] = raw
+        if errors:
+            raise ConfigError("; ".join(errors))
+        return cls(
+            sqs_failover_secs=float(values.get("SQS_FAILOVER_SECS", 300.0)),
+            local_cap_fallback=int(values.get("LOCAL_LIMIT_CAP", 20)),
+            live_pods=live_pods,
+        )
+
+
+class FailoverState:
+    """`FailoverPolicy`: per-model SQS failover windows and pod-wide Valkey-down mode.
+
+    With a `probe` (e.g. `client.ping`) Valkey-down lasts until a probe succeeds, retried with
+    exponential backoff on the running loop. Without one, or from a thread with no loop, the
+    mode lapses after the backoff window so the next call tries Valkey again.
+    """
+
+    def __init__(
+        self,
+        settings: FailoverSettings | None = None,
+        *,
+        probe: Probe | None = None,
+        metrics: MetricsSink | None = None,
+        clock: Callable[[], float] = time.monotonic,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    ) -> None:
+        """Keep settings, the optional Valkey probe, metrics sink and injected clock/sleep."""
+        self.settings = settings or FailoverSettings()
+        self._probe = probe
+        self._metrics = metrics
+        self._clock = clock
+        self._sleep = sleep
+        self._lock = threading.Lock()
+        self._sqs_until: dict[str, float] = {}
+        self._redis_down = False
+        self._redis_retry_at = 0.0
+        self._redis_backoff = self.settings.redis_retry_min_secs
+        self._probe_task: asyncio.Task[None] | None = None
+        self._limits: dict[tuple[str, str], int] = {}
+        self._local: dict[tuple[str, str], LocalLimiter] = {}
+        # Local limiters with calls still in flight when Valkey came back; reused if it fails
+        # again before they finish, so a flap never starts a fresh cap on top of busy slots.
+        self._busy: dict[tuple[str, str], LocalLimiter] = {}
+        self._live_pods: int | None = None
+
+    # ------------------------------------------------------------------ SQS
+
+    def path_override(self, model: str) -> Literal["direct"] | None:
+        """`direct` while Valkey is down or this model's SQS failover window is open, else None."""
+        if self.redis_down() or self.sqs_failover_active(model):
+            return "direct"
+        return None
+
+    def report_sqs_failure(self, model: str, exc: BaseException) -> None:
+        """Open (or extend) the model's SQS failover window after a failed or ambiguous send."""
+        alias = model.partition("@")[0]
+        with self._lock:
+            entering = alias not in self._sqs_until
+            self._sqs_until[alias] = self._clock() + self.settings.sqs_failover_secs
+        if entering:
+            logger.error(
+                "sqs_failover_active",
+                model=alias,
+                failover_secs=self.settings.sqs_failover_secs,
+                error_type=type(exc).__name__,
+            )
+            self._gauge(1.0, {"model": alias, "kind": "sqs"})
+
+    def sqs_failover_active(self, model: str) -> bool:
+        """True while `model` (alias or `alias@target`) is inside its failover window."""
+        alias = model.partition("@")[0]
+        with self._lock:
+            until = self._sqs_until.get(alias)
+            if until is None:
+                return False
+            if self._clock() < until:
+                return True
+            del self._sqs_until[alias]
+        logger.warning("sqs_failover_ended", model=alias)
+        self._gauge(0.0, {"model": alias, "kind": "sqs"})
+        return False
+
+    # ------------------------------------------------------------------ Valkey
+
+    def report_redis_failure(self, exc: BaseException) -> None:
+        """Enter (or stay in) Valkey-down mode and start the health probe if one is configured."""
+        with self._lock:
+            entering = not self._redis_down
+            if entering:
+                self._redis_backoff = self.settings.redis_retry_min_secs
+            else:
+                self._redis_backoff = min(
+                    self._redis_backoff * 2, self.settings.redis_retry_max_secs
+                )
+            self._redis_down = True
+            self._redis_retry_at = self._clock() + self._redis_backoff
+        if entering:
+            logger.error("valkey_local_mode_active", error_type=type(exc).__name__)
+            self._gauge(1.0, {"kind": "redis"})
+        self._start_probe()
+
+    def report_redis_ok(self) -> None:
+        """Leave Valkey-down mode (a probe or a real call succeeded).
+
+        Idle local limiters are dropped (rebuilt per outage); busy ones are kept aside so the
+        in-flight calls stay counted if Valkey fails again before they finish.
+        """
+        with self._lock:
+            if not self._redis_down:
+                return
+            self._redis_down = False
+            self._redis_backoff = self.settings.redis_retry_min_secs
+            self._busy = {k: lim for k, lim in self._local.items() if lim.inflight(*k) > 0}
+            self._local.clear()
+        logger.warning("valkey_local_mode_ended")
+        self._gauge(0.0, {"kind": "redis"})
+
+    def redis_down(self) -> bool:
+        """True while Valkey is considered down (probe running, or inside the backoff window)."""
+        with self._lock:
+            if not self._redis_down:
+                return False
+            probing = self._probe_task is not None and not self._probe_task.done()
+            return probing or self._clock() < self._redis_retry_at
+
+    async def probe_redis(self) -> bool:
+        """Run the probe once; success leaves Valkey-down mode. False when no probe is set."""
+        if self._probe is None:
+            return False
+        try:
+            await self._probe()
+        except Exception as exc:  # any failure of the probe means "still down"
+            logger.debug("valkey_probe_failed", error_type=type(exc).__name__)
+            return False
+        self.report_redis_ok()
+        return True
+
+    async def _probe_loop(self) -> None:
+        delay = self.settings.redis_retry_min_secs
+        while self._redis_down:
+            await self._sleep(delay)
+            if await self.probe_redis():
+                return
+            delay = min(delay * 2, self.settings.redis_retry_max_secs)
+
+    def _start_probe(self) -> None:
+        if self._probe is None:
+            return
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:  # sync caller in a thread: the backoff window covers it
+            return
+        with self._lock:
+            if self._probe_task is None or self._probe_task.done():
+                self._probe_task = loop.create_task(self._probe_loop(), name="llmq-valkey-probe")
+
+    # ------------------------------------------------------------------ local cap
+
+    def note_limit(self, model: str, target: str, limit: int | None) -> None:
+        """Remember the last shared limit seen for a model target (feeds the local cap)."""
+        if limit is not None:
+            with self._lock:
+                self._limits[(model, target)] = limit
+
+    def set_live_pods(self, count: int) -> None:
+        """Live pod count from the pod registry; feeds the next local cap (DS18)."""
+        if count >= 1:
+            with self._lock:
+                self._live_pods = count
+
+    @property
+    def live_pods(self) -> int:
+        """Pods sharing the provider limit: the registry count, else `settings.live_pods`."""
+        with self._lock:
+            live = self._live_pods
+        return self.settings.live_pods if live is None else live
+
+    def local_cap(self, model: str, target: str) -> int:
+        """Per-pod cap while Valkey is down: last known limit / live pods, else the fallback."""
+        with self._lock:
+            last = self._limits.get((model, target))
+        return local_cap(last, self.live_pods, self.settings.local_cap_fallback)
+
+    def local_limiter(self, model: str, target: str) -> LocalLimiter:
+        """The pod's `LocalLimiter` for a model target, built once per Valkey outage."""
+        key = (model, target)
+        with self._lock:
+            limiter = self._local.get(key)
+            if limiter is None and self._redis_down:
+                busy = self._busy.pop(key, None)
+                if busy is not None and busy.inflight(model, target) > 0:
+                    limiter = self._local[key] = busy
+        if limiter is None:
+            limiter = LocalLimiter(self.local_cap(model, target), sleep=self._sleep)
+            with self._lock:
+                limiter = self._local.setdefault(key, limiter)
+        return limiter
+
+    async def aclose(self) -> None:
+        """Stop the probe task; safe to call twice."""
+        task, self._probe_task = self._probe_task, None
+        if task is not None:
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+    def _gauge(self, value: float, labels: dict[str, str]) -> None:
+        if self._metrics is not None:
+            try:
+                self._metrics.gauge(FAILOVER_ACTIVE, value, labels)
+            except Exception:  # metrics must never break failover
+                logger.exception("failover_metric_failed")