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

feat(us-016): smart shard gap detection for auto-join and --model-id

Browse Source 
Tracker /v1/network/assign now accepts an optional `hf_repo` query param
to restrict assignment to a specific model, and returns `gap_found: bool`
so callers know whether they received a real gap vs a redundancy slot.

Node startup with --model-id (no explicit shard args) now queries the
tracker first for an uncovered gap for that model before defaulting to
full coverage (0..n-1). This means a second node with --model-id will
serve only the missing layers, not the whole model again.

Auto-join fallback (no --model-id) now prints why it fell through
instead of silently switching to stub-model.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 00:48:34 +03:00
parent 97eefd3d5e
commit 3286e42783
3 changed files with 104 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -112,6 +112,23 @@ def run_startup(
f"Could not read num_hidden_layers from {model_id} config. "
"Pass --shard-start and --shard-end explicitly."
)
# When no explicit shard range given, ask the tracker if there's a gap for this model.
if shard_start is None and shard_end is None:
try:
qs = urllib.parse.urlencode({
"device": device, "vram_mb": vram_mb, "hf_repo": model_id,
})
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
shard_start = net_asgn["shard_start"]
shard_end = net_asgn["shard_end"]
print(
f" Tracker found uncovered shard: "
f"layers {shard_start}–{shard_end} (of {detected})",
flush=True,
)
except Exception:
pass # No other nodes registered yet — default to full model below
shard_start = shard_start if shard_start is not None else 0
shard_end = shard_end if shard_end is not None else detected - 1
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
@@ -179,13 +196,15 @@ def run_startup(
# 3a. Auto-join: query tracker for network-wide HF model assignment.
print("Querying tracker for network assignment...", flush=True)
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb})
net_assignment: dict = {}
try:
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
except Exception:
assigned_hf_repo = None
except Exception as exc:
print(f" (auto-join unavailable: {exc})", flush=True)
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
_gap_found: bool = bool(net_assignment.get("gap_found", False))
if assigned_hf_repo:
if assigned_hf_repo and _gap_found:
assigned_shard_start: int = net_assignment["shard_start"]
assigned_shard_end: int = net_assignment["shard_end"]
assigned_num_layers: int = net_assignment["num_layers"]

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

@@ -788,9 +788,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
Query params:
vram_mb — integer VRAM in MB (0 = CPU-only node)
device — "cuda" | "cpu"
hf_repo — optional; if set, restrict search to this repo only
Returns:
{hf_repo, shard_start, shard_end, num_layers}
{hf_repo, shard_start, shard_end, num_layers, gap_found}
gap_found=true means a real uncovered gap was assigned; false means redundancy.
"""
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
params = urllib.parse.parse_qs(parsed.query)
@@ -799,6 +801,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
except ValueError:
vram_mb = 0
device = params.get("device", ["cpu"])[0]
filter_repo = params.get("hf_repo", [None])[0] # optional repo filter
with server.lock:
self._purge_expired_nodes()
@@ -811,10 +814,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
and n.shard_start is not None
and n.shard_end is not None
and n.num_layers is not None
and (filter_repo is None or n.hf_repo == filter_repo)
]
if not hf_nodes:
self._send_json(503, {"error": "no HF-model nodes registered; cannot assign shards"})
msg = (
f"no HF-model nodes registered for {filter_repo!r}"
if filter_repo
else "no HF-model nodes registered; cannot assign shards"
)
self._send_json(503, {"error": msg})
return
# Group by hf_repo; pick the one with the largest total_layers and biggest gap.
@@ -853,7 +862,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
best_repo = repo
best_num_layers = total
if best_repo is None or best_gap_size <= 0:
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]))
best_gap_start = 0
@@ -874,6 +884,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"shard_start": shard_start,
"shard_end": shard_end,
"num_layers": total_l,
"gap_found": gap_found,
})
def _handle_route(self, parsed: urllib.parse.ParseResult):

67
tests/test_node_startup.py
View File

@@ -507,6 +507,73 @@ def test_second_node_downloads_same_shard_from_peer_without_huggingface(
tracker.stop()
def test_network_assign_gap_found_field():
"""network/assign sets gap_found=True when a real gap exists, False when fully covered."""
import json as _json
import urllib.request as _ur
tracker = TrackerServer()
port = tracker.start()
try:
# Register a node covering only layers 0-11 of a 24-layer model.
data = _json.dumps({
"endpoint": "http://127.0.0.1:9200",
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"shard_start": 0,
"shard_end": 11,
"hardware_profile": {},
"score": 1.0,
}).encode()
req = _ur.Request(
f"http://127.0.0.1:{port}/v1/nodes/register",
data=data,
headers={"Content-Type": "application/json"},
method="POST",
)
with _ur.urlopen(req) as r:
r.read()
# A new node should be told there is a gap (layers 12-23).
resp = _get_json(
f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
"&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
)
assert resp["gap_found"] is True
assert resp["shard_start"] == 12, f"expected gap at 12, got {resp['shard_start']}"
assert resp["shard_end"] == 23
# Register the second node covering the gap.
data2 = _json.dumps({
"endpoint": "http://127.0.0.1:9201",
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"shard_start": 12,
"shard_end": 23,
"hardware_profile": {},
"score": 1.0,
}).encode()
req2 = _ur.Request(
f"http://127.0.0.1:{port}/v1/nodes/register",
data=data2,
headers={"Content-Type": "application/json"},
method="POST",
)
with _ur.urlopen(req2) as r:
r.read()
# Now fully covered — gap_found should be False.
resp2 = _get_json(
f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
"&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
)
assert resp2["gap_found"] is False
finally:
tracker.stop()
def test_startup_cpu_fallback(tmp_path, monkeypatch):
"""Node starts with CPU warning when no GPU is detected."""
import meshnet_node.startup as startup_mod