feat(us-023): heartbeat stats payload, request counters, dynamic reassignment response
Node now sends cumulative stats in heartbeat body: total_requests, failed_requests, queue_depth, uptime_seconds, status Stats are tracked thread-safely in _TorchHTTPServer; buffered locally during outage streak and flushed on next successful heartbeat. Tracker stores stats on _NodeEntry (total_requests, failed_requests, queue_depth, uptime_seconds, status, heartbeats_expected/received) and returns new_assignment in heartbeat response when pending_new_assignment is set. Node logs received new_assignment (hot-reload wired in US-026). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -39,8 +39,29 @@ def _start_heartbeat(
|
|||||||
node_id: str,
|
node_id: str,
|
||||||
register_payload: dict,
|
register_payload: dict,
|
||||||
interval: float = 20.0,
|
interval: float = 20.0,
|
||||||
|
node_ref: Any | None = None,
|
||||||
|
start_time: float | None = None,
|
||||||
) -> threading.Thread:
|
) -> threading.Thread:
|
||||||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts."""
|
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.
|
||||||
|
|
||||||
|
Heartbeat body carries cumulative stats (total_requests, failed_requests,
|
||||||
|
queue_depth, uptime_seconds, status). Stats are buffered locally during
|
||||||
|
outage and flushed on next successful heartbeat.
|
||||||
|
|
||||||
|
Heartbeat response may include new_assignment: {model, shard_start, shard_end}
|
||||||
|
which is logged for now (hot-reload implemented in US-026).
|
||||||
|
"""
|
||||||
|
_start_time = start_time or time.monotonic()
|
||||||
|
|
||||||
|
def _get_stats() -> dict:
|
||||||
|
uptime = time.monotonic() - _start_time
|
||||||
|
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
|
||||||
|
if node_ref is not None:
|
||||||
|
stats["total_requests"] = getattr(node_ref, "total_requests", 0)
|
||||||
|
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
|
||||||
|
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
|
||||||
|
return stats
|
||||||
|
|
||||||
def _reregister() -> bool:
|
def _reregister() -> bool:
|
||||||
nonlocal node_id
|
nonlocal node_id
|
||||||
try:
|
try:
|
||||||
@@ -76,7 +97,15 @@ def _start_heartbeat(
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
try:
|
try:
|
||||||
_post_json(hb_url, {})
|
resp = _post_json(hb_url, _get_stats())
|
||||||
|
new_asgn = resp.get("new_assignment")
|
||||||
|
if new_asgn:
|
||||||
|
print(
|
||||||
|
f" [node] tracker reassignment received: "
|
||||||
|
f"model={new_asgn.get('model')!r} "
|
||||||
|
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
except urllib.error.HTTPError as exc:
|
except urllib.error.HTTPError as exc:
|
||||||
if exc.code == 404:
|
if exc.code == 404:
|
||||||
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
||||||
@@ -205,6 +234,7 @@ def run_startup(
|
|||||||
tracker_url=tracker_url,
|
tracker_url=tracker_url,
|
||||||
route_timeout=route_timeout,
|
route_timeout=route_timeout,
|
||||||
)
|
)
|
||||||
|
_node_start_time = time.monotonic()
|
||||||
actual_port = node.start()
|
actual_port = node.start()
|
||||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||||||
if isinstance(total_layers, int) and total_layers > 0:
|
if isinstance(total_layers, int) and total_layers > 0:
|
||||||
@@ -235,7 +265,7 @@ def run_startup(
|
|||||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||||
setattr(node, "tracker_node_id", tracker_node_id)
|
setattr(node, "tracker_node_id", tracker_node_id)
|
||||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload)
|
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
setattr(node, "tracker_node_id", None)
|
setattr(node, "tracker_node_id", None)
|
||||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||||
@@ -289,6 +319,7 @@ def run_startup(
|
|||||||
tracker_url=tracker_url,
|
tracker_url=tracker_url,
|
||||||
route_timeout=route_timeout,
|
route_timeout=route_timeout,
|
||||||
)
|
)
|
||||||
|
_node_start_time = time.monotonic()
|
||||||
actual_port = node.start()
|
actual_port = node.start()
|
||||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||||
endpoint = f"http://{public_host}:{actual_port}"
|
endpoint = f"http://{public_host}:{actual_port}"
|
||||||
@@ -311,7 +342,7 @@ def run_startup(
|
|||||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||||
setattr(node, "tracker_node_id", tracker_node_id)
|
setattr(node, "tracker_node_id", tracker_node_id)
|
||||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||||
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload)
|
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
setattr(node, "tracker_node_id", None)
|
setattr(node, "tracker_node_id", None)
|
||||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||||
|
|||||||
@@ -46,6 +46,10 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
|||||||
self.tracker_mode = tracker_mode
|
self.tracker_mode = tracker_mode
|
||||||
self.tracker_url = tracker_url
|
self.tracker_url = tracker_url
|
||||||
self.route_timeout = route_timeout
|
self.route_timeout = route_timeout
|
||||||
|
self.total_requests: int = 0
|
||||||
|
self.failed_requests: int = 0
|
||||||
|
self.queue_depth: int = 0
|
||||||
|
self._stats_lock = threading.Lock()
|
||||||
|
|
||||||
|
|
||||||
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||||
@@ -219,6 +223,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
def _handle_chat_completions(self) -> None:
|
def _handle_chat_completions(self) -> None:
|
||||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||||
|
with server._stats_lock:
|
||||||
|
server.total_requests += 1
|
||||||
|
server.queue_depth += 1
|
||||||
|
try:
|
||||||
|
self._do_chat_completions(server)
|
||||||
|
finally:
|
||||||
|
with server._stats_lock:
|
||||||
|
server.queue_depth -= 1
|
||||||
|
|
||||||
|
def _record_failed_request(self) -> None:
|
||||||
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||||
|
with server._stats_lock:
|
||||||
|
server.failed_requests += 1
|
||||||
|
|
||||||
|
def _do_chat_completions(self, server: "_TorchHTTPServer") -> None:
|
||||||
body = self._read_json_body()
|
body = self._read_json_body()
|
||||||
if body is None:
|
if body is None:
|
||||||
return
|
return
|
||||||
@@ -244,6 +263,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
|
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
|
||||||
self._send_openai_response(text, model_name, False, messages)
|
self._send_openai_response(text, model_name, False, messages)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
|
self._record_failed_request()
|
||||||
self._send_json(500, {"error": f"generation failed: {exc}"})
|
self._send_json(500, {"error": f"generation failed: {exc}"})
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -620,6 +640,18 @@ class TorchNodeServer:
|
|||||||
def forward_chunk_count(self) -> int:
|
def forward_chunk_count(self) -> int:
|
||||||
return self._server.forward_chunk_count if self._server is not None else 0
|
return self._server.forward_chunk_count if self._server is not None else 0
|
||||||
|
|
||||||
|
@property
|
||||||
|
def total_requests(self) -> int:
|
||||||
|
return self._server.total_requests if self._server is not None else 0
|
||||||
|
|
||||||
|
@property
|
||||||
|
def failed_requests(self) -> int:
|
||||||
|
return self._server.failed_requests if self._server is not None else 0
|
||||||
|
|
||||||
|
@property
|
||||||
|
def queue_depth(self) -> int:
|
||||||
|
return self._server.queue_depth if self._server is not None else 0
|
||||||
|
|
||||||
def start(self) -> int:
|
def start(self) -> int:
|
||||||
if self._server is not None:
|
if self._server is not None:
|
||||||
raise RuntimeError("TorchNodeServer is already running")
|
raise RuntimeError("TorchNodeServer is already running")
|
||||||
|
|||||||
@@ -61,6 +61,12 @@ class _NodeEntry:
|
|||||||
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
||||||
"pending_directives", "last_heartbeat", "tracker_mode",
|
"pending_directives", "last_heartbeat", "tracker_mode",
|
||||||
"relay_addr", "cert_fingerprint", "peer_id",
|
"relay_addr", "cert_fingerprint", "peer_id",
|
||||||
|
# heartbeat stats (reported by node, cumulative)
|
||||||
|
"total_requests", "failed_requests", "queue_depth", "uptime_seconds",
|
||||||
|
"status", # "ready" | "loading"
|
||||||
|
"heartbeats_expected", "heartbeats_received",
|
||||||
|
# dynamic reassignment queued by the tracker
|
||||||
|
"pending_new_assignment",
|
||||||
)
|
)
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -110,6 +116,14 @@ class _NodeEntry:
|
|||||||
self.peer_id = peer_id
|
self.peer_id = peer_id
|
||||||
self.pending_directives: list[dict] = []
|
self.pending_directives: list[dict] = []
|
||||||
self.last_heartbeat: float = time.monotonic()
|
self.last_heartbeat: float = time.monotonic()
|
||||||
|
self.total_requests: int = 0
|
||||||
|
self.failed_requests: int = 0
|
||||||
|
self.queue_depth: int = 0
|
||||||
|
self.uptime_seconds: float = 0.0
|
||||||
|
self.status: str = "ready"
|
||||||
|
self.heartbeats_expected: int = 0
|
||||||
|
self.heartbeats_received: int = 0
|
||||||
|
self.pending_new_assignment: dict | None = None
|
||||||
|
|
||||||
|
|
||||||
def _select_route(
|
def _select_route(
|
||||||
@@ -1045,21 +1059,39 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._send_json(404, {"error": "node not found"})
|
self._send_json(404, {"error": "node not found"})
|
||||||
return
|
return
|
||||||
entry.last_heartbeat = time.monotonic()
|
entry.last_heartbeat = time.monotonic()
|
||||||
|
entry.heartbeats_received += 1
|
||||||
|
# P2P metadata
|
||||||
if body.get("relay_addr"):
|
if body.get("relay_addr"):
|
||||||
entry.relay_addr = body["relay_addr"]
|
entry.relay_addr = body["relay_addr"]
|
||||||
if body.get("cert_fingerprint"):
|
if body.get("cert_fingerprint"):
|
||||||
entry.cert_fingerprint = body["cert_fingerprint"]
|
entry.cert_fingerprint = body["cert_fingerprint"]
|
||||||
if body.get("peer_id"):
|
if body.get("peer_id"):
|
||||||
entry.peer_id = body["peer_id"]
|
entry.peer_id = body["peer_id"]
|
||||||
|
# Node stats (cumulative — node always sends totals, tracker replaces)
|
||||||
|
if "total_requests" in body:
|
||||||
|
entry.total_requests = int(body["total_requests"])
|
||||||
|
if "failed_requests" in body:
|
||||||
|
entry.failed_requests = int(body["failed_requests"])
|
||||||
|
if "queue_depth" in body:
|
||||||
|
entry.queue_depth = int(body["queue_depth"])
|
||||||
|
if "uptime_seconds" in body:
|
||||||
|
entry.uptime_seconds = float(body["uptime_seconds"])
|
||||||
|
if "status" in body and body["status"] in ("ready", "loading"):
|
||||||
|
entry.status = body["status"]
|
||||||
_rebalance_model_locked(server, entry.model or "stub-model")
|
_rebalance_model_locked(server, entry.model or "stub-model")
|
||||||
directives = list(entry.pending_directives)
|
directives = list(entry.pending_directives)
|
||||||
entry.pending_directives.clear()
|
entry.pending_directives.clear()
|
||||||
|
new_assignment = entry.pending_new_assignment
|
||||||
|
if new_assignment is not None:
|
||||||
|
entry.pending_new_assignment = None
|
||||||
if server.gossip is not None:
|
if server.gossip is not None:
|
||||||
server.gossip.record(node_id)
|
server.gossip.record(node_id)
|
||||||
|
resp: dict = {}
|
||||||
if directives:
|
if directives:
|
||||||
self._send_json(200, {"directives": directives})
|
resp["directives"] = directives
|
||||||
else:
|
if new_assignment is not None:
|
||||||
self._send_json(200, {})
|
resp["new_assignment"] = new_assignment
|
||||||
|
self._send_json(200, resp)
|
||||||
|
|
||||||
# ---------------------------------------------------------------- Raft handlers
|
# ---------------------------------------------------------------- Raft handlers
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user