fleet-framework 0.1.0__py3-none-any.whl

fleet/core/primitives.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, AsyncIterator, Generic, Optional, TypeVar
+
+from pydantic import BaseModel, ValidationError
+
+from fleet.core.backend import Backend, RawItem
+from fleet.core.contract import Pool, Queue, Stream, validate_tag_filter
+
+logger = logging.getLogger(__name__)
+
+P = TypeVar("P", bound=BaseModel)
+T = TypeVar("T", bound=BaseModel)
+
+
+
+@dataclass
+class PoolItem(Generic[P, T]):
+    id: str
+    payload: P
+    tags: T
+    expires_at: Optional[float]
+    _handle: "PoolHandle[P, T]"
+
+    async def release(self) -> None:
+        await self._handle._backend.pool_release(self._handle._spec.name, self.id)
+
+    async def evict(self) -> None:
+        await self._handle._backend.pool_remove(self._handle._spec.name, self.id)
+
+    async def mark_bad(self, reason: str) -> None:
+        """Evict and log. Plugins can override later via a hook for stats."""
+        logger.info("pool %s item %s marked bad: %s", self._handle._spec.name, self.id, reason)
+        await self.evict()
+
+
+class PoolFilter(Generic[P, T]):
+    def __init__(self, handle: "PoolHandle[P, T]", conditions: dict[str, Any]) -> None:
+        self._handle = handle
+        self._conditions = conditions
+
+    def where(self, **kwargs: Any) -> "PoolFilter[P, T]":
+        # validate kwargs against tag schema; chainable.
+        validated = validate_tag_filter(self._handle._spec.tags, **kwargs)
+        return PoolFilter(self._handle, {**self._conditions, **validated})
+
+    async def list(self) -> list[PoolItem[P, T]]:
+        raws = await self._handle._backend.pool_list(self._handle._spec.name)
+        out: list[PoolItem[P, T]] = []
+        for r in raws:
+            if r.in_use_until is not None:
+                continue
+            if not all(r.tags.get(k) == v for k, v in self._conditions.items()):
+                continue
+            item = self._handle._wrap(r)
+            if item is not None:
+                out.append(item)
+        return out
+
+    async def size(self) -> int:
+        return len(await self.list())
+
+    @asynccontextmanager
+    async def claim(self, *, hold_seconds: int = 30) -> AsyncIterator[Optional[PoolItem[P, T]]]:
+        raw = await self._handle._backend.pool_claim_any(
+            self._handle._spec.name, self._conditions, hold_seconds,
+        )
+        if raw is None:
+            yield None
+            return
+        item = self._handle._wrap(raw)
+        if item is None:
+            # decode failure — release the claim, treat as miss
+            await self._handle._backend.pool_release(self._handle._spec.name, raw.id)
+            yield None
+            return
+        try:
+            yield item
+        finally:
+            # if the user evicted, the claim is gone too. release is idempotent.
+            try:
+                await self._handle._backend.pool_release(self._handle._spec.name, raw.id)
+            except Exception:
+                logger.debug("pool release failed", exc_info=True)
+
+
+class PoolHandle(Generic[P, T]):
+    def __init__(self, spec: Pool[P, T], backend: Backend) -> None:
+        self._spec = spec
+        self._backend = backend
+
+    @property
+    def name(self) -> str:
+        return self._spec.name
+
+    def where(self, **kwargs: Any) -> PoolFilter[P, T]:
+        return PoolFilter(self, {}).where(**kwargs)
+
+    async def put(
+        self,
+        *,
+        payload: P,
+        tags: T,
+        ttl_seconds: Optional[int] = None,
+        id: Optional[str] = None,
+    ) -> str:
+        # local pydantic validation (re-construct enforces type) plus tags type check.
+        if not isinstance(payload, self._spec.payload):
+            raise TypeError(
+                f"pool {self._spec.name}: payload must be {self._spec.payload.__name__}, "
+                f"got {type(payload).__name__}"
+            )
+        if not isinstance(tags, self._spec.tags):
+            raise TypeError(
+                f"pool {self._spec.name}: tags must be {self._spec.tags.__name__}, "
+                f"got {type(tags).__name__}"
+            )
+        item_id = id or uuid.uuid4().hex
+        await self._backend.pool_put(
+            self._spec.name,
+            item_id,
+            payload.model_dump_json().encode("utf-8"),
+            tags.model_dump(mode="json"),
+            ttl_seconds,
+        )
+        return item_id
+
+    async def list(self) -> list[PoolItem[P, T]]:
+        return await PoolFilter(self, {}).list()
+
+    async def size(self) -> int:
+        return await PoolFilter(self, {}).size()
+
+    def _wrap(self, raw: RawItem) -> Optional[PoolItem[P, T]]:
+        try:
+            payload = self._spec.payload.model_validate_json(raw.payload)
+            tags = self._spec.tags.model_validate(raw.tags)
+        except ValidationError:
+            logger.warning("pool %s item %s failed schema validation; dropping", self._spec.name, raw.id)
+            return None
+        return PoolItem(
+            id=raw.id,
+            payload=payload,
+            tags=tags,
+            expires_at=raw.expires_at,
+            _handle=self,
+        )
+
+
+
+class KV:
+    def __init__(self, backend: Backend, namespace: str) -> None:
+        self._backend = backend
+        self._ns = namespace
+
+    async def set(self, key: str, value: Any, *, ttl_seconds: Optional[int] = None) -> None:
+        # value can be a pydantic model, dict, str, int, ..., serialized as json
+        import json
+        if hasattr(value, "model_dump_json"):
+            blob = value.model_dump_json().encode("utf-8")
+        else:
+            blob = json.dumps(value).encode("utf-8")
+        await self._backend.kv_set(self._ns, key, blob, ttl_seconds)
+
+    async def get(self, key: str) -> Any:
+        import json
+        raw = await self._backend.kv_get(self._ns, key)
+        if raw is None:
+            return None
+        try:
+            return json.loads(raw.decode("utf-8"))
+        except Exception:
+            return None
+
+    async def get_as(self, key: str, model: type[BaseModel]) -> Optional[BaseModel]:
+        raw = await self._backend.kv_get(self._ns, key)
+        if raw is None:
+            return None
+        return model.model_validate_json(raw)
+
+    async def delete(self, key: str) -> bool:
+        return await self._backend.kv_del(self._ns, key)
+
+
+
+class Lock:
+    def __init__(self, backend: Backend, name: str, hold_seconds: int, wait_seconds: float) -> None:
+        self._backend = backend
+        self._name = name
+        self._hold = hold_seconds
+        self._wait = wait_seconds
+        self._token: Optional[str] = None
+
+    async def __aenter__(self) -> bool:
+        self._token = await self._backend.lock_acquire(self._name, self._hold, self._wait)
+        return self._token is not None
+
+    async def __aexit__(self, *exc) -> None:
+        if self._token is not None:
+            try:
+                await self._backend.lock_release(self._name, self._token)
+            except Exception:
+                logger.debug("lock release failed", exc_info=True)
+            self._token = None
+
+
+
+class Counter:
+    def __init__(self, backend: Backend, name: str) -> None:
+        self._backend = backend
+        self._name = name
+
+    async def incr(self, by: int = 1) -> int:
+        return await self._backend.counter_incr(self._name, by)
+
+    async def decr(self, by: int = 1) -> int:
+        return await self._backend.counter_incr(self._name, -by)
+
+    async def get(self) -> int:
+        return await self._backend.counter_get(self._name)
+
+    async def set(self, value: int) -> None:
+        await self._backend.counter_set(self._name, value)
+
+    async def reset(self) -> None:
+        await self._backend.counter_del(self._name)
+
+
+
+class QueueHandle(Generic[P]):
+    def __init__(self, spec: Queue[P], backend: Backend) -> None:
+        self._spec = spec
+        self._backend = backend
+
+    async def push(
+        self,
+        payload: P,
+        *,
+        task_id: Optional[str] = None,
+        priority: int = 0,
+        ttl_seconds: Optional[int] = None,
+    ) -> str:
+        if not isinstance(payload, self._spec.payload):
+            raise TypeError(
+                f"queue {self._spec.name}: payload must be {self._spec.payload.__name__}, "
+                f"got {type(payload).__name__}"
+            )
+        import time
+        tid = task_id or uuid.uuid4().hex
+        envelope: dict[str, Any] = {
+            "id": tid,
+            "payload": payload.model_dump(mode="json"),
+            "attempt": 0,
+            "submitted_at": time.time(),
+        }
+        if ttl_seconds is not None:
+            envelope["ttl_seconds"] = ttl_seconds
+        body = json.dumps(envelope).encode("utf-8")
+        await self._backend.queue_push(self._spec.name, body, priority=priority)
+        return tid
+
+
+
+class StreamReader(Generic[P]):
+    def __init__(self, spec: Stream[P], backend: Backend) -> None:
+        self._spec = spec
+        self._backend = backend
+
+    async def pop(self, n: int = 1) -> list[P]:
+        raw_list = await self._backend.stream_pop(self._spec.name, n)
+        out: list[P] = []
+        for raw in raw_list:
+            try:
+                envelope = json.loads(raw.decode("utf-8") if isinstance(raw, bytes) else raw)
+                out.append(self._spec.payload.model_validate(envelope.get("payload", {})))
+            except (json.JSONDecodeError, ValidationError):
+                logger.warning("stream %s: bad item, dropping", self._spec.name)
+                continue
+        return out
+
+    async def peek(self, n: int = 10) -> list[P]:
+        raw_list = await self._backend.stream_peek(self._spec.name, n)
+        out: list[P] = []
+        for raw in raw_list:
+            try:
+                envelope = json.loads(raw.decode("utf-8") if isinstance(raw, bytes) else raw)
+                out.append(self._spec.payload.model_validate(envelope.get("payload", {})))
+            except (json.JSONDecodeError, ValidationError):
+                continue
+        return out
+
+    async def length(self) -> int:
+        return await self._backend.stream_length(self._spec.name)
+
+
+__all__ = [
+    "Counter",
+    "KV",
+    "Lock",
+    "PoolFilter",
+    "PoolHandle",
+    "PoolItem",
+    "QueueHandle",
+    "StreamReader",
+]

