from __future__ import annotations import time from typing import Any from urllib.parse import urljoin import httpx from app.core.config import get_settings def _join_url(base_url: str, path: str) -> str: return urljoin(base_url.rstrip("/") + "/", path.lstrip("/")) def _unwrap_items(payload: Any) -> list[dict[str, Any]]: if isinstance(payload, list): return [item for item in payload if isinstance(item, dict)] if isinstance(payload, dict): data = payload.get("data") if isinstance(data, dict) and isinstance(data.get("items"), list): return [item for item in data["items"] if isinstance(item, dict)] if isinstance(payload.get("items"), list): return [item for item in payload["items"] if isinstance(item, dict)] if isinstance(data, list): return [item for item in data if isinstance(item, dict)] return [] def _unwrap_dict(payload: Any) -> dict[str, Any]: if isinstance(payload, dict) and isinstance(payload.get("data"), dict): return payload["data"] return payload if isinstance(payload, dict) else {} # Inference calls are intentionally short-timeout: # - load dispatch only confirms the compute node accepted the request # (the actual model load now runs asynchronously on the node). # - status/unload must never block the platform for long when a node is # unreachable but still marked online. INFERENCE_LOAD_TIMEOUT = httpx.Timeout(30, connect=10) INFERENCE_STATUS_TIMEOUT = httpx.Timeout(30, connect=5) INFERENCE_UNLOAD_TIMEOUT = httpx.Timeout(30, connect=5) class ComputeNodeClient: """Application-side client for one compute node. The client accepts both current YG Compute API responses and common wrapper shapes such as `{code,message,data}` to make future engine/node adapters less brittle. """ def __init__(self, api_base_url: str, token: str | None = None, timeout: float | None = None) -> None: settings = get_settings() self.api_base_url = api_base_url.rstrip("/") self.token = token or settings.compute_service_token self.timeout = timeout or settings.compute_request_timeout_seconds self.route_prefix = settings.route_prefix.rstrip("/") or "/modelTF" def headers(self) -> dict[str, str]: if not self.token: return {} return {"X-Compute-Token": self.token} async def test_connection(self) -> dict[str, Any]: started = time.perf_counter() health = await self.health() gpus = await self.gpus() return { "success": True, "latency_ms": int((time.perf_counter() - started) * 1000), "health": health, "gpus": gpus, } async def health(self) -> dict[str, Any]: paths = [f"{self.route_prefix}/v1/compute/health", f"{self.route_prefix}/health", "/health"] last_error = "" async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: for path in paths: try: response = await client.get(_join_url(self.api_base_url, path)) response.raise_for_status() return _unwrap_dict(response.json()) except Exception as exc: # noqa: BLE001 - keep endpoint compatibility fallback broad last_error = str(exc) raise RuntimeError(last_error or "compute health check failed") async def gpus(self) -> list[dict[str, Any]]: paths = [ f"{self.route_prefix}/compute/resources/gpus", f"{self.route_prefix}/v1/compute/resources/gpus", "/compute/resources/gpus", ] last_error = "" async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: for path in paths: try: response = await client.get(_join_url(self.api_base_url, path)) response.raise_for_status() return _unwrap_items(response.json()) except Exception as exc: # noqa: BLE001 last_error = str(exc) raise RuntimeError(last_error or "compute gpu discovery failed") async def create_job(self, payload: dict[str, Any]) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs"), json=payload) response.raise_for_status() return _unwrap_dict(response.json()) async def preview_job(self, payload: dict[str, Any]) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.post( _join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/preview"), json=payload, ) response.raise_for_status() return _unwrap_dict(response.json()) async def validate_job(self, payload: dict[str, Any]) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.post( _join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/validate"), json=payload, ) response.raise_for_status() return _unwrap_dict(response.json()) async def check_paths(self, paths: list[dict[str, Any]]) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.post( _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/check-paths"), json={"paths": paths}, ) response.raise_for_status() return _unwrap_dict(response.json()) async def list_files( self, root: str = "data", relative_path: str = "", directories_only: bool = False, ) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.get( _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/list"), params={"root": root, "relative_path": relative_path, "directories_only": directories_only}, ) response.raise_for_status() return _unwrap_dict(response.json()) async def get_job(self, job_id: str) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.get(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}")) response.raise_for_status() return _unwrap_dict(response.json()) async def stop_job(self, job_id: str) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/stop")) response.raise_for_status() return _unwrap_dict(response.json()) async def job_logs( self, job_id: str, tail_lines: int | None = None, offset: int | None = None, limit: int | None = None, ) -> dict[str, Any]: params = { key: value for key, value in {"tail_lines": tail_lines, "offset": offset, "limit": limit}.items() if value is not None } async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.get( _join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/logs"), params=params, ) response.raise_for_status() return _unwrap_dict(response.json()) async def import_local_file(self, payload: dict[str, Any]) -> dict[str, Any]: async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.post( _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/import-local"), json=payload, ) response.raise_for_status() return _unwrap_dict(response.json()) async def prepare_cache(self, payload: dict[str, Any]) -> dict[str, Any]: async with httpx.AsyncClient(timeout=httpx.Timeout(900, connect=30), headers=self.headers()) as client: response = await client.post( _join_url(self.api_base_url, f"{self.route_prefix}/compute/cache/prepare"), json=payload, ) response.raise_for_status() return _unwrap_dict(response.json()) async def cache_status(self, resource_id: str, version_id: str | None = None) -> dict[str, Any]: params = {"resource_id": resource_id} if version_id: params["version_id"] = version_id async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client: response = await client.get( _join_url(self.api_base_url, f"{self.route_prefix}/compute/cache/status"), params=params, ) response.raise_for_status() return _unwrap_dict(response.json()) async def _request( self, method: str, path: str, json_data: dict[str, Any] | None = None, timeout: float | None = None, ) -> dict[str, Any]: """Generic request method for compute API endpoints.""" url = _join_url(self.api_base_url, f"{self.route_prefix}{path}") async with httpx.AsyncClient(timeout=timeout or 300, 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() return _unwrap_dict(response.json()) # ── Inference helpers (short timeouts — see module constants) ────────── async def inference_load(self, payload: dict[str, Any]) -> dict[str, Any]: """Dispatch a model load. Returns as soon as the node accepts the request; the node now loads asynchronously (status goes 'loading').""" return await self._request("POST", "/inference/load", json_data=payload, timeout=INFERENCE_LOAD_TIMEOUT) async def inference_status(self) -> dict[str, Any]: return await self._request("GET", "/inference/status", timeout=INFERENCE_STATUS_TIMEOUT) async def gpu_resources(self) -> list[dict[str, Any]]: """Read live per-GPU metrics from this compute node.""" return await self.gpus() async def inference_unload(self) -> dict[str, Any]: return await self._request("POST", "/inference/unload", json_data={}, timeout=INFERENCE_UNLOAD_TIMEOUT) async def upload_file( self, filename: str, content: bytes, target_relative_path: str, resource_type: str | None = None, resource_id: str | None = None, ) -> dict[str, Any]: data = { "target_relative_path": target_relative_path, "resource_type": resource_type or "", "resource_id": resource_id or "", } files = {"file": (filename, content)} timeout = httpx.Timeout(max(self.timeout, 60), connect=self.timeout) async with httpx.AsyncClient(timeout=timeout, headers=self.headers()) as client: response = await client.post( _join_url(self.api_base_url, f"{self.route_prefix}/compute/files/upload"), data=data, files=files, ) response.raise_for_status() return _unwrap_dict(response.json()) async def upload_file_to_url(self, path: str, upload_url: str, object_key: str = "", content_type: str = "application/octet-stream") -> dict[str, Any]: return await self._request("POST", "/compute/files/upload-to-url", json_data={ "path": path, "upload_url": upload_url, "object_key": object_key, "content_type": content_type, }, timeout=900)