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

wip -more responsive UI, better routing

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 09:07:54 +02:00
parent 518c259cd3
commit 3d82188dc1
14 changed files with 506 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
packages/node/meshnet_node/startup.py
View File

@@ -349,25 +349,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 +392,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 +463,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

135
packages/node/meshnet_node/torch_server.py
View File

@@ -31,6 +31,23 @@ from .server import (
)
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,
@@ -91,6 +108,26 @@ class _TorchHTTPServer(http.server.HTTPServer):
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 +150,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 +374,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 +390,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 +407,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 +418,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 +430,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 +504,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 +529,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)
@@ -849,6 +954,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())

7
packages/node/pyproject.toml
View File

@@ -15,9 +15,10 @@ dependencies = [
"bitsandbytes>=0.43",
"rich>=13",
"safetensors>=0.4",
"torch>=2.1",
"transformers>=5.12",
"websockets>=13",
"torch>=2.1",
"transformers>=5.12",
"triton-windows>=3.7; platform_system == 'Windows'",
"websockets>=13",
"zstandard>=0.22",
"kernels>=0.11.1,<0.16",
]