diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index adb4144..f15a187 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -331,6 +331,9 @@ def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHt node.stop = _stop_with_bridge # type: ignore[method-assign] +_PENDING_NODE_ID = "pending" + + def _start_heartbeat( tracker_url: str, node_id: str, @@ -421,7 +424,7 @@ def _start_heartbeat( def _loop() -> None: nonlocal node_id hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat" - outage_streak = 0 # consecutive intervals where tracker was unreachable + outage_streak = 1 if node_id == _PENDING_NODE_ID else 0 while True: time.sleep(interval) @@ -476,9 +479,6 @@ def _start_heartbeat( return t -_PENDING_NODE_ID = "pending" - - def _register_with_tracker( tracker_url: str, reg_payload: dict, @@ -782,16 +782,9 @@ def run_startup( **registration_capabilities, **relay_fields, } - tracker_node_id: str | None = None - try: - reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload) - tracker_node_id = str(reg_resp.get("node_id") or "?") - setattr(node, "tracker_node_id", tracker_node_id) - print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) - _start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time) - except Exception as exc: - setattr(node, "tracker_node_id", None) - print(f" Warning: tracker registration failed: {exc}", flush=True) + tracker_node_id = _register_with_tracker( + tracker_url, reg_payload, node, _node_start_time, + ) print( f"\n{'=' * 32}\n" @@ -917,16 +910,9 @@ def run_startup( **registration_capabilities, **relay_fields, } - tracker_node_id = None - try: - reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload) - tracker_node_id = str(reg_resp.get("node_id") or "?") - setattr(node, "tracker_node_id", tracker_node_id) - print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) - _start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time) - except Exception as exc: - setattr(node, "tracker_node_id", None) - print(f" Warning: tracker registration failed: {exc}", flush=True) + tracker_node_id = _register_with_tracker( + tracker_url, auto_reg_payload, node, _node_start_time, + ) shard_count = assigned_shard_end - assigned_shard_start + 1 print( f"\n{'=' * 32}\n" diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index d70332d..b7a67be 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -5,8 +5,10 @@ import io import os import sys import tarfile +import threading import time import types +import urllib.error import urllib.request from pathlib import Path @@ -1644,6 +1646,79 @@ def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch): tracker.stop() +def test_torch_startup_retries_registration_when_tracker_unreachable( + tmp_path, + monkeypatch, +): + """Failed initial registration should start background retry, not stay unregistered.""" + import meshnet_node.startup as startup_mod + + class FakeBackend: + total_layers = 24 + + class FakeTorchNodeServer: + def __init__(self, **kwargs): + self.backend = FakeBackend() + self.port = None + self.chat_completion_count = 0 + self.tracker_node_id = None + + def start(self): + self.port = 7000 + return self.port + + def stop(self): + pass + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cuda", "gpu_name": "Test GPU", "vram_mb": 8192, "ram_mb": 16 * 1024}, + ) + monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer) + monkeypatch.setattr( + startup_mod, + "_detect_num_layers", + lambda *_args, **_kwargs: 24, + ) + + heartbeat_calls = [] + monkeypatch.setattr( + startup_mod, + "_start_heartbeat", + lambda *args, **kwargs: heartbeat_calls.append((args, kwargs)) or threading.Thread(), + ) + + register_calls = {"count": 0} + + def flaky_register(url, payload): + register_calls["count"] += 1 + raise urllib.error.URLError("connection refused") + + monkeypatch.setattr(startup_mod, "_post_json", flaky_register) + + tracker = TrackerServer() + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + try: + node = run_startup( + tracker_url=tracker_url, + model_id="Qwen/Qwen2.5-0.5B-Instruct", + wallet_path=tmp_path / "wallet.json", + ) + try: + assert register_calls["count"] == 1 + assert node.tracker_node_id is None + assert len(heartbeat_calls) == 1 + args, kwargs = heartbeat_calls[0] + assert args[1] == startup_mod._PENDING_NODE_ID + assert kwargs["node_ref"] is node + finally: + node.stop() + finally: + tracker.stop() + + def test_real_model_startup_registers_downloaded_inventory_without_checksum( tmp_path, monkeypatch,