FabledCurator/agent/fc_agent/client.py

"""HTTP client for the FabledCurator GPU-job API.

The agent's ONLY contact with FC — lease/submit/heartbeat/fail + fetch image
bytes, all over HTTP with the bearer token. No DB/Redis.
"""
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from . import accel


class FcClient:
    def __init__(self, base_url: str, token: str, agent_id: str):
        self.base = base_url.rstrip("/")
        self.agent_id = agent_id
        # Main session: NO in-request retry — lease/fetch are cheap to redo and
        # the worker loop already backs off + re-leases on failure. (Auto-retrying
        # a lease could double-claim a batch if a response is lost.)
        self.s = self._session(token)
        # Submit session: retry in-place, because by submit time the GPU work is
        # already DONE — a momentary blip (dropped connection, gateway 5xx during
        # a curator redeploy) must not throw that work away and force a full
        # re-download + recompute on another agent. A duplicate submit after a
        # lost response is harmless: the job is already closed, so it just returns
        # 409 lease_invalid (a no-op). Idempotent enough to retry POST safely.
        retry = Retry(
            total=3, connect=3, read=3, status=3,
            backoff_factor=0.5,                       # ~0.5s, 1s, 2s between tries
            status_forcelist=(500, 502, 503, 504),    # transient server/gateway
            allowed_methods=frozenset({"POST"}),
            raise_on_status=False,                     # let raise_for_status decide
        )
        self._submit_s = self._session(token, retry)

    @staticmethod
    def _session(token: str, retry: Retry | None = None) -> requests.Session:
        s = requests.Session()
        s.headers["Authorization"] = f"Bearer {token}"
        # Many worker threads share a Session; the default pool (10) would
        # throttle them + spam "connection pool is full". Size it for the cap.
        adapter = HTTPAdapter(
            pool_connections=64, pool_maxsize=64, max_retries=retry or 0
        )
        s.mount("http://", adapter)
        s.mount("https://", adapter)
        return s

    def _submit(self, path: str, payload: dict) -> dict:
        """POST to a submit endpoint on the RETRYING session (by submit time the
        GPU work is done — a blip must not throw it away), raise on a hard error,
        and return the parsed JSON. `agent_id` is added to every body."""
        r = self._submit_s.post(
            f"{self.base}{path}",
            json={"agent_id": self.agent_id, **payload},
            timeout=120,
        )
        r.raise_for_status()
        return r.json()

    def _post_quiet(self, path: str, payload: dict) -> None:
        """Fire-and-forget POST on the main session — heartbeat/fail/release are
        best-effort, so a transport error is swallowed (the worker's own retry and
        the server's orphan-recovery cover a lost call). `agent_id` is added."""
        try:
            self.s.post(
                f"{self.base}{path}",
                json={"agent_id": self.agent_id, **payload},
                timeout=30,
            )
        except requests.RequestException:
            pass

    def lease(self, batch_size: int) -> list[dict]:
        r = self.s.post(
            f"{self.base}/api/gpu/jobs/lease",
            json={
                "agent_id": self.agent_id, "batch_size": batch_size,
                "accel": accel.summary(),
            },
            timeout=30,
        )
        r.raise_for_status()
        return r.json().get("jobs", [])

    def submit(self, job_id: int, regions: list[dict], replace_kinds: list[str]) -> dict:
        return self._submit("/api/gpu/jobs/submit", {
            "job_id": job_id, "regions": regions, "replace_kinds": replace_kinds,
        })

    def submit_embedding(self, job_id: int, embedding: list, version: str) -> dict:
        """Post a whole-image SigLIP embedding (the 'embed' task) → image_record."""
        return self._submit("/api/gpu/jobs/submit_embedding", {
            "job_id": job_id, "embedding": embedding, "embedding_version": version,
        })

    def heartbeat(self, job_ids: list[int]) -> None:
        self._post_quiet(
            "/api/gpu/jobs/heartbeat", {"job_ids": job_ids, "accel": accel.summary()},
        )

    def fail(self, job_id: int, error: str) -> None:
        self._post_quiet("/api/gpu/jobs/fail", {"job_id": job_id, "error": error})

    def release(self, job_ids: list[int]) -> None:
        # Graceful hand-back on stop so orphaned work is re-leased at once.
        if not job_ids:
            return
        self._post_quiet("/api/gpu/jobs/release", {"job_ids": job_ids})

    def fetch_image(self, image_url: str, throttle=None) -> bytes:
        # image_url is a server-relative path ("/images/...").
        # timeout=(connect, read): the read timeout is BETWEEN-BYTES, not total,
        # so a large-but-flowing download still completes — but a stuck/dead
        # connection (curator overloaded) fails in 60s instead of hanging a
        # downloader for 180s and piling up concurrent stuck requests on curator.
        # With a throttle (the worker's shared TokenBucket), the body is streamed
        # in chunks and each chunk is charged to the global bandwidth budget —
        # pausing between reads lets TCP flow control pace curator's send side.
        with self.s.get(
            f"{self.base}{image_url}", timeout=(10, 60), stream=throttle is not None
        ) as r:
            r.raise_for_status()
            if throttle is None:
                return r.content
            buf = bytearray()
            for chunk in r.iter_content(chunk_size=262_144):
                throttle.take(len(chunk))
                buf.extend(chunk)
            return bytes(buf)

    def is_reachable(self) -> bool:
        """Cheap 'is curator responding at all right now?' check. Used to decide,
        when a video can't be sampled, between a transient outage (keep retrying —
        survives a redeploy) and an unprocessable file (fail it, don't loop)."""
        try:
            r = self.s.get(f"{self.base}/api/gpu/status", timeout=5)
            return r.status_code < 500
        except requests.RequestException:
            return False

    def queue_status(self) -> dict:
        # Short timeout: this backs the UI /status poll, so a busy curator must
        # not hang the page for long (the GPU meters poll /gpu separately).
        r = self.s.get(f"{self.base}/api/gpu/status", timeout=5)
        r.raise_for_status()
        return r.json()