feat(us-016): connection/heartbeat visibility for tracker and node

Tracker now prints a line when a node registers and on every heartbeat
received. Node prints its assigned node_id after successful registration
and starts a daemon heartbeat thread (30s interval) that logs each send.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 00:41:10 +03:00
parent a7cc377d13
commit 97eefd3d5e
2 changed files with 38 additions and 2 deletions

View File

@@ -5,6 +5,8 @@ from __future__ import annotations
import json import json
import socket import socket
import sys import sys
import threading
import time
import urllib.error import urllib.error
import urllib.parse import urllib.parse
import urllib.request import urllib.request
@@ -32,6 +34,23 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
return json.loads(r.read()) return json.loads(r.read())
def _start_heartbeat(tracker_url: str, node_id: str, interval: float = 30.0) -> threading.Thread:
"""Daemon thread that sends periodic heartbeats to the tracker."""
def _loop() -> None:
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
while True:
time.sleep(interval)
try:
_post_json(hb_url, {})
print(f" [node] heartbeat sent → tracker (node {node_id[:8]})", flush=True)
except Exception as exc:
print(f" [node] WARNING: heartbeat failed: {exc}", flush=True)
t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
t.start()
return t
def run_startup( def run_startup(
tracker_url: str, tracker_url: str,
port: int = 0, port: int = 0,
@@ -119,7 +138,7 @@ def run_startup(
# Register with tracker so other nodes can auto-join this model. # Register with tracker so other nodes can auto-join this model.
total_layers = getattr(node.backend, "total_layers", None) total_layers = getattr(node.backend, "total_layers", None)
try: try:
_post_json( reg_resp = _post_json(
f"{tracker_url}/v1/nodes/register", f"{tracker_url}/v1/nodes/register",
{ {
"endpoint": endpoint, "endpoint": endpoint,
@@ -135,6 +154,9 @@ def run_startup(
"tracker_mode": (shard_start == 0), "tracker_mode": (shard_start == 0),
}, },
) )
node_id = reg_resp.get("node_id", "?")
print(f" Registered with tracker — node ID: {node_id}", flush=True)
_start_heartbeat(tracker_url, node_id)
except Exception as exc: except Exception as exc:
print(f" Warning: tracker registration failed: {exc}", flush=True) print(f" Warning: tracker registration failed: {exc}", flush=True)
@@ -187,7 +209,7 @@ def run_startup(
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}"
try: try:
_post_json( reg_resp = _post_json(
f"{tracker_url}/v1/nodes/register", f"{tracker_url}/v1/nodes/register",
{ {
"endpoint": endpoint, "endpoint": endpoint,
@@ -203,6 +225,9 @@ def run_startup(
"tracker_mode": (assigned_shard_start == 0), "tracker_mode": (assigned_shard_start == 0),
}, },
) )
node_id = reg_resp.get("node_id", "?")
print(f" Registered with tracker — node ID: {node_id}", flush=True)
_start_heartbeat(tracker_url, node_id)
except Exception as exc: except Exception as exc:
print(f" Warning: tracker registration failed: {exc}", flush=True) print(f" Warning: tracker registration failed: {exc}", flush=True)
shard_count = assigned_shard_end - assigned_shard_start + 1 shard_count = assigned_shard_end - assigned_shard_start + 1

View File

@@ -656,6 +656,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if assignment_directive is not None: if assignment_directive is not None:
entry.pending_directives.clear() entry.pending_directives.clear()
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 ""
print(
f"[tracker] node registered: {node_id[:8]} {endpoint} {model}{repo_info} {shard_info}",
flush=True,
)
payload = {"node_id": node_id} payload = {"node_id": node_id}
if assignment_directive is not None: if assignment_directive is not None:
payload["directive"] = assignment_directive payload["directive"] = assignment_directive
@@ -673,6 +680,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
_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()
print(
f"[tracker] heartbeat: {node_id[:8]} {entry.endpoint}",
flush=True,
)
if directives: if directives:
self._send_json(200, {"directives": directives}) self._send_json(200, {"directives": directives})
else: else: