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

taskiq_llm_queue/reliability/claims.py

@@ -0,0 +1,226 @@
+"""Exactly-once execution on at-least-once SQS delivery: the claim record per task (DS1-DS3).
+
+`<prefix>:{r:<id>}:job` holds `state`, `worker`, `heartbeat_at`, `deadline_at` and `failures`;
+every transition is one Lua CAS using Valkey `TIME`. Used by the claim middleware before a task
+runs, by the heartbeat loop, and by `direct` failover (DS17: claim as `direct:<pod>`).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any
+
+from redis.crc import key_slot
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.keys import KeySpace
+from taskiq_llm_queue.protocols import ClaimDecision, ClaimState
+from taskiq_llm_queue.reliability import _lua
+
+if TYPE_CHECKING:
+    import redis.asyncio
+
+    from taskiq_llm_queue.config import QueueConfig
+
+logger = get_logger(__name__)
+
+_TERMINAL = frozenset({ClaimState.DONE, ClaimState.CANCELLED, ClaimState.EXPIRED})
+
+
+class ClaimOutcome(StrEnum):
+    """Result of an owner-checked `complete_owned` / `release_owned` call.
+
+    `NOT_OWNER` means another worker holds the claim now, so nothing was written.
+    """
+
+    APPLIED = "applied"
+    NOT_RUNNING = "not_running"  # release only: the record was not running
+    NOT_OWNER = "not_owner"
+
+
+@dataclass(frozen=True, slots=True)
+class ClaimSettings:
+    """Claim timing: TTLs, worst-case execution and heartbeat staleness (DS3, DS5).
+
+    A heartbeat older than `stale_after_beats x heartbeat_interval_secs` is stale and the task
+    is reclaimed; at least 3 beats, so one missed or late beat never hands a live task over.
+    """
+
+    state_ttl_secs: int = 1800
+    exec_max_secs: int = 300
+    heartbeat_interval_secs: float = 10.0
+    max_failures: int = 3
+    stale_after_beats: float = 3.0
+
+    def __post_init__(self) -> None:
+        """Reject a staleness window that a single missed heartbeat could exceed."""
+        if self.stale_after_beats < 3:
+            raise ValueError("stale_after_beats must be >= 3 heartbeat intervals")
+
+    @property
+    def stale_after_secs(self) -> float:
+        """Heartbeat age after which a running claim may be taken over."""
+        return self.stale_after_beats * self.heartbeat_interval_secs
+
+    @classmethod
+    def from_config(cls, config: QueueConfig, **overrides: Any) -> ClaimSettings:
+        """Settings whose state TTL comes from `QueueConfig.state_ttl_secs`."""
+        return cls(state_ttl_secs=config.state_ttl_secs, **overrides)
+
+
+@dataclass(frozen=True, slots=True)
+class FailureCount:
+    """Real failures counted for a task and whether it should now go to the DLQ (DS5)."""
+
+    count: int
+    exhausted: bool
+
+
+class RedisClaimStore:
+    """`ClaimStore` on Valkey: claim, heartbeat, complete, release, failures and cancel flag.
+
+    A running task with a fresh heartbeat is never acked by another pod (DEFER, DS2); terminal
+    records outlive the message so duplicates are acked without a second provider call.
+    """
+
+    def __init__(
+        self,
+        client: redis.asyncio.Redis,
+        keys: KeySpace,
+        settings: ClaimSettings | None = None,
+    ) -> None:
+        """Register the Lua scripts on `client`; no I/O happens until the first call."""
+        self._client = client
+        self._keys = keys
+        self.settings = settings or ClaimSettings()
+        self._claim = client.register_script(_lua.CLAIM)
+        self._heartbeat = client.register_script(_lua.CLAIM_HEARTBEAT)
+        self._complete = client.register_script(_lua.CLAIM_COMPLETE)
+        self._release = client.register_script(_lua.CLAIM_RELEASE)
+
+    async def claim(self, request_id: str, *, deadline_at: float, worker: str) -> ClaimDecision:
+        """Atomically decide whether this worker runs the task (Lua CAS)."""
+        s = self.settings
+        decision = await self._claim(
+            [self._keys.job(request_id), self._keys.cancel(request_id)],
+            [
+                repr(float(deadline_at)),
+                worker,
+                s.state_ttl_secs,
+                s.exec_max_secs,
+                s.stale_after_secs,
+            ],
+        )
+        return ClaimDecision(_lua.as_text(decision))
+
+    async def heartbeat_owned(self, request_ids: Sequence[str], *, worker: str) -> list[str]:
+        """Refresh heartbeat and TTL of tasks this worker runs; return the ids it lost (DS6).
+
+        Lost means another worker owns the claim now or the record is gone; a claim this worker
+        already completed or released is skipped, not lost (its task is just finishing).
+        """
+        if not request_ids:
+            return []
+        s = self.settings
+        args: list[Any] = [worker, s.state_ttl_secs, s.exec_max_secs]
+        keys = [self._keys.job(rid) for rid in request_ids]
+        # Each job key has its own `{r:<id>}` slot and a cluster script may touch only one:
+        # one call per slot, all pipelined in a single round trip (standalone Valkey too).
+        by_slot: dict[int, list[int]] = {}
+        for index, key in enumerate(keys):
+            by_slot.setdefault(key_slot(key.encode()), []).append(index)
+        groups = list(by_slot.values())
+        replies = await _lua.run_many(
+            self._client, [(self._heartbeat, [keys[i] for i in group], args) for group in groups]
+        )
+        missing = sorted(
+            group[int(i) - 1] for group, lost in zip(groups, replies, strict=True) for i in lost
+        )
+        return [request_ids[index] for index in missing]
+
+    async def heartbeat(self, request_ids: Sequence[str], *, worker: str) -> None:
+        """Refresh heartbeats and TTLs for running tasks."""
+        await self.heartbeat_owned(request_ids, worker=worker)
+
+    async def complete(
+        self, request_id: str, state: ClaimState, *, worker: str | None = None
+    ) -> None:
+        """Mark a task done, cancelled or expired after its result is stored (DS8).
+
+        With `worker`, only the current owner may complete it (see `complete_owned`).
+        """
+        await self.complete_owned(request_id, state, worker=worker)
+
+    async def complete_owned(
+        self, request_id: str, state: ClaimState, *, worker: str | None
+    ) -> ClaimOutcome:
+        """Owner-checked `complete`: NOT_OWNER when another worker holds the claim now.
+
+        A record without an owner (released, or missing) is written; `worker=None` skips the check.
+        """
+        script, keys, args = self.complete_call(request_id, state, worker=worker)
+        return self.complete_outcome(await script(keys, args), request_id, worker=worker)
+
+    @property
+    def client(self) -> redis.asyncio.Redis:
+        """The client the scripts run on (lets a caller share one pipeline with it)."""
+        return self._client
+
+    def complete_call(
+        self, request_id: str, state: ClaimState, *, worker: str | None
+    ) -> _lua.ScriptCall:
+        """The `complete_owned` script call, to queue in a caller's pipeline (`_lua`)."""
+        if state not in _TERMINAL:
+            raise ValueError(f"complete() needs a terminal state, got {state!r}")
+        args = [state.value, self.settings.state_ttl_secs, worker or ""]
+        return self._complete, [self._keys.job(request_id)], args
+
+    def complete_outcome(self, reply: Any, request_id: str, *, worker: str | None) -> ClaimOutcome:
+        """Interpret a `complete_call` reply: APPLIED, or NOT_OWNER (logged)."""
+        if int(reply) < 0:
+            logger.warning("claim_complete_not_owner", request_id=request_id, worker=worker)
+            return ClaimOutcome.NOT_OWNER
+        return ClaimOutcome.APPLIED
+
+    async def release(self, request_id: str, *, worker: str | None = None) -> None:
+        """Hand a running task back so another pod can start it at once.
+
+        With `worker`, only the current owner may release it (see `release_owned`).
+        """
+        await self.release_owned(request_id, worker=worker)
+
+    async def release_owned(self, request_id: str, *, worker: str | None) -> ClaimOutcome:
+        """Owner-checked `release`: APPLIED, NOT_RUNNING, or NOT_OWNER (nothing written).
+
+        `worker=None` releases any running record (the old, unchecked behaviour).
+        """
+        reply = int(await self._release([self._keys.job(request_id)], [worker or ""]))
+        if reply < 0:
+            logger.warning("claim_release_not_owner", request_id=request_id, worker=worker)
+            return ClaimOutcome.NOT_OWNER
+        return ClaimOutcome.APPLIED if reply else ClaimOutcome.NOT_RUNNING
+
+    async def state(self, request_id: str) -> ClaimState | None:
+        """Current claim state, or None when there is no record."""
+        raw = await self._client.hget(self._keys.job(request_id), "state")
+        return None if raw is None else ClaimState(_lua.as_text(raw))
+
+    async def record_failure(self, request_id: str) -> FailureCount:
+        """Count one real failure (not a throttle or a release) against `max_failures` (DS5)."""
+        key = self._keys.job(request_id)
+        async with self._client.pipeline(transaction=True) as pipe:
+            pipe.hincrby(key, "failures", 1)
+            pipe.expire(key, self.settings.state_ttl_secs, nx=True)  # never leave a key without TTL
+            raw_count, _ = await pipe.execute()
+        count = int(raw_count)
+        return FailureCount(count, count >= self.settings.max_failures)
+
+    async def cancel(self, request_id: str) -> None:
+        """Set the cancel flag that workers check before every attempt and sleep (DS14)."""
+        await self._client.set(self._keys.cancel(request_id), b"1", ex=self.settings.state_ttl_secs)
+
+    async def is_cancelled(self, request_id: str) -> bool:
+        """True when the caller cancelled or timed out on this task."""
+        return bool(await self._client.exists(self._keys.cancel(request_id)))

