popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix meshnet-node model and shard flag parsing.

Browse Source 
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:
Dobromir Popov
2026-07-07 17:54:30 +02:00
parent e9a094b620
commit 7e289fef2e
4 changed files with 192 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
packages/node/meshnet_node/cli.py
View File

@@ -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")

27
packages/node/meshnet_node/startup.py
View File

@@ -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,
}