Merge branch cursor/fix-meshnet-node-param-parsing into master.

Combine shard label formatting with model/shard flag parsing and tracker registration retry. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-07 18:02:01 +02:00
parent b1f08c45cd c38e36f685
commit 1eb1e0baa2
5 changed files with 309 additions and 44 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -19,3 +19,4 @@ dist/
 !.env.testnet
 .rocm-local/*
 billing.sqlite
+.pytest-tmp/*
--- 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")
--- 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,
    }
--- 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
--- 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,