diff --git a/.gitignore b/.gitignore index 940dcd4..9816289 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ dist/ !.env.testnet .rocm-local/* billing.sqlite +.pytest-tmp/* \ No newline at end of file diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index b9632fa..6b81c63 100644 --- a/packages/node/meshnet_node/cli.py +++ b/packages/node/meshnet_node/cli.py @@ -90,6 +90,19 @@ def _run_node(cfg: dict) -> None: ) +def _resolve_model_flags( + model: str | None, + model_id: str | None, +) -> tuple[str | None, str | None]: + """Return (model_name, hf_repo_or_none) from --model / --model-id flags.""" + explicit = model_id or model + if not explicit: + return None, None + if "/" in explicit: + return explicit.split("/")[-1], explicit + return explicit, None + + def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int: """Return the first TCP port bindable on host, starting at start.""" bind_host = "" if host == "0.0.0.0" else host @@ -122,9 +135,10 @@ def _cmd_default(args) -> int: # Apply CLI overrides on top of saved config overrides: dict = {} - if args.model: - overrides["model_hf_repo"] = args.model - overrides["model_name"] = args.model.split("/")[-1] + model_name, hf_repo = _resolve_model_flags(args.model, getattr(args, "model_id", None)) + if model_name is not None: + overrides["model_name"] = model_name + overrides["model_hf_repo"] = hf_repo or "" if args.quantization: overrides["quantization"] = args.quantization if args.download_dir: @@ -215,16 +229,15 @@ def _cmd_start(args) -> int: if args.tracker: cfg["tracker_url"] = args.tracker cfg["port"] = args.port if args.port is not None else _first_available_port(args.host) - model = args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or None - if args.model_id is None and model and "/" in model: - cfg["model_hf_repo"] = model - cfg["model_name"] = model.split("/")[-1] - else: - cfg["model_name"] = model + model_name, hf_repo = _resolve_model_flags( + args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or None, + args.model_id, + ) + if model_name is not None: + cfg["model_name"] = model_name + cfg["model_hf_repo"] = hf_repo or "" cfg["quantization"] = args.quantization cfg["host"] = args.host - if args.model_id: - cfg["model_hf_repo"] = args.model_id if args.shard_start is not None: cfg["shard_start"] = args.shard_start if args.shard_end is not None: @@ -242,7 +255,7 @@ def _cmd_start(args) -> int: tracker_url=cfg["tracker_url"], port=cfg["port"], model=cfg["model_name"], - model_id=cfg.get("model_hf_repo"), + model_id=cfg.get("model_hf_repo") or None, shard_start=cfg.get("shard_start"), shard_end=cfg.get("shard_end"), quantization=cfg["quantization"].replace("bf16", "bfloat16"), @@ -288,7 +301,8 @@ def main() -> None: ) # Flags that apply to the no-subcommand (default) path - parser.add_argument("--model", metavar="HF_REPO", help="HuggingFace repo ID to serve") + parser.add_argument("--model", metavar="MODEL", help="Model name or HuggingFace repo ID to serve") + parser.add_argument("--model-id", metavar="MODEL", help="Alias for --model (catalog name or HuggingFace repo)") parser.add_argument("--quantization", "-q", choices=["bf16", "int8", "nf4", "bfloat16"], help="Quantization level") parser.add_argument("--download-dir", metavar="PATH", help="Model download directory") @@ -329,8 +343,8 @@ def main() -> None: start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)") start_cmd.add_argument("--tracker") start_cmd.add_argument("--port", type=int) - start_cmd.add_argument("--model") - start_cmd.add_argument("--model-id", help="HuggingFace repo ID") + start_cmd.add_argument("--model", help="Model name or HuggingFace repo ID") + start_cmd.add_argument("--model-id", help="Alias for --model (catalog name or HuggingFace repo)") start_cmd.add_argument("--shard-start", type=int) start_cmd.add_argument("--shard-end", type=int) start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto") diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index c31cdc3..f5a07ee 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -346,6 +346,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, @@ -383,6 +386,8 @@ def _start_heartbeat( try: resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload) node_id = resp.get("node_id", node_id) + if node_ref is not None: + setattr(node_ref, "tracker_node_id", node_id) return True except Exception: return False @@ -434,7 +439,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) @@ -489,6 +494,34 @@ def _start_heartbeat( return t +def _register_with_tracker( + tracker_url: str, + reg_payload: dict, + node: Any, + start_time: float, +) -> str | None: + """Register with the tracker, or start background retries when it is unreachable.""" + 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=start_time) + return tracker_node_id + except Exception as exc: + setattr(node, "tracker_node_id", None) + print(f" Warning: tracker registration failed: {exc}", flush=True) + print(" [node] will retry registration in the background", flush=True) + _start_heartbeat( + tracker_url, + _PENDING_NODE_ID, + reg_payload, + node_ref=node, + start_time=start_time, + ) + return None + + def _warn_virtual_network_ip(ip: str | None) -> None: """Print a warning when the auto-detected advertise IP is in a known virtual-network range. @@ -648,8 +681,11 @@ def run_startup( if probationary_line is not None: print(f" {probationary_line}", flush=True) + pinned_shard_start = shard_start + pinned_shard_end = shard_end + user_pinned_shard = pinned_shard_start is not None or pinned_shard_end is not None + if model_id: # treat "" the same as None — no explicit model given - user_pinned_shard = shard_start is not None or shard_end is not None full_sources: list[dict] = [] # Auto-detect shard range from model config if not explicitly provided if shard_start is None or shard_end is None: @@ -757,16 +793,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" @@ -784,8 +813,8 @@ def run_startup( flush=True, ) return node - if shard_start is not None or shard_end is not None: - raise ValueError("--shard-start / --shard-end require --model-id") + if user_pinned_shard and not model: + raise ValueError("--shard-start / --shard-end require --model") # 3a. Auto-join: query tracker for network-wide HF model assignment. # Skipped when the user explicitly requested a model — the shard-assignment @@ -892,16 +921,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_label = _format_shard_label( assigned_shard_start, assigned_shard_end, @@ -947,8 +969,13 @@ def run_startup( print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True) raise - shard_start: int = assignment["shard_start"] - shard_end: int = assignment["shard_end"] + shard_start = assignment["shard_start"] + shard_end = assignment["shard_end"] + if user_pinned_shard: + if pinned_shard_start is not None: + shard_start = pinned_shard_start + if pinned_shard_end is not None: + shard_end = pinned_shard_end assigned_model: str = assignment.get("model", model) hf_repo: str | None = assignment.get("hf_repo") peers: list[dict] = assignment.get("peers", []) @@ -960,10 +987,15 @@ def run_startup( if model_layers_end is not None else None ) - print( - f" Shard: {_format_shard_label(shard_start, shard_end, assigned_total_layers, model_name=assigned_model)}", - flush=True, + shard_label = _format_shard_label( + shard_start, + shard_end, + assigned_total_layers, + model_name=assigned_model, ) + if user_pinned_shard: + shard_label = f"{shard_label} (pinned)" + print(f" Shard: {shard_label}", flush=True) # 4. Download shard print("Downloading shard...", flush=True) @@ -1025,6 +1057,7 @@ def run_startup( "hardware_profile": hw, "wallet_address": address, "score": 1.0, + "managed_assignment": not user_pinned_shard, **registration_capabilities, **relay_fields, } diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py index 2161a0a..a27a098 100644 --- a/tests/test_mining_cli.py +++ b/tests/test_mining_cli.py @@ -388,6 +388,107 @@ def test_legacy_start_treats_repo_model_as_model_id(monkeypatch): assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct" +def test_legacy_start_catalog_model_with_pinned_shards(monkeypatch): + """Catalog model names accept --shard-start/--shard-end without --model-id.""" + from meshnet_node.cli import main + + captured = {} + + def fake_run_startup(*args, **kwargs): + captured.update(kwargs) + class _FakeNode: + chat_completion_count = 0 + def stop(self): pass + return _FakeNode() + + monkeypatch.setattr(sys, "argv", [ + "meshnet-node", "start", + "--tracker", "http://192.168.0.179:8080", + "--model", "Qwen3.6-35B-A3B", + "--shard-start", "0", + "--shard-end", "44", + "--port", "0", + ]) + + with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup): + with patch("time.sleep", side_effect=KeyboardInterrupt): + try: + main() + except SystemExit as exc: + assert exc.code == 0 + + assert captured["model"] == "Qwen3.6-35B-A3B" + assert captured["model_id"] is None + assert captured["shard_start"] == 0 + assert captured["shard_end"] == 44 + + +def test_legacy_start_model_id_alias_for_catalog_name(monkeypatch): + """--model-id with a catalog name routes through the tracker preset path.""" + from meshnet_node.cli import main + + captured = {} + + def fake_run_startup(*args, **kwargs): + captured.update(kwargs) + class _FakeNode: + chat_completion_count = 0 + def stop(self): pass + return _FakeNode() + + monkeypatch.setattr(sys, "argv", [ + "meshnet-node", "start", + "--tracker", "http://192.168.0.179:8080", + "--model-id", "Qwen3.6-35B-A3B", + "--port", "0", + ]) + + with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup): + with patch("time.sleep", side_effect=KeyboardInterrupt): + try: + main() + except SystemExit as exc: + assert exc.code == 0 + + assert captured["model"] == "Qwen3.6-35B-A3B" + assert captured["model_id"] is None + + +def test_legacy_start_hf_repo_with_pinned_shards(monkeypatch): + """HF repo --model with pinned shards still enters the torch startup path.""" + from meshnet_node.cli import main + + captured = {} + + def fake_run_startup(*args, **kwargs): + captured.update(kwargs) + class _FakeNode: + chat_completion_count = 0 + def stop(self): pass + return _FakeNode() + + monkeypatch.setattr(sys, "argv", [ + "meshnet-node", "start", + "--tracker", "http://192.168.0.179:8081", + "--model", "Qwen/Qwen2.5-0.5B-Instruct", + "--shard-start", "12", + "--shard-end", "23", + "--port", "0", + ]) + + with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup): + with patch("time.sleep", side_effect=KeyboardInterrupt): + try: + main() + except SystemExit as exc: + assert exc.code == 0 + + assert captured["model"] == "Qwen2.5-0.5B-Instruct" + assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert captured["shard_start"] == 12 + assert captured["shard_end"] == 23 + + def test_legacy_start_falls_back_to_env_tracker_and_model(monkeypatch): """`meshnet-node start` uses env defaults when tracker/model flags are omitted.""" import importlib diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 1c4d57b..3a22a51 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 @@ -1603,6 +1605,120 @@ def test_preset_model_startup_starts_heartbeat(tmp_path, monkeypatch): tracker.stop() +def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch): + """Explicit --shard-start/--shard-end override tracker auto-assignment.""" + import meshnet_node.startup as startup_mod + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16 * 1024}, + ) + heartbeat_calls = [] + monkeypatch.setattr( + startup_mod, + "_start_heartbeat", + lambda *args, **kwargs: heartbeat_calls.append((args, kwargs)), + ) + + tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}}) + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + try: + node = run_startup( + tracker_url=tracker_url, + model="stub-model", + shard_start=0, + shard_end=5, + wallet_path=tmp_path / "wallet.json", + cache_dir=tmp_path / "shards", + ) + try: + assert len(heartbeat_calls) == 1 + args, kwargs = heartbeat_calls[0] + reg_payload = args[2] + assert reg_payload["shard_start"] == 0 + assert reg_payload["shard_end"] == 5 + assert reg_payload["managed_assignment"] is False + finally: + node.stop() + finally: + 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,