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

taskiq_llm_queue/results/s3_backend.py

@@ -0,0 +1,105 @@
+"""Taskiq result backend on S3: JSON results at `s3://bucket/<prefix>/<task_id>.json` (P10).
+
+For results that are large or must outlive Redis. S3 has no per-object TTL, so expiry is an
+S3 lifecycle rule on `prefix` (e.g. expire after 1 day), set up outside this library.
+Adapted from taskiq-sqs `S3ResultBackend` (MIT, see NOTICE). Results are JSON, never pickle.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+from taskiq import AsyncResultBackend, TaskiqResult
+from taskiq.exceptions import ResultGetError
+
+from taskiq_llm_queue.broker.s3_offload import S3Connection, is_missing, object_key
+
+if TYPE_CHECKING:
+    from aiobotocore.config import AioConfig
+    from aiobotocore.session import AioSession
+    from types_aiobotocore_s3 import S3Client
+
+
+class S3ResultBackend(AsyncResultBackend[Any]):
+    """Stores each `TaskiqResult` as one JSON object; `get_result` may delete it after reading.
+
+    Several brokers may share one instance: `startup`/`shutdown` are refcounted and an injected
+    `client` is never closed. Expire old results with an S3 lifecycle rule on `prefix`; with
+    `delete_on_read=False` nothing here ever deletes them.
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        *,
+        prefix: str = "taskiq-results",
+        delete_on_read: bool = False,
+        declare: bool = False,
+        client: S3Client | None = None,
+        session: AioSession | None = None,
+        endpoint_url: str | None = None,
+        region_name: str | None = None,
+        aws_access_key_id: str | None = None,
+        aws_secret_access_key: str | None = None,
+        aws_session_token: str | None = None,
+        client_config: AioConfig | None = None,
+    ) -> None:
+        """Keep settings only; the S3 client is created in `startup` (no I/O here)."""
+        self.prefix = prefix
+        self.delete_on_read = delete_on_read
+        self.s3 = S3Connection(
+            bucket,
+            client=client,
+            session=session,
+            endpoint_url=endpoint_url,
+            region_name=region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
+            client_config=client_config,
+            declare=declare,
+        )
+
+    async def startup(self) -> None:
+        """Create the S3 client and check (or declare) the bucket."""
+        await self.s3.startup()
+
+    async def shutdown(self) -> None:
+        """Close the owned S3 client once its last user shuts down."""
+        await self.s3.shutdown()
+
+    async def set_result(self, task_id: str, result: TaskiqResult[Any]) -> None:
+        """Write the result as JSON (overwrites an earlier result of the same task)."""
+        await self.s3.client.put_object(
+            Bucket=self.s3.bucket,
+            Key=object_key(self.prefix, task_id),
+            Body=result.model_dump_json().encode("utf-8"),
+            ContentType="application/json",
+        )
+
+    async def is_result_ready(self, task_id: str) -> bool:
+        """True when a result object exists for `task_id` (one HEAD request)."""
+        try:
+            await self.s3.client.head_object(
+                Bucket=self.s3.bucket, Key=object_key(self.prefix, task_id)
+            )
+        except ClientError as exc:
+            if is_missing(exc):
+                return False
+            raise
+        return True
+
+    async def get_result(self, task_id: str, with_logs: bool = False) -> TaskiqResult[Any]:
+        """Read one result (then delete it when `delete_on_read`); `ResultGetError` if absent.
+
+        Read-then-delete is two requests, so two concurrent readers may both get the result.
+        """
+        key = object_key(self.prefix, task_id)
+        raw = await self.s3.get_bytes(key)
+        if raw is None:
+            raise ResultGetError
+        result = TaskiqResult[Any].model_validate_json(raw)
+        if self.delete_on_read:
+            await self.s3.client.delete_object(Bucket=self.s3.bucket, Key=key)
+        return result

taskiq_llm_queue/routing.py