taskiq_llm_queue/reliability/config_store.py

@@ -0,0 +1,148 @@
+"""Live JSON config in Valkey (routing, fallback chains) with last-good fallback and audit.
+
+Readers validate with pydantic and keep the last good value when the stored JSON is invalid or
+Valkey is unreachable, so a bad edit never breaks live calls. Every write appends who changed
+what to `KeySpace.audit` in the same Lua call, trimmed to the retention window.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Mapping
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from pydantic import BaseModel, ValidationError
+from redis.exceptions import RedisError
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.errors import ConfigError
+from taskiq_llm_queue.keys import KeySpace
+from taskiq_llm_queue.reliability import _lua
+from taskiq_llm_queue.reliability.fallback import FallbackConfig
+
+if TYPE_CHECKING:
+    import redis.asyncio
+
+logger = get_logger(__name__)
+ModelT = TypeVar("ModelT", bound=BaseModel)
+AUDIT_RETENTION_SECS = 90 * 86400
+
+
+class LiveConfigStore(Generic[ModelT]):  # PEP 695 needs 3.12; package supports 3.11
+    """One live config value of type `model_type` at `key`, with audit trail and reload loop.
+
+    `current` is the last good value (the startup default until a load succeeds); hot paths read
+    it without I/O, and `run_reload` refreshes it every `interval_secs` (`LLMQ_CFG_RELOAD_SECS`).
+    """
+
+    def __init__(
+        self,
+        client: redis.asyncio.Redis,
+        key: str,
+        model_type: type[ModelT],
+        default: ModelT,
+        *,
+        audit_key: str,
+        audit_retention_secs: int = AUDIT_RETENTION_SECS,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    ) -> None:
+        """Keep the client and the startup default; nothing is read until `load`."""
+        self._client = client
+        self.key = key
+        self._type = model_type
+        self.default = default
+        self._audit_key = audit_key
+        self._retention_ms = audit_retention_secs * 1000
+        self._sleep = sleep
+        self._current = default
+        self._warned: set[bytes] = set()
+        self.invalid_count = 0
+        self._set = client.register_script(_lua.CONFIG_SET)
+
+    @property
+    def current(self) -> ModelT:
+        """Last good value; never raises and never does I/O."""
+        return self._current
+
+    def parse(self, raw: str | bytes | Mapping[str, Any] | ModelT) -> ModelT:
+        """Validate JSON, a mapping or a model instance; raises `ConfigError` when invalid."""
+        try:
+            if isinstance(raw, self._type):
+                return raw
+            if isinstance(raw, str | bytes):
+                return self._type.model_validate_json(raw)
+            return self._type.model_validate(raw)
+        except ValidationError as exc:
+            raise ConfigError(f"invalid {self.key}: {exc}") from exc
+
+    async def load(self) -> ModelT:
+        """Read and validate the stored value; on bad data or no Valkey keep the last good one.
+
+        A missing key means the startup default. Invalid data is logged once per distinct value
+        and counted in `invalid_count` (exported as `llmq_config_invalid_total`).
+        """
+        try:
+            raw = await self._client.get(self.key)
+        except RedisError as exc:
+            logger.warning("live_config_unreachable", key=self.key, error=str(exc))
+            return self._current
+        if raw is None:
+            self._current = self.default
+            return self._current
+        raw_bytes = raw if isinstance(raw, bytes) else str(raw).encode()
+        try:
+            self._current = self.parse(raw_bytes)
+        except ConfigError as exc:
+            self.invalid_count += 1
+            if raw_bytes not in self._warned:
+                self._warned.add(raw_bytes)
+                logger.warning(
+                    "live_config_invalid",
+                    key=self.key,
+                    invalid_count=self.invalid_count,
+                    error=str(exc),
+                )
+        return self._current
+
+    async def set(self, value: str | bytes | Mapping[str, Any] | ModelT, *, actor: str) -> ModelT:
+        """Validate and store a new value, appending an audit entry; returns the stored value.
+
+        Raises `ConfigError` for invalid input, so a bad CLI edit never reaches Valkey.
+        """
+        parsed = self.parse(value)
+        await self._set(
+            [self.key, self._audit_key], [parsed.model_dump_json(), actor, self._retention_ms]
+        )
+        self._current = parsed
+        return parsed
+
+    async def reset(self, *, actor: str) -> ModelT:
+        """Delete the stored value (back to the startup default), with an audit entry."""
+        await self._set([self.key, self._audit_key], ["", actor, self._retention_ms])
+        self._current = self.default
+        return self._current
+
+    async def history(self, count: int = 20) -> list[dict[str, str]]:
+        """Latest audit entries for this key, newest first (for `llm-queue ... get`)."""
+        entries = await self._client.xrevrange(self._audit_key, count=count * 4)
+        out: list[dict[str, str]] = []
+        for _id, fields in entries or ():
+            row = {_lua.as_text(k): _lua.as_text(v) for k, v in (fields or {}).items()}
+            if row.get("key") == self.key:
+                out.append(row)
+        return out[:count]
+
+    async def run_reload(self, interval_secs: float) -> None:
+        """Reload every `interval_secs` until cancelled; errors never end the loop."""
+        while True:
+            await self.load()
+            await self._sleep(interval_secs)
+
+
+def fallback_store(
+    client: redis.asyncio.Redis, keys: KeySpace, default: FallbackConfig | None = None
+) -> LiveConfigStore[FallbackConfig]:
+    """Live store for fallback chains at `KeySpace.fallback`, audited to `KeySpace.audit`."""
+    return LiveConfigStore(
+        client, keys.fallback, FallbackConfig, default or FallbackConfig(), audit_key=keys.audit
+    )

taskiq_llm_queue/reliability/fallback.py

@@ -0,0 +1,111 @@
+"""Fallback order for one call: same model on other targets first, other models only on `chain`.
+
+Pure logic over `QueueConfig` and the live `FallbackConfig`. Embeddings never
+switch model (vectors from another model are incompatible). The worker walks the steps in order
+and checks limiter headroom itself where a step says `require_headroom`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+from taskiq_llm_queue.errors import ConfigError
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.config import QueueConfig
+    from taskiq_llm_queue.envelope import FallbackPolicy, Op
+
+
+class ChainStep(BaseModel):
+    """One other model in a chain (`alias` or `alias@target`); a plain string is accepted."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    model: str = Field(min_length=1)
+    require_headroom: bool = False
+
+
+class FallbackConfig(BaseModel):
+    """Live fallback chains stored as JSON at `KeySpace.fallback`.
+
+    `chains` maps a primary alias to the other models tried under policy `chain`;
+    `skip_targets` lists targets not usable yet per alias (e.g. APAC before sign-off).
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    chains: dict[str, tuple[ChainStep, ...]] = Field(default_factory=dict)
+    skip_targets: dict[str, frozenset[str]] = Field(default_factory=dict)
+
+    @field_validator("chains", mode="before")
+    @classmethod
+    def _strings_to_steps(cls, value: Any) -> Any:
+        if not isinstance(value, dict):
+            return value
+        return {
+            alias: [{"model": s} if isinstance(s, str) else s for s in steps]
+            if isinstance(steps, list | tuple)
+            else steps
+            for alias, steps in value.items()
+        }
+
+    @model_validator(mode="after")
+    def _no_self_reference(self) -> FallbackConfig:
+        loops = [a for a, steps in self.chains.items() if any(_alias(s.model) == a for s in steps)]
+        if loops:
+            raise ValueError(f"chains must not list the primary model itself: {loops}")
+        return self
+
+
+@dataclass(frozen=True, slots=True)
+class FallbackStep:
+    """One place to try: model alias and target, whether it changes model, headroom rule."""
+
+    alias: str
+    target: str
+    switched_model: bool = False
+    require_headroom: bool = False
+
+
+def _alias(model: str) -> str:
+    return model.partition("@")[0]
+
+
+def resolve_fallback(
+    config: QueueConfig,
+    model: str,
+    *,
+    policy: FallbackPolicy = "same_model",
+    op: Op = "converse",
+    live: FallbackConfig | None = None,
+) -> tuple[FallbackStep, ...]:
+    """Ordered steps for `model` (`alias[@target]`): the requested target, then its other targets.
+
+    With policy `chain` (and never for embeddings) the chain's models follow, each primary
+    target first. Unknown aliases or targets in the live chain are skipped, not raised, because
+    live config may lag the deployed `QueueConfig`; the requested model itself must exist.
+    """
+    live = live or FallbackConfig()
+    alias, cfg, first = config.resolve(model)
+    skip = live.skip_targets.get(alias, frozenset())
+    steps = [FallbackStep(alias, first)]
+    steps += [FallbackStep(alias, t) for t in cfg.targets if t != first and t not in skip]
+    if policy != "chain" or op == "embeddings":
+        return tuple(steps)
+    seen = {alias}
+    for step in live.chains.get(alias, ()):
+        other, _, pinned = step.model.partition("@")
+        if other in seen:
+            continue
+        try:
+            _, other_cfg, _ = config.resolve(step.model)
+        except ConfigError:
+            continue
+        seen.add(other)
+        other_skip = live.skip_targets.get(other, frozenset())
+        targets = [pinned] if pinned else [t for t in other_cfg.targets if t not in other_skip]
+        steps += [FallbackStep(other, t, True, step.require_headroom) for t in targets]
+    return tuple(steps)