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

taskiq_llm_queue/metrics/prometheus.py

@@ -0,0 +1,66 @@
+"""Optional Prometheus `MetricsSink` for local runs and the Docker stack (`[metrics]` extra).
+
+`prometheus_client` is imported when the sink is built, never at package import, and a missing
+install raises an `ImportError` naming the extra. Label names of a metric are fixed by its first
+use; later calls fill missing labels with "" and drop unknown ones.
+"""
+
+from __future__ import annotations
+
+import importlib
+import threading
+from collections.abc import Mapping
+from typing import Any
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.metrics_names import LABELS
+
+logger = get_logger(__name__)
+INSTALL_HINT = "PrometheusMetrics needs prometheus-client: pip install 'taskiq-llm-queue[metrics]'"
+
+
+class PrometheusMetrics:
+    """`MetricsSink` backed by prometheus_client Counter / Histogram / Gauge objects.
+
+    Pass a `registry` for isolation in tests; the default is prometheus_client's global one.
+    `serve(port)` starts the client's HTTP exporter for `/metrics`.
+    """
+
+    def __init__(self, registry: Any = None) -> None:
+        """Import prometheus_client now; raises `ImportError` with the extra to install."""
+        try:
+            self._prom: Any = importlib.import_module("prometheus_client")
+        except ImportError as exc:
+            raise ImportError(INSTALL_HINT) from exc
+        self._registry = registry if registry is not None else self._prom.REGISTRY
+        self._lock = threading.Lock()
+        self._metrics: dict[str, tuple[Any, tuple[str, ...]]] = {}
+
+    def incr(self, name: str, value: float = 1.0, labels: Mapping[str, str] | None = None) -> None:
+        """Increase a Counter."""
+        self._child("Counter", name, labels).inc(value)
+
+    def observe(self, name: str, value: float, labels: Mapping[str, str] | None = None) -> None:
+        """Observe into a Histogram (default buckets)."""
+        self._child("Histogram", name, labels).observe(value)
+
+    def gauge(self, name: str, value: float, labels: Mapping[str, str] | None = None) -> None:
+        """Set a Gauge."""
+        self._child("Gauge", name, labels).set(value)
+
+    def serve(self, port: int, addr: str = "0.0.0.0") -> None:
+        """Expose `/metrics` on `addr:port` from a daemon thread."""
+        self._prom.start_http_server(port, addr=addr, registry=self._registry)
+
+    def _child(self, kind: str, name: str, labels: Mapping[str, str] | None) -> Any:
+        given = {k: str(v) for k, v in (labels or {}).items() if k in LABELS}
+        with self._lock:
+            found = self._metrics.get(name)
+            if found is None:
+                names = tuple(sorted(given))
+                metric = getattr(self._prom, kind)(
+                    name, name, labelnames=names, registry=self._registry
+                )
+                found = self._metrics[name] = (metric, names)
+        metric, names = found
+        return metric.labels(**{n: given.get(n, "") for n in names}) if names else metric

taskiq_llm_queue/metrics_names.py

@@ -0,0 +1,19 @@
+"""Metric names and label sets, fixed here so dashboards and alarms don't drift from code."""
+
+from __future__ import annotations
+
+QUEUE_WAIT_SECONDS = "llmq_queue_wait_seconds"
+PROVIDER_SECONDS = "llmq_provider_seconds"
+CALLS_TOTAL = "llmq_calls_total"
+THROTTLES_TOTAL = "llmq_throttles_total"
+FALLBACK_TOTAL = "llmq_fallback_total"
+LIMIT = "llmq_limit"
+INFLIGHT = "llmq_inflight"
+BREAKER_OPEN = "llmq_breaker_open"
+SYNC_FALLBACK_TOTAL = "llmq_sync_fallback_total"
+RELEASED_UNKNOWN_TOTAL = "llmq_released_unknown_total"
+CONFIG_INVALID_TOTAL = "llmq_config_invalid_total"
+FAILOVER_ACTIVE = "llmq_failover_active"
+
+# Allowed labels; raw user/session ids are never metric labels (cardinality, privacy).
+LABELS = ("model", "target", "lane", "queue", "path", "kind", "step")

taskiq_llm_queue/protocols.py

