class SlurmWebGatewayClient:
r"""基于集中式 Slurm-web gateway 的只读 Slurm 客户端。"""
def __init__(
self,
base_url: str,
*,
cluster: str | Sequence[str],
jwt_key_file: str,
groups: Sequence[str] = (),
audience: str = "slurm-web",
token_ttl_seconds: int = 300,
timeout: float = 10.0,
transport: httpx.AsyncBaseTransport | None = None,
username: str | None = None,
) -> None:
self.base_url = base_url.rstrip("/")
self.clusters = _clusters(cluster)
self.cluster = self.clusters[0]
self.jwt_key_file = jwt_key_file
self.groups = tuple(groups) or tuple(f"cluster.{cluster_name}.user" for cluster_name in self.clusters)
self.audience = audience
self.token_ttl_seconds = token_ttl_seconds
self.timeout = timeout
self.transport = transport
self.username = _text(username)
def with_user(self, username: str) -> SlurmWebGatewayClient:
return SlurmWebGatewayClient(
self.base_url,
cluster=self.clusters,
jwt_key_file=self.jwt_key_file,
groups=self.groups,
audience=self.audience,
token_ttl_seconds=self.token_ttl_seconds,
timeout=self.timeout,
transport=self.transport,
username=username,
)
async def cluster_status(self, *, max_jobs: int = 20, max_unhealthy_nodes: int = 20) -> dict[str, Any]:
ping, nodes, partitions, jobs = await asyncio.gather(
self.ping(),
self.nodes(max_results=10000),
self.partitions(max_results=10000),
self.jobs(max_results=10000),
)
compact_nodes = nodes["nodes"]
compact_partitions = partitions["partitions"]
compact_jobs = jobs["jobs"]
return {
"ping": ping,
"nodes": {
"count": len(compact_nodes),
"by_state": _state_counts(compact_nodes),
"unhealthy": _limit([node for node in compact_nodes if _is_unhealthy_node(node)], max_unhealthy_nodes),
},
"partitions": {
"count": len(compact_partitions),
"by_state": _state_counts(compact_partitions),
"items": _limit(compact_partitions, 50),
},
"jobs": {
"count": len(compact_jobs),
"by_state": _state_counts(compact_jobs),
"items": _limit(compact_jobs, max_jobs),
},
}
async def ping(self) -> dict[str, Any]:
payloads = await self._request_all("ping")
if len(payloads) == 1:
payload = next(iter(payloads.values()))
return payload if isinstance(payload, dict) else {"payload": payload}
return {"clusters": payloads}
async def nodes(self, *, state: str | None = None, max_results: int = 100) -> dict[str, Any]:
payloads = await self._request_all("nodes")
nodes = [
{**_compact_node(item), "cluster": cluster}
for cluster, payload in payloads.items()
for item in _payload_items(payload, "nodes")
]
return _limited_collection("nodes", [node for node in nodes if _matches_state(node, state)], max_results)
async def jobs(
self,
*,
state: str | None = None,
user: str | None = None,
max_results: int = 50,
) -> dict[str, Any]:
requested_user = _text(user) or self._requester_username()
payloads = await self._request_all("jobs")
jobs = [
{**_compact_job(item), "cluster": cluster}
for cluster, payload in payloads.items()
for item in _payload_items(payload, "jobs")
]
filtered = [job for job in jobs if _matches_state(job, state) and _matches_user(job, requested_user)]
return _limited_collection("jobs", filtered, max_results)
async def partitions(self, *, max_results: int = 50) -> dict[str, Any]:
payloads = await self._request_all("partitions")
partitions = [
{**_compact_partition(item), "cluster": cluster}
for cluster, payload in payloads.items()
for item in _payload_items(payload, "partitions")
]
return _limited_collection("partitions", partitions, max_results)
async def _request_all(self, path: str) -> dict[str, Any]:
results = await asyncio.gather(*(self._request(cluster, path) for cluster in self.clusters))
return dict(zip(self.clusters, results, strict=True))
async def _request(self, cluster: str, path: str) -> Any:
token = self._token(self._requester_username())
headers = {"Accept": "application/json", "Authorization": f"Bearer {token}"}
url = f"{self.base_url}/api/agents/{cluster}/{path.lstrip('/')}"
async with httpx.AsyncClient(timeout=self.timeout, headers=headers, transport=self.transport) as client:
response = await client.get(url)
response.raise_for_status()
return response.json()
def _requester_username(self) -> str:
return self.username or _requester_username()
def _token(self, username: str) -> str:
now = int(time.time())
key = Path(self.jwt_key_file).read_text(encoding="utf-8").strip()
return _jwt_hs256(
key,
{
"aud": self.audience,
"sub": username,
"groups": list(self.groups),
"iat": now,
"exp": now + self.token_ttl_seconds,
},
)