"""HTTP server for real PyTorch-backed shard nodes.""" from __future__ import annotations import base64 import http.server import json import sys import threading import time import urllib.error import urllib.parse import urllib.request import uuid from pathlib import Path from typing import Any from .model_backend import ( InsufficientVRAMError, KVCacheMiss, MissingModelDependencyError, Quantization, TailTokenResult, TorchModelShard, validate_quantization, ) class _PipelineCacheMiss(Exception): """A downstream hop reported 409 cache_miss — head must re-prefill.""" from .server import ( _WIRE_VERSION, _compress_body, _decompress_body, _parse_shape, _validate_activation_body, ) 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, body: bytes, headers: dict[str, str], timeout: float = 120.0, ) -> tuple[int, dict[str, str], bytes]: """Send a single HTTP-shaped request through a relay RPC WebSocket. relay_addr is the wss://relay.../rpc/{peer_id} URL. Returns (status, response_headers_lower, response_body). Raises on connection failure so callers can fall back to direct. """ import websockets.sync.client as wsc # type: ignore[import] request_id = f"{time.time_ns():x}" payload = json.dumps({ "request_id": request_id, "method": "POST", "path": path, "headers": headers, "body_base64": base64.b64encode(body).decode(), }) with wsc.connect(relay_addr, open_timeout=timeout) as ws: ws.send(payload) raw = ws.recv(timeout=timeout) resp = json.loads(raw) status = int(resp.get("status", 503)) resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()} body_b64 = resp.get("body_base64") resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode() return status, resp_headers, resp_body def _is_cache_miss_body(body: bytes) -> bool: try: return json.loads(body).get("error") == "cache_miss" except (json.JSONDecodeError, AttributeError, UnicodeDecodeError): return False class _TorchHTTPServer(http.server.HTTPServer): def __init__( self, addr, handler, backend: TorchModelShard, tracker_mode: bool = False, tracker_url: str | None = None, route_timeout: float = 30.0, debug: bool = False, max_loaded_shards: int = 1, ): super().__init__(addr, handler) self.backend = backend self.backends: dict[str, TorchModelShard] = {backend.model_id: backend} self.received_activations = False self.forward_chunk_count = 0 self.tracker_mode = tracker_mode self.tracker_url = tracker_url 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]] = {} self._decode_log: dict[str, dict[str, float]] = {} def note_decode_step( self, session: str, now: float | None = None, ) -> int | None: """Count one decode forward; return the cumulative step count when a log line is due (first step of a session, then every 5s), else None.""" if now is None: now = time.monotonic() with self._stats_lock: rec = self._decode_log.get(session) if rec is None: if len(self._decode_log) >= 64: stale = [ sid for sid, r in self._decode_log.items() if now - r["seen"] > 600.0 ] for sid in stale: del self._decode_log[sid] while len(self._decode_log) >= 64: self._decode_log.pop(next(iter(self._decode_log))) self._decode_log[session] = {"steps": 1.0, "logged": now, "seen": now} return 1 rec["steps"] += 1 rec["seen"] = now if now - rec["logged"] >= 5.0: rec["logged"] = now return int(rec["steps"]) return None 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: return self.backend wanted = model_name.strip().lower() for key, shard_backend in self.backends.items(): key_l = key.lower() if key_l == wanted or key_l.rsplit("/", 1)[-1] == wanted: return shard_backend return self.backend def chat_enabled(self) -> bool: return any( shard_backend.is_head for shard_backend in self.backends.values() ) 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": self._handle_forward() elif self.path == "/v1/infer": self._handle_infer() elif self.path == "/v1/chat/completions" and server.chat_enabled(): self._handle_chat_completions() else: self.send_response(404) self.end_headers() def _handle_infer(self) -> None: body = self._read_json_body() if body is None: return messages = body.get("messages", []) prompt = "" if isinstance(messages, list) and messages: last = messages[-1] if isinstance(last, dict): prompt = str(last.get("content", "")) server: _TorchHTTPServer = self.server # type: ignore[assignment] try: payload = server.backend.encode_prompt(prompt) if server.backend.is_tail: text = server.backend.decode_tail( server.backend.torch.frombuffer( bytearray(payload.body), dtype=server.backend.torch.bfloat16, ) .reshape(payload.shape) .to(server.backend.device) ) self._send_json(200, {"text": text}) return self._send_json(200, {"activations": {"shape": payload.shape, "dtype": "bfloat16"}}) except Exception as exc: self._send_json(500, {"error": str(exc)}) def _handle_forward(self) -> None: content_type = self.headers.get("Content-Type", "") if content_type.startswith("application/json"): self._handle_prompt_forward() return self._handle_binary_forward() def _handle_prompt_forward(self) -> None: body = self._read_json_body() if body is None: return prompt = str(body.get("prompt", "")) server: _TorchHTTPServer = self.server # type: ignore[assignment] try: payload = server.backend.encode_prompt(prompt) except Exception as exc: self._send_json(400, {"error": str(exc)}) return self._send_activation(payload) def _handle_binary_forward(self) -> None: server: _TorchHTTPServer = self.server # type: ignore[assignment] try: shape = _parse_shape(self.headers.get("X-Meshnet-Shape")) dtype = self.headers.get("X-Meshnet-Dtype", "") session = self.headers["X-Meshnet-Session"] chunk_index = self.headers["X-Meshnet-Chunk-Index"] chunk_total = self.headers["X-Meshnet-Chunk-Total"] encoding = self.headers.get("X-Meshnet-Encoding") length = int(self.headers.get("Content-Length", 0)) body = self.rfile.read(length) raw_body = _decompress_body(body, encoding) _validate_activation_body(raw_body, shape, dtype) if dtype != "bfloat16": raise ValueError("real model backend requires bfloat16 activation input") chunk_index_value = int(chunk_index) chunk_total_value = int(chunk_total) if chunk_total_value <= 0 or not 0 <= chunk_index_value < chunk_total_value: raise ValueError("invalid chunk index/total") except (KeyError, ValueError, TypeError): self.send_response(400) self.send_header("X-Meshnet-Wire", _WIRE_VERSION) self.end_headers() return server.forward_chunk_count += 1 hop_index = int(self.headers.get("X-Meshnet-Hop-Index", "0")) if hop_index > 0: server.received_activations = True # Session KV-cache protocol: prefill establishes per-session state on # this node's layer range; decode reuses it. Absent header = legacy # stateless call (also the signature fake backends implement). cache_mode = self.headers.get("X-Meshnet-Cache") if chunk_index_value == 0: shard_start = getattr(server.backend, "shard_start", "?") shard_end = getattr(server.backend, "shard_end", "?") if cache_mode == "decode": # One decode forward arrives per generated token — log a # periodic per-session summary instead of one line per token. steps = server.note_decode_step(session) if steps is not None: print( f" [node] decoding layers={shard_start}-{shard_end} " f"session={session[:8]} steps={steps}" f"{self._request_log_suffix()}", flush=True, ) else: print( f" [node] forward hop={hop_index} " f"layers={shard_start}-{shard_end} " f"session={session[:8]}{self._request_log_suffix()}", flush=True, ) start_layer_header = self.headers.get("X-Meshnet-Start-Layer") start_layer = int(start_layer_header) if start_layer_header else None forward_kwargs: dict[str, object] = {} if cache_mode in ("prefill", "decode"): past_len_header = self.headers.get("X-Meshnet-Past-Len") forward_kwargs = { "session_id": session, "cache_mode": cache_mode, "past_len": int(past_len_header) if past_len_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, **forward_kwargs, ) except KVCacheMiss as exc: self._send_json(409, {"error": "cache_miss", "detail": str(exc)}) return except Exception as exc: self._send_json(500, {"error": str(exc)}) return if isinstance(result, TailTokenResult): self._send_json(200, {"text": result.text, "token_id": result.token_id}) return if isinstance(result, str): self._send_json(200, {"text": result}) return response_body = _compress_body(result.body, encoding) self.send_response(200) self.send_header("Content-Type", "application/octet-stream") self.send_header("Content-Length", str(len(response_body))) self.send_header("X-Meshnet-Wire", _WIRE_VERSION) self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in result.shape)) self.send_header("X-Meshnet-Dtype", "bfloat16") self.send_header("X-Meshnet-Session", session) self.send_header("X-Meshnet-Chunk-Index", chunk_index) self.send_header("X-Meshnet-Chunk-Total", chunk_total) if encoding: self.send_header("X-Meshnet-Encoding", encoding) if result.attention_mask_header: self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header) if result.position_ids_header: self.send_header("X-Meshnet-Position-Ids", result.position_ids_header) self.end_headers() self.wfile.write(response_body) def _send_activation(self, payload) -> None: body = payload.body self.send_response(200) self.send_header("Content-Type", "application/octet-stream") self.send_header("Content-Length", str(len(body))) self.send_header("X-Meshnet-Wire", _WIRE_VERSION) self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in payload.shape)) self.send_header("X-Meshnet-Dtype", "bfloat16") if payload.attention_mask_header: self.send_header("X-Meshnet-Attn-Mask", payload.attention_mask_header) if payload.position_ids_header: self.send_header("X-Meshnet-Position-Ids", payload.position_ids_header) self.end_headers() self.wfile.write(body) def _read_json_body(self) -> dict | None: length = int(self.headers.get("Content-Length", 0)) try: body = json.loads(self.rfile.read(length) or b"{}") except (json.JSONDecodeError, ValueError): self._send_json(400, {"error": "invalid JSON body"}) return None if not isinstance(body, dict): self._send_json(400, {"error": "JSON body must be an object"}) return None return body def _send_json(self, status: int, data: dict) -> None: payload = json.dumps(data).encode() self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(payload))) self.end_headers() try: self.wfile.write(payload) except BrokenPipeError: pass # client disconnected before we could respond — not an error 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, request_id) finally: self._track_request_end(server, request_id) 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", request_id: str) -> None: body = self._read_json_body() if body is None: return messages = body.get("messages", []) if not isinstance(messages, list): messages = [] stream = bool(body.get("stream", False)) model_name = str(body.get("model", "")) backend = server.resolve_backend(model_name) if backend is None or not backend.is_head: self._send_json(400, {"error": "model not loaded on this node"}) return max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 5120) 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()}", flush=True, ) # Fast path: this node owns the complete model — use HF generate() with KV cache. # 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 def _counting_stream(): nonlocal token_count for token_text in backend.generate_text_streaming( messages, max_tokens, temperature, top_p, ): 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) 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={elapsed:.1f} tps={tps:.2f}{self._request_log_suffix()}", final=True, ) else: text = backend.generate_text(messages, max_tokens, temperature, top_p) completion_tokens = _backend_token_count( backend, "count_text_tokens", text, fallback=len(text.split()) or 1, ) print( f" [node] chat complete tokens={completion_tokens} " f"elapsed_s={time.monotonic() - gen_started:.1f}{self._request_log_suffix()}", flush=True, ) self._send_openai_response(text, model_name, False, messages, backend=backend) except Exception as exc: self._record_failed_request() print( f" [node] chat failed after {time.monotonic() - gen_started:.1f}s: {exc}" f"{self._request_log_suffix()}", flush=True, ) self._send_json(500, {"error": f"generation failed: {exc}"}) return # Distributed path: autoregressive generation across shards with a # sharded per-node KV cache. Step 0 prefills the full prompt through the # route (each node caches state for its own layer range, keyed by a # per-generation session id); steps 1+ send only the newest token's # hidden state. A 409 cache_miss from any hop (eviction/restart/route # change) falls back to a full re-prefill — the old stateless behavior. remaining_route = self._get_remaining_route(model_name, backend=backend) print( f" [node] chat route model={model_name!r} max_tokens={max_tokens} " f"downstream={remaining_route}", flush=True, ) if not remaining_route: self._send_openai_response( "error: no downstream route — check tracker connectivity", model_name, False, messages, backend=backend, ) return # Format with chat template so the model knows it's in assistant mode. try: if hasattr(backend.tokenizer, "apply_chat_template"): prompt_text: str = backend.tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=False, ) else: raise AttributeError("no apply_chat_template") except Exception: prompt_text = " ".join( str(m.get("content", "")) for m in messages if isinstance(m, dict) and m.get("role") == "user" ) eos_token: str = getattr(backend.tokenizer, "eos_token", "") or "" generated: list[str] = [] current_text = prompt_text session_id = str(uuid.uuid4()) use_kv = bool(getattr(backend, "supports_kv_cache", False)) # EOS detection by id must work on the stateless path too: the tail # returns token_id regardless of caching, and EOS usually decodes to # "" (skip_special_tokens), so the text comparison never fires. eos_ids: set[int] = set() try: eos_ids = set(backend.eos_token_ids()) except Exception: eos_ids = set() 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] last_token_id: int | None = None def _prefill_step() -> tuple[str, int | None]: """Full-sequence prefill: initial step and cache-miss recovery.""" payload = ( backend.encode_prompt(current_text, session_id=session_id) if use_kv else backend.encode_prompt(current_text) ) return self._run_downstream_pipeline( payload, remaining_route, backend=backend, session=session_id, cache_mode="prefill" if use_kv else None, ) for step in range(max_tokens): try: if use_kv and step > 0 and last_token_id is not None: try: payload = backend.encode_next_token(last_token_id, session_id) token_str, token_id = self._run_downstream_pipeline( payload, remaining_route, backend=backend, session=session_id, cache_mode="decode", ) except (KVCacheMiss, _PipelineCacheMiss) as miss: # Evicted/restarted node or head lost its own session: # re-prefill the whole sequence once and continue cached. print( f" [node] kv cache miss at step {step} ({miss}); " f"re-prefilling {len(current_text)} chars", flush=True, ) token_str, token_id = _prefill_step() else: token_str, token_id = _prefill_step() except _PipelineCacheMiss as exc: print(f" [node] unexpected cache miss on prefill: {exc}", flush=True) break except Exception as exc: print(f" [node] distributed encode error: {exc}", flush=True) break # Stop on error responses or EOS. if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")): break if token_id is not None and token_id in eos_ids: break if eos_token and token_str == eos_token: break if not token_str and token_id is None: break last_token_id = token_id # token_str can be empty for a skipped special token that is not # EOS — keep generating from its token_id without emitting text. if token_str: generated.append(token_str) 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: 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={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}", ) last_gen_log = now if use_kv: try: backend.release_session(session_id) except Exception: pass if generated: 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) if stream_emit is not None: stream_emit(None) return self._send_openai_response(result_text, model_name, stream, messages, backend=backend) def _get_remaining_route(self, model: str, *, backend: TorchModelShard | None = None) -> list[dict]: """Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr. 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. """ 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: try: route = json.loads(injected) if isinstance(route, list): hops: list[dict] = [] for item in route: if isinstance(item, dict): hop = { "endpoint": str(item["endpoint"]), "start_layer": int(item.get("start_layer", 0)), } if item.get("relay_addr"): hop["relay_addr"] = str(item["relay_addr"]) hops.append(hop) elif isinstance(item, str): hops.append({"endpoint": item, "start_layer": 0}) 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). if server.tracker_url is None: return [] route_model = getattr(active_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=server.route_timeout) as r: route_resp = json.loads(r.read()) own_key = _own_endpoint_key(server) nodes_info = route_resp.get("nodes", []) hops: list[dict] = [] passed_self = False for node_info in nodes_info: ep = node_info.get("endpoint", "") if not ep: continue 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) 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) return [] def _run_downstream_pipeline( self, payload: object, route: list[dict], *, backend: TorchModelShard | None = None, session: str | None = None, cache_mode: str | None = None, ) -> tuple[str, int | None]: """Forward an activation through the downstream route. Returns (token_text, token_id) — token_id is None when a hop predates the KV-cache protocol. Raises _PipelineCacheMiss when a hop responds 409 cache_miss (evicted/restarted node) so the caller can re-prefill. """ server: _TorchHTTPServer = self.server # type: ignore[assignment] active_backend = backend or server.backend if not route: # Partial shard at tail: decode the activation from the previous node. # Full single-node (head+tail) is handled before entering this method. if active_backend.is_tail: try: tensor = active_backend.torch.frombuffer( bytearray(payload.body), # type: ignore[union-attr] dtype=active_backend.torch.bfloat16, ).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr] if hasattr(active_backend, "decode_tail_token"): tail = active_backend.decode_tail_token(tensor) return tail.text, tail.token_id return active_backend.decode_tail(tensor), None except Exception as exc: return f"decode error: {exc}", None return "no downstream route available for non-tail shard", None # Session is stable across all steps of one generation when the caller # provides it (KV-cache protocol); fresh per call otherwise (legacy). session = session or str(uuid.uuid4()) shape = payload.shape # type: ignore[union-attr] attn_mask = payload.attention_mask_header # type: ignore[union-attr] pos_ids = payload.position_ids_header # type: ignore[union-attr] current_body = payload.body # type: ignore[union-attr] current_shape = shape current_attn = attn_mask current_pos = pos_ids for hop_index, hop in enumerate(route): node_url = hop["endpoint"] start_layer = hop.get("start_layer", 0) relay_addr = hop.get("relay_addr") if server.debug: print( f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}" + (f" relay={relay_addr}" if relay_addr else ""), flush=True, ) headers: dict[str, str] = { "Content-Type": "application/octet-stream", "X-Meshnet-Wire": _WIRE_VERSION, "X-Meshnet-Shape": ",".join(str(d) for d in current_shape), "X-Meshnet-Dtype": "bfloat16", "X-Meshnet-Session": session, "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 cache_mode: headers["X-Meshnet-Cache"] = cache_mode past_len = getattr(payload, "past_len", None) if cache_mode == "decode" and past_len is not None: headers["X-Meshnet-Past-Len"] = str(past_len) if current_attn: headers["X-Meshnet-Attn-Mask"] = current_attn if current_pos: headers["X-Meshnet-Position-Ids"] = current_pos if relay_addr: try: status, resp_headers, resp_body = _relay_hop( relay_addr, "/forward", current_body, headers, timeout=120.0, ) if status == 409 and _is_cache_miss_body(resp_body): raise _PipelineCacheMiss(node_url) if status >= 400: print( f" [node] relay hop {hop_index} returned {status} from {relay_addr}", flush=True, ) return f"pipeline error at {node_url} via relay: status {status}", None except _PipelineCacheMiss: raise except Exception as exc: print( f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; " f"falling back to direct {node_url}", flush=True, ) relay_addr = None # fall through to direct if not relay_addr: req = urllib.request.Request( f"{node_url}/forward", data=current_body, headers=headers, method="POST", ) try: with urllib.request.urlopen(req, timeout=120.0) as r: resp_body = r.read() resp_headers = {k.lower(): v for k, v in r.headers.items()} except urllib.error.HTTPError as exc: body = exc.read() if exc.code == 409 and _is_cache_miss_body(body): raise _PipelineCacheMiss(node_url) from exc print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) return f"pipeline error at {node_url}: {exc}", None 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}", None content_type = resp_headers.get("content-type", "") if "application/json" in content_type: try: data = json.loads(resp_body) text = str(data.get("text", "")) token_id = data.get("token_id") if server.debug: print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True) return text, int(token_id) if token_id is not None else None except json.JSONDecodeError: return resp_body.decode("utf-8", errors="replace"), None # Binary activation — update and forward to next node shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape)) current_shape = _parse_shape(shape_header) current_body = resp_body current_attn = resp_headers.get("x-meshnet-attn-mask") current_pos = resp_headers.get("x-meshnet-position-ids") return "", None def _stream_openai_response(self, token_iter, model: str) -> None: """Stream tokens from an iterator as SSE chunks.""" emit = self._start_openai_stream(model) for token_text in token_iter: if not token_text: continue emit(token_text) emit(None) def _start_openai_stream(self, model: str): """Open an OpenAI-compatible SSE response and return a token emitter.""" chunk_id = "chatcmpl-node" created = int(time.time()) self.send_response(200) self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") self.end_headers() def _emit(data: str) -> None: try: self.wfile.write(f"data: {data}\n\n".encode()) self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], })) def emit_token(token_text: str | None) -> None: if token_text is None: _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], })) try: self.wfile.write(b"data: [DONE]\n\n") self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass return _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}], })) return emit_token def _send_openai_response( self, text: str, model: str, stream: bool, messages: list[dict] | None = None, backend: TorchModelShard | None = None, ) -> None: chunk_id = "chatcmpl-node" created = int(time.time()) active_backend = backend or self.server.backend # type: ignore[attr-defined] if not stream: usage = _usage_for_response(active_backend, messages or [], text) self._send_json(200, { "id": chunk_id, "object": "chat.completion", "created": created, "model": model, "choices": [{ "index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop", }], "usage": usage, }) return self.send_response(200) self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") self.end_headers() def _emit(data: str) -> None: try: self.wfile.write(f"data: {data}\n\n".encode()) self.wfile.flush() except BrokenPipeError: pass _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], })) _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}], })) _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], })) try: self.wfile.write(b"data: [DONE]\n\n") self.wfile.flush() except BrokenPipeError: pass def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]: prompt_tokens = _backend_token_count( backend, "count_prompt_tokens", messages, fallback=_fallback_message_token_count(messages), ) completion_tokens = _backend_token_count( backend, "count_text_tokens", completion_text, fallback=_fallback_text_token_count(completion_text), ) return { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens, } def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int: method: Any = getattr(backend, method_name, None) if callable(method): try: return max(0, int(method(value))) except Exception: pass return max(0, int(fallback)) def _fallback_message_token_count(messages: list[dict]) -> int: text = " ".join( str(message.get("content", "")) for message in messages if isinstance(message, dict) ) return _fallback_text_token_count(text) def _fallback_text_token_count(text: str) -> int: parts = text.split() if parts: return len(parts) return 1 if text else 0 class TorchNodeServer: """HTTP server backed by a HuggingFace causal language model shard.""" def __init__( self, host: str = "127.0.0.1", port: int = 0, model_id: str = "openai-community/gpt2", shard_start: int = 0, shard_end: int = 6, quantization: str = "bfloat16", backend: TorchModelShard | None = None, tracker_mode: bool | None = None, tracker_url: str | None = None, route_timeout: float = 30.0, cache_dir: Path | None = None, debug: bool = False, max_loaded_shards: int = 1, force_cpu: bool = False, ) -> None: self._host = host self._requested_port = port self._max_loaded_shards = max(1, max_loaded_shards) self._backend = backend or _load_backend( model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu, ) self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend} # 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._cache_dir = cache_dir self._debug = debug 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 @property def received_activations(self) -> bool: return self._server.received_activations if self._server is not None else False @property 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 @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()) def apply_tracker_directives(self, directives: list[dict]) -> dict | None: """Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more).""" add_directive = next( (directive for directive in reversed(directives) if directive.get("action") == "ADD_SHARD"), None, ) load_directive = next( (directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"), None, ) directive = add_directive or load_directive if directive is None: return None shard_start = int(directive["shard_start"]) shard_end = int(directive["shard_end"]) quantization = str(directive.get("quantization") or self._backend.quantization) model_id = str(directive.get("model") or self._backend.model_id) replacing = directive.get("action") == "LOAD_SHARD" if not replacing and len(self._backends) >= self._max_loaded_shards: print( f" [node] WARNING: ignoring ADD_SHARD for {model_id!r} — " f"loaded {len(self._backends)}/{self._max_loaded_shards} slots full", flush=True, ) return None action_label = "reassigned" if replacing else "additional" print( f" [node] loading {action_label} shard: {model_id} layers {shard_start}-{shard_end}", flush=True, ) try: new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir) except TypeError: new_backend = _load_backend(model_id, shard_start, shard_end, quantization) self._backends[model_id] = new_backend if replacing or shard_start == 0: self._backend = new_backend self._tracker_mode = shard_start == 0 print( f" [node] loaded {action_label} shard: {model_id} layers {shard_start}-{shard_end}", flush=True, ) if self._server is not None: self._server.backends = dict(self._backends) if replacing or shard_start == 0: self._server.backend = new_backend self._server.tracker_mode = self._tracker_mode return { "action": directive.get("action"), "model": model_id, "shard_start": shard_start, "shard_end": shard_end, "quantization": quantization, "tracker_mode": shard_start == 0, } def start(self) -> int: if self._server is not None: raise RuntimeError("TorchNodeServer is already running") self._server = _TorchHTTPServer( (self._host, self._requested_port), _TorchHandler, self._backend, self._tracker_mode, self._tracker_url, self._route_timeout, self._debug, self._max_loaded_shards, ) self._server.backends = dict(self._backends) self.port = self._server.server_address[1] self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) 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 self._server.shutdown() self._server.server_close() if self._thread is not None: self._thread.join(timeout=1) self._server = None self._thread = None self.port = None def _load_backend( model_id: str, shard_start: int, shard_end: int, quantization: str, cache_dir: Path | None = None, force_cpu: bool = False, ) -> TorchModelShard: from .model_backend import load_torch_shard quant = validate_quantization(quantization) try: return load_torch_shard( model_id, shard_start, shard_end, quant, cache_dir, force_cpu=force_cpu ) except MissingModelDependencyError: raise except InsufficientVRAMError as exc: print(f"ERROR: {exc}", file=sys.stderr, flush=True) raise