Spaces:
Paused
Paused
| """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() | |