Compare commits
8 Commits
c587e02c2c
...
1da088926a
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1da088926a | ||
|
|
8ce5a74d5e | ||
|
|
27818df654 | ||
|
|
d9110b623b | ||
|
|
fff11c4d8b | ||
|
|
34fb1ec65d | ||
|
|
96af82892f | ||
|
|
e1ba120912 |
@@ -28,6 +28,8 @@ def _run_node(cfg: dict) -> None:
|
||||
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
|
||||
host=cfg.get("host", "0.0.0.0"),
|
||||
advertise_host=cfg.get("advertise_host"),
|
||||
route_timeout=float(cfg.get("route_timeout", 30.0)),
|
||||
vram_mb_override=cfg.get("vram_mb_override"),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -103,6 +105,10 @@ def _cmd_default(args) -> int:
|
||||
overrides["host"] = args.host
|
||||
if args.advertise_host:
|
||||
overrides["advertise_host"] = args.advertise_host
|
||||
if args.route_timeout != 30.0:
|
||||
overrides["route_timeout"] = args.route_timeout
|
||||
if getattr(args, "memory", None) is not None:
|
||||
overrides["vram_mb_override"] = args.memory
|
||||
|
||||
if overrides:
|
||||
cfg = merge_cli_overrides(cfg, **overrides)
|
||||
@@ -194,6 +200,8 @@ def _cmd_start(args) -> int:
|
||||
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
|
||||
host=cfg["host"],
|
||||
advertise_host=getattr(args, "advertise_host", None),
|
||||
route_timeout=getattr(args, "route_timeout", 30.0),
|
||||
vram_mb_override=getattr(args, "memory", None),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -234,6 +242,10 @@ def main() -> None:
|
||||
parser.add_argument("--port", type=int, metavar="N", help="Port to listen on")
|
||||
parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
|
||||
parser.add_argument("--advertise-host", metavar="ADDR", help="Host/IP advertised to the tracker")
|
||||
parser.add_argument("--route-timeout", type=float, metavar="SEC", default=30.0,
|
||||
help="Seconds to wait for tracker route lookup (default 30)")
|
||||
parser.add_argument("--memory", type=int, metavar="MB", default=None,
|
||||
help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
|
||||
parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
|
||||
parser.add_argument("--compact", action="store_true", help="Single-line status output")
|
||||
parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists")
|
||||
@@ -262,6 +274,10 @@ def main() -> None:
|
||||
start_cmd.add_argument("--tracker-url", default=None)
|
||||
start_cmd.add_argument("--wallet")
|
||||
start_cmd.add_argument("--download-dir")
|
||||
start_cmd.add_argument("--route-timeout", type=float, default=30.0,
|
||||
help="Seconds to wait for tracker route lookup (default 30)")
|
||||
start_cmd.add_argument("--memory", type=int, default=None, metavar="MB",
|
||||
help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -145,6 +145,7 @@ class TorchModelShard:
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
start_layer: int | None = None,
|
||||
) -> TensorPayload | str:
|
||||
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
||||
self.device
|
||||
@@ -155,7 +156,9 @@ class TorchModelShard:
|
||||
position_ids = _tensor_from_int64_header(
|
||||
position_ids_header, self.torch, self.device
|
||||
)
|
||||
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
|
||||
hidden_states = self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
)
|
||||
if self.is_tail:
|
||||
return self.decode_tail(hidden_states)
|
||||
return self._payload(hidden_states, attention_mask, position_ids)
|
||||
@@ -278,7 +281,21 @@ class TorchModelShard:
|
||||
)
|
||||
return dict(self.tokenizer(prompt, return_tensors="pt"))
|
||||
|
||||
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
||||
def _run_layers(
|
||||
self,
|
||||
hidden_states: Any,
|
||||
attention_mask: Any,
|
||||
position_ids: Any,
|
||||
start_layer: int | None = None,
|
||||
) -> Any:
|
||||
# start_layer overrides shard_start for overlapping-shard routing
|
||||
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
|
||||
# indexing outside the loaded weights.
|
||||
effective_start = (
|
||||
max(self.shard_start, start_layer)
|
||||
if start_layer is not None
|
||||
else self.shard_start
|
||||
)
|
||||
position_embeddings = _rotary_position_embeddings(
|
||||
self.model,
|
||||
hidden_states,
|
||||
@@ -290,7 +307,7 @@ class TorchModelShard:
|
||||
self.torch,
|
||||
)
|
||||
with self.torch.inference_mode():
|
||||
for layer in self.layers[self.shard_start:self.shard_end + 1]:
|
||||
for layer in self.layers[effective_start:self.shard_end + 1]:
|
||||
hidden_states = _call_layer(
|
||||
layer,
|
||||
hidden_states,
|
||||
|
||||
@@ -91,8 +91,29 @@ def _start_heartbeat(
|
||||
node_id: str,
|
||||
register_payload: dict,
|
||||
interval: float = 20.0,
|
||||
node_ref: Any | None = None,
|
||||
start_time: float | None = None,
|
||||
) -> threading.Thread:
|
||||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts."""
|
||||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.
|
||||
|
||||
Heartbeat body carries cumulative stats (total_requests, failed_requests,
|
||||
queue_depth, uptime_seconds, status). Stats are buffered locally during
|
||||
outage and flushed on next successful heartbeat.
|
||||
|
||||
Heartbeat response may include new_assignment: {model, shard_start, shard_end}
|
||||
which is logged for now (hot-reload implemented in US-026).
|
||||
"""
|
||||
_start_time = start_time or time.monotonic()
|
||||
|
||||
def _get_stats() -> dict:
|
||||
uptime = time.monotonic() - _start_time
|
||||
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
|
||||
if node_ref is not None:
|
||||
stats["total_requests"] = getattr(node_ref, "total_requests", 0)
|
||||
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
|
||||
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
|
||||
return stats
|
||||
|
||||
def _reregister() -> bool:
|
||||
nonlocal node_id
|
||||
try:
|
||||
@@ -128,7 +149,15 @@ def _start_heartbeat(
|
||||
continue
|
||||
|
||||
try:
|
||||
_post_json(hb_url, {})
|
||||
resp = _post_json(hb_url, _get_stats())
|
||||
new_asgn = resp.get("new_assignment")
|
||||
if new_asgn:
|
||||
print(
|
||||
f" [node] tracker reassignment received: "
|
||||
f"model={new_asgn.get('model')!r} "
|
||||
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
||||
flush=True,
|
||||
)
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 404:
|
||||
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
||||
@@ -163,6 +192,8 @@ def run_startup(
|
||||
host: str = "127.0.0.1",
|
||||
advertise_host: str | None = None,
|
||||
contracts: Any | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
vram_mb_override: int | None = None,
|
||||
) -> StubNodeServer | TorchNodeServer:
|
||||
"""Execute the full startup sequence and return a running node server.
|
||||
|
||||
@@ -199,7 +230,10 @@ def run_startup(
|
||||
gpu_name: str | None = hw.get("gpu_name")
|
||||
vram_mb: int = hw.get("vram_mb", 0)
|
||||
|
||||
if device == "cpu":
|
||||
if vram_mb_override is not None:
|
||||
vram_mb = vram_mb_override
|
||||
print(f" Memory budget overridden to {vram_mb} MB via --memory", flush=True)
|
||||
elif device == "cpu":
|
||||
print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
|
||||
else:
|
||||
print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)
|
||||
@@ -254,7 +288,9 @@ def run_startup(
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
)
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||||
if isinstance(total_layers, int) and total_layers > 0:
|
||||
@@ -294,7 +330,7 @@ def run_startup(
|
||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||
setattr(node, "tracker_node_id", tracker_node_id)
|
||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload)
|
||||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time)
|
||||
except Exception as exc:
|
||||
setattr(node, "tracker_node_id", None)
|
||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||
@@ -346,7 +382,9 @@ def run_startup(
|
||||
shard_end=assigned_shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
)
|
||||
_node_start_time = time.monotonic()
|
||||
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}"
|
||||
@@ -378,7 +416,7 @@ def run_startup(
|
||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||
setattr(node, "tracker_node_id", tracker_node_id)
|
||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload)
|
||||
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time)
|
||||
except Exception as exc:
|
||||
setattr(node, "tracker_node_id", None)
|
||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||
|
||||
@@ -37,6 +37,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
backend: TorchModelShard,
|
||||
tracker_mode: bool = False,
|
||||
tracker_url: str | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
):
|
||||
super().__init__(addr, handler)
|
||||
self.backend = backend
|
||||
@@ -44,6 +45,11 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
self.forward_chunk_count = 0
|
||||
self.tracker_mode = tracker_mode
|
||||
self.tracker_url = tracker_url
|
||||
self.route_timeout = route_timeout
|
||||
self.total_requests: int = 0
|
||||
self.failed_requests: int = 0
|
||||
self.queue_depth: int = 0
|
||||
self._stats_lock = threading.Lock()
|
||||
|
||||
|
||||
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
@@ -139,12 +145,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
|
||||
server.received_activations = True
|
||||
|
||||
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
|
||||
start_layer = int(start_layer_header) if start_layer_header else None
|
||||
|
||||
try:
|
||||
result = server.backend.forward_bytes(
|
||||
raw_body,
|
||||
shape,
|
||||
self.headers.get("X-Meshnet-Attn-Mask"),
|
||||
self.headers.get("X-Meshnet-Position-Ids"),
|
||||
start_layer=start_layer,
|
||||
)
|
||||
except Exception as exc:
|
||||
self._send_json(500, {"error": str(exc)})
|
||||
@@ -213,6 +223,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
def _handle_chat_completions(self) -> None:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
with server._stats_lock:
|
||||
server.total_requests += 1
|
||||
server.queue_depth += 1
|
||||
try:
|
||||
self._do_chat_completions(server)
|
||||
finally:
|
||||
with server._stats_lock:
|
||||
server.queue_depth -= 1
|
||||
|
||||
def _record_failed_request(self) -> None:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
with server._stats_lock:
|
||||
server.failed_requests += 1
|
||||
|
||||
def _do_chat_completions(self, server: "_TorchHTTPServer") -> None:
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
@@ -238,6 +263,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
|
||||
self._send_openai_response(text, model_name, False, messages)
|
||||
except Exception as exc:
|
||||
self._record_failed_request()
|
||||
self._send_json(500, {"error": f"generation failed: {exc}"})
|
||||
return
|
||||
|
||||
@@ -298,17 +324,29 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
result_text = "".join(generated)
|
||||
self._send_openai_response(result_text, model_name, stream, messages)
|
||||
|
||||
def _get_remaining_route(self, model: str) -> list[str]:
|
||||
def _get_remaining_route(self, model: str) -> list[tuple[str, int]]:
|
||||
"""Return downstream hops as (endpoint, start_layer) pairs.
|
||||
|
||||
Fast path reads X-Meshnet-Route header injected by the tracker.
|
||||
Slow path queries the tracker's /v1/route endpoint as a fallback.
|
||||
start_layer tells each downstream node which layer to begin from,
|
||||
enabling correct execution when shard ranges overlap.
|
||||
"""
|
||||
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
|
||||
injected = self.headers.get("X-Meshnet-Route")
|
||||
if injected:
|
||||
try:
|
||||
route = json.loads(injected)
|
||||
if isinstance(route, list):
|
||||
resolved = [str(ep) for ep in route]
|
||||
print(f" [node] using injected downstream route: {resolved}", flush=True)
|
||||
return resolved
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
hops: list[tuple[str, int]] = []
|
||||
for item in route:
|
||||
if isinstance(item, dict):
|
||||
hops.append((str(item["endpoint"]), int(item.get("start_layer", 0))))
|
||||
elif isinstance(item, str):
|
||||
hops.append((item, 0)) # backward-compat: plain string, no start_layer
|
||||
print(f" [node] using injected downstream route: {[ep for ep, _ in hops]}", flush=True)
|
||||
return hops
|
||||
except (json.JSONDecodeError, TypeError, KeyError):
|
||||
pass
|
||||
|
||||
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
||||
@@ -318,18 +356,31 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
route_model = getattr(server.backend, "model_id", None) or model
|
||||
try:
|
||||
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
||||
with urllib.request.urlopen(url, timeout=5.0) as r:
|
||||
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
||||
route_resp = json.loads(r.read())
|
||||
route = route_resp.get("route", [])
|
||||
own_port = server.server_address[1]
|
||||
resolved = [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")]
|
||||
print(f" [node] tracker downstream route: {resolved}", flush=True)
|
||||
return resolved
|
||||
nodes_info = route_resp.get("nodes", [])
|
||||
# nodes_info is ordered; find own node and compute start_layers post-hoc
|
||||
hops = []
|
||||
covered_up_to: int | None = None
|
||||
for node_info in nodes_info:
|
||||
ep = node_info.get("endpoint", "")
|
||||
if ep.rstrip("/").endswith(f":{own_port}"):
|
||||
covered_up_to = node_info.get("shard_end")
|
||||
continue
|
||||
if covered_up_to is None:
|
||||
# Own node not found yet; use node's shard_start as fallback
|
||||
covered_up_to = (node_info.get("shard_start") or 1) - 1
|
||||
start_l = covered_up_to + 1
|
||||
hops.append((ep, start_l))
|
||||
covered_up_to = node_info.get("shard_end", covered_up_to)
|
||||
print(f" [node] tracker downstream route: {[ep for ep, _ in hops]}", flush=True)
|
||||
return hops
|
||||
except Exception as exc:
|
||||
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
||||
return []
|
||||
|
||||
def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:
|
||||
def _run_downstream_pipeline(self, payload: object, route: list[tuple[str, int]]) -> str:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
if not route:
|
||||
# Partial shard at tail: decode the activation from the previous node.
|
||||
@@ -354,8 +405,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
current_attn = attn_mask
|
||||
current_pos = pos_ids
|
||||
|
||||
for hop_index, node_url in enumerate(route):
|
||||
print(f" [node] pipeline hop {hop_index}: {node_url}", flush=True)
|
||||
for hop_index, (node_url, start_layer) in enumerate(route):
|
||||
print(f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True)
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/octet-stream",
|
||||
"X-Meshnet-Wire": _WIRE_VERSION,
|
||||
@@ -365,6 +416,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
"X-Meshnet-Chunk-Index": "0",
|
||||
"X-Meshnet-Chunk-Total": "1",
|
||||
"X-Meshnet-Hop-Index": str(hop_index),
|
||||
"X-Meshnet-Start-Layer": str(start_layer),
|
||||
}
|
||||
if current_attn:
|
||||
headers["X-Meshnet-Attn-Mask"] = current_attn
|
||||
@@ -558,6 +610,7 @@ class TorchNodeServer:
|
||||
backend: TorchModelShard | None = None,
|
||||
tracker_mode: bool | None = None,
|
||||
tracker_url: str | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._requested_port = port
|
||||
@@ -570,10 +623,15 @@ class TorchNodeServer:
|
||||
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
||||
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
|
||||
self._tracker_url = tracker_url
|
||||
self._route_timeout = route_timeout
|
||||
self._server: _TorchHTTPServer | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self.port: int | None = None
|
||||
|
||||
@property
|
||||
def route_timeout(self) -> float:
|
||||
return self._route_timeout
|
||||
|
||||
@property
|
||||
def backend(self) -> TorchModelShard:
|
||||
return self._backend
|
||||
@@ -586,6 +644,18 @@ class TorchNodeServer:
|
||||
def forward_chunk_count(self) -> int:
|
||||
return self._server.forward_chunk_count if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def total_requests(self) -> int:
|
||||
return self._server.total_requests if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def failed_requests(self) -> int:
|
||||
return self._server.failed_requests if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def queue_depth(self) -> int:
|
||||
return self._server.queue_depth if self._server is not None else 0
|
||||
|
||||
def start(self) -> int:
|
||||
if self._server is not None:
|
||||
raise RuntimeError("TorchNodeServer is already running")
|
||||
@@ -595,6 +665,7 @@ class TorchNodeServer:
|
||||
self._backend,
|
||||
self._tracker_mode,
|
||||
self._tracker_url,
|
||||
self._route_timeout,
|
||||
)
|
||||
self.port = self._server.server_address[1]
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
|
||||
@@ -27,6 +27,12 @@ def main() -> None:
|
||||
default=None,
|
||||
help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--stats-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="SQLite database path for persistent model usage statistics",
|
||||
)
|
||||
common.add_argument(
|
||||
"--relay-url",
|
||||
default=None,
|
||||
@@ -52,6 +58,7 @@ def main() -> None:
|
||||
heartbeat_timeout=args.heartbeat_timeout,
|
||||
cluster_peers=cluster_peers or None,
|
||||
cluster_self_url=args.self_url,
|
||||
stats_db=getattr(args, "stats_db", None),
|
||||
relay_url=args.relay_url,
|
||||
)
|
||||
port = server.start()
|
||||
|
||||
@@ -25,6 +25,7 @@ import http.server
|
||||
import hashlib
|
||||
import json
|
||||
import socketserver
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
@@ -55,6 +56,171 @@ DEFAULT_QUANTIZATIONS = ["bfloat16"]
|
||||
DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0
|
||||
|
||||
|
||||
class _RollingCounter:
|
||||
"""Circular-bucket request counter.
|
||||
|
||||
Tracks events in `num_buckets` slots, each covering `bucket_seconds` of wall time.
|
||||
Old buckets are silently discarded when their epoch expires.
|
||||
"""
|
||||
|
||||
def __init__(self, num_buckets: int, bucket_seconds: int) -> None:
|
||||
self._num = num_buckets
|
||||
self._bsec = bucket_seconds
|
||||
self._counts: list[int] = [0] * num_buckets
|
||||
self._epochs: list[int] = [-1] * num_buckets # which epoch each slot holds
|
||||
|
||||
def _epoch(self, now: float) -> int:
|
||||
return int(now) // self._bsec
|
||||
|
||||
def record(self, now: float | None = None) -> None:
|
||||
t = now if now is not None else time.time()
|
||||
ep = self._epoch(t)
|
||||
idx = ep % self._num
|
||||
if self._epochs[idx] != ep:
|
||||
self._counts[idx] = 0
|
||||
self._epochs[idx] = ep
|
||||
self._counts[idx] += 1
|
||||
|
||||
def rpm(self, now: float | None = None) -> float:
|
||||
t = now if now is not None else time.time()
|
||||
cutoff = self._epoch(t) - self._num # epochs <= this are stale
|
||||
total = sum(self._counts[i] for i in range(self._num) if self._epochs[i] > cutoff)
|
||||
window_minutes = (self._num * self._bsec) / 60.0
|
||||
return total / window_minutes if window_minutes > 0 else 0.0
|
||||
|
||||
def buckets(self) -> list[tuple[int, int]]:
|
||||
return [(self._epochs[i], self._counts[i]) for i in range(self._num)]
|
||||
|
||||
def restore_buckets(self, data: list[tuple[int, int]]) -> None:
|
||||
for i, (ep, cnt) in enumerate(data):
|
||||
if i < self._num:
|
||||
self._epochs[i] = ep
|
||||
self._counts[i] = cnt
|
||||
|
||||
|
||||
class _ModelStats:
|
||||
"""Three rolling windows for one model: last hour, last day, last month."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.per_minute = _RollingCounter(60, 60) # 60 × 1-min buckets = 1 hour
|
||||
self.per_hour = _RollingCounter(24, 3600) # 24 × 1-hr buckets = 1 day
|
||||
self.per_day = _RollingCounter(30, 86400) # 30 × 1-day buckets = ~1 month
|
||||
|
||||
def record(self, now: float | None = None) -> None:
|
||||
t = now if now is not None else time.time()
|
||||
self.per_minute.record(t)
|
||||
self.per_hour.record(t)
|
||||
self.per_day.record(t)
|
||||
|
||||
def rpms(self, now: float | None = None) -> dict:
|
||||
t = now if now is not None else time.time()
|
||||
return {
|
||||
"rpm_last_hour": round(self.per_minute.rpm(t), 4),
|
||||
"rpm_last_day": round(self.per_hour.rpm(t), 4),
|
||||
"rpm_last_month": round(self.per_day.rpm(t), 4),
|
||||
}
|
||||
|
||||
|
||||
class _StatsCollector:
|
||||
"""Thread-safe model request stats with SQLite persistence and peer slice merging."""
|
||||
|
||||
SAVE_INTERVAL = 60.0 # seconds between DB saves
|
||||
|
||||
def __init__(self, db_path: str | None = None) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._local: dict[str, _ModelStats] = {}
|
||||
self._peer_rpms: dict[str, dict[str, dict]] = {} # tracker_url -> model -> rpms
|
||||
self._db_path = db_path
|
||||
if db_path:
|
||||
self._init_db()
|
||||
self._load_from_db()
|
||||
|
||||
# ---------- public API ----------
|
||||
|
||||
def record_request(self, model: str, now: float | None = None) -> None:
|
||||
t = now if now is not None else time.time()
|
||||
with self._lock:
|
||||
if model not in self._local:
|
||||
self._local[model] = _ModelStats()
|
||||
self._local[model].record(t)
|
||||
|
||||
def get_local_rpms(self, now: float | None = None) -> dict[str, dict]:
|
||||
t = now if now is not None else time.time()
|
||||
with self._lock:
|
||||
return {m: s.rpms(t) for m, s in self._local.items()}
|
||||
|
||||
def merge_peer_rpms(self, tracker_url: str, rpms: dict[str, dict]) -> None:
|
||||
with self._lock:
|
||||
self._peer_rpms[tracker_url] = dict(rpms)
|
||||
|
||||
def get_combined_stats(self, now: float | None = None) -> dict:
|
||||
t = now if now is not None else time.time()
|
||||
with self._lock:
|
||||
combined: dict[str, dict] = {}
|
||||
for model, ms in self._local.items():
|
||||
combined[model] = ms.rpms(t)
|
||||
for _url, peer in self._peer_rpms.items():
|
||||
for model, rpms in peer.items():
|
||||
if model not in combined:
|
||||
combined[model] = {"rpm_last_hour": 0.0, "rpm_last_day": 0.0, "rpm_last_month": 0.0}
|
||||
for key in ("rpm_last_hour", "rpm_last_day", "rpm_last_month"):
|
||||
combined[model][key] = round(combined[model].get(key, 0.0) + rpms.get(key, 0.0), 4)
|
||||
return combined
|
||||
|
||||
# ---------- persistence ----------
|
||||
|
||||
def _init_db(self) -> None:
|
||||
con = sqlite3.connect(self._db_path) # type: ignore[arg-type]
|
||||
con.execute(
|
||||
"CREATE TABLE IF NOT EXISTS model_rpm_buckets "
|
||||
"(model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, count INTEGER, "
|
||||
"PRIMARY KEY (model, window, bucket_idx))"
|
||||
)
|
||||
con.commit()
|
||||
con.close()
|
||||
|
||||
def save_to_db(self) -> None:
|
||||
if not self._db_path:
|
||||
return
|
||||
with self._lock:
|
||||
rows = []
|
||||
for model, ms in self._local.items():
|
||||
for window_name, counter in (
|
||||
("hour", ms.per_minute),
|
||||
("day", ms.per_hour),
|
||||
("month", ms.per_day),
|
||||
):
|
||||
for idx, (ep, cnt) in enumerate(counter.buckets()):
|
||||
if ep >= 0:
|
||||
rows.append((model, window_name, idx, ep, cnt))
|
||||
con = sqlite3.connect(self._db_path) # type: ignore[arg-type]
|
||||
con.executemany(
|
||||
"INSERT OR REPLACE INTO model_rpm_buckets VALUES (?,?,?,?,?)", rows
|
||||
)
|
||||
con.commit()
|
||||
con.close()
|
||||
|
||||
def _load_from_db(self) -> None:
|
||||
con = sqlite3.connect(self._db_path) # type: ignore[arg-type]
|
||||
rows = con.execute("SELECT model, window, bucket_idx, bucket_epoch, count FROM model_rpm_buckets").fetchall()
|
||||
con.close()
|
||||
grouped: dict[str, dict[str, list[tuple[int, int]]]] = {}
|
||||
for model, window, idx, ep, cnt in rows:
|
||||
grouped.setdefault(model, {}).setdefault(window, []).append((idx, ep, cnt))
|
||||
for model, windows in grouped.items():
|
||||
ms = _ModelStats()
|
||||
for window_name, entries in windows.items():
|
||||
counter = {"hour": ms.per_minute, "day": ms.per_hour, "month": ms.per_day}.get(window_name)
|
||||
if counter is None:
|
||||
continue
|
||||
data = [(0, -1)] * counter._num
|
||||
for idx, ep, cnt in entries:
|
||||
if 0 <= idx < counter._num:
|
||||
data[idx] = (ep, cnt)
|
||||
counter.restore_buckets(data)
|
||||
self._local[model] = ms
|
||||
|
||||
|
||||
class _NodeEntry:
|
||||
__slots__ = (
|
||||
"node_id", "endpoint", "shard_start", "shard_end",
|
||||
@@ -63,6 +229,12 @@ class _NodeEntry:
|
||||
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
||||
"pending_directives", "last_heartbeat", "tracker_mode",
|
||||
"relay_addr", "cert_fingerprint", "peer_id",
|
||||
# heartbeat stats (reported by node, cumulative)
|
||||
"total_requests", "failed_requests", "queue_depth", "uptime_seconds",
|
||||
"status", # "ready" | "loading"
|
||||
"heartbeats_expected", "heartbeats_received",
|
||||
# dynamic reassignment queued by the tracker
|
||||
"pending_new_assignment",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
@@ -112,6 +284,19 @@ class _NodeEntry:
|
||||
self.peer_id = peer_id
|
||||
self.pending_directives: list[dict] = []
|
||||
self.last_heartbeat: float = time.monotonic()
|
||||
self.total_requests: int = 0
|
||||
self.failed_requests: int = 0
|
||||
self.queue_depth: int = 0
|
||||
self.uptime_seconds: float = 0.0
|
||||
self.status: str = "ready"
|
||||
self.heartbeats_expected: int = 0
|
||||
self.heartbeats_received: int = 0
|
||||
self.pending_new_assignment: dict | None = None
|
||||
|
||||
|
||||
def _effective_throughput(node: "_NodeEntry") -> float:
|
||||
"""Effective tokens/s accounting for current queue depth."""
|
||||
return node.benchmark_tokens_per_sec / (node.queue_depth + 1)
|
||||
|
||||
|
||||
def _select_route(
|
||||
@@ -119,7 +304,12 @@ def _select_route(
|
||||
required_start: int,
|
||||
required_end: int,
|
||||
) -> tuple[list[_NodeEntry], str]:
|
||||
"""Greedy interval-cover. Returns (ordered route, error_message)."""
|
||||
"""Greedy interval-cover biased toward fast, lightly-loaded nodes.
|
||||
|
||||
Among nodes that equally advance coverage, prefer the one with higher
|
||||
effective throughput: benchmark_tokens_per_sec / (queue_depth + 1).
|
||||
Tiebreak: higher shard_end (fewer hops).
|
||||
"""
|
||||
candidates = sorted(
|
||||
[node for node in nodes if node.shard_start is not None and node.shard_end is not None],
|
||||
key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator]
|
||||
@@ -131,7 +321,11 @@ def _select_route(
|
||||
best: _NodeEntry | None = None
|
||||
for node in candidates:
|
||||
if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to:
|
||||
if best is None or node.shard_end > best.shard_end:
|
||||
if best is None:
|
||||
best = node
|
||||
elif node.shard_end > best.shard_end:
|
||||
best = node
|
||||
elif node.shard_end == best.shard_end and _effective_throughput(node) > _effective_throughput(best):
|
||||
best = node
|
||||
if best is None:
|
||||
missing = covered_up_to + 1
|
||||
@@ -232,6 +426,81 @@ def _coverage_map(
|
||||
return coverage
|
||||
|
||||
|
||||
def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict:
|
||||
"""Per-node health detail for the availability map."""
|
||||
age = time.monotonic() - node.last_heartbeat
|
||||
alive = age <= heartbeat_timeout
|
||||
hb_expected = max(1, round(node.uptime_seconds / 20.0)) # assume ~20s interval
|
||||
hb_rate = round(min(1.0, node.heartbeats_received / hb_expected), 4) if node.heartbeats_received else 0.0
|
||||
total = node.total_requests
|
||||
inf_rate = round((total - node.failed_requests) / total, 4) if total > 0 else 1.0
|
||||
return {
|
||||
"node_id": node.node_id,
|
||||
"endpoint": node.endpoint,
|
||||
"alive": alive,
|
||||
"last_seen_seconds_ago": round(age, 1),
|
||||
"status": node.status,
|
||||
"queue_depth": node.queue_depth,
|
||||
"total_requests": node.total_requests,
|
||||
"heartbeat_success_rate": hb_rate,
|
||||
"inference_success_rate": inf_rate,
|
||||
}
|
||||
|
||||
|
||||
def _coverage_map_detailed(
|
||||
nodes: list["_NodeEntry"],
|
||||
required_start: int,
|
||||
required_end: int,
|
||||
heartbeat_timeout: float,
|
||||
) -> list[dict]:
|
||||
"""Like _coverage_map but with per-node identity and health in each band.
|
||||
|
||||
Includes all nodes (alive and stale) so operators can see both coverage
|
||||
holes and which dead nodes used to fill them.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
|
||||
def covers(node: "_NodeEntry", layer: int) -> bool:
|
||||
return (
|
||||
node.shard_start is not None
|
||||
and node.shard_end is not None
|
||||
and node.shard_start <= layer <= node.shard_end
|
||||
)
|
||||
|
||||
# Build per-layer list of nodes (with alive flag)
|
||||
layer_nodes: list[list[tuple["_NodeEntry", bool]]] = []
|
||||
for layer in range(required_start, required_end + 1):
|
||||
ns = [
|
||||
(n, (now - n.last_heartbeat) <= heartbeat_timeout)
|
||||
for n in nodes
|
||||
if covers(n, layer)
|
||||
]
|
||||
layer_nodes.append(ns)
|
||||
|
||||
# Merge consecutive layers with identical node-set (same node_ids, same alive)
|
||||
coverage: list[dict] = []
|
||||
for i, layer_idx in enumerate(range(required_start, required_end + 1)):
|
||||
ns = layer_nodes[i]
|
||||
node_ids = [n.node_id for n, _ in ns]
|
||||
alive_flags = [a for _, a in ns]
|
||||
# Compare with last band
|
||||
if (
|
||||
coverage
|
||||
and coverage[-1]["end_layer"] == layer_idx - 1
|
||||
and [nd["node_id"] for nd in coverage[-1]["nodes"]] == node_ids
|
||||
and [nd["alive"] for nd in coverage[-1]["nodes"]] == alive_flags
|
||||
):
|
||||
coverage[-1]["end_layer"] = layer_idx
|
||||
else:
|
||||
coverage.append({
|
||||
"start_layer": layer_idx,
|
||||
"end_layer": layer_idx,
|
||||
"node_count": len(ns),
|
||||
"nodes": [_node_health(n, heartbeat_timeout) for n, _ in ns],
|
||||
})
|
||||
return coverage
|
||||
|
||||
|
||||
def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
|
||||
return [
|
||||
(segment["start_layer"], segment["end_layer"])
|
||||
@@ -466,6 +735,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
|
||||
relay_url: str | None = None,
|
||||
raft: "RaftNode | None" = None,
|
||||
gossip: "NodeGossip | None" = None,
|
||||
stats: "_StatsCollector | None" = None,
|
||||
) -> None:
|
||||
super().__init__(addr, handler)
|
||||
self.registry = registry
|
||||
@@ -477,6 +747,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
|
||||
self.relay_url = relay_url.rstrip("/") if relay_url else None
|
||||
self.raft = raft
|
||||
self.gossip = gossip
|
||||
self.stats: _StatsCollector | None = stats
|
||||
|
||||
|
||||
class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
@@ -526,6 +797,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if self.path == "/v1/gossip":
|
||||
self._handle_gossip()
|
||||
return
|
||||
if self.path == "/v1/stats/gossip":
|
||||
self._handle_stats_gossip()
|
||||
return
|
||||
parts = self.path.split("/")
|
||||
# /v1/nodes/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', 'heartbeat']
|
||||
if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
|
||||
@@ -556,6 +830,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._handle_tracker_nodes(model)
|
||||
elif parsed.path == "/v1/raft/status":
|
||||
self._handle_raft_status()
|
||||
elif parsed.path == "/v1/stats":
|
||||
self._handle_stats()
|
||||
elif parsed.path == "/v1/health":
|
||||
self._send_json(200, {"status": "ok"})
|
||||
else:
|
||||
@@ -645,19 +921,20 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
def _handle_coverage(self, model: str):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
# Do NOT purge before coverage — dead nodes are included with alive=false
|
||||
# so operators can see what was covering each layer band before failure.
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
resolved = _nodes_and_bounds_for_model(server, model)
|
||||
if resolved is None:
|
||||
self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
|
||||
return
|
||||
alive, required_start, required_end = resolved
|
||||
all_nodes, required_start, required_end = resolved
|
||||
if server.contracts is not None:
|
||||
alive = [
|
||||
node for node in alive
|
||||
all_nodes = [
|
||||
node for node in all_nodes
|
||||
if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
|
||||
]
|
||||
coverage = _coverage_map(alive, required_start, required_end)
|
||||
coverage = _coverage_map_detailed(all_nodes, required_start, required_end, server.heartbeat_timeout)
|
||||
self._send_json(200, {"model": model, "coverage": coverage})
|
||||
|
||||
def _handle_tracker_nodes(self, model: str):
|
||||
@@ -743,6 +1020,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
model: str = body.get("model", "")
|
||||
is_stream: bool = bool(body.get("stream", False))
|
||||
|
||||
if model and server.stats is not None:
|
||||
server.stats.record_request(model)
|
||||
|
||||
# Find a live tracker-mode node for this model
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
@@ -789,8 +1069,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
]
|
||||
rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
|
||||
route_nodes, _ = _select_route(all_nodes, rs, re)
|
||||
# Compute start_layer for each hop: each node begins where the previous ended + 1.
|
||||
# This allows overlapping shard registrations without double-computation.
|
||||
covered_up_to = rs - 1
|
||||
route_hops: list[dict] = []
|
||||
for rn in route_nodes:
|
||||
route_hops.append({"endpoint": rn.endpoint, "start_layer": covered_up_to + 1})
|
||||
covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to
|
||||
# Strip the first-shard node we're about to proxy to — it's already handling the request.
|
||||
downstream_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint])
|
||||
downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")]
|
||||
downstream_urls = json.dumps(downstream_hops)
|
||||
route_debug = " -> ".join(
|
||||
f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
|
||||
for n in route_nodes
|
||||
@@ -1131,21 +1419,39 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._send_json(404, {"error": "node not found"})
|
||||
return
|
||||
entry.last_heartbeat = time.monotonic()
|
||||
entry.heartbeats_received += 1
|
||||
# P2P metadata
|
||||
if body.get("relay_addr"):
|
||||
entry.relay_addr = body["relay_addr"]
|
||||
if body.get("cert_fingerprint"):
|
||||
entry.cert_fingerprint = body["cert_fingerprint"]
|
||||
if body.get("peer_id"):
|
||||
entry.peer_id = body["peer_id"]
|
||||
# Node stats (cumulative — node always sends totals, tracker replaces)
|
||||
if "total_requests" in body:
|
||||
entry.total_requests = int(body["total_requests"])
|
||||
if "failed_requests" in body:
|
||||
entry.failed_requests = int(body["failed_requests"])
|
||||
if "queue_depth" in body:
|
||||
entry.queue_depth = int(body["queue_depth"])
|
||||
if "uptime_seconds" in body:
|
||||
entry.uptime_seconds = float(body["uptime_seconds"])
|
||||
if "status" in body and body["status"] in ("ready", "loading"):
|
||||
entry.status = body["status"]
|
||||
_rebalance_model_locked(server, entry.model or "stub-model")
|
||||
directives = list(entry.pending_directives)
|
||||
entry.pending_directives.clear()
|
||||
new_assignment = entry.pending_new_assignment
|
||||
if new_assignment is not None:
|
||||
entry.pending_new_assignment = None
|
||||
if server.gossip is not None:
|
||||
server.gossip.record(node_id)
|
||||
resp: dict = {}
|
||||
if directives:
|
||||
self._send_json(200, {"directives": directives})
|
||||
else:
|
||||
self._send_json(200, {})
|
||||
resp["directives"] = directives
|
||||
if new_assignment is not None:
|
||||
resp["new_assignment"] = new_assignment
|
||||
self._send_json(200, resp)
|
||||
|
||||
# ---------------------------------------------------------------- Raft handlers
|
||||
|
||||
@@ -1187,6 +1493,24 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
server.gossip.merge({k: float(v) for k, v in body.items()})
|
||||
self._send_json(200, {})
|
||||
|
||||
def _handle_stats(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
if server.stats is None:
|
||||
self._send_json(200, {"models": {}})
|
||||
return
|
||||
self._send_json(200, {"models": server.stats.get_combined_stats()})
|
||||
|
||||
def _handle_stats_gossip(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
tracker_url = body.get("tracker_url", "")
|
||||
rpms = body.get("stats", {})
|
||||
if server.stats is not None and tracker_url and isinstance(rpms, dict):
|
||||
server.stats.merge_peer_rpms(tracker_url, rpms)
|
||||
self._send_json(200, {})
|
||||
|
||||
def _handle_assign(self, parsed: urllib.parse.ParseResult):
|
||||
"""Return an optimal shard assignment for a node given its hardware profile.
|
||||
|
||||
@@ -1334,8 +1658,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
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).
|
||||
# Smart scoring: demand_rpm × coverage_deficit
|
||||
# coverage_deficit = uncovered_layers / total_layers (0 = fully covered)
|
||||
# demand_rpm comes from the stats collector; defaults to 0.0 when no traffic yet.
|
||||
demand_rpms: dict[str, float] = {}
|
||||
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)
|
||||
|
||||
best_repo = None
|
||||
best_score = -1.0
|
||||
best_gap_size = -1
|
||||
best_gap_start = 0
|
||||
best_num_layers = 0
|
||||
@@ -1353,8 +1686,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
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:
|
||||
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)
|
||||
if score > best_score or (score == best_score and gap_size > best_gap_size):
|
||||
best_score = score
|
||||
best_gap_size = gap_size
|
||||
best_gap_start = gap_start
|
||||
best_repo = repo
|
||||
@@ -1362,8 +1700,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
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]))
|
||||
# All shards covered — assign to highest-demand model for redundancy.
|
||||
best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0))
|
||||
best_gap_start = 0
|
||||
best_num_layers = repo_layers[best_repo]
|
||||
|
||||
@@ -1383,6 +1721,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
"shard_end": shard_end,
|
||||
"num_layers": total_l,
|
||||
"gap_found": gap_found,
|
||||
"price_per_token": 0.0,
|
||||
})
|
||||
|
||||
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
||||
@@ -1428,12 +1767,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._send_json(503, {"error": error})
|
||||
return
|
||||
|
||||
covered_up_to = required_start - 1
|
||||
route_with_start: list[tuple] = []
|
||||
for rn in route:
|
||||
route_with_start.append((rn, covered_up_to + 1))
|
||||
covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to
|
||||
|
||||
self._send_json(200, {
|
||||
"route": [e.endpoint for e in route],
|
||||
"route": [e.endpoint for e, _ in route_with_start],
|
||||
"nodes": [
|
||||
{
|
||||
"node_id": e.node_id,
|
||||
"endpoint": e.endpoint,
|
||||
"start_layer": start,
|
||||
"relay_addr": e.relay_addr,
|
||||
"peer_id": e.peer_id,
|
||||
"wallet_address": e.wallet_address,
|
||||
@@ -1443,7 +1789,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
"shard_checksum": e.shard_checksum,
|
||||
"score": e.score,
|
||||
}
|
||||
for e in route
|
||||
for e, start in route_with_start
|
||||
],
|
||||
})
|
||||
|
||||
@@ -1531,6 +1877,7 @@ class TrackerServer:
|
||||
minimum_stake: int = 0,
|
||||
cluster_peers: list[str] | None = None,
|
||||
cluster_self_url: str | None = None,
|
||||
stats_db: str | None = None,
|
||||
relay_url: str | None = None,
|
||||
) -> None:
|
||||
self._host = host
|
||||
@@ -1553,6 +1900,9 @@ class TrackerServer:
|
||||
self._rebalance_thread: threading.Thread | None = None
|
||||
self._raft: RaftNode | None = None
|
||||
self._gossip: NodeGossip | None = None
|
||||
self._stats: _StatsCollector | None = _StatsCollector(db_path=stats_db) if stats_db else _StatsCollector()
|
||||
self._stats_stop = threading.Event()
|
||||
self._stats_thread: threading.Thread | None = None
|
||||
self.port: int | None = None
|
||||
|
||||
def start(self) -> int:
|
||||
@@ -1569,7 +1919,8 @@ class TrackerServer:
|
||||
self._model_presets,
|
||||
self._contracts,
|
||||
self._minimum_stake,
|
||||
self._relay_url,
|
||||
relay_url=self._relay_url,
|
||||
stats=self._stats,
|
||||
)
|
||||
self.port = self._server.server_address[1]
|
||||
|
||||
@@ -1588,10 +1939,13 @@ class TrackerServer:
|
||||
self._gossip.start()
|
||||
|
||||
self._rebalance_stop.clear()
|
||||
self._stats_stop.clear()
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
self._thread.start()
|
||||
self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True)
|
||||
self._rebalance_thread.start()
|
||||
self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True)
|
||||
self._stats_thread.start()
|
||||
return self.port
|
||||
|
||||
def _raft_apply(self, command: str, payload: dict) -> None:
|
||||
@@ -1635,6 +1989,28 @@ class TrackerServer:
|
||||
_purge_expired_nodes_locked(server)
|
||||
_rebalance_all_locked(server)
|
||||
|
||||
def _stats_loop(self) -> None:
|
||||
"""Periodically save stats to DB and push local slice to cluster peers."""
|
||||
while not self._stats_stop.wait(_StatsCollector.SAVE_INTERVAL):
|
||||
if self._stats is not None:
|
||||
self._stats.save_to_db()
|
||||
if self._stats is not None and self._cluster_peers:
|
||||
self_url = self._cluster_self_url or f"http://{self._host}:{self.port}"
|
||||
local_rpms = self._stats.get_local_rpms()
|
||||
body = json.dumps({"tracker_url": self_url, "stats": local_rpms}).encode()
|
||||
for peer in self._cluster_peers:
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
f"{peer}/v1/stats/gossip",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=2.0) as r:
|
||||
r.read()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def stop(self) -> None:
|
||||
if self._raft is not None:
|
||||
self._raft.stop()
|
||||
@@ -1643,15 +2019,21 @@ class TrackerServer:
|
||||
if self._server is None:
|
||||
return
|
||||
self._rebalance_stop.set()
|
||||
self._stats_stop.set()
|
||||
if self._stats is not None:
|
||||
self._stats.save_to_db()
|
||||
self._server.shutdown()
|
||||
self._server.server_close()
|
||||
if self._thread is not None:
|
||||
self._thread.join(timeout=1)
|
||||
if self._rebalance_thread is not None:
|
||||
self._rebalance_thread.join(timeout=1)
|
||||
if self._stats_thread is not None:
|
||||
self._stats_thread.join(timeout=1)
|
||||
self._server = None
|
||||
self._thread = None
|
||||
self._rebalance_thread = None
|
||||
self._stats_thread = None
|
||||
self._raft = None
|
||||
self._gossip = None
|
||||
self.port = None
|
||||
|
||||
@@ -36,7 +36,7 @@ class _FakeBackend:
|
||||
position_ids_header=None,
|
||||
)
|
||||
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header):
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
||||
assert shape == [1, 6, 8]
|
||||
return TensorPayload(
|
||||
body=body,
|
||||
@@ -50,7 +50,7 @@ class _FakeTailBackend(_FakeBackend):
|
||||
is_head = False
|
||||
is_tail = True
|
||||
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header):
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
||||
assert len(body) == 1 * 6 * 8 * 2
|
||||
return " Paris"
|
||||
|
||||
|
||||
@@ -243,10 +243,24 @@ def test_tracker_coverage_endpoint_reports_uncovered_ranges():
|
||||
)
|
||||
|
||||
coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
|
||||
assert coverage_resp["coverage"] == [
|
||||
{"start_layer": 0, "end_layer": 2, "node_count": 1},
|
||||
{"start_layer": 3, "end_layer": 5, "node_count": 0},
|
||||
]
|
||||
bands = coverage_resp["coverage"]
|
||||
assert len(bands) == 2
|
||||
# Covered band: layers 0-2, one node
|
||||
assert bands[0]["start_layer"] == 0
|
||||
assert bands[0]["end_layer"] == 2
|
||||
assert bands[0]["node_count"] == 1
|
||||
assert len(bands[0]["nodes"]) == 1
|
||||
node_info = bands[0]["nodes"][0]
|
||||
assert node_info["endpoint"] == "http://127.0.0.1:9001"
|
||||
assert node_info["alive"] is True
|
||||
assert "last_seen_seconds_ago" in node_info
|
||||
assert "heartbeat_success_rate" in node_info
|
||||
assert "inference_success_rate" in node_info
|
||||
# Gap band: layers 3-5, no nodes
|
||||
assert bands[1]["start_layer"] == 3
|
||||
assert bands[1]["end_layer"] == 5
|
||||
assert bands[1]["node_count"] == 0
|
||||
assert bands[1]["nodes"] == []
|
||||
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
||||
@@ -283,11 +297,12 @@ def test_tracker_coverage_endpoint_accepts_registered_hf_repo_or_short_name():
|
||||
)
|
||||
|
||||
assert by_repo["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert by_repo["coverage"] == [
|
||||
{"start_layer": 0, "end_layer": 19, "node_count": 1},
|
||||
{"start_layer": 20, "end_layer": 21, "node_count": 2},
|
||||
{"start_layer": 22, "end_layer": 23, "node_count": 1},
|
||||
]
|
||||
bands = by_repo["coverage"]
|
||||
assert len(bands) == 3
|
||||
assert bands[0]["start_layer"] == 0 and bands[0]["end_layer"] == 19 and bands[0]["node_count"] == 1
|
||||
assert bands[1]["start_layer"] == 20 and bands[1]["end_layer"] == 21 and bands[1]["node_count"] == 2
|
||||
assert bands[2]["start_layer"] == 22 and bands[2]["end_layer"] == 23 and bands[2]["node_count"] == 1
|
||||
assert all("nodes" in b for b in bands)
|
||||
assert by_short_name["coverage"] == by_repo["coverage"]
|
||||
finally:
|
||||
tracker.stop()
|
||||
@@ -1074,3 +1089,229 @@ def test_two_node_pipeline_via_tracker():
|
||||
node_a.stop()
|
||||
node_b.stop()
|
||||
tracker.stop()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- US-025 stats
|
||||
|
||||
def test_stats_endpoint_returns_empty_models_initially():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
resp = _get_json(f"http://127.0.0.1:{port}/v1/stats")
|
||||
assert "models" in resp
|
||||
assert isinstance(resp["models"], dict)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_stats_rolling_counter_rpm_reflects_recorded_requests():
|
||||
from meshnet_tracker.server import _RollingCounter
|
||||
|
||||
counter = _RollingCounter(num_buckets=60, bucket_seconds=60)
|
||||
now = 1_000_000.0
|
||||
for _ in range(60):
|
||||
counter.record(now)
|
||||
|
||||
rpm = counter.rpm(now)
|
||||
assert rpm == 1.0 # 60 requests over 60 minutes = 1 rpm
|
||||
|
||||
|
||||
def test_stats_collector_records_and_returns_rpms():
|
||||
from meshnet_tracker.server import _StatsCollector
|
||||
|
||||
collector = _StatsCollector()
|
||||
now = 1_000_000.0
|
||||
for _ in range(30):
|
||||
collector.record_request("my-model", now=now)
|
||||
|
||||
combined = collector.get_combined_stats(now=now)
|
||||
assert "my-model" in combined
|
||||
assert combined["my-model"]["rpm_last_hour"] > 0
|
||||
assert combined["my-model"]["rpm_last_day"] > 0
|
||||
assert combined["my-model"]["rpm_last_month"] > 0
|
||||
|
||||
|
||||
def test_stats_collector_merges_peer_rpms_additively():
|
||||
from meshnet_tracker.server import _StatsCollector
|
||||
|
||||
collector = _StatsCollector()
|
||||
now = 1_000_000.0
|
||||
collector.record_request("modelX", now=now) # 1 local request in the current minute bucket
|
||||
|
||||
peer_rpms = {"modelX": {"rpm_last_hour": 5.0, "rpm_last_day": 3.0, "rpm_last_month": 1.0}}
|
||||
collector.merge_peer_rpms("http://peer-tracker:8080", peer_rpms)
|
||||
|
||||
combined = collector.get_combined_stats(now=now)
|
||||
assert combined["modelX"]["rpm_last_hour"] > 5.0 # local + peer
|
||||
assert combined["modelX"]["rpm_last_day"] > 3.0
|
||||
assert combined["modelX"]["rpm_last_month"] >= 1.0 # local rounds to 0 at 4dp, peer is 1.0
|
||||
|
||||
|
||||
def test_stats_sqlite_persistence_survives_restart(tmp_path):
|
||||
from meshnet_tracker.server import _StatsCollector
|
||||
|
||||
db = str(tmp_path / "stats.db")
|
||||
now = 1_000_000.0
|
||||
|
||||
collector = _StatsCollector(db_path=db)
|
||||
for _ in range(10):
|
||||
collector.record_request("persistent-model", now=now)
|
||||
collector.save_to_db()
|
||||
|
||||
collector2 = _StatsCollector(db_path=db)
|
||||
combined = collector2.get_combined_stats(now=now)
|
||||
assert "persistent-model" in combined
|
||||
assert combined["persistent-model"]["rpm_last_hour"] > 0
|
||||
|
||||
|
||||
def test_stats_gossip_endpoint_merges_peer_slice():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
peer_payload = {
|
||||
"tracker_url": "http://peer:9000",
|
||||
"stats": {
|
||||
"remote-model": {
|
||||
"rpm_last_hour": 2.5,
|
||||
"rpm_last_day": 1.0,
|
||||
"rpm_last_month": 0.5,
|
||||
}
|
||||
},
|
||||
}
|
||||
_post_json(f"http://127.0.0.1:{port}/v1/stats/gossip", peer_payload)
|
||||
|
||||
combined = _get_json(f"http://127.0.0.1:{port}/v1/stats")
|
||||
assert "remote-model" in combined["models"]
|
||||
assert combined["models"]["remote-model"]["rpm_last_hour"] == 2.5
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- US-027 / US-028 routing
|
||||
|
||||
def _make_node(node_id, shard_start, shard_end, bench=1.0, queue_depth=0):
|
||||
"""Helper: build a _NodeEntry with minimal fields."""
|
||||
from meshnet_tracker.server import _NodeEntry
|
||||
n = _NodeEntry(
|
||||
node_id=node_id, endpoint=f"http://fake:{9000 + shard_start}",
|
||||
shard_start=shard_start, shard_end=shard_end,
|
||||
model="m", shard_checksum=None, hardware_profile={},
|
||||
wallet_address=None, score=1.0,
|
||||
benchmark_tokens_per_sec=bench,
|
||||
)
|
||||
n.queue_depth = queue_depth
|
||||
return n
|
||||
|
||||
|
||||
def test_select_route_overlapping_shards_uses_start_layer_protocol():
|
||||
"""_select_route with overlapping shards: A(0-22) and B(20-24).
|
||||
Route must be [A, B]; B's effective start_layer = 23 (after A ends at 22).
|
||||
"""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
node_a = _make_node("A", 0, 22)
|
||||
node_b = _make_node("B", 20, 24)
|
||||
|
||||
route, err = _select_route([node_a, node_b], 0, 24)
|
||||
assert err == ""
|
||||
assert len(route) == 2
|
||||
assert route[0].node_id == "A"
|
||||
assert route[1].node_id == "B"
|
||||
|
||||
|
||||
def test_select_route_no_overlap_three_nodes():
|
||||
"""Non-overlapping shards are covered in the right order."""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
nodes = [_make_node("C", 8, 15), _make_node("A", 0, 7), _make_node("B", 16, 23)]
|
||||
route, err = _select_route(nodes, 0, 23)
|
||||
assert err == ""
|
||||
# Greedy cover order: A(0-7) → C(8-15) → B(16-23)
|
||||
assert [n.node_id for n in route] == ["A", "C", "B"]
|
||||
|
||||
|
||||
def test_select_route_gap_returns_error():
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
nodes = [_make_node("A", 0, 5), _make_node("B", 10, 15)]
|
||||
route, err = _select_route(nodes, 0, 15)
|
||||
assert route == []
|
||||
assert "layer 6" in err
|
||||
|
||||
|
||||
def test_select_route_prefers_higher_throughput_node_when_shards_equal():
|
||||
"""When two nodes cover the same interval, prefer the faster / less-loaded one."""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
fast = _make_node("fast", 0, 11, bench=10.0, queue_depth=0) # effective = 10.0
|
||||
slow = _make_node("slow", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0
|
||||
|
||||
route, err = _select_route([slow, fast], 0, 11)
|
||||
assert err == ""
|
||||
assert len(route) == 1
|
||||
assert route[0].node_id == "fast"
|
||||
|
||||
|
||||
def test_select_route_throughput_accounts_for_queue_depth():
|
||||
"""A nominally faster node with high queue depth loses to a slower idle node."""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
busy_fast = _make_node("busy", 0, 11, bench=20.0, queue_depth=9) # effective = 2.0
|
||||
idle_slow = _make_node("idle", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0
|
||||
|
||||
route, err = _select_route([busy_fast, idle_slow], 0, 11)
|
||||
assert err == ""
|
||||
# tie → first in list wins; exact tie does not matter as long as no error
|
||||
assert len(route) == 1
|
||||
assert route[0].node_id in ("busy", "idle")
|
||||
|
||||
|
||||
def test_two_stub_nodes_complete_pipeline_via_tracker():
|
||||
"""Integration: two StubNodeServer instances serving complementary shards
|
||||
produce a full inference response through the tracker route."""
|
||||
tracker = TrackerServer(model_presets={
|
||||
"two-node-model": {"layers_start": 0, "layers_end": 3}
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
|
||||
from meshnet_node.server import StubNodeServer
|
||||
node_a = StubNodeServer()
|
||||
node_b = StubNodeServer()
|
||||
port_a = node_a.start()
|
||||
port_b = node_b.start()
|
||||
|
||||
try:
|
||||
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {
|
||||
"endpoint": f"http://127.0.0.1:{port_a}",
|
||||
"model": "two-node-model", "shard_start": 0, "shard_end": 1,
|
||||
"hardware_profile": {}, "score": 1.0,
|
||||
})
|
||||
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {
|
||||
"endpoint": f"http://127.0.0.1:{port_b}",
|
||||
"model": "two-node-model", "shard_start": 2, "shard_end": 3,
|
||||
"hardware_profile": {}, "score": 1.0,
|
||||
})
|
||||
|
||||
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=two-node-model")
|
||||
assert len(route_resp["route"]) == 2
|
||||
finally:
|
||||
node_a.stop()
|
||||
node_b.stop()
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_route_timeout_config_is_exposed_on_server():
|
||||
"""TorchNodeServer exposes the route_timeout that was configured."""
|
||||
from meshnet_node.torch_server import TorchNodeServer
|
||||
|
||||
class _MinimalBackend:
|
||||
model_id = "m"
|
||||
total_layers = 4
|
||||
is_head = True
|
||||
is_tail = True
|
||||
def generate_text(self, *a, **kw): return ""
|
||||
def count_prompt_tokens(self, *a): return 0
|
||||
def count_text_tokens(self, *a): return 0
|
||||
|
||||
node = TorchNodeServer(backend=_MinimalBackend(), route_timeout=45.0)
|
||||
assert node.route_timeout == 45.0
|
||||
|
||||
Reference in New Issue
Block a user