diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index 9d45754..bee0779 100644 --- a/packages/node/meshnet_node/cli.py +++ b/packages/node/meshnet_node/cli.py @@ -27,6 +27,7 @@ def _run_node(cfg: dict) -> None: cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None, host=cfg.get("host", "0.0.0.0"), route_timeout=float(cfg.get("route_timeout", 30.0)), + vram_mb_override=cfg.get("vram_mb_override"), ) except Exception as exc: print(f"\nERROR: {exc}", file=sys.stderr, flush=True) @@ -89,6 +90,8 @@ def _cmd_default(args) -> int: overrides["host"] = args.host if args.route_timeout != 30.0: overrides["route_timeout"] = args.route_timeout + if getattr(args, "memory", None) is not None: + overrides["vram_mb_override"] = args.memory if overrides: cfg = merge_cli_overrides(cfg, **overrides) @@ -177,6 +180,7 @@ def _cmd_start(args) -> int: host=cfg["host"], advertise_host=getattr(args, "advertise_host", None), route_timeout=getattr(args, "route_timeout", 30.0), + vram_mb_override=getattr(args, "memory", None), ) except Exception as exc: print(f"ERROR: {exc}", file=sys.stderr, flush=True) @@ -218,6 +222,8 @@ def main() -> None: parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)") parser.add_argument("--route-timeout", type=float, metavar="SEC", default=30.0, help="Seconds to wait for tracker route lookup (default 30)") + parser.add_argument("--memory", type=int, metavar="MB", default=None, + help="Override autodetected VRAM/RAM budget in MB used for shard assignment") parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)") parser.add_argument("--compact", action="store_true", help="Single-line status output") parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists") @@ -248,6 +254,8 @@ def main() -> None: start_cmd.add_argument("--download-dir") start_cmd.add_argument("--route-timeout", type=float, default=30.0, help="Seconds to wait for tracker route lookup (default 30)") + start_cmd.add_argument("--memory", type=int, default=None, metavar="MB", + help="Override autodetected VRAM/RAM budget in MB used for shard assignment") args = parser.parse_args() diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index ae67d1b..876c7ce 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -141,6 +141,7 @@ def run_startup( advertise_host: str | None = None, contracts: Any | None = None, route_timeout: float = 30.0, + vram_mb_override: int | None = None, ) -> StubNodeServer | TorchNodeServer: """Execute the full startup sequence and return a running node server. @@ -177,7 +178,10 @@ def run_startup( gpu_name: str | None = hw.get("gpu_name") vram_mb: int = hw.get("vram_mb", 0) - if device == "cpu": + if vram_mb_override is not None: + vram_mb = vram_mb_override + print(f" Memory budget overridden to {vram_mb} MB via --memory", flush=True) + elif device == "cpu": print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True) else: print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True) diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index e96aa72..a63e8f5 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -1550,8 +1550,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]: repo_layers[n.hf_repo] = n.num_layers - # Pick the repo where the gap is largest (most work to do). + # Smart scoring: demand_rpm × coverage_deficit + # coverage_deficit = uncovered_layers / total_layers (0 = fully covered) + # demand_rpm comes from the stats collector; defaults to 0.0 when no traffic yet. + demand_rpms: dict[str, float] = {} + if server.stats is not None: + local_rpms = server.stats.get_local_rpms() + for repo in repo_groups: + demand_rpms[repo] = local_rpms.get(repo, {}).get("rpm_last_hour", 0.0) + best_repo = None + best_score = -1.0 best_gap_size = -1 best_gap_start = 0 best_num_layers = 0 @@ -1569,8 +1578,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): gap_start = max(gap_start, e + 1) else: break - gap_size = max(0, (total - 1) - gap_start + 1) # layers remaining uncovered - if gap_size > best_gap_size: + gap_size = max(0, (total - 1) - gap_start + 1) + coverage_deficit = gap_size / max(total, 1) + demand = demand_rpms.get(repo, 0.0) + # +1.0 floor on demand so models with no recorded traffic still compete by coverage. + score = (demand + 1.0) * (coverage_deficit + 0.01) + if score > best_score or (score == best_score and gap_size > best_gap_size): + best_score = score best_gap_size = gap_size best_gap_start = gap_start best_repo = repo @@ -1578,8 +1592,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): gap_found = best_gap_size > 0 if not gap_found: - # All shards are covered — still assign to the model with most nodes for redundancy. - best_repo = max(repo_groups, key=lambda r: len(repo_groups[r])) + # All shards covered — assign to highest-demand model for redundancy. + best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0)) best_gap_start = 0 best_num_layers = repo_layers[best_repo] @@ -1599,6 +1613,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "shard_end": shard_end, "num_layers": total_l, "gap_found": gap_found, + "price_per_token": 0.0, }) def _handle_route(self, parsed: urllib.parse.ParseResult):