This commit is contained in:
Dobromir Popov
2026-07-08 20:00:56 +03:00
28 changed files with 14182 additions and 12368 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -197,12 +197,43 @@ def _max_assignable_layers(
memory_mb: int,
total_layers: int | None,
bytes_per_layer: int | None = None,
*,
safety_fraction: float = 0.8,
) -> int:
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
return 0
budget_bytes = memory_mb * 1024 * 1024
layer_bytes = bytes_per_layer or _DEFAULT_BYTES_PER_LAYER
return min(total_layers, int((budget_bytes * 0.8) // layer_bytes))
return min(total_layers, int((budget_bytes * safety_fraction) // layer_bytes))
def _runtime_shard_safety_fraction(device: str) -> float:
"""CPU partial loads need room for model skeletons, tokenizer, and allocator peaks."""
return 0.55 if device != "cuda" else 0.8
def _cap_auto_assigned_shard(
shard_start: int,
shard_end: int,
total_layers: int | None,
memory_mb: int,
bytes_per_layer: int | None,
device: str,
) -> tuple[int, bool, int]:
if bytes_per_layer is None or total_layers is None:
return shard_end, False, 0
max_layers = _max_assignable_layers(
memory_mb,
total_layers,
bytes_per_layer=bytes_per_layer,
safety_fraction=_runtime_shard_safety_fraction(device),
)
if max_layers <= 0:
return shard_end, False, max_layers
assigned_layers = shard_end - shard_start + 1
if assigned_layers <= max_layers:
return shard_end, False, max_layers
return min(total_layers - 1, shard_start + max_layers - 1), True, max_layers
def _format_shard_label(
@@ -226,13 +257,19 @@ def _shard_budget_line(
total_layers: int | None,
quantization: str,
bytes_per_layer: int | None = None,
safety_fraction: float = 0.8,
) -> str:
memory_gb = memory_mb / 1024
gb_str = f"{memory_gb:.1f} GB"
budget_quantization = "bfloat16" if quantization == "auto" else quantization
if total_layers is None or total_layers <= 0:
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
max_layers = _max_assignable_layers(memory_mb, total_layers, bytes_per_layer=bytes_per_layer)
max_layers = _max_assignable_layers(
memory_mb,
total_layers,
bytes_per_layer=bytes_per_layer,
safety_fraction=safety_fraction,
)
# Remaining capacity after one full model load (rough estimate)
shard_bytes = max_layers * (bytes_per_layer or _DEFAULT_BYTES_PER_LAYER)
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
@@ -349,25 +386,38 @@ def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHt
_PENDING_NODE_ID = "pending"
_HEARTBEAT_INTERVAL_IDLE = 20.0
_HEARTBEAT_INTERVAL_BUSY = 3.0
def _start_heartbeat(
tracker_url: str,
node_id: str,
register_payload: dict,
interval: float = 20.0,
interval: float = _HEARTBEAT_INTERVAL_IDLE,
node_ref: Any | None = None,
start_time: float | None = None,
) -> threading.Thread:
"""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.
queue_depth, current_requests, 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 _current_requests_snapshot() -> list[dict]:
if node_ref is None:
return []
getter = getattr(node_ref, "current_requests", None)
if getter is None:
return []
current = getter() if callable(getter) else getter
return list(current) if isinstance(current, list) else []
def _get_stats() -> dict:
uptime = time.monotonic() - _start_time
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
@@ -379,8 +429,16 @@ def _start_heartbeat(
)
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
current_requests = _current_requests_snapshot()
if current_requests:
stats["current_requests"] = current_requests
return stats
def _sleep_interval() -> float:
if _current_requests_snapshot() or (node_ref is not None and getattr(node_ref, "queue_depth", 0) > 0):
return _HEARTBEAT_INTERVAL_BUSY
return interval
def _reregister() -> bool:
nonlocal node_id
try:
@@ -442,7 +500,7 @@ def _start_heartbeat(
outage_streak = 1 if node_id == _PENDING_NODE_ID else 0
while True:
time.sleep(interval)
time.sleep(_sleep_interval())
if outage_streak > 0:
# Tracker was down — attempt re-registration first (it may have restarted
@@ -708,6 +766,25 @@ def run_startup(
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"]
asgn_total_layers = int(net_asgn.get("num_layers") or detected)
asgn_bytes_per_layer = _assignment_bytes_per_layer(net_asgn, quantization)
capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
shard_start,
shard_end,
asgn_total_layers,
memory_budget_mb,
asgn_bytes_per_layer,
device,
)
if was_capped:
original_end = shard_end
shard_end = capped_shard_end
print(
f" WARNING: tracker assigned layers {shard_start}-{original_end}, "
f"but CPU-safe runtime budget fits {max_runtime_layers}/{asgn_total_layers} layers; "
f"loading layers {shard_start}-{shard_end} instead.",
flush=True,
)
full_sources = (
[] if tracker_source_disabled
else _full_model_sources(net_asgn.get("model_sources", []))
@@ -723,7 +800,7 @@ def run_startup(
)
print(
f" Tracker found uncovered shard: "
f"layers {shard_start}{shard_end} (of {detected})",
f"layers {shard_start}-{shard_end} (of {detected})",
flush=True,
)
except Exception:
@@ -752,6 +829,7 @@ def run_startup(
shard_label = _format_shard_label(shard_start, shard_end, total_layers)
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
node.set_advertised_endpoint(endpoint)
local_base_url = f"http://127.0.0.1:{actual_port}"
relay_bridge, relay_fields = _start_relay_bridge_if_available(
tracker_url,
@@ -838,18 +916,36 @@ def run_startup(
assigned_shard_start: int = net_assignment["shard_start"]
assigned_shard_end: int = net_assignment["shard_end"]
assigned_num_layers: int = net_assignment["num_layers"]
assigned_bytes_per_layer = _assignment_bytes_per_layer(net_assignment, quantization)
capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard(
assigned_shard_start,
assigned_shard_end,
assigned_num_layers,
memory_budget_mb,
assigned_bytes_per_layer,
device,
)
if was_capped:
original_end = assigned_shard_end
assigned_shard_end = capped_shard_end
print(
f" WARNING: tracker assigned layers {assigned_shard_start}-{original_end}, "
f"but CPU-safe runtime budget fits {max_runtime_layers}/{assigned_num_layers} layers; "
f"loading layers {assigned_shard_start}-{assigned_shard_end} instead.",
flush=True,
)
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
if _gap_found:
print(
f" Assigned gap: {assigned_hf_repo} "
f"layers {assigned_shard_start}{assigned_shard_end} "
f"layers {assigned_shard_start}-{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
else:
print(
f" Assigned redundant copy: {assigned_hf_repo} "
f"layers {assigned_shard_start}{assigned_shard_end} "
f"layers {assigned_shard_start}-{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
@@ -882,6 +978,7 @@ def run_startup(
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}"
node.set_advertised_endpoint(endpoint)
local_base_url = f"http://127.0.0.1:{actual_port}"
relay_bridge, relay_fields = _start_relay_bridge_if_available(
tracker_url,
@@ -935,7 +1032,7 @@ def run_startup(
f" Wallet: {address}\n"
f" Model ID: {assigned_hf_repo}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization, bytes_per_layer=assigned_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n"
@@ -1001,6 +1098,7 @@ def run_startup(
memory_budget_mb,
assigned_total_layers,
assignment_bytes_per_layer,
safety_fraction=_runtime_shard_safety_fraction(device),
)
if pinned_layers > max_layers:
raise ValueError(
@@ -1058,6 +1156,7 @@ def run_startup(
shard_label = f"{shard_label} (pinned)"
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
node.set_advertised_endpoint(endpoint)
local_base_url = f"http://127.0.0.1:{actual_port}"
relay_bridge, relay_fields = _start_relay_bridge_if_available(
tracker_url,
@@ -1094,7 +1193,7 @@ def run_startup(
f" Wallet: {address}\n"
f" Model ID: {hf_repo}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n"
@@ -1171,7 +1270,7 @@ def run_startup(
f"meshnet-node ready\n"
f" Wallet: {address}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {node_id}\n"
f" Hardware: {hw_str}\n"

View File

@@ -31,6 +31,69 @@ from .server import (
)
def _endpoint_key(url: str) -> str:
"""Normalize http(s) endpoints for host:port comparison."""
parsed = urllib.parse.urlparse(url.rstrip("/"))
host = (parsed.hostname or "").lower()
if not host:
return url.rstrip("/").lower()
port = parsed.port
if port is None:
port = 443 if parsed.scheme == "https" else 80
return f"{host}:{port}"
def _own_endpoint_key(server: _TorchHTTPServer) -> str:
advertised = getattr(server, "advertised_endpoint", None)
if advertised:
return _endpoint_key(advertised)
host, port = server.server_address
return _endpoint_key(f"http://{host}:{port}")
def _clamp_downstream_hops(
hops: list[dict],
backend: TorchModelShard | None,
) -> list[dict]:
"""Ensure downstream start_layer continues after this shard's layers."""
if not hops or backend is None:
return hops
shard_end = getattr(backend, "shard_end", None)
if shard_end is None:
return hops
min_start = int(shard_end) + 1
clamped: list[dict] = []
for hop in hops:
adjusted = dict(hop)
if int(adjusted.get("start_layer", 0)) < min_start:
adjusted["start_layer"] = min_start
clamped.append(adjusted)
return clamped
def _format_downstream_route(hops: list[dict]) -> str:
return ", ".join(
f"{h['endpoint']}@{h.get('start_layer', 0)}" for h in hops
)
def _write_progress_line(state: list[bool], message: str, *, final: bool = False) -> None:
"""Rewrite one in-place progress line (\\r) or finish with a newline."""
if final:
if state[0]:
sys.stdout.write("\r" + message + "\n")
state[0] = False
else:
print(message, flush=True)
return
if state[0]:
sys.stdout.write("\r" + message)
else:
sys.stdout.write(message)
state[0] = True
sys.stdout.flush()
def _relay_hop(
relay_addr: str,
path: str,
@@ -87,10 +150,31 @@ class _TorchHTTPServer(http.server.HTTPServer):
self.route_timeout = route_timeout
self.debug = debug
self.max_loaded_shards = max(1, max_loaded_shards)
self.advertised_endpoint: str | None = None
self.total_requests: int = 0
self.failed_requests: int = 0
self.queue_depth: int = 0
self._stats_lock = threading.Lock()
self._active_requests: dict[str, dict[str, Any]] = {}
def snapshot_current_requests(self) -> list[dict[str, Any]]:
"""In-flight request snapshots for tracker heartbeats."""
now = time.monotonic()
with self._stats_lock:
out: list[dict[str, Any]] = []
for rec in self._active_requests.values():
elapsed = max(now - float(rec["started"]), 1e-6)
tokens = int(rec.get("tokens") or 0)
out.append({
"request_id": str(rec["request_id"]),
"model": str(rec.get("model") or ""),
"kind": str(rec.get("kind") or "chat"),
"tokens": tokens,
"elapsed_seconds": round(elapsed, 1),
"tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
"routing_complete": bool(rec.get("routing_complete")),
})
return out
def resolve_backend(self, model_name: str | None) -> TorchModelShard | None:
if not model_name:
@@ -113,10 +197,53 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
pass
def _request_id(self) -> str:
return (
self.headers.get("X-Meshnet-Request-Id")
or self.headers.get("X-Request-Id")
or f"local-{time.time_ns():x}"
)
def _request_log_suffix(self) -> str:
req_id = self.headers.get("X-Meshnet-Request-Id") or self.headers.get("X-Request-Id")
return f" request_id={req_id}" if req_id else ""
def _track_request_begin(
self,
server: "_TorchHTTPServer",
request_id: str,
model: str,
) -> None:
with server._stats_lock:
server._active_requests[request_id] = {
"request_id": request_id,
"model": model,
"kind": "chat",
"started": time.monotonic(),
"tokens": 0,
"routing_complete": False,
}
def _track_request_progress(
self,
server: "_TorchHTTPServer",
request_id: str,
*,
tokens: int,
routing_complete: bool = False,
) -> None:
with server._stats_lock:
rec = server._active_requests.get(request_id)
if rec is None:
return
rec["tokens"] = tokens
if routing_complete:
rec["routing_complete"] = True
def _track_request_end(self, server: "_TorchHTTPServer", request_id: str) -> None:
with server._stats_lock:
server._active_requests.pop(request_id, None)
def do_POST(self):
server: _TorchHTTPServer = self.server # type: ignore[assignment]
if self.path == "/forward":
@@ -294,12 +421,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
def _handle_chat_completions(self) -> None:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
request_id = self._request_id()
with server._stats_lock:
server.total_requests += 1
server.queue_depth += 1
try:
self._do_chat_completions(server)
self._do_chat_completions(server, request_id)
finally:
self._track_request_end(server, request_id)
with server._stats_lock:
server.queue_depth -= 1
@@ -308,7 +437,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
with server._stats_lock:
server.failed_requests += 1
def _do_chat_completions(self, server: "_TorchHTTPServer") -> None:
def _do_chat_completions(self, server: "_TorchHTTPServer", request_id: str) -> None:
body = self._read_json_body()
if body is None:
return
@@ -325,6 +454,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
temperature = float(body.get("temperature") or 1.0)
top_p = float(body.get("top_p") or 1.0)
self._track_request_begin(server, request_id, model_name)
print(
f" [node] processing chat model={model_name!r} stream={stream} "
f"max_tokens={max_tokens}{self._request_log_suffix()}",
@@ -335,6 +465,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# Avoids the single-token-per-forward-pass limitation of the distributed path.
if backend.is_head and backend.is_tail:
gen_started = time.monotonic()
progress_line = [False]
try:
if stream:
token_count = 0
@@ -346,13 +477,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
):
if token_text:
token_count += 1
self._track_request_progress(
server, request_id, tokens=token_count, routing_complete=True,
)
yield token_text
self._stream_openai_response(_counting_stream(), model_name)
print(
elapsed = time.monotonic() - gen_started
tps = token_count / max(elapsed, 1e-6)
_write_progress_line(
progress_line,
f" [node] chat complete (stream) tokens={token_count} "
f"elapsed_s={time.monotonic() - gen_started:.1f}{self._request_log_suffix()}",
flush=True,
f"elapsed_s={elapsed:.1f} tps={tps:.2f}{self._request_log_suffix()}",
final=True,
)
else:
text = backend.generate_text(messages, max_tokens, temperature, top_p)
@@ -414,10 +551,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
stream_emit = None
if stream:
stream_emit = self._start_openai_stream(model_name)
self._track_request_progress(server, request_id, tokens=0, routing_complete=True)
_GENERATION_LOG_INTERVAL = 5.0
gen_started = time.monotonic()
last_gen_log = gen_started
progress_line = [False]
for step in range(max_tokens):
try:
@@ -437,20 +576,33 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if stream_emit is not None:
stream_emit(token_str)
current_text = current_text + token_str
self._track_request_progress(
server,
request_id,
tokens=len(generated),
routing_complete=True,
)
now = time.monotonic()
if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL:
print(
elapsed = now - gen_started
token_count = len(generated)
tps = token_count / max(elapsed, 1e-6)
_write_progress_line(
progress_line,
f" [node] generating step={step + 1}/{max_tokens} "
f"tokens={len(generated)} elapsed_s={now - gen_started:.1f}",
flush=True,
f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
)
last_gen_log = now
if generated:
print(
f" [node] generation complete tokens={len(generated)} "
f"elapsed_s={time.monotonic() - gen_started:.1f}",
flush=True,
elapsed = time.monotonic() - gen_started
token_count = len(generated)
tps = token_count / max(elapsed, 1e-6)
_write_progress_line(
progress_line,
f" [node] generation complete tokens={token_count} "
f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
final=True,
)
result_text = "".join(generated)
@@ -467,6 +619,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
start_layer tells each downstream node which layer to begin from,
enabling correct execution when shard ranges overlap.
"""
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
injected = self.headers.get("X-Meshnet-Route")
if injected:
@@ -485,14 +640,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
hops.append(hop)
elif isinstance(item, str):
hops.append({"endpoint": item, "start_layer": 0})
print(f" [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True)
hops = _clamp_downstream_hops(hops, active_backend)
print(
f" [node] using injected downstream route: {_format_downstream_route(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).
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
if server.tracker_url is None:
return []
route_model = getattr(active_backend, "model_id", None) or model
@@ -500,23 +657,31 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
route_resp = json.loads(r.read())
own_port = server.server_address[1]
own_key = _own_endpoint_key(server)
nodes_info = route_resp.get("nodes", [])
hops = []
covered_up_to: int | None = None
hops: list[dict] = []
passed_self = False
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")
if not ep:
continue
if covered_up_to is None:
covered_up_to = (node_info.get("shard_start") or 1) - 1
hop = {"endpoint": ep, "start_layer": covered_up_to + 1}
if _endpoint_key(ep) == own_key:
passed_self = True
continue
if not passed_self:
continue
hop = {
"endpoint": ep,
"start_layer": int(node_info.get("start_layer", 0)),
}
if node_info.get("relay_addr"):
hop["relay_addr"] = str(node_info["relay_addr"])
hops.append(hop)
covered_up_to = node_info.get("shard_end", covered_up_to)
print(f" [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True)
hops = _clamp_downstream_hops(hops, active_backend)
print(
f" [node] tracker downstream route: {_format_downstream_route(hops)}",
flush=True,
)
return hops
except Exception as exc:
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
@@ -849,6 +1014,12 @@ class TorchNodeServer:
def queue_depth(self) -> int:
return self._server.queue_depth if self._server is not None else 0
@property
def current_requests(self) -> list[dict[str, Any]]:
if self._server is None:
return []
return self._server.snapshot_current_requests()
@property
def loaded_model_ids(self) -> list[str]:
return list(self._backends.keys())
@@ -928,6 +1099,11 @@ class TorchNodeServer:
self._thread.start()
return self.port
def set_advertised_endpoint(self, endpoint: str) -> None:
"""Set the LAN-facing endpoint used for route self-detection."""
if self._server is not None:
self._server.advertised_endpoint = endpoint
def stop(self) -> None:
if self._server is None:
return

View File

@@ -1,33 +1,34 @@
[build-system]
requires = ["setuptools>=64"]
build-backend = "setuptools.build_meta"
[project]
name = "meshnet-node"
version = "0.1.0"
description = "Distributed Inference Network node client"
requires-python = ">=3.10"
dependencies = [
"cryptography>=41",
"huggingface-hub>=0.20",
"accelerate>=0.28",
"bitsandbytes>=0.43",
"rich>=13",
"safetensors>=0.4",
"torch>=2.1",
"transformers>=5.12",
"websockets>=13",
"zstandard>=0.22",
"kernels>=0.11.1,<0.16",
]
[project.scripts]
meshnet-node = "meshnet_node.cli:main"
[tool.setuptools.packages.find]
where = ["."]
include = ["meshnet_node*"]
[tool.setuptools.package-data]
meshnet_node = ["*.json"]
[build-system]
requires = ["setuptools>=64"]
build-backend = "setuptools.build_meta"
[project]
name = "meshnet-node"
version = "0.1.0"
description = "Distributed Inference Network node client"
requires-python = ">=3.10"
dependencies = [
"cryptography>=41",
"huggingface-hub>=0.20",
"accelerate>=0.28",
"bitsandbytes>=0.43",
"rich>=13",
"safetensors>=0.4",
"torch>=2.1",
"transformers>=5.12",
"triton-windows>=3.7; platform_system == 'Windows'",
"websockets>=13",
"zstandard>=0.22",
"kernels>=0.11.1,<0.16",
]
[project.scripts]
meshnet-node = "meshnet_node.cli:main"
[tool.setuptools.packages.find]
where = ["."]
include = ["meshnet_node*"]
[tool.setuptools.package-data]
meshnet_node = ["*.json"]

View File

@@ -1,13 +1,13 @@
"""meshnet-tracker CLI entry point."""
import argparse
import os
import sys
import time
from pathlib import Path
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
from .billing import DEFAULT_BILLING_DB_PATH
"""meshnet-tracker CLI entry point."""
import argparse
import os
import sys
import time
from pathlib import Path
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
from .billing import DEFAULT_BILLING_DB_PATH
from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH
from .logging_setup import (
DEFAULT_LOG_BACKUP_COUNT,
@@ -15,258 +15,299 @@ from .logging_setup import (
DEFAULT_LOG_MAX_BYTES,
configure_tracker_file_logging,
)
from .routing_stats import RoutingConfig
from .server import (
DEFAULT_CALLER_CREDIT_USDT,
DEFAULT_DEVNET_TOPUP_USDT,
TrackerServer,
derive_relay_url_from_public_tracker_url,
)
DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3"
def _load_env_file(path: Path) -> None:
"""Load simple KEY=VALUE pairs from an env file without overriding env vars."""
if not path.exists():
return
try:
lines = path.read_text().splitlines()
except OSError:
return
for line in lines:
text = line.strip()
if not text or text.startswith("#"):
continue
if text.startswith("export "):
text = text[len("export "):].strip()
if "=" not in text:
continue
key, value = text.split("=", 1)
key = key.strip()
if not key or key in os.environ:
continue
value = value.strip()
if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
value = value[1:-1]
os.environ[key] = value
def _load_env_defaults() -> None:
"""Load local and user-level tracker env defaults before parsing arguments."""
_load_env_file(Path.cwd() / ".env")
_load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env")
def main() -> None:
_load_env_defaults()
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on")
common.add_argument("--port", type=int, default=8080, help="Port to listen on")
common.add_argument(
"--heartbeat-timeout",
type=float,
default=30.0,
help="Seconds before a node is removed from the registry after missed heartbeat",
)
common.add_argument(
"--cluster-peers",
default="",
help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)",
)
common.add_argument(
"--self-url",
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,
help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
)
common.add_argument(
"--billing-db",
default=DEFAULT_BILLING_DB_PATH,
metavar="PATH",
help=(
"SQLite database path for the USDT billing ledger "
f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)"
),
)
common.add_argument(
"--no-billing",
action="store_true",
help="Disable the USDT billing ledger",
)
common.add_argument(
"--max-charge-per-request",
type=float,
default=None,
help=(
"Reject chat completion requests whose prompt plus requested completion "
"token bound would cost more than this many USDT"
),
)
common.add_argument(
"--starting-credit",
type=float,
default=DEFAULT_CALLER_CREDIT_USDT,
metavar="USDT",
help=(
"One-time Caller Credit granted when an account creates its first "
f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require "
"deposits before inference)"
),
)
common.add_argument(
"--devnet-topup",
type=float,
default=DEFAULT_DEVNET_TOPUP_USDT,
metavar="USDT",
help=(
"Dashboard devnet top-up faucet: each click credits this many USDT "
f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; "
"MUST be 0 on mainnet deployments)"
),
)
common.add_argument(
"--registry-db",
default=DEFAULT_REGISTRY_DB_PATH,
metavar="PATH",
help=(
"SQLite database path for persisted strike/ban/reputation registry "
f"state (default: {DEFAULT_REGISTRY_DB_PATH})"
),
)
common.add_argument(
"--no-registry-contracts",
action="store_true",
help="Disable the local contract registry used for strike/ban/reputation enforcement",
)
common.add_argument(
"--accounts-db",
default=DEFAULT_ACCOUNTS_DB_PATH,
metavar="PATH",
help=(
"SQLite database path for dashboard user accounts "
f"(default: {DEFAULT_ACCOUNTS_DB_PATH})"
),
)
common.add_argument(
"--no-accounts",
action="store_true",
help="Disable dashboard user accounts (registration/login)",
)
common.add_argument(
"--solana-rpc-url",
default=None,
help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury",
)
common.add_argument(
"--usdt-mint",
default=None,
help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py",
)
common.add_argument(
"--treasury-keypair",
default=None,
metavar="PATH",
help="Treasury keypair JSON path (only on settlement-capable trackers)",
)
common.add_argument(
"--settle-period",
type=float,
default=86400.0,
help="Max seconds between payouts to a node (dev: 60, prod: 86400)",
)
common.add_argument(
"--payout-threshold",
type=float,
default=5.0,
help="Pending USDT that triggers an immediate payout (dev: 0)",
)
common.add_argument(
"--payout-dust-floor",
type=float,
default=0.01,
help="Never pay out less than this many USDT",
)
common.add_argument(
"--validator-service-token",
default=None,
help=(
"Service token the validator uses on POST /v1/billing/forfeit "
"(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)"
),
)
common.add_argument(
"--hive-secret",
default=None,
help=(
"Shared secret authenticating gossip between tracker peers "
"(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)"
),
)
common.add_argument(
"--toploc-calibration-db",
default=None,
metavar="PATH",
help=(
"SQLite path for the AH-021 honest-noise TOPLOC calibration corpus "
"(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)"
),
)
common.add_argument(
"--toploc-reference-node-url",
default=None,
help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)",
)
common.add_argument(
"--toploc-calibration-gate-min-hardware-profiles",
type=int,
default=1,
help=(
"Distinct (GPU model, dtype) profiles the corpus must cover before "
"gate_status.ready is true (alpha exception: fleet size is acceptable)"
),
)
common.add_argument(
"--enable-hf-pricing",
action="store_true",
help=(
"Enable the daily dynamic pricing refresh (issue 23): for presets with a "
"curated hf_aliases list, sets the client price to 80%% of the cheapest "
"matching HuggingFace inference-marketplace rate. Presets without "
"hf_aliases are unaffected and keep their static price."
),
)
common.add_argument(
"--hf-pricing-log-db",
default=None,
metavar="PATH",
help=(
"SQLite database path for the dynamic pricing change log "
f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; "
"enables GET /v1/pricing/hf/history)"
),
)
common.add_argument(
"--hf-pricing-refresh-interval",
type=float,
default=86400.0,
help="Seconds between dynamic pricing refresh passes (default: daily)",
)
DEFAULT_CALLER_CREDIT_USDT,
DEFAULT_DEVNET_TOPUP_USDT,
TrackerServer,
derive_relay_url_from_public_tracker_url,
)
DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3"
def _load_env_file(path: Path) -> None:
"""Load simple KEY=VALUE pairs from an env file without overriding env vars."""
if not path.exists():
return
try:
lines = path.read_text().splitlines()
except OSError:
return
for line in lines:
text = line.strip()
if not text or text.startswith("#"):
continue
if text.startswith("export "):
text = text[len("export "):].strip()
if "=" not in text:
continue
key, value = text.split("=", 1)
key = key.strip()
if not key or key in os.environ:
continue
value = value.strip()
if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
value = value[1:-1]
os.environ[key] = value
def _load_env_defaults() -> None:
"""Load local and user-level tracker env defaults before parsing arguments."""
_load_env_file(Path.cwd() / ".env")
_load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env")
def _routing_config_from_args(args: argparse.Namespace) -> RoutingConfig | None:
"""Build a RoutingConfig from CLI flags; None keeps env-var/server defaults."""
overrides = {
"explore_share": args.route_explore_share,
"weight_alpha": args.route_weight_alpha,
"stats_half_life_seconds": args.route_stats_half_life,
}
set_values = {key: value for key, value in overrides.items() if value is not None}
if not set_values:
return None
return RoutingConfig(**set_values)
def main() -> None:
_load_env_defaults()
common = argparse.ArgumentParser(add_help=False)
common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on")
common.add_argument("--port", type=int, default=8080, help="Port to listen on")
common.add_argument(
"--heartbeat-timeout",
type=float,
default=30.0,
help="Seconds before a node is removed from the registry after missed heartbeat",
)
common.add_argument(
"--cluster-peers",
default="",
help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)",
)
common.add_argument(
"--self-url",
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,
help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
)
common.add_argument(
"--billing-db",
default=DEFAULT_BILLING_DB_PATH,
metavar="PATH",
help=(
"SQLite database path for the USDT billing ledger "
f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)"
),
)
common.add_argument(
"--no-billing",
action="store_true",
help="Disable the USDT billing ledger",
)
common.add_argument(
"--max-charge-per-request",
type=float,
default=None,
help=(
"Reject chat completion requests whose prompt plus requested completion "
"token bound would cost more than this many USDT"
),
)
common.add_argument(
"--starting-credit",
type=float,
default=DEFAULT_CALLER_CREDIT_USDT,
metavar="USDT",
help=(
"One-time Caller Credit granted when an account creates its first "
f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require "
"deposits before inference)"
),
)
common.add_argument(
"--devnet-topup",
type=float,
default=DEFAULT_DEVNET_TOPUP_USDT,
metavar="USDT",
help=(
"Dashboard devnet top-up faucet: each click credits this many USDT "
f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; "
"MUST be 0 on mainnet deployments)"
),
)
common.add_argument(
"--registry-db",
default=DEFAULT_REGISTRY_DB_PATH,
metavar="PATH",
help=(
"SQLite database path for persisted strike/ban/reputation registry "
f"state (default: {DEFAULT_REGISTRY_DB_PATH})"
),
)
common.add_argument(
"--no-registry-contracts",
action="store_true",
help="Disable the local contract registry used for strike/ban/reputation enforcement",
)
common.add_argument(
"--accounts-db",
default=DEFAULT_ACCOUNTS_DB_PATH,
metavar="PATH",
help=(
"SQLite database path for dashboard user accounts "
f"(default: {DEFAULT_ACCOUNTS_DB_PATH})"
),
)
common.add_argument(
"--no-accounts",
action="store_true",
help="Disable dashboard user accounts (registration/login)",
)
common.add_argument(
"--solana-rpc-url",
default=None,
help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury",
)
common.add_argument(
"--usdt-mint",
default=None,
help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py",
)
common.add_argument(
"--treasury-keypair",
default=None,
metavar="PATH",
help="Treasury keypair JSON path (only on settlement-capable trackers)",
)
common.add_argument(
"--settle-period",
type=float,
default=86400.0,
help="Max seconds between payouts to a node (dev: 60, prod: 86400)",
)
common.add_argument(
"--payout-threshold",
type=float,
default=5.0,
help="Pending USDT that triggers an immediate payout (dev: 0)",
)
common.add_argument(
"--payout-dust-floor",
type=float,
default=0.01,
help="Never pay out less than this many USDT",
)
common.add_argument(
"--validator-service-token",
default=None,
help=(
"Service token the validator uses on POST /v1/billing/forfeit "
"(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)"
),
)
common.add_argument(
"--hive-secret",
default=None,
help=(
"Shared secret authenticating gossip between tracker peers "
"(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)"
),
)
common.add_argument(
"--toploc-calibration-db",
default=None,
metavar="PATH",
help=(
"SQLite path for the AH-021 honest-noise TOPLOC calibration corpus "
"(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)"
),
)
common.add_argument(
"--toploc-reference-node-url",
default=None,
help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)",
)
common.add_argument(
"--toploc-calibration-gate-min-hardware-profiles",
type=int,
default=1,
help=(
"Distinct (GPU model, dtype) profiles the corpus must cover before "
"gate_status.ready is true (alpha exception: fleet size is acceptable)"
),
)
common.add_argument(
"--enable-hf-pricing",
action="store_true",
help=(
"Enable the daily dynamic pricing refresh (issue 23): for presets with a "
"curated hf_aliases list, sets the client price to 80%% of the cheapest "
"matching HuggingFace inference-marketplace rate. Presets without "
"hf_aliases are unaffected and keep their static price."
),
)
common.add_argument(
"--hf-pricing-log-db",
default=None,
metavar="PATH",
help=(
"SQLite database path for the dynamic pricing change log "
f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; "
"enables GET /v1/pricing/hf/history)"
),
)
common.add_argument(
"--hf-pricing-refresh-interval",
type=float,
default=86400.0,
help="Seconds between dynamic pricing refresh passes (default: daily)",
)
common.add_argument(
"--models-dir",
default=None,
metavar="PATH",
help="Local HuggingFace snapshot root advertised as tracker model-file source (default: MESHNET_MODELS_DIR)",
)
common.add_argument(
"--route-explore-share",
type=float,
default=None,
metavar="FRACTION",
help=(
"Fraction of requests routed down unproven/stale routes to gather "
"throughput statistics (ADR-0021; default 0.3, lower once traffic grows)"
),
)
common.add_argument(
"--route-weight-alpha",
type=float,
default=None,
metavar="ALPHA",
help=(
"Traffic weight exponent among proven routes: share ∝ tps^alpha "
"(default 1.0 — a 1.5x-faster route gets 1.5x the traffic)"
),
)
common.add_argument(
"--route-stats-half-life",
type=float,
default=None,
metavar="SECONDS",
help="Half-life for decaying route throughput observations (default 600)",
)
common.add_argument(
"--log-dir",
default=DEFAULT_LOG_DIR,
@@ -295,18 +336,18 @@ def main() -> None:
action="store_true",
help="Disable rotating tracker log files and only write to the terminal",
)
parser = argparse.ArgumentParser(
prog="meshnet-tracker",
description="Distributed Inference Network node registry and route selection",
parents=[common],
)
subparsers = parser.add_subparsers(dest="command")
subparsers.add_parser("start", help="Start the tracker server", parents=[common])
args = parser.parse_args()
parser = argparse.ArgumentParser(
prog="meshnet-tracker",
description="Distributed Inference Network node registry and route selection",
parents=[common],
)
subparsers = parser.add_subparsers(dest="command")
subparsers.add_parser("start", help="Start the tracker server", parents=[common])
args = parser.parse_args()
if args.command in {None, "start"}:
if not args.no_file_logs:
log_dir = configure_tracker_file_logging(
@@ -316,62 +357,63 @@ def main() -> None:
)
print(f"meshnet-tracker logs: {log_dir}", flush=True)
cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()]
relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url)
treasury = None
if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair:
from meshnet_contracts.solana_adapter import SolanaCustodialTreasury
treasury = SolanaCustodialTreasury(
args.solana_rpc_url, args.usdt_mint, args.treasury_keypair,
)
contracts = None
if not args.no_registry_contracts:
from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found]
contracts = LocalSolanaContracts(registry_db=args.registry_db)
server = TrackerServer(
host=args.host,
port=args.port,
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=relay_url,
enable_billing=not args.no_billing,
billing_db=None if args.no_billing else args.billing_db,
max_charge_per_request=args.max_charge_per_request,
starting_credit=args.starting_credit,
devnet_topup_amount=args.devnet_topup,
contracts=contracts,
accounts_db=None if args.no_accounts else args.accounts_db,
treasury=treasury,
settle_period=args.settle_period,
payout_threshold=args.payout_threshold,
payout_dust_floor=args.payout_dust_floor,
validator_service_token=args.validator_service_token,
hive_secret=args.hive_secret,
toploc_calibration_db=args.toploc_calibration_db,
toploc_reference_node_url=args.toploc_reference_node_url,
toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles,
enable_hf_pricing=args.enable_hf_pricing,
hf_pricing_log_db=(
args.hf_pricing_log_db
or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None)
),
hf_pricing_refresh_interval=args.hf_pricing_refresh_interval,
models_dir=args.models_dir,
)
port = server.start()
print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
server.stop()
sys.exit(0)
else:
parser.print_help()
if __name__ == "__main__":
main()
relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url)
treasury = None
if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair:
from meshnet_contracts.solana_adapter import SolanaCustodialTreasury
treasury = SolanaCustodialTreasury(
args.solana_rpc_url, args.usdt_mint, args.treasury_keypair,
)
contracts = None
if not args.no_registry_contracts:
from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found]
contracts = LocalSolanaContracts(registry_db=args.registry_db)
server = TrackerServer(
host=args.host,
port=args.port,
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=relay_url,
enable_billing=not args.no_billing,
billing_db=None if args.no_billing else args.billing_db,
max_charge_per_request=args.max_charge_per_request,
starting_credit=args.starting_credit,
devnet_topup_amount=args.devnet_topup,
contracts=contracts,
accounts_db=None if args.no_accounts else args.accounts_db,
treasury=treasury,
settle_period=args.settle_period,
payout_threshold=args.payout_threshold,
payout_dust_floor=args.payout_dust_floor,
validator_service_token=args.validator_service_token,
hive_secret=args.hive_secret,
toploc_calibration_db=args.toploc_calibration_db,
toploc_reference_node_url=args.toploc_reference_node_url,
toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles,
enable_hf_pricing=args.enable_hf_pricing,
hf_pricing_log_db=(
args.hf_pricing_log_db
or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None)
),
hf_pricing_refresh_interval=args.hf_pricing_refresh_interval,
models_dir=args.models_dir,
routing_config=_routing_config_from_args(args),
)
port = server.start()
print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
server.stop()
sys.exit(0)
else:
parser.print_help()
if __name__ == "__main__":
main()

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,257 @@
"""Learned route statistics for dynamic bandit-style route selection (ADR-0021).
The tracker treats each viable route (ordered chain of node shards covering a
model) as a bandit arm. Observed end-to-end tokens/sec per route is kept as a
time-decayed EWMA. Selection splits traffic between:
- **exploit**: weighted-random among *proven* routes, weight ∝ tps ** alpha
(alpha=1.0 → a 1.5x-faster route gets 1.5x the traffic);
- **scout**: with probability `explore_share`, the least-measured unproven or
stale route is chosen so the tracker keeps learning as the network morphs.
Staleness has two mechanisms:
- continuous: sample mass decays with `stats_half_life_seconds`, so old
observations fade;
- abrupt: every node join/leave bumps the model's *topology epoch*; stats from
an older epoch keep their EWMA as a prior but drop back into the scout pool
until re-measured.
Route signatures embed node ids and shard ranges, so a node re-registering
with a different shard produces a new arm automatically.
"""
from __future__ import annotations
import random
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Iterable
@dataclass(frozen=True)
class RoutingConfig:
explore_share: float = 0.3
weight_alpha: float = 1.0
stats_half_life_seconds: float = 600.0
min_sample_tokens: int = 8
# One fresh sample has mass 1.0 and decays from there; 0.5 keeps a single
# observation "proven" for one half-life before demoting it to the scout pool.
min_proven_weight: float = 0.5
max_candidate_routes: int = 8
prune_after_seconds: float = 86400.0
@dataclass
class RouteStat:
ewma_tps: float = 0.0
weight: float = 0.0 # decayed effective sample mass
last_sample_ts: float = 0.0
epoch: int = 0
samples: int = 0 # lifetime raw sample count (display only)
def decayed_weight(self, now: float, half_life: float) -> float:
if self.weight <= 0.0:
return 0.0
age = max(0.0, now - self.last_sample_ts)
return self.weight * 0.5 ** (age / half_life)
@dataclass
class RouteCandidate:
nodes: list[Any]
signature: str
prior_tps: float = 0.0
def route_signature(model_key: str, nodes: Iterable[Any]) -> str:
hops = "->".join(
f"{getattr(n, 'node_id', '?')}[{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}]"
for n in nodes
)
return f"{model_key}|{hops}"
class RouteStatsStore:
"""Thread-safe per-route decayed throughput statistics."""
def __init__(self, config: RoutingConfig | None = None) -> None:
self.config = config or RoutingConfig()
self._lock = threading.Lock()
self._stats: dict[str, RouteStat] = {}
self._epochs: dict[str, int] = {}
def epoch(self, model_key: str) -> int:
with self._lock:
return self._epochs.get(model_key, 0)
def bump_epoch(self, model_keys: Iterable[str | None]) -> None:
"""Mark the topology changed for the given model keys (node join/leave)."""
with self._lock:
for key in model_keys:
if key:
self._epochs[key] = self._epochs.get(key, 0) + 1
def record_sample(
self,
model_key: str,
signature: str,
tokens: int,
elapsed_seconds: float,
now: float | None = None,
) -> bool:
"""Fold one completed request into the route's EWMA.
Returns False (and records nothing) for samples below
`min_sample_tokens` — near-empty completions come from broken routes
and would poison the arm with meaningless throughput values.
"""
cfg = self.config
if tokens < cfg.min_sample_tokens or elapsed_seconds <= 0.0:
return False
tps = tokens / elapsed_seconds
ts = time.time() if now is None else now
with self._lock:
stat = self._stats.get(signature)
if stat is None:
stat = RouteStat()
self._stats[signature] = stat
carried = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
total = carried + 1.0
stat.ewma_tps = (stat.ewma_tps * carried + tps) / total
stat.weight = total
stat.last_sample_ts = ts
stat.epoch = self._epochs.get(model_key, 0)
stat.samples += 1
return True
def snapshot(self, signature: str, model_key: str, now: float | None = None) -> dict:
"""Point-in-time view of one route's learned state."""
ts = time.time() if now is None else now
cfg = self.config
with self._lock:
stat = self._stats.get(signature)
current_epoch = self._epochs.get(model_key, 0)
if stat is None:
return {"tps": None, "weight": 0.0, "samples": 0, "status": "unsampled"}
weight = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
if stat.epoch != current_epoch:
status = "stale"
elif weight < cfg.min_proven_weight:
status = "decayed" if stat.samples else "unsampled"
else:
status = "proven"
return {
"tps": round(stat.ewma_tps, 4) if stat.samples else None,
"weight": round(weight, 4),
"samples": stat.samples,
"status": status,
}
def prune(self, now: float | None = None) -> int:
"""Drop routes with no samples for `prune_after_seconds`."""
ts = time.time() if now is None else now
cutoff = ts - self.config.prune_after_seconds
with self._lock:
dead = [sig for sig, stat in self._stats.items() if stat.last_sample_ts < cutoff]
for sig in dead:
del self._stats[sig]
return len(dead)
def choose_route(
candidates: list[RouteCandidate],
store: RouteStatsStore,
model_key: str,
rng: random.Random | None = None,
now: float | None = None,
) -> tuple[RouteCandidate | None, dict]:
"""Pick a route: ε-scout among unproven arms, else weighted ∝ tps**alpha.
Returns (candidate, decision) where decision explains the pick for logs
and diagnostics: {"mode": "scout"|"exploit"|"prior", ...}.
"""
if not candidates:
return None, {"mode": "none"}
rng = rng or random
cfg = store.config
proven: list[tuple[RouteCandidate, float]] = []
scouts: list[tuple[RouteCandidate, float]] = []
for cand in candidates:
snap = store.snapshot(cand.signature, model_key, now=now)
if snap["status"] == "proven":
proven.append((cand, max(float(snap["tps"] or 0.0), 1e-6)))
else:
scouts.append((cand, float(snap["weight"])))
if scouts and (not proven or rng.random() < cfg.explore_share):
# Least-measured first so new/stale arms accumulate samples fastest;
# tiebreak on prior estimate so plausible routes get scouted first.
scouts.sort(key=lambda item: (item[1], -item[0].prior_tps))
pick = scouts[0][0]
return pick, {"mode": "scout", "signature": pick.signature}
if proven:
weights = [tps ** cfg.weight_alpha for _, tps in proven]
pick = rng.choices([cand for cand, _ in proven], weights=weights, k=1)[0]
return pick, {
"mode": "exploit",
"signature": pick.signature,
"candidates": len(proven),
}
# No stats anywhere yet — fall back to the prior (benchmark-derived) estimate.
weights = [max(cand.prior_tps, 1e-6) ** cfg.weight_alpha for cand in candidates]
pick = rng.choices(candidates, weights=weights, k=1)[0]
return pick, {"mode": "prior", "signature": pick.signature}
def route_table(
candidates: list[RouteCandidate],
store: RouteStatsStore,
model_key: str,
now: float | None = None,
) -> list[dict]:
"""Diagnostics rows: learned tps, coefficient vs best, expected traffic share."""
cfg = store.config
rows = []
for cand in candidates:
snap = store.snapshot(cand.signature, model_key, now=now)
rows.append({"candidate": cand, **snap})
proven = [r for r in rows if r["status"] == "proven"]
scouts = [r for r in rows if r["status"] != "proven"]
best_tps = max((float(r["tps"]) for r in proven), default=0.0)
exploit_budget = 1.0 - (cfg.explore_share if scouts and proven else 0.0)
if not proven:
exploit_budget = 0.0
weight_sum = sum(float(r["tps"]) ** cfg.weight_alpha for r in proven) or 1.0
out = []
for r in rows:
cand: RouteCandidate = r["candidate"]
if r["status"] == "proven":
share = exploit_budget * (float(r["tps"]) ** cfg.weight_alpha) / weight_sum
coefficient = round(float(r["tps"]) / best_tps, 3) if best_tps else None
else:
share = (
(cfg.explore_share if proven else 1.0) / len(scouts)
if scouts
else 0.0
)
coefficient = None
out.append({
"signature": cand.signature,
"hops": [
{
"node_id": getattr(n, "node_id", "?"),
"shard": f"{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}",
"endpoint": getattr(n, "endpoint", "?"),
}
for n in cand.nodes
],
"tps": r["tps"],
"coefficient": coefficient,
"expected_share": round(share, 4),
"samples": r["samples"],
"weight": r["weight"],
"status": r["status"],
"prior_tps": round(cand.prior_tps, 4),
})
out.sort(key=lambda r: (-(r["tps"] or 0.0), -r["prior_tps"]))
return out

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff