Compare commits
4 Commits
1bdfce657d
...
c75e9708ae
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c75e9708ae | ||
|
|
3286e42783 | ||
|
|
97eefd3d5e | ||
|
|
a7cc377d13 |
@@ -5,6 +5,8 @@ from __future__ import annotations
|
|||||||
import json
|
import json
|
||||||
import socket
|
import socket
|
||||||
import sys
|
import sys
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
import urllib.error
|
import urllib.error
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
import urllib.request
|
import urllib.request
|
||||||
@@ -32,6 +34,23 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
|
|||||||
return json.loads(r.read())
|
return json.loads(r.read())
|
||||||
|
|
||||||
|
|
||||||
|
def _start_heartbeat(tracker_url: str, node_id: str, interval: float = 20.0) -> threading.Thread:
|
||||||
|
"""Daemon thread that sends periodic heartbeats to the tracker."""
|
||||||
|
def _loop() -> None:
|
||||||
|
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||||||
|
while True:
|
||||||
|
time.sleep(interval)
|
||||||
|
try:
|
||||||
|
_post_json(hb_url, {})
|
||||||
|
print(f" [node] heartbeat sent → tracker (node {node_id[:8]})", flush=True)
|
||||||
|
except Exception as exc:
|
||||||
|
print(f" [node] WARNING: heartbeat failed: {exc}", flush=True)
|
||||||
|
|
||||||
|
t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
|
||||||
|
t.start()
|
||||||
|
return t
|
||||||
|
|
||||||
|
|
||||||
def run_startup(
|
def run_startup(
|
||||||
tracker_url: str,
|
tracker_url: str,
|
||||||
port: int = 0,
|
port: int = 0,
|
||||||
@@ -84,7 +103,7 @@ def run_startup(
|
|||||||
if probationary_line is not None:
|
if probationary_line is not None:
|
||||||
print(f" {probationary_line}", flush=True)
|
print(f" {probationary_line}", flush=True)
|
||||||
|
|
||||||
if model_id is not None:
|
if model_id: # treat "" the same as None — no explicit model given
|
||||||
# Auto-detect shard range from model config if not explicitly provided
|
# Auto-detect shard range from model config if not explicitly provided
|
||||||
if shard_start is None or shard_end is None:
|
if shard_start is None or shard_end is None:
|
||||||
detected = _detect_num_layers(model_id)
|
detected = _detect_num_layers(model_id)
|
||||||
@@ -93,6 +112,23 @@ def run_startup(
|
|||||||
f"Could not read num_hidden_layers from {model_id} config. "
|
f"Could not read num_hidden_layers from {model_id} config. "
|
||||||
"Pass --shard-start and --shard-end explicitly."
|
"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_start = shard_start if shard_start is not None else 0
|
||||||
shard_end = shard_end if shard_end is not None else detected - 1
|
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)
|
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
|
||||||
@@ -116,6 +152,31 @@ def run_startup(
|
|||||||
shard_label = f"layers {shard_start}–{shard_end}"
|
shard_label = f"layers {shard_start}–{shard_end}"
|
||||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||||
endpoint = f"http://{public_host}:{actual_port}"
|
endpoint = f"http://{public_host}:{actual_port}"
|
||||||
|
# Register with tracker so other nodes can auto-join this model.
|
||||||
|
total_layers = getattr(node.backend, "total_layers", None)
|
||||||
|
try:
|
||||||
|
reg_resp = _post_json(
|
||||||
|
f"{tracker_url}/v1/nodes/register",
|
||||||
|
{
|
||||||
|
"endpoint": endpoint,
|
||||||
|
"model": model_id.split("/")[-1],
|
||||||
|
"hf_repo": model_id,
|
||||||
|
"num_layers": total_layers,
|
||||||
|
"shard_start": shard_start,
|
||||||
|
"shard_end": shard_end,
|
||||||
|
"hardware_profile": hw,
|
||||||
|
"wallet_address": address,
|
||||||
|
"quantization": quantization,
|
||||||
|
"score": 1.0,
|
||||||
|
"tracker_mode": (shard_start == 0),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
node_id = reg_resp.get("node_id", "?")
|
||||||
|
print(f" Registered with tracker — node ID: {node_id}", flush=True)
|
||||||
|
_start_heartbeat(tracker_url, node_id)
|
||||||
|
except Exception as exc:
|
||||||
|
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||||
|
|
||||||
print(
|
print(
|
||||||
f"\n{'=' * 32}\n"
|
f"\n{'=' * 32}\n"
|
||||||
f"meshnet-node ready\n"
|
f"meshnet-node ready\n"
|
||||||
@@ -132,7 +193,79 @@ def run_startup(
|
|||||||
if shard_start is not None or shard_end is not None:
|
if shard_start is not None or shard_end is not None:
|
||||||
raise ValueError("--shard-start / --shard-end require --model-id")
|
raise ValueError("--shard-start / --shard-end require --model-id")
|
||||||
|
|
||||||
# 3. Shard assignment from tracker
|
# 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}")
|
||||||
|
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 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"]
|
||||||
|
print(
|
||||||
|
f" Assigned: {assigned_hf_repo} "
|
||||||
|
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||||||
|
f"(of {assigned_num_layers})",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
print("Loading real PyTorch model shard...", flush=True)
|
||||||
|
node = TorchNodeServer(
|
||||||
|
host=host,
|
||||||
|
port=port,
|
||||||
|
model_id=assigned_hf_repo,
|
||||||
|
shard_start=assigned_shard_start,
|
||||||
|
shard_end=assigned_shard_end,
|
||||||
|
quantization=quantization,
|
||||||
|
tracker_url=tracker_url,
|
||||||
|
)
|
||||||
|
actual_port = node.start()
|
||||||
|
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||||
|
endpoint = f"http://{public_host}:{actual_port}"
|
||||||
|
try:
|
||||||
|
reg_resp = _post_json(
|
||||||
|
f"{tracker_url}/v1/nodes/register",
|
||||||
|
{
|
||||||
|
"endpoint": endpoint,
|
||||||
|
"model": assigned_hf_repo.split("/")[-1],
|
||||||
|
"hf_repo": assigned_hf_repo,
|
||||||
|
"num_layers": assigned_num_layers,
|
||||||
|
"shard_start": assigned_shard_start,
|
||||||
|
"shard_end": assigned_shard_end,
|
||||||
|
"hardware_profile": hw,
|
||||||
|
"wallet_address": address,
|
||||||
|
"quantization": quantization,
|
||||||
|
"score": 1.0,
|
||||||
|
"tracker_mode": (assigned_shard_start == 0),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
node_id = reg_resp.get("node_id", "?")
|
||||||
|
print(f" Registered with tracker — node ID: {node_id}", flush=True)
|
||||||
|
_start_heartbeat(tracker_url, node_id)
|
||||||
|
except Exception as exc:
|
||||||
|
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||||
|
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||||||
|
print(
|
||||||
|
f"\n{'=' * 32}\n"
|
||||||
|
f"meshnet-node ready (auto-joined)\n"
|
||||||
|
f" Wallet: {address}\n"
|
||||||
|
f" Model ID: {assigned_hf_repo}\n"
|
||||||
|
f" Shard: layers {assigned_shard_start}–{assigned_shard_end} "
|
||||||
|
f"({shard_count} of {assigned_num_layers})\n"
|
||||||
|
f" Quantization: {quantization}\n"
|
||||||
|
f" Endpoint: {endpoint}\n"
|
||||||
|
f" Hardware: {device.upper()}\n"
|
||||||
|
f"{'=' * 32}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
return node
|
||||||
|
|
||||||
|
# 3b. Shard assignment from tracker (stub-model / preset-based path)
|
||||||
print("Querying tracker for shard assignment...", flush=True)
|
print("Querying tracker for shard assignment...", flush=True)
|
||||||
assign_qs = urllib.parse.urlencode({
|
assign_qs = urllib.parse.urlencode({
|
||||||
"model": model,
|
"model": model,
|
||||||
|
|||||||
@@ -52,7 +52,7 @@ DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0
|
|||||||
class _NodeEntry:
|
class _NodeEntry:
|
||||||
__slots__ = (
|
__slots__ = (
|
||||||
"node_id", "endpoint", "shard_start", "shard_end",
|
"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",
|
"score", "vram_bytes", "ram_bytes", "quantizations",
|
||||||
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
||||||
"pending_directives", "last_heartbeat", "tracker_mode",
|
"pending_directives", "last_heartbeat", "tracker_mode",
|
||||||
@@ -76,6 +76,8 @@ class _NodeEntry:
|
|||||||
quantization: str | None = None,
|
quantization: str | None = None,
|
||||||
managed_assignment: bool = False,
|
managed_assignment: bool = False,
|
||||||
tracker_mode: bool = False,
|
tracker_mode: bool = False,
|
||||||
|
hf_repo: str | None = None,
|
||||||
|
num_layers: int | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.node_id = node_id
|
self.node_id = node_id
|
||||||
self.endpoint = endpoint
|
self.endpoint = endpoint
|
||||||
@@ -93,6 +95,8 @@ class _NodeEntry:
|
|||||||
self.quantization = quantization
|
self.quantization = quantization
|
||||||
self.managed_assignment = managed_assignment
|
self.managed_assignment = managed_assignment
|
||||||
self.tracker_mode = tracker_mode
|
self.tracker_mode = tracker_mode
|
||||||
|
self.hf_repo = hf_repo
|
||||||
|
self.num_layers = num_layers
|
||||||
self.pending_directives: list[dict] = []
|
self.pending_directives: list[dict] = []
|
||||||
self.last_heartbeat: float = time.monotonic()
|
self.last_heartbeat: float = time.monotonic()
|
||||||
|
|
||||||
@@ -254,7 +258,12 @@ def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
|
|||||||
if (now - entry.last_heartbeat) > server.heartbeat_timeout
|
if (now - entry.last_heartbeat) > server.heartbeat_timeout
|
||||||
]
|
]
|
||||||
for node_id in expired_ids:
|
for node_id in expired_ids:
|
||||||
del server.registry[node_id]
|
entry = server.registry.pop(node_id)
|
||||||
|
print(
|
||||||
|
f"[tracker] node expired: {node_id[:8]} {entry.endpoint} "
|
||||||
|
f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
if expired_ids:
|
if expired_ids:
|
||||||
_rebalance_all_locked(server)
|
_rebalance_all_locked(server)
|
||||||
return expired_ids
|
return expired_ids
|
||||||
@@ -426,6 +435,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._handle_routes(parsed)
|
self._handle_routes(parsed)
|
||||||
elif parsed.path == "/v1/nodes/assign":
|
elif parsed.path == "/v1/nodes/assign":
|
||||||
self._handle_assign(parsed)
|
self._handle_assign(parsed)
|
||||||
|
elif parsed.path == "/v1/network/assign":
|
||||||
|
self._handle_network_assign(parsed)
|
||||||
elif parsed.path == "/v1/models":
|
elif parsed.path == "/v1/models":
|
||||||
self._handle_models()
|
self._handle_models()
|
||||||
elif parsed.path.startswith("/v1/coverage/"):
|
elif parsed.path.startswith("/v1/coverage/"):
|
||||||
@@ -607,6 +618,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
return
|
return
|
||||||
|
|
||||||
tracker_mode = bool(body.get("tracker_mode", False))
|
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())
|
node_id = str(uuid.uuid4())
|
||||||
entry = _NodeEntry(
|
entry = _NodeEntry(
|
||||||
@@ -626,9 +649,23 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
quantization=quantization,
|
quantization=quantization,
|
||||||
managed_assignment=not explicit_shard,
|
managed_assignment=not explicit_shard,
|
||||||
tracker_mode=tracker_mode,
|
tracker_mode=tracker_mode,
|
||||||
|
hf_repo=hf_repo,
|
||||||
|
num_layers=num_layers,
|
||||||
)
|
)
|
||||||
with server.lock:
|
with server.lock:
|
||||||
self._purge_expired_nodes()
|
self._purge_expired_nodes()
|
||||||
|
# Dedup: if this endpoint is already registered, remove the old entry first.
|
||||||
|
stale_ids = [
|
||||||
|
eid for eid, e in server.registry.items()
|
||||||
|
if e.endpoint == entry.endpoint.rstrip("/")
|
||||||
|
]
|
||||||
|
for eid in stale_ids:
|
||||||
|
old = server.registry.pop(eid)
|
||||||
|
print(
|
||||||
|
f"[tracker] node re-registered: replaced {eid[:8]} with {node_id[:8]}"
|
||||||
|
f" {old.endpoint}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
server.registry[node_id] = entry
|
server.registry[node_id] = entry
|
||||||
if entry.managed_assignment:
|
if entry.managed_assignment:
|
||||||
_rebalance_model_locked(server, model)
|
_rebalance_model_locked(server, model)
|
||||||
@@ -636,6 +673,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if assignment_directive is not None:
|
if assignment_directive is not None:
|
||||||
entry.pending_directives.clear()
|
entry.pending_directives.clear()
|
||||||
|
|
||||||
|
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
|
||||||
|
repo_info = f" [{hf_repo}]" if hf_repo else ""
|
||||||
|
print(
|
||||||
|
f"[tracker] node registered: {node_id[:8]} {endpoint} {model}{repo_info} {shard_info}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
|
||||||
payload = {"node_id": node_id}
|
payload = {"node_id": node_id}
|
||||||
if assignment_directive is not None:
|
if assignment_directive is not None:
|
||||||
payload["directive"] = assignment_directive
|
payload["directive"] = assignment_directive
|
||||||
@@ -653,6 +697,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
_rebalance_model_locked(server, entry.model or "stub-model")
|
_rebalance_model_locked(server, entry.model or "stub-model")
|
||||||
directives = list(entry.pending_directives)
|
directives = list(entry.pending_directives)
|
||||||
entry.pending_directives.clear()
|
entry.pending_directives.clear()
|
||||||
|
print(
|
||||||
|
f"[tracker] heartbeat: {node_id[:8]} {entry.endpoint}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
if directives:
|
if directives:
|
||||||
self._send_json(200, {"directives": directives})
|
self._send_json(200, {"directives": directives})
|
||||||
else:
|
else:
|
||||||
@@ -751,6 +799,111 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
|
**({"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"
|
||||||
|
hf_repo — optional; if set, restrict search to this repo only
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
{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)
|
||||||
|
try:
|
||||||
|
vram_mb = int(params.get("vram_mb", ["0"])[0])
|
||||||
|
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()
|
||||||
|
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
|
||||||
|
and (filter_repo is None or n.hf_repo == filter_repo)
|
||||||
|
]
|
||||||
|
|
||||||
|
if not hf_nodes:
|
||||||
|
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.
|
||||||
|
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
|
||||||
|
|
||||||
|
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
|
||||||
|
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,
|
||||||
|
"gap_found": gap_found,
|
||||||
|
})
|
||||||
|
|
||||||
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
||||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||||
params = urllib.parse.parse_qs(parsed.query)
|
params = urllib.parse.parse_qs(parsed.query)
|
||||||
@@ -761,15 +914,28 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
model = model_list[0]
|
model = model_list[0]
|
||||||
preset = server.model_presets.get(model)
|
preset = server.model_presets.get(model)
|
||||||
if preset is None:
|
|
||||||
self._send_json(404, {"error": f"unknown model preset: {model!r}"})
|
|
||||||
return
|
|
||||||
|
|
||||||
required_start, required_end = _preset_layer_bounds(preset)
|
|
||||||
|
|
||||||
with server.lock:
|
with server.lock:
|
||||||
self._purge_expired_nodes()
|
self._purge_expired_nodes()
|
||||||
alive = [node for node in server.registry.values() if node.model == model]
|
if preset is not None:
|
||||||
|
# Preset-based routing (stub-model system).
|
||||||
|
alive = [node for node in server.registry.values() if node.model == model]
|
||||||
|
required_start, required_end = _preset_layer_bounds(preset)
|
||||||
|
else:
|
||||||
|
# HF model routing: match by hf_repo (full) or model short name.
|
||||||
|
alive = [
|
||||||
|
node for node in server.registry.values()
|
||||||
|
if (node.hf_repo == model or node.model == model)
|
||||||
|
and node.shard_start is not None
|
||||||
|
and node.shard_end is not None
|
||||||
|
and node.num_layers is not None
|
||||||
|
]
|
||||||
|
if not alive:
|
||||||
|
self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
|
||||||
|
return
|
||||||
|
required_start = 0
|
||||||
|
required_end = max(n.num_layers for n in alive) - 1 # type: ignore[type-var]
|
||||||
|
|
||||||
if server.contracts is not None:
|
if server.contracts is not None:
|
||||||
alive = [
|
alive = [
|
||||||
node for node in alive
|
node for node in alive
|
||||||
@@ -871,7 +1037,7 @@ class TrackerServer:
|
|||||||
self,
|
self,
|
||||||
host: str = "127.0.0.1",
|
host: str = "127.0.0.1",
|
||||||
port: int = 0,
|
port: int = 0,
|
||||||
heartbeat_timeout: float = 30.0,
|
heartbeat_timeout: float = 90.0,
|
||||||
rebalance_interval: float = 30.0,
|
rebalance_interval: float = 30.0,
|
||||||
model_presets: dict | None = None,
|
model_presets: dict | None = None,
|
||||||
contracts: Any | None = None,
|
contracts: Any | None = None,
|
||||||
|
|||||||
@@ -507,6 +507,163 @@ def test_second_node_downloads_same_shard_from_peer_without_huggingface(
|
|||||||
tracker.stop()
|
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_route_finds_hf_model_across_two_nodes():
|
||||||
|
"""Tracker /v1/route returns ordered route for HF model even without a preset."""
|
||||||
|
import json as _json
|
||||||
|
import urllib.request as _ur
|
||||||
|
|
||||||
|
tracker = TrackerServer()
|
||||||
|
port = tracker.start()
|
||||||
|
try:
|
||||||
|
def register(endpoint, shard_start, shard_end):
|
||||||
|
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,
|
||||||
|
"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()
|
||||||
|
|
||||||
|
register("http://127.0.0.1:9300", 0, 11)
|
||||||
|
register("http://127.0.0.1:9301", 12, 23)
|
||||||
|
|
||||||
|
# Route by hf_repo (full identifier).
|
||||||
|
resp = _get_json(
|
||||||
|
f"http://127.0.0.1:{port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
|
||||||
|
)
|
||||||
|
assert resp["route"] == ["http://127.0.0.1:9300", "http://127.0.0.1:9301"]
|
||||||
|
|
||||||
|
# Route also works by short model name.
|
||||||
|
resp2 = _get_json(
|
||||||
|
f"http://127.0.0.1:{port}/v1/route?model=Qwen2.5-0.5B-Instruct"
|
||||||
|
)
|
||||||
|
assert resp2["route"] == ["http://127.0.0.1:9300", "http://127.0.0.1:9301"]
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_register_deduplicates_same_endpoint():
|
||||||
|
"""Re-registering the same endpoint replaces the old entry, not duplicates it."""
|
||||||
|
import json as _json
|
||||||
|
import urllib.request as _ur
|
||||||
|
|
||||||
|
tracker = TrackerServer()
|
||||||
|
port = tracker.start()
|
||||||
|
try:
|
||||||
|
def register(shard_start, shard_end):
|
||||||
|
data = _json.dumps({
|
||||||
|
"endpoint": "http://127.0.0.1:9400",
|
||||||
|
"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,
|
||||||
|
"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:
|
||||||
|
return _json.loads(r.read())
|
||||||
|
|
||||||
|
register(0, 23) # initial full-model registration
|
||||||
|
register(12, 23) # re-register with corrected shard range
|
||||||
|
|
||||||
|
# After re-register, tracker should see only one node at 12-23 for this endpoint.
|
||||||
|
# If both were still registered, the gap scan would find no gap (0-23 still covers).
|
||||||
|
# With dedup, the old 0-23 is gone and a real gap 0-11 exists.
|
||||||
|
assign_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 assign_resp["gap_found"] is True
|
||||||
|
assert assign_resp["shard_start"] == 0, "old 0-23 entry should have been replaced"
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
def test_startup_cpu_fallback(tmp_path, monkeypatch):
|
def test_startup_cpu_fallback(tmp_path, monkeypatch):
|
||||||
"""Node starts with CPU warning when no GPU is detected."""
|
"""Node starts with CPU warning when no GPU is detected."""
|
||||||
import meshnet_node.startup as startup_mod
|
import meshnet_node.startup as startup_mod
|
||||||
|
|||||||
Reference in New Issue
Block a user