Expose node shard capacity limits

This commit is contained in:
Dobromir Popov
2026-06-30 21:29:30 +02:00
parent 94d2216a28
commit b06b86a6f0
3 changed files with 84 additions and 0 deletions

View File

@@ -438,6 +438,21 @@ def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
return int((node.vram_bytes * 0.8) // layer_bytes)
def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict:
"""Operator-facing capacity fields for inspection endpoints."""
summary = {
"vram_bytes": node.vram_bytes,
"ram_bytes": node.ram_bytes,
"max_loaded_shards": node.max_loaded_shards,
"quantizations": list(node.quantizations),
"quantization": node.quantization,
"benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
}
if preset is not None:
summary["max_assignable_layers"] = _node_layer_capacity(node, preset)
return summary
def _coverage_map(
nodes: list[_NodeEntry],
required_start: int,
@@ -480,6 +495,7 @@ def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict:
"total_requests": node.total_requests,
"heartbeat_success_rate": hb_rate,
"inference_success_rate": inf_rate,
"capacity": _node_capacity_summary(node),
}
@@ -1114,6 +1130,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
with server.lock:
self._purge_expired_nodes()
nodes = list(server.registry.values())
def capacity_for(node: _NodeEntry) -> dict:
preset = None
if node.model:
preset = server.model_presets.get(node.model)
if preset is None and node.hf_repo and node.num_layers:
preset = _hf_rebalance_preset([node])
return _node_capacity_summary(node, preset)
self._send_json(200, {
"relay_url": server.relay_url,
"nodes": [
@@ -1128,6 +1153,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"shard_end": node.shard_end,
"tracker_mode": node.tracker_mode,
"last_heartbeat": node.last_heartbeat,
"capacity": capacity_for(node),
}
for node in nodes
],