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

distributd cache

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 22:53:03 +02:00
parent 436e872abe
commit daddbaa4a3
8 changed files with 966 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -17,11 +17,17 @@ 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,
@@ -128,6 +134,13 @@ def _relay_hop(
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,
@@ -346,6 +359,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
start_layer = int(start_layer_header) if start_layer_header else None
# 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")
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,
@@ -353,11 +379,18 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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
@@ -512,10 +545,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._send_json(500, {"error": f"generation failed: {exc}"})
return
# Distributed path: autoregressive generation across shards.
# We do N single-step forward passes (no cross-node KV cache), which is slow
# but correct. Each step: head encodes current sequence → forwards through route
# → tail returns the next token string → append → repeat.
# 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} "
@@ -548,6 +583,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
generated: list[str] = []
current_text = prompt_text
session_id = str(uuid.uuid4())
use_kv = bool(getattr(backend, "supports_kv_cache", False))
eos_ids: set[int] = set()
if use_kv:
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)
@@ -557,25 +601,63 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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:
payload = backend.encode_prompt(current_text)
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
token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
if not token_str:
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
generated.append(token_str)
if stream_emit is not None:
stream_emit(token_str)
current_text = current_text + token_str
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,
@@ -594,6 +676,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
)
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)
@@ -687,7 +775,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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) -> str:
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:
@@ -699,12 +801,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
bytearray(payload.body), # type: ignore[union-attr]
dtype=active_backend.torch.bfloat16,
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
return active_backend.decode_tail(tensor)
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}"
return "no downstream route available for non-tail shard"
return f"decode error: {exc}", None
return "no downstream route available for non-tail shard", None
session = str(uuid.uuid4())
# 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]
@@ -734,6 +841,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"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:
@@ -743,12 +855,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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}"
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}; "
@@ -767,26 +883,33 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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}"
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
return text, int(token_id) if token_id is not None else None
except json.JSONDecodeError:
return resp_body.decode("utf-8", errors="replace")
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 ""
return "", None
def _stream_openai_response(self, token_iter, model: str) -> None:
"""Stream tokens from an iterator as SSE chunks."""