diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index d5863dd..127cd7b 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -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 "" diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index b36a7d9..52fec68 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -1,6 +1,7 @@ """US-035: tracker web dashboard — served from any tracker, embedded asset.""" import json +import time import urllib.request from meshnet_contracts import LocalSolanaContracts @@ -91,3 +92,45 @@ def test_console_endpoint_exposes_tracker_events(): tracker.stop() assert any(event["message"] == "node registered" for event in data["events"]) + + +def test_console_node_lifecycle_events_include_model_health(): + tracker = TrackerServer(heartbeat_timeout=0.05) + port = tracker.start() + try: + body = json.dumps({ + "endpoint": "http://127.0.0.1:9002", + "model": "console-health-test", + "hf_repo": "example/console-health-test", + "num_layers": 4, + "shard_start": 0, + "shard_end": 1, + "hardware_profile": {}, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/nodes/register", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + urllib.request.urlopen(req).read() + + registered = json.loads(urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/console").read()) + registered_event = next( + event for event in registered["events"] + if event["message"] == "node registered" + ) + assert registered_event["fields"]["model_health"]["served_model_copies"] == 0.5 + assert registered_event["fields"]["model_health"]["coverage_percentage"] == 50.0 + + time.sleep(0.06) + urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/network/map").read() + expired = json.loads(urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/console").read()) + expired_event = next( + event for event in expired["events"] + if event["message"] == "node expired" + ) + assert expired_event["fields"]["model_health"]["served_model_copies"] == 0.0 + assert expired_event["fields"]["model_health"]["coverage_percentage"] == 0.0 + finally: + tracker.stop()