@@ -0,0 +1,173 @@
+"""Canary router: picks `sqs` or `direct` per call from live, audited routing config (§15.1).
+
+`RoutingConfig` lives as JSON at `KeySpace.routing` and is reloaded every 30 s, so an operator
+can promote, roll back or kill the queue path with `llm-queue routing ...` and no restart.
+`LLMQueue` calls `CanaryRouter.route` on every request; it never does I/O.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import random
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.config import DEFAULT_ID_HASH_SALT, Lane
+from taskiq_llm_queue.envelope import keyed_digest
+from taskiq_llm_queue.metrics_names import CONFIG_INVALID_TOTAL
+from taskiq_llm_queue.reliability.config_store import LiveConfigStore
+
+if TYPE_CHECKING:
+    import redis.asyncio
+
+    from taskiq_llm_queue.keys import KeySpace
+    from taskiq_llm_queue.metrics import MetricsSink
+
+logger = get_logger(__name__)
+
+RoutePath = Literal["sqs", "direct"]
+Pct = Annotated[int, Field(ge=0, le=100)]
+DEFAULT_RELOAD_SECS = 30.0
+
+
+class RoutingConfig(BaseModel):
+    """Live canary config: kill switch, sqs percentage per model alias and lane, allowlist.
+
+    A model or lane that is not listed routes `direct` (0 %), so the empty default sends
+    everything direct. `allow_users` always go `sqs` unless the kill switch is on.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    kill_switch: bool = False
+    models: dict[str, dict[Lane, Pct]] = Field(default_factory=dict)
+    allow_users: frozenset[str] = frozenset()
+
+    def pct(self, model: str, lane: Lane) -> int:
+        """Share of calls (0-100) for `model` (alias or `alias@target`) and lane sent to sqs."""
+        return self.models.get(model.partition("@")[0], {}).get(lane, 0)
+
+    def with_pct(self, model: str, lane: Lane, pct: int) -> RoutingConfig:
+        """Copy with one model x lane percentage changed (promote / roll back)."""
+        models = {alias: dict(lanes) for alias, lanes in self.models.items()}
+        models.setdefault(model.partition("@")[0], {})[lane] = pct
+        return RoutingConfig.model_validate({**self.model_dump(), "models": models})
+
+    def with_kill(self, on: bool) -> RoutingConfig:
+        """Copy with the kill switch set (`on` routes every call `direct`)."""
+        return self.model_copy(update={"kill_switch": on})
+
+    def with_user(self, user_id: str, *, allowed: bool) -> RoutingConfig:
+        """Copy with `user_id` added to or removed from the allowlist."""
+        users = self.allow_users | {user_id} if allowed else self.allow_users - {user_id}
+        return self.model_copy(update={"allow_users": users})
+
+
+def bucket(key: str, salt: str | bytes = DEFAULT_ID_HASH_SALT) -> int:
+    """Stable bucket 0-99 for a session or user id; equal on every pod sharing `salt`."""
+    return int.from_bytes(keyed_digest(key, salt), "big") % 100
+
+
+def routing_store(
+    client: redis.asyncio.Redis, keys: KeySpace, default: RoutingConfig | None = None
+) -> LiveConfigStore[RoutingConfig]:
+    """Live store for routing at `KeySpace.routing`, audited to `KeySpace.audit`."""
+    return LiveConfigStore(
+        client, keys.routing, RoutingConfig, default or RoutingConfig(), audit_key=keys.audit
+    )
+
+
+class CanaryRouter:
+    """`PathRouter` over a `LiveConfigStore[RoutingConfig]`, sticky per session (fallback user).
+
+    `route` reads the last good config only. `start` runs the reload loop as a task; an invalid
+    stored config keeps the last good one and emits `routing_config_invalid` (log + metric).
+    """
+
+    def __init__(
+        self,
+        store: LiveConfigStore[RoutingConfig],
+        *,
+        metrics: MetricsSink | None = None,
+        reload_secs: float = DEFAULT_RELOAD_SECS,
+        rand: Callable[[], float] = random.random,
+        salt: str | bytes = DEFAULT_ID_HASH_SALT,
+    ) -> None:
+        """Keep the store and metrics sink; `rand` buckets calls that carry no id at all.
+
+        `salt` keys the sticky bucket hash (`QueueConfig.id_hash_salt`); all pods must share it.
+        """
+        self.store = store
+        self._salt = salt
+        self._metrics = metrics
+        self._reload_secs = reload_secs
+        self._rand = rand
+        self._seen_invalid = store.invalid_count
+        self._task: asyncio.Task[None] | None = None
+
+    @property
+    def config(self) -> RoutingConfig:
+        """The config `route` uses now (last good)."""
+        return self.store.current
+
+    def route(
+        self, model: str, lane: Lane, session_id: str | None, user_id: str | None
+    ) -> RoutePath:
+        """Decide the path of one call: kill switch, then allowlist, then sticky percentage."""
+        cfg = self.store.current
+        if cfg.kill_switch:
+            return "direct"
+        if user_id is not None and user_id in cfg.allow_users:
+            return "sqs"
+        pct = cfg.pct(model, lane)
+        if pct <= 0:
+            return "direct"
+        if pct >= 100:
+            return "sqs"
+        key = session_id or user_id
+        # No id at all: not sticky, but still honours the percentage on average.
+        point = bucket(key, self._salt) if key else int(self._rand() * 100)
+        return "sqs" if point < pct else "direct"
+
+    async def reload(self) -> RoutingConfig:
+        """Load the stored config once; report `routing_config_invalid` when it was rejected."""
+        cfg = await self.store.load()
+        invalid = self.store.invalid_count
+        if invalid > self._seen_invalid:
+            # The store warns once per bad value; this event is what the alarm (§13) keys on.
+            logger.error("routing_config_invalid", key=self.store.key, invalid_count=invalid)
+            if self._metrics is not None:
+                self._metrics.incr(
+                    CONFIG_INVALID_TOTAL, invalid - self._seen_invalid, {"kind": "routing"}
+                )
+            self._seen_invalid = invalid
+        return cfg
+
+    async def run(self) -> None:
+        """Reload every `reload_secs` until cancelled; a failed reload never ends the loop."""
+        while True:
+            try:
+                await self.reload()
+            except Exception:  # the loop must survive anything; route() keeps last good
+                logger.exception("routing_reload_failed")
+            await asyncio.sleep(self._reload_secs)
+
+    def start(self) -> asyncio.Task[None]:
+        """Start the reload loop on the running event loop (idempotent)."""
+        if self._task is None or self._task.done():
+            self._task = asyncio.get_running_loop().create_task(
+                self.run(), name="llmq-routing-reload"
+            )
+        return self._task
+
+    async def aclose(self) -> None:
+        """Stop the reload loop; safe to call twice."""
+        task, self._task = self._task, None
+        if task is not None:
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task

