Spaces:
Running
Running
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| import httpx | |
| class WorkerClient: | |
| def __init__(self, server_url: str, api_key: str, timeout: float = 60.0): | |
| # trust_env=False: the worker talks to the local backend directly. | |
| # On Windows httpx would otherwise honor the system proxy | |
| # (e.g. Clash on 127.0.0.1:7897), which cannot reach localhost and | |
| # turns every claim into a 502 Bad Gateway. | |
| self.client = httpx.Client( | |
| base_url=server_url.rstrip("/"), | |
| headers={"X-API-Key": api_key}, | |
| timeout=httpx.Timeout(timeout, connect=10.0), | |
| trust_env=False, | |
| ) | |
| def close(self): | |
| self.client.close() | |
| def claim(self, worker_id: str, capabilities: list[str], device: str) -> dict | None: | |
| response = self.client.post( | |
| "/api/v1/worker/jobs/claim", | |
| json={"worker_id": worker_id, "capabilities": capabilities, "device": device}, | |
| ) | |
| if response.status_code == 204: | |
| return None | |
| response.raise_for_status() | |
| return response.json() | |
| def download_input(self, run_id: str, destination: Path) -> None: | |
| # A cold input package may need to scan the NAS/SSHFS mount. The | |
| # normal 60-second polling timeout must not abort a valid upload. | |
| download_timeout = httpx.Timeout(None, connect=10.0) | |
| with self.client.stream( | |
| "GET", f"/api/v1/worker/jobs/{run_id}/input", | |
| timeout=download_timeout, | |
| ) as response: | |
| response.raise_for_status() | |
| with destination.open("wb") as output: | |
| for chunk in response.iter_bytes(1024 * 1024): | |
| output.write(chunk) | |
| def heartbeat(self, run_id: str, worker_id: str, lease_token: str, | |
| progress: float, node_states: dict): | |
| response = self.client.post( | |
| f"/api/v1/worker/jobs/{run_id}/heartbeat", | |
| json={"worker_id": worker_id, "lease_token": lease_token, | |
| "progress": progress, "node_states": node_states}, | |
| ) | |
| response.raise_for_status() | |
| def complete(self, run_id: str, worker_id: str, lease_token: str, | |
| result_zip: Path, node_states: dict, outputs: dict): | |
| with result_zip.open("rb") as source: | |
| response = self.client.post( | |
| f"/api/v1/worker/jobs/{run_id}/complete", | |
| data={ | |
| "worker_id": worker_id, | |
| "lease_token": lease_token, | |
| "node_states": json.dumps(node_states), | |
| "outputs": json.dumps(outputs), | |
| }, | |
| files={"result_zip": (result_zip.name, source, "application/zip")}, | |
| timeout=None, | |
| ) | |
| response.raise_for_status() | |
| return response.json() | |
| def fail(self, run_id: str, worker_id: str, lease_token: str, | |
| error: str, retry: bool = True): | |
| response = self.client.post( | |
| f"/api/v1/worker/jobs/{run_id}/fail", | |
| json={"worker_id": worker_id, "lease_token": lease_token, | |
| "error": error, "retry": retry}, | |
| ) | |
| response.raise_for_status() | |