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

taskiq_llm_queue/__init__.py

@@ -0,0 +1,138 @@
+"""taskiq-llm-queue: an LLM request client with its own durable queue (taskiq + Amazon SQS).
+
+This is the public front door: `__all__` lists the client surface (the client and its
+results, config, request context, errors and extension points). Wire and plumbing names live
+in their submodules (`envelope`, `keys`, `protocols`, `errors`, `routing_hooks`). Importing
+the package has no side effects: no clients, threads or network I/O.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+from taskiq_llm_queue._version import __version__
+from taskiq_llm_queue.config import (
+    BreakerTuning,
+    Lane,
+    LimiterTuning,
+    ModelConfig,
+    QueueConfig,
+    Target,
+)
+from taskiq_llm_queue.context import RequestContext, current_context, request_context
+from taskiq_llm_queue.envelope import CallMeta, FallbackPolicy
+from taskiq_llm_queue.errors import (
+    ConfigError,
+    LLMQueueCancelled,
+    LLMQueueError,
+    LLMQueueTimeout,
+    ProviderError,
+    QueueUnavailable,
+    UnsupportedEnvelopeError,
+)
+from taskiq_llm_queue.routing_hooks import FailoverPolicy, PathRouter, QueueHooks
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.client import (
+        ChatResult,
+        ConverseResult,
+        EmbedResult,
+        InvokeResult,
+        LLMQueue,
+    )
+
+    # Plumbing kept importable for compatibility (typed for mypy; not in __all__).
+    from taskiq_llm_queue.envelope import ENVELOPE_VERSION as ENVELOPE_VERSION
+    from taskiq_llm_queue.envelope import TASK_CALL as TASK_CALL
+    from taskiq_llm_queue.envelope import TASK_EMBED as TASK_EMBED
+    from taskiq_llm_queue.envelope import TASK_JOB as TASK_JOB
+    from taskiq_llm_queue.envelope import Envelope as Envelope
+    from taskiq_llm_queue.envelope import LLMReply as LLMReply
+    from taskiq_llm_queue.envelope import parse_envelope as parse_envelope
+    from taskiq_llm_queue.errors import ErrorInfo as ErrorInfo
+    from taskiq_llm_queue.errors import ErrorKind as ErrorKind
+    from taskiq_llm_queue.keys import KeySpace as KeySpace
+    from taskiq_llm_queue.protocols import ClaimDecision as ClaimDecision
+    from taskiq_llm_queue.protocols import ClaimState as ClaimState
+    from taskiq_llm_queue.protocols import ClaimStore as ClaimStore
+    from taskiq_llm_queue.protocols import Lease as Lease
+    from taskiq_llm_queue.protocols import Limiter as Limiter
+    from taskiq_llm_queue.protocols import Provider as Provider
+    from taskiq_llm_queue.protocols import ReceiveGate as ReceiveGate
+    from taskiq_llm_queue.protocols import ResultStore as ResultStore
+    from taskiq_llm_queue.routing_hooks import CallPath as CallPath
+
+# Loaded on first access: the client pulls in taskiq, redis and aiobotocore, and importing the
+# package must stay free of them (library rule; see tests/unit/test_contracts.py).
+_CLIENT = "taskiq_llm_queue.client"
+_LAZY = {
+    "LLMQueue": _CLIENT,
+    "ConverseResult": _CLIENT,
+    "ChatResult": _CLIENT,
+    "InvokeResult": _CLIENT,
+    "EmbedResult": _CLIENT,
+}
+# Wire / plumbing names that used to be re-exported here: still importable from the package
+# for compatibility, but not public API; import them from their submodules.
+_PLUMBING = {
+    "ENVELOPE_VERSION": "taskiq_llm_queue.envelope",
+    "TASK_CALL": "taskiq_llm_queue.envelope",
+    "TASK_EMBED": "taskiq_llm_queue.envelope",
+    "TASK_JOB": "taskiq_llm_queue.envelope",
+    "Envelope": "taskiq_llm_queue.envelope",
+    "LLMReply": "taskiq_llm_queue.envelope",
+    "parse_envelope": "taskiq_llm_queue.envelope",
+    "KeySpace": "taskiq_llm_queue.keys",
+    "ErrorInfo": "taskiq_llm_queue.errors",
+    "ErrorKind": "taskiq_llm_queue.errors",
+    "CallPath": "taskiq_llm_queue.routing_hooks",
+    "ClaimDecision": "taskiq_llm_queue.protocols",
+    "ClaimState": "taskiq_llm_queue.protocols",
+    "ClaimStore": "taskiq_llm_queue.protocols",
+    "Lease": "taskiq_llm_queue.protocols",
+    "Limiter": "taskiq_llm_queue.protocols",
+    "Provider": "taskiq_llm_queue.protocols",
+    "ReceiveGate": "taskiq_llm_queue.protocols",
+    "ResultStore": "taskiq_llm_queue.protocols",
+}
+
+__all__ = [
+    "BreakerTuning",
+    "CallMeta",
+    "ChatResult",
+    "ConfigError",
+    "ConverseResult",
+    "EmbedResult",
+    "FailoverPolicy",
+    "FallbackPolicy",
+    "InvokeResult",
+    "LLMQueue",
+    "LLMQueueCancelled",
+    "LLMQueueError",
+    "LLMQueueTimeout",
+    "Lane",
+    "LimiterTuning",
+    "ModelConfig",
+    "PathRouter",
+    "ProviderError",
+    "QueueConfig",
+    "QueueHooks",
+    "QueueUnavailable",
+    "RequestContext",
+    "Target",
+    "UnsupportedEnvelopeError",
+    "__version__",
+    "current_context",
+    "request_context",
+]
+
+
+def __getattr__(name: str) -> Any:
+    """Import the client on first use of `LLMQueue` / a result type; plumbing names likewise."""
+    module = _LAZY.get(name) or _PLUMBING.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(module), name)
+    globals()[name] = value
+    return value

taskiq_llm_queue/_logging.py

@@ -0,0 +1,40 @@
+"""Structured loggers that always go through stdlib `logging` (no global structlog config).
+
+Unconfigured structlog prints to stdout and ignores the app's levels and handlers, so every
+module logs via `get_logger(__name__)`: structlog's key-value API on top of a stdlib logger.
+Apps keep full control (levels, handlers, loguru's InterceptHandler, structlog's
+ProcessorFormatter); the fields ride on `record.llmq` as a dict.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping, MutableMapping
+from typing import Any
+
+import structlog
+
+_PASSTHROUGH = ("exc_info", "stack_info")
+
+
+def _to_stdlib(
+    _logger: Any, _method: str, event_dict: MutableMapping[str, Any]
+) -> tuple[tuple[str], dict[str, Any]]:
+    # Message is readable on its own ("event k=v ..."); the structured copy lives on one
+    # `llmq` attribute so field names can never clash with LogRecord's reserved attributes.
+    event = str(event_dict.pop("event", ""))
+    kwargs: dict[str, Any] = {k: event_dict.pop(k) for k in _PASSTHROUGH if k in event_dict}
+    fields: Mapping[str, Any] = dict(event_dict)
+    message = " ".join([event, *(f"{k}={v}" for k, v in fields.items())])
+    kwargs["extra"] = {"llmq": {"event": event, **fields}}
+    kwargs["stacklevel"] = 4  # point %(funcName)s/%(lineno)d at the library call site
+    return (message,), kwargs
+
+
+def get_logger(name: str) -> structlog.stdlib.BoundLogger:
+    """Key-value logger bound to `logging.getLogger(name)`; used by every library module."""
+    return structlog.wrap_logger(  # type: ignore[no-any-return]
+        logging.getLogger(name),
+        processors=[_to_stdlib],
+        wrapper_class=structlog.stdlib.BoundLogger,
+    )