taskiq_llm_queue/routing_hooks.py

@@ -0,0 +1,90 @@
+"""Extension points of `LLMQueue`: path routing, SQS/Valkey failover and observability hooks.
+
+`PathRouter` (canary routing) and `FailoverPolicy` are implemented in `routing.py` /
+`failover.py`; `QueueHooks` is subclassed by the metrics layer. This module has no heavy
+imports, so implementations can depend on it freely.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
+
+from taskiq_llm_queue._logging import get_logger
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.config import Lane
+
+logger = get_logger(__name__)
+
+CallPath = Literal["sqs", "direct"]
+
+
+@runtime_checkable
+class PathRouter(Protocol):
+    """Chooses how one call runs: through its model's SQS queue or `direct` in this process.
+
+    Without a router `LLMQueue` uses `sqs` when the model has queues, else `direct`.
+    """
+
+    def route(
+        self, model: str, lane: Lane, session_id: str | None, user_id: str | None
+    ) -> CallPath:
+        """Path for one call of model alias `model` in `lane` (sticky per session)."""
+        ...
+
+
+@runtime_checkable
+class FailoverPolicy(Protocol):
+    """Sender-side failover state: SQS outages per model and a pod-wide Valkey outage.
+
+    `LLMQueue` reports failures here and asks `path_override` / `redis_down` before each call;
+    implementations must never block or raise.
+    """
+
+    def path_override(self, model: str) -> Literal["direct"] | None:
+        """`direct` while `model` is failed over from SQS, else None (use the router)."""
+        ...
+
+    def report_sqs_failure(self, model: str, exc: BaseException) -> None:
+        """An SQS send for `model` failed; may start a failover window."""
+        ...
+
+    def report_redis_failure(self, exc: BaseException) -> None:
+        """A Valkey call failed; may switch the pod to Valkey-down mode (DS18)."""
+        ...
+
+    def redis_down(self) -> bool:
+        """True while the pod treats Valkey as unreachable."""
+        ...
+
+
+class QueueHooks:
+    """No-op callbacks for client and worker events; the metrics layer overrides what it needs.
+
+    Every hook is called through `safe_hook`, so a failing hook never breaks a call. Labels are
+    low-cardinality only (model alias, target, lane, path, kind, step).
+    """
+
+    def on_call(self, model: str, lane: str, path: str, ok: bool, secs: float) -> None:
+        """One client call finished (`ok` False for errors and timeouts)."""
+
+    def on_attempt(self, model: str, target: str, kind: str, secs: float) -> None:
+        """One provider attempt finished; `kind` is `ok` or an `ErrorKind` value."""
+
+    def on_fallback(self, model: str, step: str) -> None:
+        """A call was served by fallback step `step` (`alias@target`) instead of the first."""
+
+    def on_queue_wait_ms(self, model: str, lane: str, wait_ms: int) -> None:
+        """A worker started a task `wait_ms` after it was enqueued."""
+
+    def on_sync_fallback(self, model: str) -> None:
+        """A sync call ran `direct` under the shared limiter (`llmq_sync_fallback_total`)."""
+
+
+def safe_hook(hook: Callable[..., Any], *args: Any) -> None:
+    """Run one hook, logging (never raising) whatever it throws."""
+    try:
+        hook(*args)
+    except Exception:
+        logger.exception("routing_hook_failed", hook=getattr(hook, "__name__", repr(hook)))