fleet-framework 0.1.0__py3-none-any.whl

fleet/core/memory_backend.py

@@ -0,0 +1,492 @@
+"""In-process Backend implementation. Drop-in for tests + single-host dev.
+
+No external dependencies. Same semantics as RedisBackend within a single
+event loop. Not safe across processes or threads — use the Redis backend
+for any deployment that has more than one process touching the same state.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import secrets
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from fleet.core.backend import RawItem
+
+
+@dataclass
+class _PoolEntry:
+    payload: bytes
+    tags: dict[str, Any]
+    expires_at: Optional[float]
+
+
+class InMemoryBackend:
+    def __init__(self, url: str = "memory://") -> None:
+        self._url = url
+        self._kv: dict[tuple[str, str], tuple[bytes, Optional[float]]] = {}
+        self._pools: dict[str, dict[str, _PoolEntry]] = defaultdict(dict)
+        self._claims: dict[tuple[str, str], float] = {}
+        self._locks: dict[str, tuple[str, float]] = {}
+        self._counters: dict[str, int] = defaultdict(int)
+        # Three priority tiers per queue: high, normal, low. Reservation
+        # drains them in that order.
+        self._queues: dict[tuple[str, str], asyncio.Queue[bytes]] = defaultdict(asyncio.Queue)
+        # task_id -> (queue_name, body, deadline). Keyed by (name, tid) elsewhere.
+        self._processing: dict[tuple[str, str], tuple[bytes, float]] = {}
+        # name -> list[(ready_at, body)] sorted by ready_at
+        self._scheduled: dict[str, list[tuple[float, bytes]]] = defaultdict(list)
+        self._dead: dict[str, list[bytes]] = defaultdict(list)
+        self._streams: dict[str, list[tuple[float, bytes]]] = defaultdict(list)
+        self._workers: dict[tuple[str, str], dict[str, Any]] = {}
+        self._catalog: dict[str, bytes] = {}
+        # list[(ts, frame_dict)] sorted by ts, capped at max_entries on push.
+        self._stats_history: dict[tuple[str, str], list[tuple[float, dict[str, Any]]]] = defaultdict(list)
+
+    async def aclose(self) -> None:
+        return None
+
+    def _expired_kv(self, key: tuple[str, str]) -> bool:
+        rec = self._kv.get(key)
+        return rec is not None and rec[1] is not None and rec[1] < time.time()
+
+    async def kv_set(self, ns: str, key: str, value: bytes, ttl_seconds: Optional[int]) -> None:
+        exp = time.time() + ttl_seconds if (ttl_seconds and ttl_seconds > 0) else None
+        self._kv[(ns, key)] = (value, exp)
+
+    async def kv_setnx(
+        self, ns: str, key: str, value: bytes, ttl_seconds: Optional[int]
+    ) -> bool:
+        k = (ns, key)
+        if self._expired_kv(k):
+            self._kv.pop(k, None)
+        if k in self._kv:
+            return False
+        exp = time.time() + ttl_seconds if (ttl_seconds and ttl_seconds > 0) else None
+        self._kv[k] = (value, exp)
+        return True
+
+    async def kv_get(self, ns: str, key: str) -> Optional[bytes]:
+        k = (ns, key)
+        if self._expired_kv(k):
+            self._kv.pop(k, None)
+            return None
+        rec = self._kv.get(k)
+        return rec[0] if rec else None
+
+    async def kv_del(self, ns: str, key: str) -> bool:
+        return self._kv.pop((ns, key), None) is not None
+
+    def _prune_pool(self, name: str) -> None:
+        now = time.time()
+        bucket = self._pools[name]
+        to_del = [iid for iid, e in bucket.items() if e.expires_at is not None and e.expires_at < now]
+        for iid in to_del:
+            bucket.pop(iid, None)
+        for (n, iid) in list(self._claims):
+            if n == name and self._claims[(n, iid)] < now:
+                self._claims.pop((n, iid), None)
+
+    async def pool_put(
+        self,
+        name: str,
+        item_id: str,
+        payload: bytes,
+        tags: dict[str, Any],
+        ttl_seconds: Optional[int],
+    ) -> None:
+        exp = time.time() + ttl_seconds if (ttl_seconds and ttl_seconds > 0) else None
+        self._pools[name][item_id] = _PoolEntry(payload=payload, tags=dict(tags), expires_at=exp)
+        self._prune_pool(name)
+
+    async def pool_list(self, name: str) -> list[RawItem]:
+        self._prune_pool(name)
+        now = time.time()
+        out: list[RawItem] = []
+        for iid, e in self._pools[name].items():
+            claim_exp = self._claims.get((name, iid))
+            in_use = claim_exp if (claim_exp is not None and claim_exp > now) else None
+            out.append(RawItem(
+                id=iid,
+                payload=e.payload,
+                tags=e.tags,
+                expires_at=e.expires_at,
+                in_use_until=in_use,
+            ))
+        return out
+
+    async def pool_remove(self, name: str, item_id: str) -> bool:
+        removed = self._pools[name].pop(item_id, None) is not None
+        self._claims.pop((name, item_id), None)
+        return removed
+
+    async def pool_release(self, name: str, item_id: str) -> bool:
+        return self._claims.pop((name, item_id), None) is not None
+
+    async def pool_claim_any(
+        self,
+        name: str,
+        where: dict[str, Any],
+        hold_seconds: int,
+    ) -> Optional[RawItem]:
+        self._prune_pool(name)
+        now = time.time()
+        # newest items first, like RedisBackend
+        candidates = sorted(
+            self._pools[name].items(),
+            key=lambda kv: (kv[1].expires_at or 0),
+            reverse=True,
+        )
+        for iid, entry in candidates:
+            claim_exp = self._claims.get((name, iid))
+            if claim_exp is not None and claim_exp > now:
+                continue
+            if not all(entry.tags.get(k) == v for k, v in where.items()):
+                continue
+            self._claims[(name, iid)] = now + hold_seconds
+            return RawItem(
+                id=iid,
+                payload=entry.payload,
+                tags=entry.tags,
+                expires_at=entry.expires_at,
+                in_use_until=now + hold_seconds,
+            )
+        return None
+
+    async def lock_acquire(self, name: str, hold_seconds: int, wait_seconds: float) -> Optional[str]:
+        deadline = time.monotonic() + max(0.0, wait_seconds)
+        backoff = 0.05
+        while True:
+            now = time.time()
+            current = self._locks.get(name)
+            if current is None or current[1] < now:
+                token = secrets.token_hex(16)
+                self._locks[name] = (token, now + max(1, hold_seconds))
+                return token
+            if time.monotonic() >= deadline:
+                return None
+            await asyncio.sleep(min(backoff, deadline - time.monotonic()))
+            backoff = min(backoff * 2, 0.5)
+
+    async def lock_release(self, name: str, token: str) -> bool:
+        current = self._locks.get(name)
+        if current is None:
+            return False
+        if current[0] != token:
+            return False
+        self._locks.pop(name, None)
+        return True
+
+    async def counter_incr(self, name: str, by: int) -> int:
+        self._counters[name] += by
+        return self._counters[name]
+
+    async def counter_get(self, name: str) -> int:
+        return self._counters.get(name, 0)
+
+    async def counter_set(self, name: str, value: int) -> None:
+        self._counters[name] = value
+
+    async def counter_del(self, name: str) -> None:
+        self._counters.pop(name, None)
+
+    def _tier(self, priority: int) -> str:
+        if priority > 0:
+            return "high"
+        if priority < 0:
+            return "low"
+        return "normal"
+
+    def _tiers_in_order(self) -> list[str]:
+        return ["high", "normal", "low"]
+
+    async def queue_push(self, name: str, body: bytes, *, priority: int = 0) -> None:
+        await self._queues[(name, self._tier(priority))].put(body)
+
+    async def queue_size(self, name: str) -> int:
+        return sum(
+            self._queues[(name, t)].qsize()
+            for t in self._tiers_in_order()
+            if (name, t) in self._queues
+        )
+
+    async def queue_peek(self, name: str, n: int) -> list[bytes]:
+        if n <= 0:
+            return []
+        out: list[bytes] = []
+        for t in self._tiers_in_order():
+            key = (name, t)
+            if key not in self._queues:
+                continue
+            q = self._queues[key]
+            for body in list(q._queue):  # type: ignore[attr-defined]
+                if len(out) >= n:
+                    return out
+                out.append(body)
+        return out
+
+    async def pool_size(self, name: str) -> int:
+        return len(self._pools.get(name, {}))
+
+    def _move_due_scheduled(self, name: str) -> None:
+        now = time.time()
+        due = [b for (rt, b) in self._scheduled[name] if rt <= now]
+        if not due:
+            return
+        self._scheduled[name] = [(rt, b) for (rt, b) in self._scheduled[name] if rt > now]
+        # Re-pushed scheduled tasks go to the normal tier.
+        q = self._queues[(name, "normal")]
+        for body in due:
+            q.put_nowait(body)
+
+    async def queue_reserve(
+        self, name: str, timeout_seconds: float, visibility_seconds: int
+    ) -> Optional[bytes]:
+        self._move_due_scheduled(name)
+        body: Optional[bytes] = None
+        # Try each tier non-blockingly, then poll if nothing immediately available.
+        deadline = time.monotonic() + max(0.0, timeout_seconds)
+        while True:
+            for tier in self._tiers_in_order():
+                q = self._queues[(name, tier)]
+                if not q.empty():
+                    body = q.get_nowait()
+                    break
+            if body is not None or time.monotonic() >= deadline:
+                break
+            await asyncio.sleep(min(0.05, max(0.0, deadline - time.monotonic())))
+        if body is None:
+            return None
+        try:
+            envelope = json.loads(body.decode("utf-8"))
+            tid = str(envelope.get("id", ""))
+        except Exception:
+            tid = ""
+        if tid:
+            d = time.time() + max(1, visibility_seconds)
+            self._processing[(name, tid)] = (body, d)
+        return body
+
+    async def queue_ack(self, name: str, task_id: str) -> bool:
+        return self._processing.pop((name, task_id), None) is not None
+
+    async def queue_nack(
+        self, name: str, task_id: str, body: bytes, *, delay_seconds: float
+    ) -> bool:
+        had = self._processing.pop((name, task_id), None) is not None
+        if delay_seconds > 0:
+            self._scheduled[name].append((time.time() + delay_seconds, body))
+            self._scheduled[name].sort(key=lambda x: x[0])
+        else:
+            await self._queues[(name, "normal")].put(body)
+        return had
+
+    async def queue_sweep_expired(self, name: str, now: Optional[float] = None) -> int:
+        now = now if now is not None else time.time()
+        moved = 0
+        for key in list(self._processing.keys()):
+            if key[0] != name:
+                continue
+            body, deadline = self._processing[key]
+            if deadline <= now:
+                self._processing.pop(key, None)
+                await self._queues[(name, "normal")].put(body)
+                moved += 1
+        return moved
+
+    async def queue_dead(self, name: str, body: bytes) -> None:
+        self._dead[name].append(body)
+
+    async def dlq_peek(self, name: str, n: int) -> list[bytes]:
+        return list(self._dead.get(name, [])[:n])
+
+    async def dlq_pop(self, name: str, n: int) -> list[bytes]:
+        bucket = self._dead.get(name, [])
+        taken = bucket[:n]
+        self._dead[name] = bucket[n:]
+        return taken
+
+    async def dlq_length(self, name: str) -> int:
+        return len(self._dead.get(name, []))
+
+    async def dlq_clear(self, name: str) -> int:
+        n = len(self._dead.get(name, []))
+        self._dead[name] = []
+        return n
+
+    async def stream_push(
+        self,
+        name: str,
+        envelope: bytes,
+        *,
+        score: float,
+        max_len: int,
+        ttl_seconds: int,
+    ) -> None:
+        bucket = self._streams[name]
+        bucket.append((score, envelope))
+        bucket.sort(key=lambda x: x[0])
+        if ttl_seconds > 0:
+            cutoff = score - ttl_seconds
+            self._streams[name] = [(s, m) for (s, m) in bucket if s >= cutoff]
+        if max_len > 0 and len(self._streams[name]) > max_len:
+            self._streams[name] = self._streams[name][-max_len:]
+
+    async def stream_pop(self, name: str, n: int) -> list[bytes]:
+        bucket = self._streams[name]
+        if n <= 0:
+            return []
+        taken = bucket[:n]
+        self._streams[name] = bucket[n:]
+        return [m for _s, m in taken]
+
+    async def stream_peek(self, name: str, n: int) -> list[bytes]:
+        return [m for _s, m in self._streams[name][:n]]
+
+    async def stream_length(self, name: str) -> int:
+        return len(self._streams.get(name, []))
+
+    async def stream_trim(self, name: str, *, max_len: int, ttl_seconds: int) -> None:
+        if name not in self._streams:
+            return
+        bucket = self._streams[name]
+        if ttl_seconds > 0:
+            cutoff = time.time() - ttl_seconds
+            bucket = [(s, m) for (s, m) in bucket if s >= cutoff]
+        if max_len > 0 and len(bucket) > max_len:
+            bucket = bucket[-max_len:]
+        self._streams[name] = bucket
+
+    async def event_publish(self, topic: str, body: bytes) -> None:
+        # No subscribers in-memory. Plugins that use ctx.events.subscribe must
+        # use a real broker; the memory backend is for tests + single-process.
+        return None
+
+    async def worker_register(
+        self, automation_type: str, worker_id: str, hardware: dict[str, Any]
+    ) -> None:
+        key = (automation_type, worker_id)
+        existing = self._workers.get(key, {})
+        existing["hardware"] = dict(hardware)
+        existing.setdefault("config", {})
+        existing.setdefault("config_gen", 0)
+        existing.setdefault("state", "IDLE")
+        existing.setdefault("last_seen", 0.0)
+        existing.setdefault("last_error", None)
+        existing.setdefault("stats", {})
+        self._workers[key] = existing
+
+    async def worker_get_config(
+        self, automation_type: str, worker_id: str
+    ) -> tuple[dict[str, Any], int]:
+        rec = self._workers.get((automation_type, worker_id))
+        if rec is None:
+            return {}, 0
+        return dict(rec["config"]), rec["config_gen"]
+
+    async def worker_set_config(
+        self, automation_type: str, worker_id: str, config: dict[str, Any]
+    ) -> tuple[dict[str, Any], int]:
+        key = (automation_type, worker_id)
+        rec = self._workers.setdefault(key, {
+            "hardware": {}, "config": {}, "config_gen": 0,
+            "state": "IDLE", "last_seen": 0.0, "last_error": None, "stats": {},
+        })
+        rec["config"] = dict(config)
+        rec["config_gen"] = int(rec.get("config_gen", 0)) + 1
+        return rec["config"], rec["config_gen"]
+
+    async def worker_set_state(
+        self,
+        automation_type: str,
+        worker_id: str,
+        state: str,
+        last_error: Optional[str],
+        last_seen: float,
+    ) -> None:
+        key = (automation_type, worker_id)
+        rec = self._workers.setdefault(key, {
+            "hardware": {}, "config": {}, "config_gen": 0,
+            "state": "IDLE", "last_seen": 0.0, "last_error": None, "stats": {},
+        })
+        rec["state"] = state
+        rec["last_seen"] = last_seen
+        rec["last_error"] = last_error
+
+    async def worker_set_stats(
+        self, automation_type: str, worker_id: str, stats: dict[str, Any]
+    ) -> None:
+        key = (automation_type, worker_id)
+        rec = self._workers.setdefault(key, {
+            "hardware": {}, "config": {}, "config_gen": 0,
+            "state": "IDLE", "last_seen": 0.0, "last_error": None, "stats": {},
+        })
+        rec["stats"] = dict(stats)
+
+    async def worker_stats_push(
+        self,
+        automation_type: str,
+        worker_id: str,
+        ts: float,
+        frame: dict[str, Any],
+        *,
+        max_entries: int = 240,
+    ) -> None:
+        bucket = self._stats_history[(automation_type, worker_id)]
+        bucket.append((ts, dict(frame)))
+        bucket.sort(key=lambda x: x[0])
+        if max_entries > 0 and len(bucket) > max_entries:
+            del bucket[: len(bucket) - max_entries]
+
+    async def worker_stats_history(
+        self,
+        automation_type: str,
+        worker_id: str,
+        *,
+        since: float = 0,
+        limit: int = 240,
+    ) -> list[dict[str, Any]]:
+        bucket = self._stats_history.get((automation_type, worker_id), [])
+        return [dict(f) for (t, f) in bucket if t >= since][:limit]
+
+    async def worker_get(
+        self, automation_type: str, worker_id: str
+    ) -> Optional[dict[str, Any]]:
+        rec = self._workers.get((automation_type, worker_id))
+        if rec is None:
+            return None
+        return {
+            "hardware": dict(rec["hardware"]),
+            "config": dict(rec["config"]),
+            "config_gen": rec["config_gen"],
+            "state": rec["state"],
+            "last_seen": rec["last_seen"],
+            "last_error": rec["last_error"],
+            "stats": dict(rec["stats"]),
+        }
+
+    async def worker_list(
+        self, automation_type: Optional[str] = None
+    ) -> list[tuple[str, str]]:
+        if automation_type is None:
+            return list(self._workers.keys())
+        return [(at, wid) for (at, wid) in self._workers if at == automation_type]
+
+    async def worker_remove(self, automation_type: str, worker_id: str) -> bool:
+        return self._workers.pop((automation_type, worker_id), None) is not None
+
+    async def catalog_set(self, automation_type: str, doc: bytes) -> None:
+        self._catalog[automation_type] = doc
+
+    async def catalog_get(self, automation_type: str) -> Optional[bytes]:
+        return self._catalog.get(automation_type)
+
+    async def catalog_all(self) -> dict[str, bytes]:
+        return dict(self._catalog)
+
+
+__all__ = ["InMemoryBackend"]

fleet/core/metrics.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import time
+from collections import defaultdict, deque
+from dataclasses import dataclass, field
+from threading import Lock
+from typing import Any
+
+
+@dataclass
+class _Counter:
+    n: int = 0
+
+
+@dataclass
+class _Gauge:
+    v: float = 0.0
+
+
+@dataclass
+class _Histogram:
+    samples: deque = field(default_factory=lambda: deque(maxlen=512))
+
+
+class SlotMetrics:
+    # per-slot counters. dumped via WS Stats frames and surfaced on the dashboard.
+    def __init__(self) -> None:
+        self._lock = Lock()
+        self._counters: dict[str, _Counter] = defaultdict(_Counter)
+        self._gauges: dict[str, _Gauge] = defaultdict(_Gauge)
+        self._hists: dict[str, _Histogram] = defaultdict(_Histogram)
+        self._created_at = time.monotonic()
+
+    def incr(self, name: str, by: int = 1) -> None:
+        with self._lock:
+            self._counters[name].n += by
+
+    def set(self, name: str, value: float) -> None:
+        with self._lock:
+            self._gauges[name].v = value
+
+    def observe(self, name: str, value: float) -> None:
+        with self._lock:
+            self._hists[name].samples.append(value)
+
+    def snapshot(self) -> dict[str, Any]:
+        with self._lock:
+            counters = {k: v.n for k, v in self._counters.items()}
+            gauges = {k: v.v for k, v in self._gauges.items()}
+            hists = {}
+            for k, h in self._hists.items():
+                if not h.samples:
+                    continue
+                s = sorted(h.samples)
+                hists[k] = {
+                    "n": len(s),
+                    "p50": s[len(s) // 2],
+                    "p95": s[int(len(s) * 0.95)],
+                    "p99": s[int(len(s) * 0.99)],
+                }
+        return {"counters": counters, "gauges": gauges, "histograms": hists}

fleet/core/otel.py

@@ -0,0 +1,97 @@
+"""Optional OpenTelemetry integration.
+
+Importing this module never fails: if `opentelemetry-api` is not installed
+or `OTEL_EXPORTER_OTLP_ENDPOINT` (or `FLEET_OTEL_ENABLED`) is not set, the
+tracer falls back to no-op context managers so wrapping a hot path costs
+roughly one Python function call.
+
+Activation:
+
+    pip install fleet-core[otel]
+    export OTEL_EXPORTER_OTLP_ENDPOINT=http://collector:4318
+    export OTEL_SERVICE_NAME=fleet-worker
+
+Then spans appear in your collector for `slot.run_continuous`,
+`slot.run_one`, `slot.reserve`, and `api.submit_task`. All spans set
+`fleet.automation_type`, `fleet.worker_id`, `fleet.slot_id`, and (where
+relevant) `fleet.task_id` attributes.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from contextlib import contextmanager, nullcontext
+from typing import Any, Iterator
+
+logger = logging.getLogger(__name__)
+
+_tracer: Any = None
+_disabled = False
+
+
+def _build_tracer() -> Any:
+    global _tracer, _disabled
+    if _disabled:
+        return None
+    if _tracer is not None:
+        return _tracer
+    if not _otel_requested():
+        _disabled = True
+        return None
+    try:
+        from opentelemetry import trace
+    except ImportError:
+        logger.info("opentelemetry-api not installed; OTel hooks disabled")
+        _disabled = True
+        return None
+    _tracer = trace.get_tracer("fleet-core")
+    return _tracer
+
+
+def _otel_requested() -> bool:
+    if os.environ.get("FLEET_OTEL_ENABLED", "").lower() in ("1", "true", "yes"):
+        return True
+    return any(k.startswith("OTEL_EXPORTER_") for k in os.environ)
+
+
+@contextmanager
+def span(name: str, **attrs: Any) -> Iterator[Any]:
+    """Start an OTel span if OTel is configured, otherwise a no-op.
+
+    Attributes are passed in as kwargs and are added to the span. Any
+    None-valued attributes are skipped (OTel rejects them)."""
+    tracer = _build_tracer()
+    if tracer is None:
+        with nullcontext() as ctx:
+            yield ctx
+        return
+    cleaned = {f"fleet.{k}": v for k, v in attrs.items() if v is not None}
+    with tracer.start_as_current_span(name, attributes=cleaned) as s:
+        yield s
+
+
+def record_exception(exc: BaseException) -> None:
+    """Attach an exception to the active span if one exists."""
+    tracer = _build_tracer()
+    if tracer is None:
+        return
+    try:
+        from opentelemetry import trace
+        current = trace.get_current_span()
+        if current is not None:
+            current.record_exception(exc)
+    except Exception:
+        return
+
+
+def reset_for_tests() -> None:
+    """Test helper. Clears cached tracer so monkeypatched env reloads."""
+    global _tracer, _disabled
+    _tracer = None
+    _disabled = False
+
+
+def is_active() -> bool:
+    """True if OTel is wired and ready to emit spans."""
+    return _build_tracer() is not None