181 lines
6.7 KiB
Python
181 lines
6.7 KiB
Python
|
"""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 {}
|