taskiq_llm_queue/_version.py

@@ -0,0 +1,5 @@
+"""Package version, read by hatchling at build time and re-exported as `__version__`."""
+
+from __future__ import annotations
+
+__version__ = "0.1.0a1"

taskiq_llm_queue/broker/__init__.py

@@ -0,0 +1,86 @@
+"""SQS broker for taskiq with priority, capacity-aware receive and batching (see docs/DESIGN.md).
+
+`HybridSQSBroker` plus its settings (`SQSQueue`, `Batching`, `Heartbeat`), event hooks and
+the SQS error-family classifier. Importing this package creates no client and does no I/O.
+"""
+
+from __future__ import annotations
+
+from taskiq_llm_queue.broker.acker import VisibilityResult
+from taskiq_llm_queue.broker.broker import (
+    DEDUPLICATION_ID_LABEL,
+    DELAY_LABEL,
+    EXPIRY_ATTRIBUTE,
+    EXPIRY_LABEL,
+    GROUP_ID_LABEL,
+    QUEUE_LABEL,
+    SKIP_BATCHING_LABEL,
+    SQS_MESSAGE_ID_LABEL,
+    SQS_QUEUE_LABEL,
+    SQS_RECEIPT_HANDLE_LABEL,
+    SQS_RECEIVE_COUNT_LABEL,
+    DeliveryFormatter,
+    HybridSQSBroker,
+    ReceiveBudget,
+)
+from taskiq_llm_queue.broker.delivery import Delivery
+from taskiq_llm_queue.broker.errors import (
+    BrokerConfigError,
+    ErrorFamily,
+    InvalidLabelError,
+    MessageTooLargeError,
+    SQSSendError,
+    backoff_delay,
+    classify_code,
+    classify_error,
+)
+from taskiq_llm_queue.broker.events import BrokerHooks
+from taskiq_llm_queue.broker.filters import MessageFilter, envelope_supported
+from taskiq_llm_queue.broker.queue import Batching, Heartbeat, SQSQueue
+from taskiq_llm_queue.broker.s3_offload import (
+    OFFLOAD_LABEL,
+    OffloadedPayloadMissingError,
+    S3OffloadError,
+    S3OffloadMiddleware,
+)
+from taskiq_llm_queue.broker.scheduler import SlotScheduler
+from taskiq_llm_queue.broker.serializer import WireSerializer
+
+__all__ = [
+    "DEDUPLICATION_ID_LABEL",
+    "DELAY_LABEL",
+    "EXPIRY_ATTRIBUTE",
+    "EXPIRY_LABEL",
+    "GROUP_ID_LABEL",
+    "OFFLOAD_LABEL",
+    "QUEUE_LABEL",
+    "SKIP_BATCHING_LABEL",
+    "SQS_MESSAGE_ID_LABEL",
+    "SQS_QUEUE_LABEL",
+    "SQS_RECEIPT_HANDLE_LABEL",
+    "SQS_RECEIVE_COUNT_LABEL",
+    "Batching",
+    "BrokerConfigError",
+    "BrokerHooks",
+    "Delivery",
+    "DeliveryFormatter",
+    "ErrorFamily",
+    "Heartbeat",
+    "HybridSQSBroker",
+    "InvalidLabelError",
+    "MessageFilter",
+    "MessageTooLargeError",
+    "OffloadedPayloadMissingError",
+    "ReceiveBudget",
+    "S3OffloadError",
+    "S3OffloadMiddleware",
+    "SQSQueue",
+    "SQSSendError",
+    "SlotScheduler",
+    "VisibilityResult",
+    "WireSerializer",
+    "backoff_delay",
+    "classify_code",
+    "classify_error",
+    "envelope_supported",
+]

