Fix meshnet-node model and shard flag parsing.
Unify --model and --model-id so catalog names use the tracker path, and allow --shard-start/--shard-end with --model instead of requiring --model-id. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -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:
|
def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int:
|
||||||
"""Return the first TCP port bindable on host, starting at start."""
|
"""Return the first TCP port bindable on host, starting at start."""
|
||||||
bind_host = "" if host == "0.0.0.0" else host
|
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
|
# Apply CLI overrides on top of saved config
|
||||||
overrides: dict = {}
|
overrides: dict = {}
|
||||||
if args.model:
|
model_name, hf_repo = _resolve_model_flags(args.model, getattr(args, "model_id", None))
|
||||||
overrides["model_hf_repo"] = args.model
|
if model_name is not None:
|
||||||
overrides["model_name"] = args.model.split("/")[-1]
|
overrides["model_name"] = model_name
|
||||||
|
overrides["model_hf_repo"] = hf_repo or ""
|
||||||
if args.quantization:
|
if args.quantization:
|
||||||
overrides["quantization"] = args.quantization
|
overrides["quantization"] = args.quantization
|
||||||
if args.download_dir:
|
if args.download_dir:
|
||||||
@@ -215,16 +229,15 @@ def _cmd_start(args) -> int:
|
|||||||
if args.tracker:
|
if args.tracker:
|
||||||
cfg["tracker_url"] = args.tracker
|
cfg["tracker_url"] = args.tracker
|
||||||
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
|
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
|
model_name, hf_repo = _resolve_model_flags(
|
||||||
if args.model_id is None and model and "/" in model:
|
args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or None,
|
||||||
cfg["model_hf_repo"] = model
|
args.model_id,
|
||||||
cfg["model_name"] = model.split("/")[-1]
|
)
|
||||||
else:
|
if model_name is not None:
|
||||||
cfg["model_name"] = model
|
cfg["model_name"] = model_name
|
||||||
|
cfg["model_hf_repo"] = hf_repo or ""
|
||||||
cfg["quantization"] = args.quantization
|
cfg["quantization"] = args.quantization
|
||||||
cfg["host"] = args.host
|
cfg["host"] = args.host
|
||||||
if args.model_id:
|
|
||||||
cfg["model_hf_repo"] = args.model_id
|
|
||||||
if args.shard_start is not None:
|
if args.shard_start is not None:
|
||||||
cfg["shard_start"] = args.shard_start
|
cfg["shard_start"] = args.shard_start
|
||||||
if args.shard_end is not None:
|
if args.shard_end is not None:
|
||||||
@@ -242,7 +255,7 @@ def _cmd_start(args) -> int:
|
|||||||
tracker_url=cfg["tracker_url"],
|
tracker_url=cfg["tracker_url"],
|
||||||
port=cfg["port"],
|
port=cfg["port"],
|
||||||
model=cfg["model_name"],
|
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_start=cfg.get("shard_start"),
|
||||||
shard_end=cfg.get("shard_end"),
|
shard_end=cfg.get("shard_end"),
|
||||||
quantization=cfg["quantization"].replace("bf16", "bfloat16"),
|
quantization=cfg["quantization"].replace("bf16", "bfloat16"),
|
||||||
@@ -288,7 +301,8 @@ def main() -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Flags that apply to the no-subcommand (default) path
|
# 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"],
|
parser.add_argument("--quantization", "-q", choices=["bf16", "int8", "nf4", "bfloat16"],
|
||||||
help="Quantization level")
|
help="Quantization level")
|
||||||
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
|
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 = subparsers.add_parser("start", help="Start node (legacy flags)")
|
||||||
start_cmd.add_argument("--tracker")
|
start_cmd.add_argument("--tracker")
|
||||||
start_cmd.add_argument("--port", type=int)
|
start_cmd.add_argument("--port", type=int)
|
||||||
start_cmd.add_argument("--model")
|
start_cmd.add_argument("--model", help="Model name or HuggingFace repo ID")
|
||||||
start_cmd.add_argument("--model-id", help="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-start", type=int)
|
||||||
start_cmd.add_argument("--shard-end", type=int)
|
start_cmd.add_argument("--shard-end", type=int)
|
||||||
start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
|
start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
|
||||||
|
|||||||
@@ -633,8 +633,11 @@ def run_startup(
|
|||||||
if probationary_line is not None:
|
if probationary_line is not None:
|
||||||
print(f" {probationary_line}", flush=True)
|
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
|
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] = []
|
full_sources: list[dict] = []
|
||||||
# Auto-detect shard range from model config if not explicitly provided
|
# Auto-detect shard range from model config if not explicitly provided
|
||||||
if shard_start is None or shard_end is None:
|
if shard_start is None or shard_end is None:
|
||||||
@@ -773,8 +776,8 @@ def run_startup(
|
|||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
return node
|
return node
|
||||||
if shard_start is not None or shard_end is not None:
|
if user_pinned_shard and not model:
|
||||||
raise ValueError("--shard-start / --shard-end require --model-id")
|
raise ValueError("--shard-start / --shard-end require --model")
|
||||||
|
|
||||||
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
||||||
# Skipped when the user explicitly requested a model — the shard-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)
|
print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
shard_start: int = assignment["shard_start"]
|
shard_start = assignment["shard_start"]
|
||||||
shard_end: int = assignment["shard_end"]
|
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)
|
assigned_model: str = assignment.get("model", model)
|
||||||
hf_repo: str | None = assignment.get("hf_repo")
|
hf_repo: str | None = assignment.get("hf_repo")
|
||||||
peers: list[dict] = assignment.get("peers", [])
|
peers: list[dict] = assignment.get("peers", [])
|
||||||
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
||||||
assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
|
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
|
# 4. Download shard
|
||||||
print("Downloading shard...", flush=True)
|
print("Downloading shard...", flush=True)
|
||||||
@@ -1002,6 +1016,7 @@ def run_startup(
|
|||||||
"hardware_profile": hw,
|
"hardware_profile": hw,
|
||||||
"wallet_address": address,
|
"wallet_address": address,
|
||||||
"score": 1.0,
|
"score": 1.0,
|
||||||
|
"managed_assignment": not user_pinned_shard,
|
||||||
**registration_capabilities,
|
**registration_capabilities,
|
||||||
**relay_fields,
|
**relay_fields,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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"
|
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):
|
def test_legacy_start_falls_back_to_env_tracker_and_model(monkeypatch):
|
||||||
"""`meshnet-node start` uses env defaults when tracker/model flags are omitted."""
|
"""`meshnet-node start` uses env defaults when tracker/model flags are omitted."""
|
||||||
import importlib
|
import importlib
|
||||||
|
|||||||
@@ -1603,6 +1603,47 @@ def test_preset_model_startup_starts_heartbeat(tmp_path, monkeypatch):
|
|||||||
tracker.stop()
|
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(
|
def test_real_model_startup_registers_downloaded_inventory_without_checksum(
|
||||||
tmp_path,
|
tmp_path,
|
||||||
monkeypatch,
|
monkeypatch,
|
||||||
|
|||||||
Reference in New Issue
Block a user