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 61e950a..6484c25 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -633,8 +633,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: @@ -773,8 +776,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 @@ -933,14 +936,25 @@ 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", []) model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", []) assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization) - print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True) + if user_pinned_shard: + print( + f" Shard: layers {shard_start}-{shard_end} of {assigned_model} (pinned)", + flush=True, + ) + else: + print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True) # 4. Download shard print("Downloading shard...", flush=True) @@ -1002,6 +1016,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 269d0c4..d70332d 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -1603,6 +1603,47 @@ 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_real_model_startup_registers_downloaded_inventory_without_checksum( tmp_path, monkeypatch,