Spaces:
Running
Running
File size: 3,252 Bytes
d646dd7 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | 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()
|