fleet/core/protocol.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import json
+from typing import Annotated, Any, Literal, Union
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+PROTOCOL_VERSION: int = 1
+
+WorkerState = Literal["IDLE", "RUNNING", "ERROR", "DRAINING"]
+
+
+class ProtocolError(Exception):
+    pass
+
+
+class _Frame(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class Hardware(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    cores: int = Field(ge=1)
+    ram_gb: int = Field(ge=0)
+    hostname: str
+    kernel: str
+
+
+# worker -> master
+
+
+class Register(_Frame):
+    type: Literal["register"] = "register"
+    worker_id: str
+    automation_type: str
+    protocol_version: int
+    hardware: Hardware
+    fleet_core_version: str = ""
+    """Worker's installed fleet-core version. Master logs a mismatch warning
+    when the major differs from its own. Empty for backwards compat with
+    pre-0.1 workers."""
+
+
+class State(_Frame):
+    type: Literal["state"] = "state"
+    state: WorkerState
+    config_gen: int = Field(ge=0)
+    ts: int = Field(ge=0)
+    last_error: str | None = None
+
+
+class SlotStats(_Frame):
+    """Per-slot snapshot inside a Stats frame."""
+    slot_id: int = Field(ge=0)
+    age_seconds: float = Field(ge=0.0)
+    cpu_pct: float = Field(ge=0.0, default=0.0)
+    rss_mb: int = Field(ge=0, default=0)
+    counters: dict[str, int] = Field(default_factory=dict)
+
+
+class Stats(_Frame):
+    type: Literal["stats"] = "stats"
+    slots: int = Field(ge=0)
+    cpu_pct: float = Field(ge=0.0)
+    rss_mb: int = Field(ge=0)
+    ts: int = Field(ge=0)
+    metrics: dict[str, Any] = Field(default_factory=dict)
+    per_slot: list[SlotStats] = Field(default_factory=list)
+
+
+class Output(_Frame):
+    # generic worker -> master output frame. payload is plugin-typed (validated server-side).
+    type: Literal["output"] = "output"
+    slot_id: int = Field(ge=0)
+    ts: int = Field(ge=0)
+    payload: dict[str, Any]
+
+
+# master -> worker
+
+
+class ConfigChanged(_Frame):
+    type: Literal["config_changed"] = "config_changed"
+    config_gen: int = Field(ge=0)
+
+
+class TaskAssign(_Frame):
+    # used by BatchAutomation. master pushes a single task to a worker slot.
+    type: Literal["task_assign"] = "task_assign"
+    task_id: str
+    payload: dict[str, Any]
+
+
+class Ping(_Frame):
+    type: Literal["ping"] = "ping"
+    id: str
+
+
+class Drain(_Frame):
+    """Master → worker: drain in-flight slots and exit cleanly.
+
+    After receiving, the worker sets shutdown on every slot, waits for
+    them to complete (no hard cancel), reports DRAINING, then exits.
+    """
+    type: Literal["drain"] = "drain"
+
+
+Message = Annotated[
+    Union[Register, State, Stats, Output, ConfigChanged, TaskAssign, Ping, Drain],
+    Field(discriminator="type"),
+]
+
+
+_TYPE_REGISTRY: dict[str, type[BaseModel]] = {
+    "register": Register,
+    "state": State,
+    "stats": Stats,
+    "output": Output,
+    "config_changed": ConfigChanged,
+    "task_assign": TaskAssign,
+    "ping": Ping,
+    "drain": Drain,
+}
+
+
+def encode(msg: BaseModel) -> str:
+    if not isinstance(msg, BaseModel):
+        raise TypeError(f"encode() expected pydantic BaseModel, got {type(msg).__name__}")
+    return msg.model_dump_json()
+
+
+def decode(raw: str | bytes | bytearray) -> BaseModel:
+    if not isinstance(raw, (str, bytes, bytearray)):
+        raise ProtocolError(f"decode() expects str/bytes, got {type(raw).__name__}")
+    try:
+        data: Any = json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise ProtocolError(f"malformed JSON: {e}") from e
+    if not isinstance(data, dict):
+        raise ProtocolError(f"frame root must be object, got {type(data).__name__}")
+    if "type" not in data:
+        raise ProtocolError("frame missing required field 'type'")
+    tag = data["type"]
+    if not isinstance(tag, str):
+        raise ProtocolError(f"frame 'type' must be string, got {type(tag).__name__}")
+    cls = _TYPE_REGISTRY.get(tag)
+    if cls is None:
+        raise ProtocolError(f"unknown message type: {tag!r}")
+    try:
+        return cls.model_validate(data)
+    except ValidationError as e:
+        raise ProtocolError(f"validation failed for type={tag!r}: {e}") from e
+
+
+__all__ = [
+    "ConfigChanged",
+    "Drain",
+    "Hardware",
+    "Message",
+    "Output",
+    "PROTOCOL_VERSION",
+    "Ping",
+    "ProtocolError",
+    "Register",
+    "State",
+    "Stats",
+    "TaskAssign",
+    "WorkerState",
+    "decode",
+    "encode",
+]

fleet/core/proxy.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import importlib.metadata as md
+import logging
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator, ClassVar, Optional, Protocol, runtime_checkable
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ProxySession:
+    url: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+    _release: Optional[Any] = None
+    _report_bad: Optional[Any] = None
+
+    async def release(self) -> None:
+        if self._release is not None:
+            await self._release(self)
+
+    async def report_bad(self, reason: str) -> None:
+        if self._report_bad is not None:
+            await self._report_bad(self, reason)
+
+
+@runtime_checkable
+class ProxyProvider(Protocol):
+    name: ClassVar[str]
+
+    @classmethod
+    def from_config(cls, cfg: dict[str, Any]) -> "ProxyProvider": ...
+
+    async def acquire(
+        self,
+        *,
+        slot_id: int,
+        sticky: bool = True,
+        country: Optional[str | list[str]] = None,
+    ) -> ProxySession: ...
+
+
+class ProxyHandle:
+    # what ctx.proxy returns. wraps a provider and the slot's identity.
+    def __init__(self, provider: Optional[ProxyProvider], slot_id: int) -> None:
+        self._provider = provider
+        self._slot_id = slot_id
+
+    @property
+    def provider(self) -> Optional[ProxyProvider]:
+        return self._provider
+
+    @asynccontextmanager
+    async def session(
+        self,
+        *,
+        sticky: bool = True,
+        country: Optional[str | list[str]] = None,
+    ) -> AsyncIterator[ProxySession]:
+        if self._provider is None:
+            yield ProxySession(url="")
+            return
+        sess = await self._provider.acquire(
+            slot_id=self._slot_id, sticky=sticky, country=country
+        )
+        try:
+            yield sess
+        finally:
+            try:
+                await sess.release()
+            except Exception:
+                logger.debug("proxy session release failed", exc_info=True)
+
+
+_REGISTRY: dict[str, type] = {}
+
+
+def register_provider(name: str):
+    def _wrap(cls):
+        if not hasattr(cls, "from_config"):
+            raise TypeError(f"{cls.__name__} must define classmethod from_config")
+        cls.name = name
+        if name in _REGISTRY and _REGISTRY[name] is not cls:
+            raise ValueError(f"proxy provider '{name}' already registered")
+        _REGISTRY[name] = cls
+        return cls
+
+    return _wrap
+
+
+def get_provider_registry() -> dict[str, type]:
+    return dict(_REGISTRY)
+
+
+def load_provider_entry_points(group: str = "fleet.providers.proxy") -> dict[str, type]:
+    try:
+        eps = md.entry_points(group=group)
+    except TypeError:
+        eps = md.entry_points().get(group, [])  # type: ignore[assignment]
+    for ep in eps:
+        try:
+            ep.load()
+        except Exception:
+            logger.exception("failed to load proxy provider entry-point %s", ep.name)
+    return get_provider_registry()
+
+
+def build_provider(name: str, cfg: dict[str, Any]) -> ProxyProvider:
+    load_provider_entry_points()
+    cls = _REGISTRY.get(name)
+    if cls is None:
+        known = ", ".join(sorted(_REGISTRY.keys())) or "(none installed)"
+        raise ValueError(f"unknown proxy provider '{name}'. installed: {known}")
+    return cls.from_config(cfg)
+
+
+
+@register_provider("static")
+class StaticProvider:
+    # the original behavior: one URL, always the same.
+    def __init__(self, url: Optional[str]) -> None:
+        self._url = url
+
+    @classmethod
+    def from_config(cls, cfg: dict[str, Any]) -> "StaticProvider":
+        return cls(url=cfg.get("url"))
+
+    async def acquire(
+        self,
+        *,
+        slot_id: int,
+        sticky: bool = True,
+        country: Optional[str | list[str]] = None,
+    ) -> ProxySession:
+        return ProxySession(url=self._url or "")
+
+
+
+@register_provider("rotating-list")
+class RotatingListProvider:
+    def __init__(self, urls: list[str]) -> None:
+        self._urls = urls
+        self._next = 0
+
+    @classmethod
+    def from_config(cls, cfg: dict[str, Any]) -> "RotatingListProvider":
+        urls = list(cfg.get("urls") or [])
+        if not urls:
+            raise ValueError("rotating-list provider requires non-empty 'urls' list")
+        return cls(urls=urls)
+
+    async def acquire(
+        self,
+        *,
+        slot_id: int,
+        sticky: bool = True,
+        country: Optional[str | list[str]] = None,
+    ) -> ProxySession:
+        if sticky:
+            # deterministic by slot_id so the same slot keeps the same URL.
+            url = self._urls[slot_id % len(self._urls)]
+        else:
+            url = self._urls[self._next % len(self._urls)]
+            self._next += 1
+        return ProxySession(url=url, metadata={"sticky": sticky})