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:
@@ -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,26 +171,24 @@ 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)
|
||||
reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": model_id.split("/")[-1],
|
||||
"hf_repo": model_id,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
}
|
||||
try:
|
||||
reg_resp = _post_json(
|
||||
f"{tracker_url}/v1/nodes/register",
|
||||
{
|
||||
"endpoint": endpoint,
|
||||
"model": model_id.split("/")[-1],
|
||||
"hf_repo": model_id,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
},
|
||||
)
|
||||
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,26 +242,24 @@ 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}"
|
||||
auto_reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_hf_repo.split("/")[-1],
|
||||
"hf_repo": assigned_hf_repo,
|
||||
"num_layers": assigned_num_layers,
|
||||
"shard_start": assigned_shard_start,
|
||||
"shard_end": assigned_shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (assigned_shard_start == 0),
|
||||
}
|
||||
try:
|
||||
reg_resp = _post_json(
|
||||
f"{tracker_url}/v1/nodes/register",
|
||||
{
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_hf_repo.split("/")[-1],
|
||||
"hf_repo": assigned_hf_repo,
|
||||
"num_layers": assigned_num_layers,
|
||||
"shard_start": assigned_shard_start,
|
||||
"shard_end": assigned_shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (assigned_shard_start == 0),
|
||||
},
|
||||
)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user