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

[verified] feat: complete Ralph task workstreams

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-12 11:17:03 +03:00
parent 9a1b15c020
commit 377346c301
37 changed files with 5862 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,17 +3,17 @@
from __future__ import annotations
import base64
import http.client
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 typing import Any, Mapping
from .model_backend import (
InsufficientVRAMError,
@@ -22,16 +22,32 @@ from .model_backend import (
Quantization,
TailTokenResult,
TorchModelShard,
_tensor_from_bfloat16_bytes,
validate_quantization,
)
from .seam_telemetry import GenerationTelemetry
from .activation_compression import (
CompressionPolicies,
CompressionPolicy,
compress_activation,
decompress_activation,
)
class _PipelineCacheMiss(Exception):
"""A downstream hop reported 409 cache_miss — head must re-prefill."""
class _RelayRequestUncertainError(ConnectionError):
"""A relay request may have reached the peer but produced no response."""
class _DirectRequestUncertainError(ConnectionError):
"""A direct request may have reached the downstream node but did not finish."""
from .server import (
_WIRE_VERSION,
_compress_body,
_decompress_body,
_parse_shape,
_validate_activation_body,
)
@@ -107,6 +123,7 @@ class _RelayHopClient:
self.relay_addr = relay_addr
self.timeout = timeout
self._ws = None
self._lock = threading.RLock()
def request(
self,
@@ -118,42 +135,116 @@ class _RelayHopClient:
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
if self._ws is None:
self._ws = wsc.connect(
self.relay_addr,
open_timeout=self.timeout,
max_size=ws_max_size(),
compression=None,
)
request_id = f"{time.time_ns():x}"
frame = encode_binary_frame({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
}, body)
self._ws.send(frame)
raw = self._ws.recv(timeout=self.timeout)
if isinstance(raw, (bytes, bytearray)):
resp_header, resp_body = decode_binary_frame(bytes(raw))
status = int(resp_header.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()}
return status, resp_headers, resp_body
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
with self._lock:
if self._ws is None:
self._ws = wsc.connect(
self.relay_addr,
open_timeout=self.timeout,
max_size=ws_max_size(),
compression=None,
)
request_id = headers.get("X-Meshnet-Activation-Id") or uuid.uuid4().hex
frame = encode_binary_frame({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
}, body)
try:
# A send failure is uncertain too: bytes may already have been
# accepted by the kernel or peer before the exception surfaced.
self._ws.send(frame)
raw = self._ws.recv(timeout=self.timeout)
if isinstance(raw, (bytes, bytearray)):
resp_header, resp_body = decode_binary_frame(bytes(raw))
response_id = str(resp_header.get("request_id") or "")
status = int(resp_header.get("status", 503))
resp_headers = {
k.lower(): v for k, v in (resp_header.get("headers") or {}).items()
}
else:
resp = json.loads(raw)
response_id = str(resp.get("request_id") or "")
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()
)
if response_id and response_id != request_id:
raise ValueError("relay response request_id did not match request")
return status, resp_headers, resp_body
except Exception as exc:
self.close()
raise _RelayRequestUncertainError(
"relay connection failed after forwarding request; refusing replay"
) from exc
def close(self) -> None:
if self._ws is not None:
with self._lock:
if self._ws is not None:
try:
self._ws.close()
except Exception:
pass
finally:
self._ws = None
class _DirectHopClient:
"""One serialized HTTP/1.1 connection to one downstream hop.
Generation handlers own these clients, so a cached Route Session reuses a
TCP connection without sharing it between concurrent Route Sessions.
"""
def __init__(self, endpoint: str, timeout: float = 120.0) -> None:
parsed = urllib.parse.urlsplit(endpoint.rstrip("/"))
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
raise ValueError(f"invalid downstream endpoint: {endpoint!r}")
self.timeout = timeout
self._scheme = parsed.scheme
self._host = parsed.hostname
self._port = parsed.port
self._base_path = parsed.path.rstrip("/")
self._connection: http.client.HTTPConnection | None = None
self._lock = threading.RLock()
def _connect(self) -> http.client.HTTPConnection:
connection_type = (
http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection
)
return connection_type(self._host, self._port, timeout=self.timeout)
def request(
self, path: str, body: bytes, headers: Mapping[str, str],
) -> tuple[int, dict[str, str], bytes]:
request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}"
with self._lock:
if self._connection is None:
self._connection = self._connect()
try:
self._ws.close()
except Exception:
pass
finally:
self._ws = None
self._connection.request("POST", request_path, body=body, headers=dict(headers))
response = self._connection.getresponse()
response_body = response.read()
response_headers = {key.lower(): value for key, value in response.headers.items()}
return response.status, response_headers, response_body
except Exception as exc:
self.close()
raise _DirectRequestUncertainError(
"direct connection failed after forwarding request; refusing replay"
) from exc
def close(self) -> None:
with self._lock:
if self._connection is not None:
try:
self._connection.close()
finally:
self._connection = None
def _relay_hop(
@@ -178,18 +269,15 @@ def _relay_hop(
client.close()
# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs).
_COMPRESS_MIN_BYTES = 64 * 1024
_COMPRESSION_POLICIES = CompressionPolicies()
def _maybe_compress_activation(body: bytes) -> tuple[bytes, str | None]:
"""zstd-compress large activation bodies; returns (wire_body, encoding)."""
if len(body) < _COMPRESS_MIN_BYTES:
return body, None
try:
return _compress_body(body, "zstd"), "zstd"
except Exception:
return body, None
def _maybe_compress_activation(
body: bytes, policy: CompressionPolicy | None = None,
) -> tuple[bytes, str | None]:
"""Compatibility wrapper for callers that only need wire body and encoding."""
result = compress_activation(body, policy or _COMPRESSION_POLICIES.for_condition("lan", "prefill"))
return result.body, result.encoding
def _is_cache_miss_body(body: bytes) -> bool:
@@ -285,6 +373,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
"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")),
"telemetry": rec["telemetry"].snapshot(now=now) if rec.get("telemetry") else None,
})
return out
@@ -306,6 +395,10 @@ class _TorchHTTPServer(http.server.HTTPServer):
class _TorchHandler(http.server.BaseHTTPRequestHandler):
# HTTP/1.1 is required for Route Session-owned downstream connections.
# Finite responses below provide Content-Length; streams are chunked.
protocol_version = "HTTP/1.1"
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
pass
@@ -317,6 +410,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
)
def _request_log_suffix(self) -> str:
activation_id = self.headers.get("X-Meshnet-Activation-Id")
if activation_id:
return f" activation_id={activation_id}"
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 ""
@@ -334,6 +430,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"started": time.monotonic(),
"tokens": 0,
"routing_complete": False,
"telemetry": None,
}
def _track_request_progress(
@@ -366,6 +463,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._handle_chat_completions()
else:
self.send_response(404)
self.send_header("Content-Length", "0")
self.end_headers()
def _handle_infer(self) -> None:
@@ -427,7 +525,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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)
raw_body = decompress_activation(body, encoding).body
_validate_activation_body(raw_body, shape, dtype)
if dtype != "bfloat16":
raise ValueError("real model backend requires bfloat16 activation input")
@@ -438,6 +536,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
except (KeyError, ValueError, TypeError):
self.send_response(400)
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.send_header("Content-Length", "0")
self.end_headers()
return
@@ -511,7 +610,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._send_json(200, {"text": result})
return
response_body = _compress_body(result.body, encoding)
route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan")
phase_condition = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
response_compression = compress_activation(
result.body, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition),
)
response_body = response_compression.body
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(response_body)))
@@ -521,8 +625,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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 response_compression.encoding:
self.send_header("X-Meshnet-Encoding", response_compression.encoding)
if result.attention_mask_header:
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
if result.position_ids_header:
@@ -700,6 +804,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
current_text = prompt_text
session_id = str(uuid.uuid4())
telemetry = GenerationTelemetry(session_id)
with server._stats_lock:
current = server._active_requests.get(request_id)
if current is not None:
current["telemetry"] = telemetry
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
@@ -722,6 +831,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
last_token_id: int | None = None
failure_reason: str | None = None
relay_clients: dict[str, _RelayHopClient] = {}
direct_clients: dict[str, _DirectHopClient] = {}
def _prefill_step() -> tuple[str, int | None]:
"""Full-sequence prefill: initial step and cache-miss recovery."""
@@ -734,6 +844,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
payload, remaining_route, backend=backend,
session=session_id, cache_mode="prefill" if use_kv else None,
relay_clients=relay_clients,
direct_clients=direct_clients if use_kv else None,
telemetry=telemetry,
)
for step in range(max_tokens):
@@ -745,6 +857,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
payload, remaining_route, backend=backend,
session=session_id, cache_mode="decode",
relay_clients=relay_clients,
direct_clients=direct_clients,
telemetry=telemetry,
)
except (KVCacheMiss, _PipelineCacheMiss) as miss:
# Evicted/restarted node or head lost its own session:
@@ -758,11 +872,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
else:
token_str, token_id = _prefill_step()
except _PipelineCacheMiss as exc:
print(f" [node] unexpected cache miss on prefill: {exc}", flush=True)
print(f" [node] unexpected cache miss on prefill session={session_id[:8]}: {exc}", flush=True)
failure_reason = f"cache miss on prefill: {exc}"
break
except Exception as exc:
print(f" [node] distributed encode error: {exc}", flush=True)
print(f" [node] distributed encode error session={session_id[:8]}: {exc}", flush=True)
failure_reason = f"distributed encode error: {exc}"
break
# Stop on error responses or EOS.
@@ -789,14 +903,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
tokens=len(generated),
routing_complete=True,
)
telemetry.note_tokens(len(generated))
now = time.monotonic()
if telemetry.report_due:
summary = telemetry.snapshot(now=now)
seams = summary["seams"]
if seams:
latest = seams[-1]
print(
f" [node] seam telemetry session={session_id[:8]} "
f"phase={latest['phase']} hop={latest['hop']} "
f"activations={latest['activations']} "
f"avg_ms={latest['avg_latency_ms']:.2f} "
f"wire_bytes={latest['wire_bytes']} "
f"response_bytes={latest['response_bytes']} "
f"tps={summary['tokens_per_sec']:.2f} "
f"activation_id={latest['last_activation_id']}",
flush=True,
)
telemetry.mark_reported(now=now)
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" [node] generating step={step + 1}/{max_tokens} session={session_id[:8]} "
f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
)
last_gen_log = now
@@ -808,6 +940,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
pass
for relay_client in relay_clients.values():
relay_client.close()
for direct_client in direct_clients.values():
direct_client.close()
if generated:
elapsed = time.monotonic() - gen_started
@@ -815,10 +949,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
tps = token_count / max(elapsed, 1e-6)
_write_progress_line(
progress_line,
f" [node] generation complete tokens={token_count} "
f" [node] generation complete session={session_id[:8]} tokens={token_count} "
f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
final=True,
)
telemetry.close()
result_text = "".join(generated)
# A failure before the first token is an upstream error, not an empty
@@ -921,6 +1056,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
session: str | None = None,
cache_mode: str | None = None,
relay_clients: dict[str, _RelayHopClient] | None = None,
direct_clients: dict[str, _DirectHopClient] | None = None,
telemetry: GenerationTelemetry | None = None,
) -> tuple[str, int | None]:
"""Forward an activation through the downstream route.
@@ -935,10 +1072,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# 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]
tensor = _tensor_from_bfloat16_bytes(
payload.body, payload.shape, active_backend.torch, # type: ignore[union-attr]
).to(active_backend.device)
if hasattr(active_backend, "decode_tail_token"):
tail = active_backend.decode_tail_token(tensor)
return tail.text, tail.token_id
@@ -968,7 +1104,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
+ (f" relay={relay_addr}" if relay_addr else ""),
flush=True,
)
wire_body, wire_encoding = _maybe_compress_activation(current_body)
phase = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
route_condition = "relay" if relay_addr else "lan"
compression = compress_activation(
current_body, _COMPRESSION_POLICIES.for_condition(route_condition, phase),
)
wire_body, wire_encoding = compression.body, compression.encoding
headers: dict[str, str] = {
"Content-Type": "application/octet-stream",
"X-Meshnet-Wire": _WIRE_VERSION,
@@ -979,9 +1120,17 @@ 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-Activation-Id": uuid.uuid4().hex,
"X-Meshnet-Compression-Route": route_condition,
}
if wire_encoding:
headers["X-Meshnet-Encoding"] = wire_encoding
if telemetry is not None:
telemetry.record_compression(
phase=phase, hop=hop_index, node=node_url,
input_bytes=compression.input_bytes, output_bytes=compression.output_bytes,
elapsed_seconds=compression.elapsed_seconds,
)
if cache_mode:
headers["X-Meshnet-Cache"] = cache_mode
past_len = getattr(payload, "past_len", None)
@@ -992,6 +1141,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if current_pos:
headers["X-Meshnet-Position-Ids"] = current_pos
if relay_addr:
connection_reused = bool(
relay_clients is not None and relay_addr in relay_clients
)
seam_started = time.monotonic()
try:
if relay_clients is None:
status, resp_headers, resp_body = _relay_hop(
@@ -1004,44 +1157,100 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
status, resp_headers, resp_body = relay_client.request(
"/forward", wire_body, headers,
)
if telemetry is not None:
telemetry.record_seam(
activation_id=headers["X-Meshnet-Activation-Id"],
phase=cache_mode or "prefill",
hop=hop_index,
node=node_url,
latency_seconds=time.monotonic() - seam_started,
wire_bytes=len(wire_body),
response_bytes=len(resp_body),
connection_reused=connection_reused,
)
if status == 409 and _is_cache_miss_body(resp_body):
raise _PipelineCacheMiss(node_url)
if status >= 400:
detail = _response_error_snippet(resp_body)
print(
f" [node] relay hop {hop_index} returned {status} from {relay_addr}: {detail}",
f" [node] relay hop {hop_index} session={session[:8]} returned "
f"{status} from {relay_addr}: {detail}",
flush=True,
)
return f"pipeline error at {node_url} via relay: status {status}: {detail}", None
except _PipelineCacheMiss:
raise
except _RelayRequestUncertainError as exc:
# The activation may already have mutated the downstream
# KV cache. Do not replay it on a direct connection.
print(
f" [node] relay hop {hop_index} session={session[:8]} outcome is uncertain at "
f"{relay_addr}: {exc}",
flush=True,
)
return f"pipeline relay outcome uncertain at {node_url}: {exc}", None
except Exception as exc:
print(
f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
f" [node] relay hop {hop_index} session={session[:8]} 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=wire_body,
headers=headers,
method="POST",
connection_reused = bool(
direct_clients is not None and node_url in direct_clients
)
seam_started = time.monotonic()
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
detail = _response_error_snippet(body)
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}: {detail}", flush=True)
return f"pipeline error at {node_url}: {exc}: {detail}", None
if direct_clients is None:
direct_client = _DirectHopClient(node_url, timeout=120.0)
try:
status, resp_headers, resp_body = direct_client.request(
"/forward", wire_body, headers,
)
finally:
direct_client.close()
else:
direct_client = direct_clients.setdefault(
node_url, _DirectHopClient(node_url, timeout=120.0),
)
status, resp_headers, resp_body = direct_client.request(
"/forward", wire_body, headers,
)
if telemetry is not None:
telemetry.record_seam(
activation_id=headers["X-Meshnet-Activation-Id"],
phase=cache_mode or "prefill",
hop=hop_index,
node=node_url,
latency_seconds=time.monotonic() - seam_started,
wire_bytes=len(wire_body),
response_bytes=len(resp_body),
connection_reused=connection_reused,
)
if status == 409 and _is_cache_miss_body(resp_body):
raise _PipelineCacheMiss(node_url)
if status >= 400:
detail = _response_error_snippet(resp_body)
print(
f" [node] pipeline hop {hop_index} session={session[:8]} failed at {node_url}: "
f"status {status}: {detail}", flush=True,
)
return f"pipeline error at {node_url}: status {status}: {detail}", None
except _PipelineCacheMiss:
raise
except _DirectRequestUncertainError as exc:
print(
f" [node] pipeline hop {hop_index} session={session[:8]} outcome is uncertain "
f"at {node_url}: {exc}",
flush=True,
)
return f"pipeline direct outcome uncertain at {node_url}: {exc}", None
except Exception as exc:
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
print(
f" [node] pipeline hop {hop_index} session={session[:8]} "
f"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:
@@ -1058,11 +1267,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
current_shape = _parse_shape(shape_header)
try:
current_body = _decompress_body(
decompression = decompress_activation(
resp_body, resp_headers.get("x-meshnet-encoding")
)
current_body = decompression.body
if telemetry is not None:
telemetry.record_compression(
phase=phase, hop=hop_index, node=node_url,
input_bytes=decompression.input_bytes, output_bytes=decompression.output_bytes,
elapsed_seconds=decompression.elapsed_seconds, decompression=True,
)
except ValueError as exc:
print(f" [node] pipeline hop {hop_index} bad response encoding: {exc}", flush=True)
print(
f" [node] pipeline hop {hop_index} session={session[:8]} "
f"bad response encoding: {exc}", flush=True,
)
return f"pipeline error at {node_url}: {exc}", None
current_attn = resp_headers.get("x-meshnet-attn-mask")
current_pos = resp_headers.get("x-meshnet-position-ids")
@@ -1084,14 +1303,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.send_header("Transfer-Encoding", "chunked")
self.end_headers()
def _emit(data: str) -> None:
def _emit(data: str) -> bool:
try:
self.wfile.write(f"data: {data}\n\n".encode())
body = f"data: {data}\n\n".encode()
self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
self.wfile.flush()
return True
except (BrokenPipeError, ConnectionResetError):
pass
return False
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
@@ -1105,7 +1327,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# instead of showing an empty completion.
_emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
try:
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass
@@ -1117,7 +1339,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}))
try:
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass
@@ -1159,14 +1381,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.send_header("Transfer-Encoding", "chunked")
self.end_headers()
def _emit(data: str) -> None:
def _emit(data: str) -> bool:
try:
self.wfile.write(f"data: {data}\n\n".encode())
body = f"data: {data}\n\n".encode()
self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
self.wfile.flush()
except BrokenPipeError:
pass
return True
except (BrokenPipeError, ConnectionResetError):
return False
_emit(json.dumps({
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
@@ -1184,9 +1409,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}))
try:
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
self.wfile.flush()
except BrokenPipeError:
except (BrokenPipeError, ConnectionResetError):
pass
@@ -1255,6 +1480,7 @@ class TorchNodeServer:
debug: bool = False,
max_loaded_shards: int = 1,
force_cpu: bool = False,
recipe_params: Mapping[str, Any] | None = None,
) -> None:
self._host = host
self._requested_port = port
@@ -1266,6 +1492,7 @@ class TorchNodeServer:
quantization,
cache_dir,
force_cpu=force_cpu,
recipe_params=recipe_params,
)
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
@@ -1415,13 +1642,20 @@ def _load_backend(
quantization: str,
cache_dir: Path | None = None,
force_cpu: bool = False,
recipe_params: Mapping[str, Any] | None = None,
) -> 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
model_id,
shard_start,
shard_end,
quant,
cache_dir,
force_cpu=force_cpu,
recipe_params=recipe_params,
)
except MissingModelDependencyError:
raise