fix: node auto-re-registers with tracker after 404 heartbeat (tracker restart)

When the tracker restarts, nodes' registrations are lost. The heartbeat loop
was catching the 404 and printing a warning but never re-registering, leaving
the node permanently invisible until manually restarted.

Fix: on HTTP 404 heartbeat response, the loop re-posts the original
registration payload to /v1/nodes/register and updates the node_id and
heartbeat URL for subsequent beats. This also handles tracker expiry races.

The register_payload is now passed into _start_heartbeat so the thread has
everything needed for re-registration without reaching back into run_startup.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 01:15:15 +03:00
parent 1e6016e76f
commit d701ae9ba2

View File

@@ -34,15 +34,32 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
return json.loads(r.read())
def _start_heartbeat(tracker_url: str, node_id: str, interval: float = 20.0) -> threading.Thread:
"""Daemon thread that sends periodic heartbeats to the tracker."""
def _start_heartbeat(
tracker_url: str,
node_id: str,
register_payload: dict,
interval: float = 20.0,
) -> threading.Thread:
"""Daemon thread: sends heartbeats; re-registers automatically on 404 (tracker restart)."""
def _loop() -> None:
nonlocal node_id
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 urllib.error.HTTPError as exc:
if exc.code == 404:
print(" [node] tracker lost registration — re-registering...", flush=True)
try:
resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
node_id = resp.get("node_id", node_id)
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
print(f" [node] re-registered — node ID: {node_id}", flush=True)
except Exception as re_exc:
print(f" [node] WARNING: re-registration failed: {re_exc}", flush=True)
else:
print(f" [node] WARNING: heartbeat failed: {exc}", flush=True)
except Exception as exc:
print(f" [node] WARNING: heartbeat failed: {exc}", flush=True)
@@ -154,10 +171,7 @@ def run_startup(
endpoint = f"http://{public_host}:{actual_port}"
# Register with tracker so other nodes can auto-join this model.
total_layers = getattr(node.backend, "total_layers", None)
try:
reg_resp = _post_json(
f"{tracker_url}/v1/nodes/register",
{
reg_payload = {
"endpoint": endpoint,
"model": model_id.split("/")[-1],
"hf_repo": model_id,
@@ -169,11 +183,12 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (shard_start == 0),
},
)
}
try:
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
node_id = reg_resp.get("node_id", "?")
print(f" Registered with tracker — node ID: {node_id}", flush=True)
_start_heartbeat(tracker_url, node_id)
_start_heartbeat(tracker_url, node_id, reg_payload)
except Exception as exc:
print(f" Warning: tracker registration failed: {exc}", flush=True)
@@ -227,10 +242,7 @@ def run_startup(
actual_port = node.start()
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
try:
reg_resp = _post_json(
f"{tracker_url}/v1/nodes/register",
{
auto_reg_payload = {
"endpoint": endpoint,
"model": assigned_hf_repo.split("/")[-1],
"hf_repo": assigned_hf_repo,
@@ -242,11 +254,12 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (assigned_shard_start == 0),
},
)
}
try:
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
node_id = reg_resp.get("node_id", "?")
print(f" Registered with tracker — node ID: {node_id}", flush=True)
_start_heartbeat(tracker_url, node_id)
_start_heartbeat(tracker_url, node_id, auto_reg_payload)
except Exception as exc:
print(f" Warning: tracker registration failed: {exc}", flush=True)
shard_count = assigned_shard_end - assigned_shard_start + 1