feat: auto-join network — node discovers missing shards from tracker
Tracker:
- _NodeEntry gains hf_repo + num_layers fields (parsed from register body)
- GET /v1/network/assign — finds the biggest uncovered shard gap across
registered HF-model nodes; returns {hf_repo, shard_start, shard_end, num_layers}
- Returns 503 when no HF-model nodes are registered yet
Node startup:
- When model_id is set: registers with tracker including hf_repo + num_layers
so other nodes can auto-join this model
- When model_id is empty/None: queries /v1/network/assign, gets assigned the
missing layers, loads TorchNodeServer with the assigned shard automatically
- Fixes empty-string model_id leaking from DEFAULTS (treats "" same as None)
Usage: `meshnet-node start --tracker http://localhost:8080 --quantization bfloat16`
Node discovers what to serve and joins the network without any model flags.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -52,7 +52,7 @@ DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0
|
||||
class _NodeEntry:
|
||||
__slots__ = (
|
||||
"node_id", "endpoint", "shard_start", "shard_end",
|
||||
"model", "shard_checksum", "hardware_profile", "wallet_address",
|
||||
"model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
|
||||
"score", "vram_bytes", "ram_bytes", "quantizations",
|
||||
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
||||
"pending_directives", "last_heartbeat", "tracker_mode",
|
||||
@@ -76,6 +76,8 @@ class _NodeEntry:
|
||||
quantization: str | None = None,
|
||||
managed_assignment: bool = False,
|
||||
tracker_mode: bool = False,
|
||||
hf_repo: str | None = None,
|
||||
num_layers: int | None = None,
|
||||
) -> None:
|
||||
self.node_id = node_id
|
||||
self.endpoint = endpoint
|
||||
@@ -93,6 +95,8 @@ class _NodeEntry:
|
||||
self.quantization = quantization
|
||||
self.managed_assignment = managed_assignment
|
||||
self.tracker_mode = tracker_mode
|
||||
self.hf_repo = hf_repo
|
||||
self.num_layers = num_layers
|
||||
self.pending_directives: list[dict] = []
|
||||
self.last_heartbeat: float = time.monotonic()
|
||||
|
||||
@@ -426,6 +430,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._handle_routes(parsed)
|
||||
elif parsed.path == "/v1/nodes/assign":
|
||||
self._handle_assign(parsed)
|
||||
elif parsed.path == "/v1/network/assign":
|
||||
self._handle_network_assign(parsed)
|
||||
elif parsed.path == "/v1/models":
|
||||
self._handle_models()
|
||||
elif parsed.path.startswith("/v1/coverage/"):
|
||||
@@ -607,6 +613,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
return
|
||||
|
||||
tracker_mode = bool(body.get("tracker_mode", False))
|
||||
hf_repo = body.get("hf_repo")
|
||||
if hf_repo is not None and not isinstance(hf_repo, str):
|
||||
self._send_json(400, {"error": "hf_repo must be a string"})
|
||||
return
|
||||
num_layers_body = body.get("num_layers")
|
||||
num_layers: int | None = None
|
||||
if num_layers_body is not None:
|
||||
try:
|
||||
num_layers = int(num_layers_body)
|
||||
except (TypeError, ValueError):
|
||||
self._send_json(400, {"error": "num_layers must be an integer"})
|
||||
return
|
||||
|
||||
node_id = str(uuid.uuid4())
|
||||
entry = _NodeEntry(
|
||||
@@ -626,6 +644,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
quantization=quantization,
|
||||
managed_assignment=not explicit_shard,
|
||||
tracker_mode=tracker_mode,
|
||||
hf_repo=hf_repo,
|
||||
num_layers=num_layers,
|
||||
)
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
@@ -751,6 +771,100 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
|
||||
})
|
||||
|
||||
def _handle_network_assign(self, parsed: urllib.parse.ParseResult):
|
||||
"""Assign a new node to fill the biggest uncovered shard gap across HF-model nodes.
|
||||
|
||||
Query params:
|
||||
vram_mb — integer VRAM in MB (0 = CPU-only node)
|
||||
device — "cuda" | "cpu"
|
||||
|
||||
Returns:
|
||||
{hf_repo, shard_start, shard_end, num_layers}
|
||||
"""
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
params = urllib.parse.parse_qs(parsed.query)
|
||||
try:
|
||||
vram_mb = int(params.get("vram_mb", ["0"])[0])
|
||||
except ValueError:
|
||||
vram_mb = 0
|
||||
device = params.get("device", ["cpu"])[0]
|
||||
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
all_nodes = list(server.registry.values())
|
||||
|
||||
# Collect only nodes that registered a real HF model (have hf_repo + shard bounds).
|
||||
hf_nodes = [
|
||||
n for n in all_nodes
|
||||
if n.hf_repo
|
||||
and n.shard_start is not None
|
||||
and n.shard_end is not None
|
||||
and n.num_layers is not None
|
||||
]
|
||||
|
||||
if not hf_nodes:
|
||||
self._send_json(503, {"error": "no HF-model nodes registered; cannot assign shards"})
|
||||
return
|
||||
|
||||
# Group by hf_repo; pick the one with the largest total_layers and biggest gap.
|
||||
from collections import defaultdict
|
||||
repo_groups: dict = defaultdict(list)
|
||||
repo_layers: dict = {}
|
||||
for n in hf_nodes:
|
||||
repo_groups[n.hf_repo].append(n)
|
||||
# Use the largest num_layers seen for this repo.
|
||||
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).
|
||||
best_repo = None
|
||||
best_gap_size = -1
|
||||
best_gap_start = 0
|
||||
best_num_layers = 0
|
||||
|
||||
for repo, nodes in repo_groups.items():
|
||||
total = repo_layers[repo]
|
||||
covered = sorted(
|
||||
[(n.shard_start, n.shard_end) for n in nodes],
|
||||
key=lambda t: t[0],
|
||||
)
|
||||
# Walk from 0 to find first uncovered layer.
|
||||
gap_start = 0
|
||||
for s, e in covered:
|
||||
if s <= gap_start:
|
||||
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:
|
||||
best_gap_size = gap_size
|
||||
best_gap_start = gap_start
|
||||
best_repo = repo
|
||||
best_num_layers = total
|
||||
|
||||
if best_repo is None or best_gap_size <= 0:
|
||||
# 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
|
||||
best_num_layers = repo_layers[best_repo]
|
||||
|
||||
# Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM.
|
||||
total_l = best_num_layers
|
||||
if device == "cuda" and vram_mb >= 8192:
|
||||
max_layers = total_l
|
||||
else:
|
||||
max_layers = max(1, total_l // 2)
|
||||
|
||||
shard_start = best_gap_start
|
||||
shard_end = min(total_l - 1, shard_start + max_layers - 1)
|
||||
|
||||
self._send_json(200, {
|
||||
"hf_repo": best_repo,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"num_layers": total_l,
|
||||
})
|
||||
|
||||
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
params = urllib.parse.parse_qs(parsed.query)
|
||||
|
||||
Reference in New Issue
Block a user