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


class FcClient:
    def __init__(self, base_url: str, token: str, agent_id: str):
        self.base = base_url.rstrip("/")
        self.agent_id = agent_id
        self.s = requests.Session()
        self.s.headers["Authorization"] = f"Bearer {token}"
        # Many worker threads share this 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)
        self.s.mount("http://", adapter)
        self.s.mount("https://", adapter)

    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},
            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:
        r = self.s.post(
            f"{self.base}/api/gpu/jobs/submit",
            json={
                "agent_id": self.agent_id, "job_id": job_id,
                "regions": regions, "replace_kinds": replace_kinds,
            },
            timeout=120,
        )
        r.raise_for_status()
        return r.json()

    def submit_embedding(self, job_id: int, embedding: list, version: str) -> dict:
        """Post a whole-image SigLIP embedding (the 'embed' task) → image_record."""
        r = self.s.post(
            f"{self.base}/api/gpu/jobs/submit_embedding",
            json={
                "agent_id": self.agent_id, "job_id": job_id,
                "embedding": embedding, "embedding_version": version,
            },
            timeout=120,
        )
        r.raise_for_status()
        return r.json()

    def heartbeat(self, job_ids: list[int]) -> None:
        try:
            self.s.post(
                f"{self.base}/api/gpu/jobs/heartbeat",
                json={"agent_id": self.agent_id, "job_ids": job_ids},
                timeout=30,
            )
        except requests.RequestException:
            pass

    def fail(self, job_id: int, error: str) -> None:
        try:
            self.s.post(
                f"{self.base}/api/gpu/jobs/fail",
                json={"agent_id": self.agent_id, "job_id": job_id, "error": error},
                timeout=30,
            )
        except requests.RequestException:
            pass

    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
        try:
            self.s.post(
                f"{self.base}/api/gpu/jobs/release",
                json={"agent_id": self.agent_id, "job_ids": job_ids},
                timeout=30,
            )
        except requests.RequestException:
            pass

    def fetch_image(self, image_url: str) -> bytes:
        # image_url is a server-relative path ("/images/...").
        r = self.s.get(f"{self.base}{image_url}", timeout=180)
        r.raise_for_status()
        return r.content

    def queue_status(self) -> dict:
        r = self.s.get(f"{self.base}/api/gpu/status", timeout=15)
        r.raise_for_status()
        return r.json()