Spaces:
Paused
Paused
File size: 2,304 Bytes
da839c2 5dbd0bf da839c2 5dbd0bf da839c2 5dbd0bf da839c2 1a36c6b da839c2 5dbd0bf 1a36c6b 5dbd0bf 1a36c6b 5dbd0bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """RunPod serverless backend: submit a job, read its status."""
from __future__ import annotations
import httpx
from .sim_backend import BackendError
class RunPodError(BackendError):
pass
class RunPodClient:
name = "runpod"
FUNCTIONS = {"run_simulation", "sweep"} # what worker/Dockerfile's handler can run
def __init__(
self, api_key: str, endpoint_id: str, transport: httpx.BaseTransport | None = None, timeout: float = 30.0
):
if not api_key or not endpoint_id:
raise RunPodError("RUNPOD_API_KEY and RUNPOD_ENDPOINT_ID must be set")
self._client = httpx.Client(
base_url=f"https://api.runpod.ai/v2/{endpoint_id}",
headers={"Authorization": f"Bearer {api_key}"},
timeout=timeout,
transport=transport,
)
def supports(self, function: str) -> bool:
return function in self.FUNCTIONS
def submit(self, job_input: dict, execution_timeout_s: int, function: str = "run_simulation") -> dict:
"""POST /run. Returns {"id": ..., "status": "IN_QUEUE"}."""
if not self.supports(function):
raise RunPodError(f"{function} is not available on RunPod (use SIM_BACKEND=modal)")
payload = {
"input": {**job_input, "function": function},
"policy": {"executionTimeout": int(execution_timeout_s * 1000)},
}
try:
r = self._client.post("/run", json=payload)
except httpx.HTTPError as e:
raise RunPodError(f"RunPod unreachable: {e}") from e
if r.status_code != 200:
raise RunPodError(f"RunPod /run returned HTTP {r.status_code}: {r.text[:300]}")
data = r.json()
if "id" not in data:
raise RunPodError(f"unexpected RunPod response: {str(data)[:300]}")
return data
def status(self, job_id: str) -> dict:
"""GET /status/{id}: {"status", "output"?, "error"?, "executionTime"?, "delayTime"?}."""
try:
r = self._client.get(f"/status/{job_id}")
except httpx.HTTPError as e:
raise RunPodError(f"RunPod unreachable: {e}") from e
if r.status_code != 200:
raise RunPodError(f"RunPod /status returned HTTP {r.status_code}: {r.text[:300]}")
return r.json()
|