From 1299a6bb1c65c0864382e47e314f5f865681c88d Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 7 Jul 2026 18:30:30 +0300 Subject: [PATCH] balancing improvements --- billing.sqlite | Bin 24576 -> 24576 bytes meshnet_registry.sqlite3 | Bin 12288 -> 12288 bytes packages/node/meshnet_node/cli.py | 6 +- packages/node/meshnet_node/startup.py | 39 +++- packages/tracker/meshnet_tracker/server.py | 255 ++++++++++++++++++++- tests/test_node_startup.py | 78 +++++++ 6 files changed, 358 insertions(+), 20 deletions(-) diff --git a/billing.sqlite b/billing.sqlite index 3452dccd8f28a2dc80fef64fd3b6567ddef4ef04..eec52b48119325239da5cdb198f48228ddf7ec9e 100644 GIT binary patch delta 267 zcmZoTz}Rqrae_3X_(U0JM)8daOXQg)xF1etS5Rf*ez;jjVKet+UL9`5RP!VwV@opw z!xVEv6JvABBr}8L)Ra^ca|08Tv{b`niOsw^-<3I-cx4#)Ch*E^EEM34Ybumrjc#RU zPtHgzN>9}#WJI-6W{Q%Pf)bI+CqFdFnmo^4mCZ&HjK;e>NprE0-g|UgTg`uI5o`Hp#nVGS9t>dl9ZYNl~pKYEHSHK7W*Pu&v delta 56 zcmV-80LTA;zyW~30gxL37Lgo70T!`fq#pwp3=NY9ASJUdAhQdRfG)EQEcYV@0S=G= O4v?{Nh7PlkUT^{56%j!I diff --git a/meshnet_registry.sqlite3 b/meshnet_registry.sqlite3 index 4cfeef0d1dc72c27f44098d5c81468c5349505c8..d5788e01d0f5907b24baec52c54a5d5ad1f799c4 100644 GIT binary patch delta 244 zcmZojXh@hK%_uZc#+gxQW5N=CW+8@Uli3AS7?({ftcqw17h(-Q9G)$~k%1lwRQcxmB?PNRYY;hA~3qwOAJp&6f zGh-8@TE| None: node = run_startup( tracker_url=cfg["tracker_url"], port=cfg.get("port", 7000), - model=cfg.get("model_name") or "stub-model", + model=cfg.get("model_name") or None, model_id=cfg.get("model_hf_repo") or None, shard_start=cfg.get("shard_start"), shard_end=cfg.get("shard_end"), @@ -215,8 +215,8 @@ 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 "stub-model" - if args.model_id is None and "/" in model: + 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: diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index e71414f..a184e9b 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -482,7 +482,7 @@ def _warn_virtual_network_ip(ip: str | None) -> None: def run_startup( tracker_url: str, port: int = 0, - model: str = "stub-model", + model: str | None = None, model_id: str | None = None, shard_start: int | None = None, shard_end: int | None = None, @@ -755,8 +755,9 @@ def run_startup( # query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh # tracker would otherwise print a scary 503 for the model-less auto-join. net_assignment: dict = {} - if model and model != "stub-model": - print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True) + if model_id or (model and model != "stub-model"): + if model: + print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True) else: print("Querying tracker for network assignment...", flush=True) assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb}) @@ -767,17 +768,25 @@ def run_startup( assigned_hf_repo: str | None = net_assignment.get("hf_repo") _gap_found: bool = bool(net_assignment.get("gap_found", False)) - if assigned_hf_repo and _gap_found: + if assigned_hf_repo: assigned_shard_start: int = net_assignment["shard_start"] assigned_shard_end: int = net_assignment["shard_end"] assigned_num_layers: int = net_assignment["num_layers"] assigned_model_sources: list[dict] = net_assignment.get("model_sources", []) - print( - f" Assigned: {assigned_hf_repo} " - f"layers {assigned_shard_start}–{assigned_shard_end} " - f"(of {assigned_num_layers})", - flush=True, - ) + if _gap_found: + print( + f" Assigned gap: {assigned_hf_repo} " + f"layers {assigned_shard_start}–{assigned_shard_end} " + f"(of {assigned_num_layers})", + flush=True, + ) + else: + print( + f" Assigned redundant copy: {assigned_hf_repo} " + f"layers {assigned_shard_start}–{assigned_shard_end} " + f"(of {assigned_num_layers})", + flush=True, + ) full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources) if full_sources: print("Downloading assigned model snapshot...", flush=True) @@ -874,10 +883,16 @@ def run_startup( ) return node - # 3b. Shard assignment from tracker (stub-model / preset-based path) + if model != "stub-model": + raise RuntimeError( + "Tracker did not assign a model. Join a network that already serves one, " + "or start with --model ." + ) + + # 3b. Stub preset path (tests / explicit stub-model only) print("Querying tracker for shard assignment...", flush=True) assign_qs = urllib.parse.urlencode({ - "model": model, + "model": model or "stub-model", "device": device, # CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's # VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 6e61b25..69b77cc 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -944,6 +944,229 @@ def _pool_summary(nodes: list[_NodeEntry], preset: dict | None = None) -> dict: } +def _preset_for_node(server: "_TrackerHTTPServer", node: _NodeEntry) -> dict | None: + if node.hf_repo: + _resolved, preset = _resolve_model_preset(server.model_presets, node.hf_repo) + if preset is not None: + return preset + if node.model: + preset = server.model_presets.get(node.model) + if preset is not None: + return preset + if node.hf_repo and node.num_layers: + return _hf_rebalance_preset([node]) + return None + + +def _assignment_memory_bytes(node: _NodeEntry, preset: dict | None) -> int: + if preset is None or node.shard_start is None or node.shard_end is None: + return 0 + layers = node.shard_end - node.shard_start + 1 + if layers <= 0: + return 0 + quantization = _node_quantization(node, preset) + layer_bytes = _preset_bytes_per_layer(preset).get(quantization) + if not layer_bytes: + layer_bytes = next(iter(_preset_bytes_per_layer(preset).values())) + return layers * int(layer_bytes) + + +def _endpoint_memory_pool( + server: "_TrackerHTTPServer", + nodes: list[_NodeEntry], +) -> dict: + """Per-host RAM/VRAM budget, usage, spare capacity, and loaded assignments.""" + if not nodes: + return {} + host = nodes[0] + budget_bytes, budget_source = _node_memory_budget_bytes(host) + reserve_bytes = int(budget_bytes * 0.8) + loaded: list[dict] = [] + used_bytes = 0 + for node in nodes: + preset = _preset_for_node(server, node) + assignment_bytes = _assignment_memory_bytes(node, preset) + used_bytes += assignment_bytes + loaded.append({ + "node_id": node.node_id, + "model": node.model, + "hf_repo": node.hf_repo, + "shard_start": node.shard_start, + "shard_end": node.shard_end, + "memory_bytes": assignment_bytes, + "status": node.status, + }) + max_slots = max(node.max_loaded_shards for node in nodes) + loaded_slots = len([node for node in nodes if node.shard_start is not None]) + spare_bytes = max(0, reserve_bytes - used_bytes) + return { + "endpoint": host.endpoint, + "memory_budget_bytes": budget_bytes, + "memory_budget_source": budget_source, + "memory_reserve_bytes": reserve_bytes, + "memory_used_bytes": used_bytes, + "memory_spare_bytes": spare_bytes, + "loaded_slots": loaded_slots, + "max_loaded_shards": max_slots, + "spare_slots": max(0, max_slots - loaded_slots), + "loaded": loaded, + } + + +def _memory_pool_map(server: "_TrackerHTTPServer") -> dict: + """Aggregate and per-endpoint view of assignable RAM across the hive.""" + from collections import defaultdict + + groups: dict[str, list[_NodeEntry]] = defaultdict(list) + for node in server.registry.values(): + groups[node.endpoint.rstrip("/")].append(node) + + hosts = [_endpoint_memory_pool(server, group) for group in groups.values()] + hosts.sort(key=lambda item: (-item.get("memory_spare_bytes", 0), item.get("endpoint", ""))) + total_budget = sum(item["memory_budget_bytes"] for item in hosts) + total_used = sum(item["memory_used_bytes"] for item in hosts) + total_spare = sum(item["memory_spare_bytes"] for item in hosts) + total_slots = sum(item["max_loaded_shards"] for item in hosts) + loaded_slots = sum(item["loaded_slots"] for item in hosts) + return { + "total_memory_budget_bytes": total_budget, + "total_memory_used_bytes": total_used, + "total_memory_spare_bytes": total_spare, + "total_loaded_slots": loaded_slots, + "total_max_loaded_shards": total_slots, + "total_spare_slots": max(0, total_slots - loaded_slots), + "hosts": hosts, + } + + +def _add_shard_directive( + node: _NodeEntry, + model: str, + start: int, + end: int, + quantization: str, + *, + model_sources: list[dict] | None = None, +) -> dict: + directive = { + "action": "ADD_SHARD", + "model": model, + "start_layer": start, + "end_layer": end, + "shard_start": start, + "shard_end": end, + "quantization": quantization, + } + if model_sources: + directive["model_sources"] = model_sources + return directive + + +def _model_demand_and_supply( + server: "_TrackerHTTPServer", + model_key: str, + preset: dict, +) -> dict: + resolved_name, _ = _resolve_model_preset(server.model_presets, model_key) + lookup = preset.get("hf_repo") or resolved_name or model_key + required_start, required_end = _preset_layer_bounds(preset) + model_nodes = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name or model_key, preset) # type: ignore[arg-type] + ] + health = _model_health_summary(server, model_key, preset.get("hf_repo")) + local_rpms = server.stats.get_local_rpms() if server.stats is not None else {} + demand_rpm = _repo_demand_rpm(server, str(lookup), local_rpms) + return { + "model": resolved_name or model_key, + "hf_repo": preset.get("hf_repo"), + "preset": preset, + "demand_rpm": demand_rpm, + "served_model_copies": float(health.get("served_model_copies") or 0.0), + "coverage_percentage": float(health.get("coverage_percentage") or 0.0), + "required_start": required_start, + "required_end": required_end, + "model_nodes": model_nodes, + } + + +def _scale_demanded_models_locked(server: "_TrackerHTTPServer") -> None: + """Use spare host RAM/slots to add copies of high-demand models.""" + pool = _memory_pool_map(server) + candidates: list[tuple[float, str, dict]] = [] + for model_key, preset in server.model_presets.items(): + if not isinstance(preset, dict) or not preset.get("hf_repo"): + continue + info = _model_demand_and_supply(server, model_key, preset) + demand = info["demand_rpm"] + copies = info["served_model_copies"] + coverage = info["coverage_percentage"] + if coverage < 100.0: + continue + target_copies = max(1.0, demand / 5.0) + if copies >= target_copies: + continue + candidates.append((demand, model_key, info)) + + if not candidates: + return + + candidates.sort(key=lambda item: (-item[0], item[2]["served_model_copies"])) + for _demand, model_key, info in candidates: + preset = info["preset"] + hf_repo = str(info["hf_repo"]) + required_start = info["required_start"] + required_end = info["required_end"] + total_layers = required_end - required_start + 1 + layer_bytes = next(iter(_preset_bytes_per_layer(preset).values())) + full_copy_bytes = total_layers * layer_bytes + + for host in pool["hosts"]: + if host["spare_slots"] <= 0 or host["memory_spare_bytes"] < full_copy_bytes: + continue + host_nodes = [ + server.registry[node["node_id"]] + for node in host["loaded"] + if node["node_id"] in server.registry + ] + if not host_nodes: + continue + if any((n.hf_repo or n.model) == hf_repo for n in host_nodes): + continue + anchor = max(host_nodes, key=lambda n: n.benchmark_tokens_per_sec) + if anchor.status != "ready" or anchor.pending_new_assignment is not None: + continue + capacity = min(_node_layer_capacity(anchor, preset), total_layers) + if capacity <= 0: + continue + quantization = _node_quantization(anchor, preset) + shard_end = min(required_end, required_start + capacity - 1) + assignment = _add_shard_directive( + anchor, + hf_repo, + required_start, + shard_end, + quantization, + ) + anchor.pending_new_assignment = assignment + anchor.pending_directives.append(assignment) + _tracker_log( + server, + "info", + "scale demanded model copy", + endpoint=anchor.endpoint, + node_id=anchor.node_id, + model=model_key, + hf_repo=hf_repo, + shard=f"{required_start}-{shard_end}", + demand_rpm=round(info["demand_rpm"], 4), + served_model_copies=info["served_model_copies"], + target_copies=round(max(1.0, info["demand_rpm"] / 5.0), 2), + host_spare_bytes=host["memory_spare_bytes"], + ) + break + + def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict: if preset is None: return {"recommended": False} @@ -1169,6 +1392,20 @@ def _assign_redundant_managed_nodes( _emit_shard_change_directives(node, model, previous_range, preset) +def _repo_demand_rpm( + server: "_TrackerHTTPServer", + repo: str, + local_rpms: dict[str, dict], +) -> float: + """Look up recent request rate by hf_repo or preset alias.""" + if repo in local_rpms: + return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0) + resolved_name, _preset = _resolve_model_preset(server.model_presets, repo) + if resolved_name and resolved_name in local_rpms: + return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0) + return 0.0 + + def _relay_http_request_frames( relay_addr: str, path: str, @@ -4537,7 +4774,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): 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) + demand_rpms[repo] = _repo_demand_rpm(server, repo, local_rpms) best_repo = None best_score = -1.0 @@ -4547,6 +4784,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): for repo, nodes in repo_groups.items(): total = repo_layers[repo] + served_copies = _served_model_copies(nodes, 0, total - 1) covered = sorted( [(n.shard_start, n.shard_end) for n in nodes], key=lambda t: t[0], @@ -4561,8 +4799,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): 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) + # Prefer demanded models already being served; coverage gaps still dominate. + score = (demand + 1.0) * (coverage_deficit + 0.01) + served_copies * 0.01 if score > best_score or (score == best_score and gap_size > best_gap_size): best_score = score best_gap_size = gap_size @@ -4572,8 +4810,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): gap_found = best_gap_size > 0 if not gap_found: - # All shards covered — assign to highest-demand model for redundancy. - best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0)) + # All shards covered — add redundancy on the most demanded/served model. + best_repo = max( + repo_groups, + key=lambda repo: ( + demand_rpms.get(repo, 0.0), + _served_model_copies(repo_groups[repo], 0, repo_layers[repo] - 1), + len(repo_groups[repo]), + ), + ) best_gap_start = 0 best_num_layers = repo_layers[best_repo] diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 59a6e6a..269d0c4 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -1431,6 +1431,84 @@ def test_later_node_auto_joins_existing_public_hf_model_with_only_tracker_url( assert route_resp["route"] == ["http://203.0.113.20:8001", "http://203.0.113.21:8002"] +def test_later_node_auto_joins_redundant_copy_when_model_is_fully_covered( + tmp_path, + monkeypatch, +): + """Model-less joins should load the served HF model even when gap_found=false.""" + import meshnet_node.startup as startup_mod + + captured = {} + + class FakeBackend: + total_layers = 24 + + class FakeTorchNodeServer: + def __init__(self, **kwargs): + captured.update(kwargs) + self.backend = FakeBackend() + self.port = None + self.chat_completion_count = 0 + self.total_requests = 0 + self.failed_requests = 0 + self.queue_depth = 0 + + def start(self): + self.port = 8003 + return self.port + + def stop(self): + pass + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0}, + ) + monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer) + + tracker = TrackerServer() + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + try: + for endpoint, shard_start, shard_end in ( + ("http://203.0.113.30:8001", 0, 11), + ("http://203.0.113.31:8001", 12, 23), + ): + data = json.dumps({ + "endpoint": endpoint, + "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", + "num_layers": 24, + "shard_start": shard_start, + "shard_end": shard_end, + "tracker_mode": shard_start == 0, + "hardware_profile": {}, + "score": 1.0, + }).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/nodes/register", + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as resp: + resp.read() + + node = run_startup( + tracker_url=tracker_url, + advertise_host="203.0.113.32", + wallet_path=tmp_path / "wallet.json", + ) + try: + assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert captured["shard_start"] == 0 + finally: + node.stop() + finally: + tracker.stop() + + # --------------------------------------------------------------------------- # Full startup integration test # ---------------------------------------------------------------------------