taskiq_llm_queue/broker/acker.py

@@ -0,0 +1,239 @@
+"""Batched acks (`DeleteMessageBatch`) and visibility changes (`ChangeMessageVisibilityBatch`).
+
+At 10K messages in flight one call per message is too many requests, so acks, heartbeats
+and releases go through per-queue micro-batches (A6). Failed batch entries are retried one
+by one; an invalid receipt handle on delete is fine (the message is already gone or owned
+elsewhere and its claim is `done`, DS7), while on a visibility change it means "lost" (DS6).
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Sequence
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any
+
+from taskiq_llm_queue._logging import get_logger
+from taskiq_llm_queue.broker.batching import MicroBatcher
+from taskiq_llm_queue.broker.errors import (
+    ErrorFamily,
+    backoff_delay,
+    classify_code,
+    classify_error,
+    error_code,
+)
+from taskiq_llm_queue.broker.events import BrokerHooks, safe_call
+from taskiq_llm_queue.broker.queue import MAX_BATCH_ENTRIES
+
+if TYPE_CHECKING:
+    from types_aiobotocore_sqs.client import SQSClient
+
+logger = get_logger(__name__)
+
+
+class VisibilityResult(StrEnum):
+    """Outcome of one visibility change: done, handle invalid (lost), or failed (retry later)."""
+
+    OK = "ok"
+    LOST = "lost"
+    FAILED = "failed"
+
+
+class Acker:
+    """Owns the delete and change-visibility batchers for every queue of one broker.
+
+    `delete` and `change_visibility` return futures immediately; `flush` and `close` push
+    out whatever is queued (used on shutdown and drain, A5).
+    """
+
+    def __init__(
+        self,
+        client: SQSClient,
+        *,
+        hooks: BrokerHooks,
+        ack_flush_secs: float = 0.5,
+        visibility_flush_secs: float = 0.5,
+        retry_attempts: int = 3,
+        backoff_base_secs: float = 0.2,
+        backoff_cap_secs: float = 2.0,
+    ) -> None:
+        """Create no batchers yet; one pair is made per queue on first use.
+
+        `retry_attempts` is the only throttle/5xx retry layer: the broker's SQS client makes at
+        most one quick botocore retry (standard mode), so one call costs <= 2 x this many.
+        """
+        self._client = client
+        self._hooks = hooks
+        self._ack_flush_secs = ack_flush_secs
+        self._visibility_flush_secs = visibility_flush_secs
+        self._retry_attempts = retry_attempts
+        self._backoff_base_secs = backoff_base_secs
+        self._backoff_cap_secs = backoff_cap_secs
+        self._deleters: dict[str, MicroBatcher[str, bool]] = {}
+        self._changers: dict[str, MicroBatcher[tuple[str, int], VisibilityResult]] = {}
+
+    def delete(self, queue: str, url: str, receipt_handle: str) -> asyncio.Future[bool]:
+        """Queue one ack; resolves True when deleted or already gone, False on failure."""
+        batcher = self._deleters.get(url)
+        if batcher is None:
+            batcher = MicroBatcher[str, bool](
+                lambda items: self._delete_batch(queue, url, items),
+                max_items=MAX_BATCH_ENTRIES,
+                max_delay_secs=self._ack_flush_secs,
+                name=f"ack:{queue}",
+            )
+            self._deleters[url] = batcher
+        return batcher.submit(receipt_handle)
+
+    def change_visibility(
+        self, queue: str, url: str, receipt_handle: str, visibility_secs: int
+    ) -> asyncio.Future[VisibilityResult]:
+        """Queue one visibility change (heartbeat extend, or 0 / N for a release)."""
+        batcher = self._changers.get(url)
+        if batcher is None:
+            batcher = MicroBatcher[tuple[str, int], VisibilityResult](
+                lambda items: self._visibility_batch(queue, url, items),
+                max_items=MAX_BATCH_ENTRIES,
+                max_delay_secs=self._visibility_flush_secs,
+                name=f"visibility:{queue}",
+            )
+            self._changers[url] = batcher
+        return batcher.submit((receipt_handle, visibility_secs))
+
+    async def flush(self) -> None:
+        """Send every queued ack and visibility change now and wait for the results."""
+        batchers: list[MicroBatcher[Any, Any]] = [
+            *self._changers.values(),
+            *self._deleters.values(),
+        ]
+        await asyncio.gather(*(b.flush_now() for b in batchers))
+
+    async def close(self) -> None:
+        """Flush and stop all batchers; idempotent."""
+        batchers: list[MicroBatcher[Any, Any]] = [
+            *self._changers.values(),
+            *self._deleters.values(),
+        ]
+        await asyncio.gather(*(b.close() for b in batchers))
+        self._deleters.clear()
+        self._changers.clear()
+
+    def _batchers(self) -> list[MicroBatcher[Any, Any]]:
+        return [*self._changers.values(), *self._deleters.values()]
+
+    async def _delete_batch(self, queue: str, url: str, handles: Sequence[str]) -> list[bool]:
+        entries: list[Any] = [{"Id": str(i), "ReceiptHandle": h} for i, h in enumerate(handles)]
+        resp = await self.call(
+            queue,
+            "delete_batch",
+            lambda: self._client.delete_message_batch(QueueUrl=url, Entries=entries),
+        )
+        if resp is None:
+            return [False] * len(handles)
+        results = [True] * len(handles)
+        for failed in resp.get("Failed", []):
+            index = int(failed["Id"])
+            family = classify_code(failed.get("Code", ""), failed.get("Message", ""))
+            if family is ErrorFamily.INVALID_HANDLE:
+                logger.debug("sqs_ack_handle_invalid", queue=queue)
+                continue
+            results[index] = await self._delete_one(queue, url, handles[index])
+        acked = sum(results)
+        if acked:
+            safe_call(self._hooks.on_ack, queue, acked)
+        return results
+
+    async def _delete_one(self, queue: str, url: str, handle: str) -> bool:
+        try:
+            await self.call(
+                queue,
+                "delete",
+                lambda: self._client.delete_message(QueueUrl=url, ReceiptHandle=handle),
+                raise_final=True,
+            )
+        except Exception as exc:
+            if classify_error(exc) is ErrorFamily.INVALID_HANDLE:
+                return True
+            logger.warning("sqs_ack_failed", queue=queue, error=str(exc))
+            return False
+        return True
+
+    async def _visibility_batch(
+        self, queue: str, url: str, items: Sequence[tuple[str, int]]
+    ) -> list[VisibilityResult]:
+        entries: list[Any] = [
+            {"Id": str(i), "ReceiptHandle": h, "VisibilityTimeout": secs}
+            for i, (h, secs) in enumerate(items)
+        ]
+        resp = await self.call(
+            queue,
+            "visibility_batch",
+            lambda: self._client.change_message_visibility_batch(QueueUrl=url, Entries=entries),
+        )
+        if resp is None:
+            return [VisibilityResult.FAILED] * len(items)
+        results = [VisibilityResult.OK] * len(items)
+        for failed in resp.get("Failed", []):
+            index = int(failed["Id"])
+            family = classify_code(failed.get("Code", ""), failed.get("Message", ""))
+            if family is ErrorFamily.INVALID_HANDLE:
+                results[index] = VisibilityResult.LOST
+            else:
+                results[index] = await self._visibility_one(queue, url, *items[index])
+        return results
+
+    async def _visibility_one(
+        self, queue: str, url: str, handle: str, secs: int
+    ) -> VisibilityResult:
+        try:
+            await self.call(
+                queue,
+                "visibility",
+                lambda: self._client.change_message_visibility(
+                    QueueUrl=url, ReceiptHandle=handle, VisibilityTimeout=secs
+                ),
+                raise_final=True,
+            )
+        except Exception as exc:
+            if classify_error(exc) is ErrorFamily.INVALID_HANDLE:
+                return VisibilityResult.LOST
+            return VisibilityResult.FAILED
+        return VisibilityResult.OK
+
+    async def call(
+        self,
+        queue: str,
+        operation: str,
+        call: Callable[[], Awaitable[Any]],
+        *,
+        raise_final: bool = False,
+    ) -> Any:
+        """Run one SQS call, retrying throttles and 5xx with jittered backoff (DS16)."""
+        attempt = 0
+        while True:
+            try:
+                return await call()
+            except Exception as exc:
+                family = classify_error(exc)
+                safe_call(self._hooks.on_error, queue, operation, family.value)
+                if family.retryable and attempt + 1 < self._retry_attempts:
+                    await asyncio.sleep(
+                        backoff_delay(
+                            attempt,
+                            base_secs=self._backoff_base_secs,
+                            cap_secs=self._backoff_cap_secs,
+                        )
+                    )
+                    attempt += 1
+                    continue
+                if raise_final:
+                    raise
+                logger.warning(
+                    "sqs_call_failed",
+                    operation=operation,
+                    queue=queue,
+                    error_code=error_code(exc),
+                    attempt=attempt + 1,
+                    error=str(exc),
+                )
+                return None