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

Track observed node throughput

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-02 23:28:20 +02:00
parent a938c19a82
commit 83b042d94b
6 changed files with 322 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -180,6 +180,73 @@ class _RollingCounter:
self._counts[i] = cnt
class _RollingThroughput:
"""Circular buckets for observed tokens/sec.
Each bucket stores total output tokens and total elapsed seconds. TPS for a
window is the ratio across non-stale buckets, which avoids over-weighting
small fast requests.
"""
def __init__(self, num_buckets: int, bucket_seconds: int) -> None:
self._num = num_buckets
self._bsec = bucket_seconds
self._tokens: list[int] = [0] * num_buckets
self._seconds: list[float] = [0.0] * num_buckets
self._samples: list[int] = [0] * num_buckets
self._epochs: list[int] = [-1] * num_buckets
def _epoch(self, now: float) -> int:
return int(now) // self._bsec
def record(self, total_tokens: int, elapsed_seconds: float, now: float | None = None) -> None:
if total_tokens <= 0 or elapsed_seconds <= 0:
return
t = now if now is not None else time.time()
ep = self._epoch(t)
idx = ep % self._num
if self._epochs[idx] != ep:
self._tokens[idx] = 0
self._seconds[idx] = 0.0
self._samples[idx] = 0
self._epochs[idx] = ep
self._tokens[idx] += int(total_tokens)
self._seconds[idx] += float(elapsed_seconds)
self._samples[idx] += 1
def stats(self, now: float | None = None) -> dict:
t = now if now is not None else time.time()
cutoff = self._epoch(t) - self._num
tokens = 0
seconds = 0.0
samples = 0
for i in range(self._num):
if self._epochs[i] > cutoff:
tokens += self._tokens[i]
seconds += self._seconds[i]
samples += self._samples[i]
return {
"tokens_per_sec": round(tokens / seconds, 4) if seconds > 0 else None,
"tokens": tokens,
"seconds": round(seconds, 6),
"sample_count": samples,
}
def buckets(self) -> list[tuple[int, int, float, int]]:
return [
(self._epochs[i], self._tokens[i], self._seconds[i], self._samples[i])
for i in range(self._num)
]
def restore_buckets(self, data: list[tuple[int, int, float, int]]) -> None:
for i, (ep, tokens, seconds, samples) in enumerate(data):
if i < self._num:
self._epochs[i] = ep
self._tokens[i] = tokens
self._seconds[i] = seconds
self._samples[i] = samples
class _ModelStats:
"""Three rolling windows for one model: last hour, last day, last month."""
@@ -203,6 +270,34 @@ class _ModelStats:
}
class _NodeModelThroughput:
"""Observed throughput windows for one node/model pair."""
def __init__(self) -> None:
self.per_minute = _RollingThroughput(60, 60)
self.per_hour = _RollingThroughput(24, 3600)
self.per_day = _RollingThroughput(30, 86400)
def record(self, total_tokens: int, elapsed_seconds: float, now: float | None = None) -> None:
t = now if now is not None else time.time()
self.per_minute.record(total_tokens, elapsed_seconds, t)
self.per_hour.record(total_tokens, elapsed_seconds, t)
self.per_day.record(total_tokens, elapsed_seconds, t)
def stats(self, now: float | None = None) -> dict:
hour = self.per_minute.stats(now)
day = self.per_hour.stats(now)
month = self.per_day.stats(now)
return {
"tokens_per_sec_last_hour": hour["tokens_per_sec"],
"tokens_per_sec_last_day": day["tokens_per_sec"],
"tokens_per_sec_last_month": month["tokens_per_sec"],
"sample_count_last_hour": hour["sample_count"],
"tokens_last_hour": hour["tokens"],
"seconds_last_hour": hour["seconds"],
}
class _StatsCollector:
"""Thread-safe model request stats with SQLite persistence and peer slice merging."""
@@ -211,6 +306,7 @@ class _StatsCollector:
def __init__(self, db_path: str | None = None) -> None:
self._lock = threading.Lock()
self._local: dict[str, _ModelStats] = {}
self._node_model: dict[tuple[str, str], _NodeModelThroughput] = {}
self._peer_rpms: dict[str, dict[str, dict]] = {} # tracker_url -> model -> rpms
self._db_path = db_path
if db_path:
@@ -226,11 +322,50 @@ class _StatsCollector:
self._local[model] = _ModelStats()
self._local[model].record(t)
def record_node_throughput(
self,
node_id: str,
model: str,
total_tokens: int,
elapsed_seconds: float,
now: float | None = None,
) -> None:
if not node_id or not model or total_tokens <= 0 or elapsed_seconds <= 0:
return
t = now if now is not None else time.time()
with self._lock:
key = (node_id, model)
if key not in self._node_model:
self._node_model[key] = _NodeModelThroughput()
self._node_model[key].record(total_tokens, elapsed_seconds, 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 get_node_model_stats(self, node_id: str, model: str, now: float | None = None) -> dict:
t = now if now is not None else time.time()
empty = {
"tokens_per_sec_last_hour": None,
"tokens_per_sec_last_day": None,
"tokens_per_sec_last_month": None,
"sample_count_last_hour": 0,
"tokens_last_hour": 0,
"seconds_last_hour": 0.0,
}
with self._lock:
stat = self._node_model.get((node_id, model))
return stat.stats(t) if stat is not None else dict(empty)
def get_node_throughput_stats(self, now: float | None = None) -> dict[str, dict]:
t = now if now is not None else time.time()
with self._lock:
result: dict[str, dict] = {}
for (node_id, model), stat in self._node_model.items():
result.setdefault(node_id, {"models": {}})["models"][model] = stat.stats(t)
return result
def merge_peer_rpms(self, tracker_url: str, rpms: dict[str, dict]) -> None:
with self._lock:
self._peer_rpms[tracker_url] = dict(rpms)
@@ -258,6 +393,12 @@ class _StatsCollector:
"(model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, count INTEGER, "
"PRIMARY KEY (model, window, bucket_idx))"
)
con.execute(
"CREATE TABLE IF NOT EXISTS node_model_tps_buckets "
"(node_id TEXT, model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, "
"tokens INTEGER, seconds REAL, samples INTEGER, "
"PRIMARY KEY (node_id, model, window, bucket_idx))"
)
con.commit()
con.close()
@@ -275,16 +416,34 @@ class _StatsCollector:
for idx, (ep, cnt) in enumerate(counter.buckets()):
if ep >= 0:
rows.append((model, window_name, idx, ep, cnt))
tps_rows = []
for (node_id, model), stat in self._node_model.items():
for window_name, counter in (
("hour", stat.per_minute),
("day", stat.per_hour),
("month", stat.per_day),
):
for idx, (ep, tokens, seconds, samples) in enumerate(counter.buckets()):
if ep >= 0:
tps_rows.append((node_id, model, window_name, idx, ep, tokens, seconds, samples))
con = sqlite3.connect(self._db_path) # type: ignore[arg-type]
con.executemany(
"INSERT OR REPLACE INTO model_rpm_buckets VALUES (?,?,?,?,?)", rows
)
con.executemany(
"INSERT OR REPLACE INTO node_model_tps_buckets VALUES (?,?,?,?,?,?,?,?)",
tps_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()
tps_rows = con.execute(
"SELECT node_id, model, window, bucket_idx, bucket_epoch, tokens, seconds, samples "
"FROM node_model_tps_buckets"
).fetchall()
con.close()
grouped: dict[str, dict[str, list[tuple[int, int]]]] = {}
for model, window, idx, ep, cnt in rows:
@@ -301,6 +460,23 @@ class _StatsCollector:
data[idx] = (ep, cnt)
counter.restore_buckets(data)
self._local[model] = ms
tps_grouped: dict[tuple[str, str], dict[str, list[tuple[int, int, int, float, int]]]] = {}
for node_id, model, window, idx, ep, tokens, seconds, samples in tps_rows:
tps_grouped.setdefault((node_id, model), {}).setdefault(window, []).append(
(idx, ep, tokens, seconds, samples)
)
for key, windows in tps_grouped.items():
stat = _NodeModelThroughput()
for window_name, entries in windows.items():
counter = {"hour": stat.per_minute, "day": stat.per_hour, "month": stat.per_day}.get(window_name)
if counter is None:
continue
data = [(0, 0, 0.0, 0)] * counter._num
for idx, ep, tokens, seconds, samples in entries:
if 0 <= idx < counter._num:
data[idx] = (ep, int(tokens), float(seconds), int(samples))
counter.restore_buckets(data)
self._node_model[key] = stat
class _NodeEntry:
@@ -309,6 +485,7 @@ class _NodeEntry:
"model", "hf_repo", "num_layers", "model_metadata", "shard_checksum", "hardware_profile", "wallet_address",
"score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards",
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
"model_tokens_per_sec",
"pending_directives", "last_heartbeat", "tracker_mode",
"relay_addr", "cert_fingerprint", "peer_id",
# heartbeat stats (reported by node, cumulative)
@@ -361,6 +538,7 @@ class _NodeEntry:
self.benchmark_tokens_per_sec = benchmark_tokens_per_sec
self.quantization = quantization
self.managed_assignment = managed_assignment
self.model_tokens_per_sec: dict[str, float] = {}
self.tracker_mode = tracker_mode
self.hf_repo = hf_repo
self.num_layers = num_layers
@@ -380,20 +558,31 @@ class _NodeEntry:
self.pending_new_assignment: dict | None = None
def _effective_throughput(node: "_NodeEntry") -> float:
def _effective_throughput(node: "_NodeEntry", model: str | None = None) -> float:
"""Effective tokens/s accounting for current queue depth."""
return node.benchmark_tokens_per_sec / (node.queue_depth + 1)
observed = None
if model:
observed = node.model_tokens_per_sec.get(model)
if observed is None:
for alias in _model_aliases(model):
observed = node.model_tokens_per_sec.get(alias)
if observed is not None:
break
base = observed if observed is not None and observed > 0 else node.benchmark_tokens_per_sec
return base / (node.queue_depth + 1)
def _select_route(
nodes: list[_NodeEntry],
required_start: int,
required_end: int,
model: str | None = None,
) -> tuple[list[_NodeEntry], str]:
"""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).
effective throughput: observed per-model tokens/sec / (queue_depth + 1),
falling back to startup benchmark_tokens_per_sec until observations exist.
Tiebreak: higher shard_end (fewer hops).
"""
candidates = sorted(
@@ -411,7 +600,7 @@ def _select_route(
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):
elif node.shard_end == best.shard_end and _effective_throughput(node, model) > _effective_throughput(best, model):
best = node
if best is None:
missing = covered_up_to + 1
@@ -1379,6 +1568,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
preset = _hf_rebalance_preset([node])
return _node_capacity_summary(node, preset)
def throughput_for(node: _NodeEntry) -> dict:
if server.stats is None:
return {}
models = [m for m in (node.hf_repo, node.model) if m]
result = {}
for model in models:
result[model] = server.stats.get_node_model_stats(node.node_id, model)
return result
self._send_json(200, {
"relay_url": server.relay_url,
"pool": _pool_summary(nodes),
@@ -1408,6 +1606,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"tracker_mode": node.tracker_mode,
"last_heartbeat": node.last_heartbeat,
"capacity": capacity_for(node),
"throughput": throughput_for(node),
"stats": _node_health(node, server.heartbeat_timeout),
}
for node in nodes
@@ -1513,8 +1712,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
}})
return
# Simple round-robin via list length modulo (stateless, good enough)
node = candidates[int(time.time() * 1000) % len(candidates)]
node = max(candidates, key=lambda n: _effective_throughput(n, model))
target_url = f"{node.endpoint}/v1/chat/completions"
request_id = str(body.get("id") or f"req-{time.time_ns():x}")
@@ -1537,7 +1735,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if pinned_nodes is not None:
route_nodes = pinned_nodes
else:
route_nodes, _ = _select_route(all_nodes, rs, re)
route_nodes, _ = _select_route(all_nodes, rs, re, model=route_model)
# 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
@@ -1591,12 +1789,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
f"(direct endpoint {target_url})",
flush=True,
)
started = time.monotonic()
relayed = _relay_http_request(
node.relay_addr,
path="/v1/chat/completions",
body=raw_body,
headers=relay_headers,
)
elapsed = time.monotonic() - started
if relayed is not None:
self._send_relayed_response(relayed)
if int(relayed.get("status", 503)) < 400:
@@ -1605,6 +1805,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
tokens = _usage_total_tokens(json.loads(body_text)) or 0
except (json.JSONDecodeError, TypeError):
tokens = 0
self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes)
self._bill_completed(api_key, model, tokens, node_work)
return
print(
@@ -1614,6 +1815,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
)
try:
started = time.monotonic()
upstream = urllib.request.urlopen(req, timeout=300.0)
print(f"[tracker] proxy connected {request_id}: {target_url}", flush=True)
except urllib.error.HTTPError as exc:
@@ -1674,17 +1876,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
pass
except BrokenPipeError:
pass
elapsed = time.monotonic() - started
# Bill even on client disconnect — the nodes did the work.
# Chunk count approximates generated tokens when the stream
# carries no usage record.
observed_tokens = stream_tokens if stream_tokens is not None else chunk_count
self._record_observed_throughput(model, route_model, observed_tokens, elapsed, route_nodes)
self._bill_completed(
api_key, model,
stream_tokens if stream_tokens is not None else chunk_count,
observed_tokens,
node_work,
)
else:
# Non-streaming: buffer and relay
resp_body = upstream.read()
elapsed = time.monotonic() - started
print(
f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}",
flush=True,
@@ -1701,8 +1907,42 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
tokens = _usage_total_tokens(json.loads(resp_body)) or 0
except json.JSONDecodeError:
tokens = 0
self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes)
self._bill_completed(api_key, model, tokens, node_work)
def _record_observed_throughput(
self,
requested_model: str,
route_model: str,
total_tokens: int,
elapsed_seconds: float,
route_nodes: list[_NodeEntry],
) -> None:
"""Record observed route TPS for participating nodes.
The tracker sees end-to-end request duration, not per-hop timings, so
each hop gets the same route-level observation for now. Per-hop telemetry
can refine this later without changing the external stats shape.
"""
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.stats is None or total_tokens <= 0 or elapsed_seconds <= 0:
return
models = [m for m in (requested_model, route_model) if m]
if len(models) == 2 and models[0] == models[1]:
models = [models[0]]
for node in route_nodes:
for model in models:
server.stats.record_node_throughput(
node.node_id,
model,
total_tokens=total_tokens,
elapsed_seconds=elapsed_seconds,
)
stats = server.stats.get_node_model_stats(node.node_id, model)
observed = stats.get("tokens_per_sec_last_hour")
if observed is not None:
node.model_tokens_per_sec[model] = float(observed)
def _bill_completed(
self,
api_key: str | None,
@@ -2073,9 +2313,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
def _handle_stats(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.stats is None:
self._send_json(200, {"models": {}})
self._send_json(200, {"models": {}, "nodes": {}})
return
self._send_json(200, {"models": server.stats.get_combined_stats()})
self._send_json(200, {
"models": server.stats.get_combined_stats(),
"nodes": server.stats.get_node_throughput_stats(),
})
def _handle_stats_gossip(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]