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

balancing improvements

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 18:30:30 +03:00
parent f220fd2210
commit 1299a6bb1c
6 changed files with 358 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
billing.sqlite
View File

Binary file not shown.

BIN
meshnet_registry.sqlite3
View File

Binary file not shown.

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

@@ -52,7 +52,7 @@ def _run_node(cfg: dict) -> 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:

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

@@ -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 <HF_REPO>."
)
# 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

255
packages/tracker/meshnet_tracker/server.py
View File

@@ -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]

78
tests/test_node_startup.py
View File

@@ -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
# ---------------------------------------------------------------------------