@@ -0,0 +1,140 @@
+"""Interfaces between the layers: providers, limiter, claims, results and receive gates.
+
+Each workstream implements one of these and codes against the others, so pieces can be built
+and tested in parallel with in-memory fakes before the real implementations land.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.config import Lane, Target
+    from taskiq_llm_queue.envelope import LLMReply, Op
+    from taskiq_llm_queue.errors import ErrorInfo
+
+
+class ClaimState(StrEnum):
+    """Lifecycle of one task's claim record in Redis."""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    DONE = "done"
+    CANCELLED = "cancelled"
+    EXPIRED = "expired"
+    RELEASED = "released"
+
+
+class ClaimDecision(StrEnum):
+    """What a worker does with a received message after trying to claim it."""
+
+    RUN = "run"
+    DEFER = "defer"  # another pod is running it: extend visibility, never delete
+    ACK_DONE = "ack_done"
+    ACK_CANCELLED = "ack_cancelled"
+    ACK_EXPIRED = "ack_expired"
+
+
+@dataclass(frozen=True, slots=True)
+class Lease:
+    """One in-flight slot held for one provider attempt; released in `finally`."""
+
+    model: str
+    target: str
+    token: str
+    expires_at: float
+
+
+@runtime_checkable
+class Provider(Protocol):
+    """An LLM backend: performs one attempt and classifies its failures."""
+
+    name: str
+
+    async def call(self, op: Op, params: Mapping[str, Any], target: Target) -> dict[str, Any]:
+        """Run one attempt and return the provider-native response body."""
+        ...
+
+    def classify(self, exc: BaseException) -> ErrorInfo:
+        """Map a provider exception to throttle / retryable / fatal."""
+        ...
+
+
+@runtime_checkable
+class ReceiveGate(Protocol):
+    """Lets the broker stop receiving for a queue, e.g. while its model is throttled."""
+
+    def is_open(self) -> bool:
+        """True when the broker may receive now."""
+        ...
+
+    async def wait_open(self) -> None:
+        """Return once receiving is allowed again."""
+        ...
+
+
+@runtime_checkable
+class Limiter(Protocol):
+    """Cluster-wide adaptive concurrency limit per model target, learned from throttles."""
+
+    async def acquire(
+        self,
+        model: str,
+        target: str,
+        lane: Lane,
+        timeout_secs: float,
+        *,
+        should_abort: Callable[[], Awaitable[bool]] | None = None,
+    ) -> Lease | None:
+        """Take a slot for one attempt, or None if none frees up within the timeout.
+
+        `should_abort` (e.g. the cancel flag, DS14) is checked while waiting; True gives up.
+        """
+        ...
+
+    async def release(self, lease: Lease, *, throttled: bool) -> None:
+        """Return the slot and feed the outcome into the limit (grow or shrink)."""
+        ...
+
+    def gate(self, model: str, target: str) -> ReceiveGate:
+        """Gate that closes while this model target is paused after a throttle."""
+        ...
+
+
+@runtime_checkable
+class ClaimStore(Protocol):
+    """Exactly-once execution on top of at-least-once delivery."""
+
+    async def claim(self, request_id: str, *, deadline_at: float, worker: str) -> ClaimDecision:
+        """Atomically decide whether this worker runs the task (Lua CAS)."""
+        ...
+
+    async def heartbeat(self, request_ids: Sequence[str], *, worker: str) -> None:
+        """Refresh heartbeats and TTLs for running tasks."""
+        ...
+
+    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 (owner `worker`)."""
+        ...
+
+    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 (owner `worker`)."""
+        ...
+
+
+@runtime_checkable
+class ResultStore(Protocol):
+    """Stores task results and hands them to waiting callers in batches."""
+
+    async def put(self, reply: LLMReply, *, ttl_secs: int) -> None:
+        """Store one result (JSON) with a TTL."""
+        ...
+
+    async def take_many(self, task_ids: Sequence[str]) -> dict[str, LLMReply]:
+        """Fetch and delete whichever of these results are ready."""
+        ...

taskiq_llm_queue/providers/__init__.py

