This commit is contained in:
Dobromir Popov
2026-07-07 15:26:12 +02:00
parent 2a0d414593
commit 938a0a721b
2 changed files with 170 additions and 98 deletions

View File

@@ -728,11 +728,11 @@ def _coverage_percentage(
return round((covered / required_layers) * 100, 2)
def _served_model_copies(
nodes: list[_NodeEntry],
required_start: int,
required_end: int,
) -> float:
def _served_model_copies(
nodes: list[_NodeEntry],
required_start: int,
required_end: int,
) -> float:
required_layers = required_end - required_start + 1
if required_layers <= 0:
return 0.0
@@ -751,14 +751,60 @@ def _served_model_copies(
return 0.0
complete_copies = min(layer_counts)
residual_layers = sum(1 for count in layer_counts if count > complete_copies)
return round(complete_copies + (residual_layers / required_layers), 2)
def _preset_layer_bounds(preset: dict) -> tuple[int, int]:
start = int(preset.get("layers_start", 0))
if "layers_end" in preset:
return start, int(preset["layers_end"])
residual_layers = sum(1 for count in layer_counts if count > complete_copies)
return round(complete_copies + (residual_layers / required_layers), 2)
def _model_health_summary(
server: "_TrackerHTTPServer",
model: str | None,
hf_repo: str | None = None,
) -> dict:
lookup = hf_repo or model
if not lookup:
return {}
resolved_name, preset = _resolve_model_preset(server.model_presets, lookup)
if preset is not None:
required_start, required_end = _preset_layer_bounds(preset)
model_nodes = [
node for node in server.registry.values()
if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type]
]
model_id = resolved_name or lookup
else:
model_nodes = [
node for node in server.registry.values()
if _node_matches_model(node, lookup)
and node.shard_start is not None
and node.shard_end is not None
and node.num_layers is not None
]
if not model_nodes:
return {
"model": lookup,
"served_model_copies": 0.0,
"coverage_percentage": 0.0,
"node_count": 0,
}
required_start = 0
required_end = max(node.num_layers for node in model_nodes if node.num_layers is not None) - 1
model_id = hf_repo or model or lookup
return {
"model": model_id,
"required_start": required_start,
"required_end": required_end,
"served_model_copies": _served_model_copies(model_nodes, required_start, required_end),
"coverage_percentage": _coverage_percentage(model_nodes, required_start, required_end),
"node_count": len(model_nodes),
}
def _preset_layer_bounds(preset: dict) -> tuple[int, int]:
start = int(preset.get("layers_start", 0))
if "layers_end" in preset:
return start, int(preset["layers_end"])
return start, start + int(preset["total_layers"]) - 1
@@ -1288,27 +1334,32 @@ def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantiza
}
def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
now = time.monotonic()
expired_ids = [
node_id for node_id, entry in server.registry.items()
if (now - entry.last_heartbeat) > server.heartbeat_timeout
]
for node_id in expired_ids:
entry = server.registry.pop(node_id)
_tracker_log(
server,
"warn",
"node expired",
node_id=node_id,
endpoint=entry.endpoint,
model=entry.model,
shard=f"{entry.shard_start}-{entry.shard_end}",
heartbeat_timeout_seconds=server.heartbeat_timeout,
)
if expired_ids:
_rebalance_all_locked(server)
return expired_ids
def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
now = time.monotonic()
expired_ids = [
node_id for node_id, entry in server.registry.items()
if (now - entry.last_heartbeat) > server.heartbeat_timeout
]
expired_entries: list[tuple[str, _NodeEntry]] = []
for node_id in expired_ids:
entry = server.registry.pop(node_id)
expired_entries.append((node_id, entry))
if expired_ids:
_rebalance_all_locked(server)
for node_id, entry in expired_entries:
_tracker_log(
server,
"warn",
"node expired",
node_id=node_id,
endpoint=entry.endpoint,
model=entry.model,
hf_repo=entry.hf_repo,
shard=f"{entry.shard_start}-{entry.shard_end}",
heartbeat_timeout_seconds=server.heartbeat_timeout,
model_health=_model_health_summary(server, entry.model, entry.hf_repo),
)
return expired_ids
def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
@@ -2170,39 +2221,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
result[model] = server.stats.get_node_model_stats(node.node_id, model)
return result
def model_supply_for(node: _NodeEntry) -> dict:
model = node.model or node.hf_repo
if not model:
return {}
resolved_name, preset = _resolve_model_preset(server.model_presets, model)
if preset is not None:
required_start, required_end = _preset_layer_bounds(preset)
model_nodes = [
n for n in nodes
if _node_matches_preset(n, resolved_name, preset) # type: ignore[arg-type]
]
else:
model_nodes = [
n for n in nodes
if _node_matches_model(n, model)
and n.shard_start is not None
and n.shard_end is not None
and n.num_layers is not None
]
if not model_nodes:
return {}
required_start = 0
required_end = max(n.num_layers for n in model_nodes if n.num_layers is not None) - 1
return {
"model": resolved_name or model,
"required_start": required_start,
"required_end": required_end,
"served_model_copies": _served_model_copies(
model_nodes,
required_start,
required_end,
),
}
def model_supply_for(node: _NodeEntry) -> dict:
return _model_health_summary(server, node.model, node.hf_repo)
self._send_json(200, {
"relay_url": server.relay_url,
@@ -3088,41 +3108,50 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
with server.lock:
self._purge_expired_nodes()
# Dedup: if this endpoint is already registered, remove the old entry first.
stale_ids = [
eid for eid, e in server.registry.items()
if e.endpoint == entry.endpoint.rstrip("/")
]
for eid in stale_ids:
old = server.registry.pop(eid)
_tracker_log(
server,
"info",
"node re-registered",
old_node_id=eid,
node_id=node_id,
endpoint=old.endpoint,
)
server.registry[node_id] = entry
if entry.managed_assignment:
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
stale_ids = [
eid for eid, e in server.registry.items()
if e.endpoint == entry.endpoint.rstrip("/")
]
stale_entries: list[tuple[str, _NodeEntry]] = []
for eid in stale_ids:
old = server.registry.pop(eid)
stale_entries.append((eid, old))
server.registry[node_id] = entry
if entry.managed_assignment:
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
_rebalance_model_locked(server, model)
assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None
if assignment_directive is not None:
entry.pending_directives.clear()
_tracker_log(
server,
"info",
"node registered",
node_id=node_id,
if assignment_directive is not None:
entry.pending_directives.clear()
model_health = _model_health_summary(server, entry.model, entry.hf_repo)
for eid, old in stale_entries:
_tracker_log(
server,
"info",
"node re-registered",
old_node_id=eid,
node_id=node_id,
endpoint=old.endpoint,
old_model=old.model,
old_hf_repo=old.hf_repo,
old_model_health=_model_health_summary(server, old.model, old.hf_repo),
model_health=model_health,
)
_tracker_log(
server,
"info",
"node registered",
node_id=node_id,
endpoint=entry.endpoint,
model=entry.model,
hf_repo=entry.hf_repo,
shard=f"{entry.shard_start}-{entry.shard_end}",
tracker_mode=entry.tracker_mode,
)
hf_repo=entry.hf_repo,
shard=f"{entry.shard_start}-{entry.shard_end}",
tracker_mode=entry.tracker_mode,
model_health=model_health,
)
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
repo_info = f" [{hf_repo}]" if hf_repo else ""