fix inference

This commit is contained in:
Dobromir Popov
2026-06-30 13:01:29 +03:00
parent dade97ee67
commit c691e8d5d1
4 changed files with 105 additions and 9 deletions

View File

@@ -204,7 +204,7 @@ def run_startup(
tracker_url=tracker_url, tracker_url=tracker_url,
) )
actual_port = node.start() actual_port = node.start()
total_layers = getattr(node.backend, "total_layers", None) total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
if isinstance(total_layers, int) and total_layers > 0: if isinstance(total_layers, int) and total_layers > 0:
layer_count = shard_end - shard_start + 1 layer_count = shard_end - shard_start + 1
shard_label = f"layers {shard_start}{shard_end}; {layer_count} of {total_layers}" shard_label = f"layers {shard_start}{shard_end}; {layer_count} of {total_layers}"
@@ -213,7 +213,7 @@ def run_startup(
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}" endpoint = f"http://{public_host}:{actual_port}"
# Register with tracker so other nodes can auto-join this model. # Register with tracker so other nodes can auto-join this model.
total_layers = getattr(node.backend, "total_layers", None) total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
reg_payload = { reg_payload = {
"endpoint": endpoint, "endpoint": endpoint,
"model": model_id.split("/")[-1], "model": model_id.split("/")[-1],

View File

@@ -246,6 +246,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# but correct. Each step: head encodes current sequence → forwards through route # but correct. Each step: head encodes current sequence → forwards through route
# → tail returns the next token string → append → repeat. # → tail returns the next token string → append → repeat.
remaining_route = self._get_remaining_route(model_name) remaining_route = self._get_remaining_route(model_name)
print(
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
f"downstream={remaining_route}",
flush=True,
)
if not remaining_route: if not remaining_route:
self._send_openai_response( self._send_openai_response(
"error: no downstream route — check tracker connectivity", "error: no downstream route — check tracker connectivity",
@@ -300,7 +305,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
try: try:
route = json.loads(injected) route = json.loads(injected)
if isinstance(route, list): if isinstance(route, list):
return [str(ep) for ep in route] resolved = [str(ep) for ep in route]
print(f" [node] using injected downstream route: {resolved}", flush=True)
return resolved
except (json.JSONDecodeError, TypeError): except (json.JSONDecodeError, TypeError):
pass pass
@@ -315,7 +322,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
route_resp = json.loads(r.read()) route_resp = json.loads(r.read())
route = route_resp.get("route", []) route = route_resp.get("route", [])
own_port = server.server_address[1] own_port = server.server_address[1]
return [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")] 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
except Exception as exc: except Exception as exc:
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True) print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
return [] return []
@@ -346,6 +355,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
current_pos = pos_ids current_pos = pos_ids
for hop_index, node_url in enumerate(route): for hop_index, node_url in enumerate(route):
print(f" [node] pipeline hop {hop_index}: {node_url}", flush=True)
headers: dict[str, str] = { headers: dict[str, str] = {
"Content-Type": "application/octet-stream", "Content-Type": "application/octet-stream",
"X-Meshnet-Wire": _WIRE_VERSION, "X-Meshnet-Wire": _WIRE_VERSION,
@@ -371,12 +381,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
resp_body = r.read() resp_body = r.read()
resp_headers = {k.lower(): v for k, v in r.headers.items()} resp_headers = {k.lower(): v for k, v in r.headers.items()}
except Exception as exc: except Exception as exc:
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
return f"pipeline error at {node_url}: {exc}" return f"pipeline error at {node_url}: {exc}"
content_type = resp_headers.get("content-type", "") content_type = resp_headers.get("content-type", "")
if "application/json" in content_type: if "application/json" in content_type:
try: try:
data = json.loads(resp_body) data = json.loads(resp_body)
return str(data.get("text", "")) text = str(data.get("text", ""))
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
return text
except json.JSONDecodeError: except json.JSONDecodeError:
return resp_body.decode("utf-8", errors="replace") return resp_body.decode("utf-8", errors="replace")
# Binary activation — update and forward to next node # Binary activation — update and forward to next node

View File

@@ -24,6 +24,7 @@ HTTP API contract:
import http.server import http.server
import hashlib import hashlib
import json import json
import socketserver
import threading import threading
import time import time
import urllib.parse import urllib.parse
@@ -291,7 +292,7 @@ def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
for node_id in expired_ids: for node_id in expired_ids:
entry = server.registry.pop(node_id) entry = server.registry.pop(node_id)
print( print(
f"[tracker] node expired: {node_id[:8]} {entry.endpoint} " f"[tracker] node expired: {node_id} {entry.endpoint} "
f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)", f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)",
flush=True, flush=True,
) )
@@ -418,7 +419,9 @@ def _node_id_for_registration(
return f"{wallet_prefix}-{digest}" return f"{wallet_prefix}-{digest}"
class _TrackerHTTPServer(http.server.HTTPServer): class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
daemon_threads = True
def __init__( def __init__(
self, self,
addr: tuple, addr: tuple,
@@ -706,6 +709,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
# Simple round-robin via list length modulo (stateless, good enough) # Simple round-robin via list length modulo (stateless, good enough)
node = candidates[int(time.time() * 1000) % len(candidates)] node = candidates[int(time.time() * 1000) % len(candidates)]
target_url = f"{node.endpoint}/v1/chat/completions" target_url = f"{node.endpoint}/v1/chat/completions"
request_id = str(body.get("id") or f"req-{time.time_ns():x}")
# Pre-resolve the downstream route so the first-shard node skips its own # Pre-resolve the downstream route so the first-shard node skips its own
# tracker query. We already hold the full registry picture — no need for # tracker query. We already hold the full registry picture — no need for
@@ -726,6 +730,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
route_nodes, _ = _select_route(all_nodes, rs, re) route_nodes, _ = _select_route(all_nodes, rs, re)
# Strip the first-shard node we're about to proxy to — it's already handling the request. # 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_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint])
route_debug = " -> ".join(
f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
for n in route_nodes
)
print(
f"[tracker] proxy route {request_id}: model={model!r} "
f"head={node.node_id}@{node.endpoint} downstream={downstream_urls} "
f"route={route_debug or '<empty>'}",
flush=True,
)
req = urllib.request.Request( req = urllib.request.Request(
target_url, target_url,
@@ -743,6 +757,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
try: try:
upstream = urllib.request.urlopen(req, timeout=300.0) 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: except urllib.error.HTTPError as exc:
# Relay error status + body from node # Relay error status + body from node
err_body = exc.read() err_body = exc.read()
@@ -756,6 +771,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
pass pass
return return
except Exception as exc: except Exception as exc:
print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True)
self._send_json(503, {"error": { self._send_json(503, {"error": {
"message": f"upstream node unreachable: {exc}", "message": f"upstream node unreachable: {exc}",
"type": "service_unavailable", "type": "service_unavailable",
@@ -783,6 +799,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
else: else:
# Non-streaming: buffer and relay # Non-streaming: buffer and relay
resp_body = upstream.read() resp_body = upstream.read()
print(
f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}",
flush=True,
)
self.send_response(200) self.send_response(200)
self.send_header("Content-Type", content_type) self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(resp_body))) self.send_header("Content-Length", str(len(resp_body)))
@@ -966,7 +986,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
for eid in stale_ids: for eid in stale_ids:
old = server.registry.pop(eid) old = server.registry.pop(eid)
print( print(
f"[tracker] node re-registered: replaced {eid[:8]} with {node_id[:8]}" f"[tracker] node re-registered: replaced {eid} with {node_id}"
f" {old.endpoint}", f" {old.endpoint}",
flush=True, flush=True,
) )
@@ -980,7 +1000,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
repo_info = f" [{hf_repo}]" if hf_repo else "" repo_info = f" [{hf_repo}]" if hf_repo else ""
print( print(
f"[tracker] node registered: {node_id[:8]} {endpoint} {model}{repo_info} {shard_info}", f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info}",
flush=True, flush=True,
) )

View File

@@ -50,6 +50,69 @@ def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect():
_TrackerHandler._send_json(DummyHandler(), 200, {"ok": True}) _TrackerHandler._send_json(DummyHandler(), 200, {"ok": True})
def test_tracker_serves_health_while_proxy_request_is_in_flight():
"""Long inference proxy requests must not block heartbeats/health checks."""
class SlowChatHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args):
pass
def do_POST(self):
if self.path != "/v1/chat/completions":
self.send_response(404)
self.end_headers()
return
length = int(self.headers.get("Content-Length", 0))
self.rfile.read(length)
time.sleep(2.0)
body = json.dumps({"choices": [{"message": {"content": "ok"}}]}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
slow_node = http.server.HTTPServer(("127.0.0.1", 0), SlowChatHandler)
slow_thread = threading.Thread(target=slow_node.serve_forever, daemon=True)
slow_thread.start()
tracker = TrackerServer(heartbeat_timeout=60.0)
tracker_port = tracker.start()
proxy_error = []
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{slow_node.server_address[1]}",
"model": "slow-model", "hf_repo": "org/slow-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0, "tracker_mode": True,
"hardware_profile": {}, "score": 1.0},
)
def call_proxy():
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "slow-model", "messages": [{"role": "user", "content": "hi"}]},
)
except Exception as exc:
proxy_error.append(exc)
proxy_thread = threading.Thread(target=call_proxy)
proxy_thread.start()
time.sleep(0.2)
with urllib.request.urlopen(f"http://127.0.0.1:{tracker_port}/v1/health", timeout=1.0) as resp:
assert resp.status == 200
proxy_thread.join(timeout=3.0)
assert not proxy_thread.is_alive()
assert not proxy_error
finally:
tracker.stop()
slow_node.shutdown()
slow_node.server_close()
slow_thread.join(timeout=1.0)
def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix(): def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix():
tracker = TrackerServer() tracker = TrackerServer()
tracker_port = tracker.start() tracker_port = tracker.start()