@@ -0,0 +1,62 @@
+"""LLM providers: Amazon Bedrock and OpenAI-compatible servers behind the `Provider` protocol.
+
+Workers call `ProviderRegistry.for_model(alias).call(op, params, target)` and classify
+failures with `.classify(exc)`. Importing this package never imports aiobotocore or openai;
+each SDK is imported on a provider's first call.
+"""
+
+from __future__ import annotations
+
+from taskiq_llm_queue.providers.bedrock import BedrockProvider, decode_invoke_body
+from taskiq_llm_queue.providers.classify import (
+    classify_bedrock,
+    classify_openai,
+    error_family,
+    parse_retry_after,
+)
+from taskiq_llm_queue.providers.embeddings import (
+    COHERE_MAX_BATCH,
+    OPENAI_DEFAULT_CHUNK_SIZE,
+    TITAN_MAX_BATCH,
+    batches,
+    cohere_body,
+    embedding_body,
+    max_batch_size,
+    parse_embeddings,
+    titan_body,
+)
+from taskiq_llm_queue.providers.hosts import check_url, denied_pattern
+from taskiq_llm_queue.providers.openai_compat import OpenAICompatProvider
+from taskiq_llm_queue.providers.registry import (
+    BUILTIN_PROVIDERS,
+    ENTRY_POINT_GROUP,
+    ProviderFactory,
+    ProviderRegistry,
+    load_object,
+)
+
+__all__ = [
+    "BUILTIN_PROVIDERS",
+    "COHERE_MAX_BATCH",
+    "ENTRY_POINT_GROUP",
+    "OPENAI_DEFAULT_CHUNK_SIZE",
+    "TITAN_MAX_BATCH",
+    "BedrockProvider",
+    "OpenAICompatProvider",
+    "ProviderFactory",
+    "ProviderRegistry",
+    "batches",
+    "check_url",
+    "classify_bedrock",
+    "classify_openai",
+    "cohere_body",
+    "decode_invoke_body",
+    "denied_pattern",
+    "embedding_body",
+    "error_family",
+    "load_object",
+    "max_batch_size",
+    "parse_embeddings",
+    "parse_retry_after",
+    "titan_body",
+]

taskiq_llm_queue/providers/bedrock.py

