"""Client for talking to a compute node's REST API. This is the application-side bridge: the application backend never runs GPU workloads itself; it dispatches them to a compute node and reads back status, logs and metrics through this client. """ from __future__ import annotations import httpx from app.core.config import get_settings def _compute_timeout() -> float: return float(get_settings().compute_request_timeout_seconds or 30.0) def _auth_headers() -> dict: token = get_settings().compute_service_token if not token: return {} return {"Authorization": f"Bearer {token}"} class ComputeNodeClient: """Thin wrapper over a single compute node's HTTP API.""" def __init__(self, base_url: str, timeout: float | None = None): self.base_url = base_url.rstrip("/") self.timeout = timeout or _compute_timeout() # ---- health ------------------------------------------------------- def health(self) -> dict: last_err = None for path in ("/v1/compute/health", "/health"): try: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.get(f"{self.base_url}{path}", headers=_auth_headers()) if r.status_code == 200: return r.json() except Exception as exc: # noqa: BLE001 last_err = str(exc) raise RuntimeError(f"compute node unhealthy: {last_err}") # ---- gpus --------------------------------------------------------- def gpus(self) -> list: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.get( f"{self.base_url}/compute/resources/gpus", headers=_auth_headers(), ) r.raise_for_status() return r.json().get("data", []) # ---- jobs --------------------------------------------------------- def create_job(self, payload: dict) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/jobs", json=payload, headers=_auth_headers(), ) r.raise_for_status() return r.json() def preview_job(self, payload: dict) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/jobs/preview", json=payload, headers=_auth_headers(), ) r.raise_for_status() return r.json() def validate_job(self, payload: dict) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/jobs/validate", json=payload, headers=_auth_headers(), ) r.raise_for_status() return r.json() def get_job(self, job_id: str) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.get( f"{self.base_url}/compute/jobs/{job_id}", headers=_auth_headers(), ) r.raise_for_status() return r.json() def stop_job(self, job_id: str) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/jobs/{job_id}/stop", headers=_auth_headers(), ) r.raise_for_status() return r.json() def job_logs(self, job_id: str, cursor: int = 0, limit: int = 200) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.get( f"{self.base_url}/compute/jobs/{job_id}/logs", params={"cursor": cursor, "limit": limit}, headers=_auth_headers(), ) r.raise_for_status() return r.json() # ---- files -------------------------------------------------------- def check_paths(self, paths: list) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/files/check-paths", json={"paths": paths}, headers=_auth_headers(), ) r.raise_for_status() return r.json() def list_files(self, path: str = "/") -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.get( f"{self.base_url}/compute/files/list", params={"path": path}, headers=_auth_headers(), ) r.raise_for_status() return r.json() def import_local_file(self, src_path: str, dest_name: str | None = None) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/files/import-local", json={"src_path": src_path, "dest_name": dest_name}, headers=_auth_headers(), ) r.raise_for_status() return r.json() def upload_file(self, filename: str, content: bytes, content_type: str | None = None) -> dict: with httpx.Client(timeout=self.timeout, verify=False) as c: r = c.post( f"{self.base_url}/compute/files/upload", files={"file": (filename, content, content_type)}, headers=_auth_headers(), ) r.raise_for_status() return r.json() # ---- inference (async) ------------------------------------------- def headers(self) -> dict[str, str]: """Return auth headers for compute node requests.""" token = get_settings().compute_service_token if not token: return {} return {"X-Compute-Token": token} @property def route_prefix(self) -> str: return get_settings().route_prefix.rstrip("/") or "/modelTF" async def _request(self, method: str, path: str, json_data: dict | None = None) -> dict: """Generic async request method for compute API endpoints.""" prefix = self.route_prefix url = f"{self.base_url.rstrip('/')}{prefix}{path}" async with httpx.AsyncClient(timeout=300, verify=False, headers=self.headers()) as client: if method.upper() == "GET": response = await client.get(url) else: response = await client.post(url, json=json_data) response.raise_for_status() data = response.json() if isinstance(data, dict) and isinstance(data.get("data"), dict): return data["data"] return data if isinstance(data, dict) else {}