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

misc

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 23:32:51 +03:00
parent d648da3344
commit 94046f1102
6 changed files with 644 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,6 +19,8 @@ from .model_backend import (
InsufficientVRAMError,
MissingModelDependencyError,
Quantization,
ShardCacheMiss,
TensorPayload,
TorchModelShard,
validate_quantization,
)
@@ -31,6 +33,16 @@ from .server import (
)
class _PipelineCacheMiss(RuntimeError):
"""Downstream shard reported that its session-local cache was unavailable."""
class _PipelineResult:
def __init__(self, text: str, token_id: int | None = None):
self.text = text
self.token_id = token_id
def _endpoint_key(url: str) -> str:
"""Normalize http(s) endpoints for host:port comparison."""
parsed = urllib.parse.urlparse(url.rstrip("/"))
@@ -94,6 +106,48 @@ def _write_progress_line(state: list[bool], message: str, *, final: bool = False
sys.stdout.flush()
def _int_header(value: str | None) -> int | None:
if value is None or value == "":
return None
return int(value)
def _cache_mode_header(value: str | None) -> str:
return value if value in {"prefill", "decode"} else "stateless"
def _encode_prompt_for_session(backend: TorchModelShard, prompt: str, session_id: str) -> TensorPayload:
method = getattr(backend, "encode_prompt_cached", None)
if callable(method):
return method(prompt, session_id)
return backend.encode_prompt(prompt)
def _token_id_from_text(backend: TorchModelShard, text: str) -> int | None:
tokenizer = getattr(backend, "tokenizer", None)
if tokenizer is None or not callable(tokenizer):
return None
try:
encoded = tokenizer(text, return_tensors="pt", add_special_tokens=False)
except TypeError:
try:
encoded = tokenizer(text, return_tensors="pt")
except Exception:
return None
except Exception:
return None
input_ids = encoded.get("input_ids") if isinstance(encoded, dict) else getattr(encoded, "input_ids", None)
if input_ids is None:
return None
try:
return int(input_ids[0, -1].item())
except Exception:
try:
return int(input_ids[0][-1])
except Exception:
return None
def _relay_hop(
relay_addr: str,
path: str,
@@ -353,13 +407,28 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.headers.get("X-Meshnet-Attn-Mask"),
self.headers.get("X-Meshnet-Position-Ids"),
start_layer=start_layer,
session_id=session,
cache_mode=_cache_mode_header(self.headers.get("X-Meshnet-Cache-Mode")),
seq_len=_int_header(self.headers.get("X-Meshnet-Seq-Len")),
)
except ShardCacheMiss 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, str):
self._send_json(200, {"text": result})
token_id = None
if hasattr(server.backend, "_last_decoded_token_id"):
try:
token_id = int(getattr(server.backend, "_last_decoded_token_id"))
except Exception:
token_id = None
data: dict[str, Any] = {"text": result}
if token_id is not None:
data["token_id"] = token_id
self._send_json(200, data)
return
response_body = _compress_body(result.body, encoding)
@@ -513,9 +582,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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.
# Step 0 prefills the full prompt and creates shard-local caches. Later
# cached steps send only the previous token's activation through the route.
remaining_route = self._get_remaining_route(model_name, backend=backend)
print(
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
@@ -547,6 +615,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
generated: list[str] = []
current_text = prompt_text
session_id = str(uuid.uuid4())
last_token_id: int | None = None
current_seq_len: int | None = None
stream_emit = None
if stream:
@@ -560,11 +631,49 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
for step in range(max_tokens):
try:
payload = backend.encode_prompt(current_text)
if step == 0 or last_token_id is None or current_seq_len is None:
payload = _encode_prompt_for_session(backend, current_text, session_id)
current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else None
cache_mode = "prefill"
seq_len = current_seq_len
else:
seq_len = current_seq_len
try:
payload = backend.encode_token_cached(last_token_id, seq_len, session_id)
cache_mode = "decode"
except ShardCacheMiss:
payload = _encode_prompt_for_session(backend, current_text, session_id)
current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else current_seq_len
cache_mode = "prefill"
seq_len = current_seq_len
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)
try:
result = self._run_downstream_pipeline(
payload,
remaining_route,
backend=backend,
session_id=session_id,
cache_mode=cache_mode,
seq_len=seq_len,
)
except _PipelineCacheMiss:
try:
payload = _encode_prompt_for_session(backend, current_text, session_id)
current_seq_len = int(payload.shape[1]) if len(payload.shape) > 1 else current_seq_len
result = self._run_downstream_pipeline(
payload,
remaining_route,
backend=backend,
session_id=session_id,
cache_mode="prefill",
seq_len=current_seq_len,
)
except Exception as exc:
print(f" [node] distributed cache-miss recovery failed: {exc}", flush=True)
break
token_str = result.text
if not token_str:
break
# Stop on error responses or EOS.
@@ -573,6 +682,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if eos_token and token_str == eos_token:
break
generated.append(token_str)
last_token_id = result.token_id if result.token_id is not None else _token_id_from_text(backend, token_str)
if last_token_id is not None and current_seq_len is not None:
current_seq_len += 1
if stream_emit is not None:
stream_emit(token_str)
current_text = current_text + token_str
@@ -687,7 +799,16 @@ 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_id: str | None = None,
cache_mode: str = "stateless",
seq_len: int | None = None,
) -> _PipelineResult:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
if not route:
@@ -699,12 +820,14 @@ 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)
token_id = active_backend.decode_tail_token_id(tensor)
text = active_backend.tokenizer.decode([token_id], skip_special_tokens=True)
return _PipelineResult(text, token_id)
except Exception as exc:
return f"decode error: {exc}"
return "no downstream route available for non-tail shard"
return _PipelineResult(f"decode error: {exc}")
return _PipelineResult("no downstream route available for non-tail shard")
session = str(uuid.uuid4())
session = session_id 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]
@@ -733,7 +856,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"X-Meshnet-Chunk-Total": "1",
"X-Meshnet-Hop-Index": str(hop_index),
"X-Meshnet-Start-Layer": str(start_layer),
"X-Meshnet-Cache-Mode": cache_mode,
}
if seq_len is not None:
headers["X-Meshnet-Seq-Len"] = str(seq_len)
if current_attn:
headers["X-Meshnet-Attn-Mask"] = current_attn
if current_pos:
@@ -744,11 +870,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
relay_addr, "/forward", current_body, headers, timeout=120.0,
)
if status >= 400:
if status == 409:
raise _PipelineCacheMiss(f"cache miss at {node_url}")
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 _PipelineResult(f"pipeline error at {node_url} via relay: status {status}")
except _PipelineCacheMiss:
raise
except Exception as exc:
print(
f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
@@ -767,26 +897,34 @@ 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:
if exc.code == 409:
raise _PipelineCacheMiss(f"cache miss at {node_url}") from exc
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
return _PipelineResult(f"pipeline error at {node_url}: {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 _PipelineResult(f"pipeline error at {node_url}: {exc}")
content_type = resp_headers.get("content-type", "")
if "application/json" in content_type:
try:
data = json.loads(resp_body)
if data.get("error") == "cache_miss":
raise _PipelineCacheMiss(f"cache miss at {node_url}")
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 _PipelineResult(text, int(token_id) if token_id is not None else None)
except json.JSONDecodeError:
return resp_body.decode("utf-8", errors="replace")
return _PipelineResult(resp_body.decode("utf-8", errors="replace"))
# 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 _PipelineResult("")
def _stream_openai_response(self, token_iter, model: str) -> None:
"""Stream tokens from an iterator as SSE chunks."""