@@ -0,0 +1,188 @@
+"""Amazon Bedrock runtime provider: Converse and InvokeModel through aiobotocore.
+
+Payloads are provider-native and pass through unchanged, so responses are exactly what
+`ChatBedrockConverse` / `BedrockEmbeddings` parse. SDK retries are off (one attempt per call)
+so throttles reach the adaptive limiter instead of being hidden by botocore.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+from collections.abc import Iterable, Mapping
+from contextlib import AsyncExitStack
+from typing import TYPE_CHECKING, Any
+
+import orjson
+
+from taskiq_llm_queue.errors import ConfigError, ErrorInfo
+from taskiq_llm_queue.providers.classify import classify_bedrock
+from taskiq_llm_queue.providers.hosts import check_url
+
+if TYPE_CHECKING:
+    from aiobotocore.session import AioSession
+    from types_aiobotocore_bedrock_runtime import BedrockRuntimeClient
+
+    from taskiq_llm_queue.config import QueueConfig, Target
+    from taskiq_llm_queue.envelope import Op
+
+# Unverified starting values: Bedrock replies can run minutes, and the pool must not be the
+# bottleneck below the limiter (botocore's default pool is 10 connections per client).
+DEFAULT_READ_TIMEOUT_SECS = 300.0
+DEFAULT_CONNECT_TIMEOUT_SECS = 10.0
+DEFAULT_MAX_POOL_CONNECTIONS = 100
+
+
+class BedrockProvider:
+    """`Provider` for the Bedrock runtime: ops `converse` and `invoke_model`.
+
+    One aiobotocore client per (region, endpoint), created on first use and closed by
+    `aclose()`. The model id is `Target.profile_id` when set, else `params["modelId"]`.
+    """
+
+    name = "bedrock"
+
+    def __init__(
+        self,
+        *,
+        session: AioSession | None = None,
+        region: str | None = None,
+        endpoint_url: str | None = None,
+        denied_hosts: Iterable[str] = (),
+        read_timeout_secs: float = DEFAULT_READ_TIMEOUT_SECS,
+        connect_timeout_secs: float = DEFAULT_CONNECT_TIMEOUT_SECS,
+        max_pool_connections: int = DEFAULT_MAX_POOL_CONNECTIONS,
+        client_kwargs: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Keep settings only; no SDK import and no client until the first call.
+
+        `region=None` (and no `Target.region`) lets the AWS chain pick the region.
+        `client_kwargs` go to `create_client` (e.g. explicit credentials for a fake endpoint).
+        """
+        self._denied_hosts = tuple(denied_hosts)
+        check_url(endpoint_url, self._denied_hosts)
+        self._session = session
+        self._region = region
+        self._endpoint_url = endpoint_url
+        self._config_kwargs: dict[str, Any] = {
+            "read_timeout": read_timeout_secs,
+            "connect_timeout": connect_timeout_secs,
+            "max_pool_connections": max_pool_connections,
+            # `max_attempts: 1` would mean one *retry* (botocore adds 1); total is exact.
+            "retries": {"total_max_attempts": 1, "mode": "standard"},
+        }
+        self._client_kwargs = dict(client_kwargs or {})
+        self._clients: dict[tuple[str | None, str | None], BedrockRuntimeClient] = {}
+        self._stack = AsyncExitStack()
+        self._lock = asyncio.Lock()
+
+    @classmethod
+    def from_config(
+        cls, config: QueueConfig, *, session: AioSession | None = None, **kwargs: Any
+    ) -> BedrockProvider:
+        """Build from `QueueConfig` (region, endpoint, deny list); checks every Bedrock target URL.
+
+        `session` is the pod's shared `AioSession` (`LLMQueue.start` passes it); extra keyword
+        arguments go to the constructor (e.g. `client_kwargs`) and win over the config's
+        `provider_read_timeout_secs` / `provider_connect_timeout_secs`.
+        """
+        for model in config.models.values():
+            if model.provider == "bedrock":
+                for target in model.targets.values():
+                    check_url(target.base_url, config.denied_hosts)
+        return cls(
+            region=config.aws_region,
+            endpoint_url=config.bedrock_endpoint_url,
+            denied_hosts=config.denied_hosts,
+            session=session,
+            **{
+                "read_timeout_secs": config.provider_read_timeout_secs,
+                "connect_timeout_secs": config.provider_connect_timeout_secs,
+                **kwargs,
+            },
+        )
+
+    async def call(self, op: Op, params: Mapping[str, Any], target: Target) -> dict[str, Any]:
+        """Run one Converse or InvokeModel attempt and return a JSON-safe response dict.
+
+        `converse` returns the Converse body without `ResponseMetadata`; `invoke_model` returns
+        `{"body_b64", "contentType"}` (see `decode_invoke_body`).
+        """
+        request = dict(params)
+        model_id = target.profile_id or request.get("modelId")
+        if not model_id:
+            raise ValueError("Bedrock call needs Target.profile_id or params['modelId']")
+        request["modelId"] = model_id
+        if op == "converse":
+            client = await self._client(target)
+            response: dict[str, Any] = dict(await client.converse(**request))
+            response.pop("ResponseMetadata", None)
+            return response
+        if op == "invoke_model":
+            body = request.get("body")
+            if isinstance(body, Mapping | list):
+                request["body"] = orjson.dumps(body)
+            request.setdefault("contentType", "application/json")
+            request.setdefault("accept", "application/json")
+            client = await self._client(target)
+            raw = await client.invoke_model(**request)
+            async with raw["body"] as stream:
+                data = await stream.read()
+            return {
+                "body_b64": base64.b64encode(data).decode("ascii"),
+                "contentType": raw.get("contentType", ""),
+            }
+        raise ValueError(f"Bedrock provider does not support op {op!r}")
+
+    def classify(self, exc: BaseException) -> ErrorInfo:
+        """Throttle / retryable / fatal by Bedrock error family (see `classify_bedrock`)."""
+        return classify_bedrock(exc)
+
+    async def aclose(self) -> None:
+        """Close every client this provider created; safe to call more than once."""
+        async with self._lock:
+            self._clients.clear()
+            stack, self._stack = self._stack, AsyncExitStack()
+            await stack.aclose()
+
+    async def _client(self, target: Target) -> BedrockRuntimeClient:
+        """Return the cached client for this target's region and endpoint, creating it once."""
+        endpoint = target.base_url or self._endpoint_url
+        region = target.region or self._region
+        key = (region, endpoint)
+        if (client := self._clients.get(key)) is not None:
+            return client
+        check_url(endpoint, self._denied_hosts)
+        async with self._lock:
+            if (client := self._clients.get(key)) is not None:
+                return client
+            try:
+                from aiobotocore.config import AioConfig
+                from aiobotocore.session import get_session
+            except ImportError as exc:  # pragma: no cover - aiobotocore is a core dependency
+                raise ConfigError(
+                    "BedrockProvider needs aiobotocore: pip install 'taskiq-llm-queue'"
+                ) from exc
+            if self._session is None:
+                self._session = get_session()
+            context = self._session.create_client(
+                "bedrock-runtime",
+                region_name=region,
+                endpoint_url=endpoint,
+                config=AioConfig(**self._config_kwargs),
+                **self._client_kwargs,
+            )
+            client = await self._stack.enter_async_context(context)
+            self._clients[key] = client
+            return client
+
+
+def decode_invoke_body(response: Mapping[str, Any]) -> Any:
+    """Decode an `invoke_model` result: parsed JSON for JSON content types, else raw bytes.
+
+    Used by embedding callers to turn `{"body_b64", "contentType"}` back into the model body.
+    """
+    data = base64.b64decode(response["body_b64"])
+    if "json" in str(response.get("contentType", "")).lower():
+        return orjson.loads(data)
+    return data