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

ram pool map

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 18:35:36 +03:00
parent 1299a6bb1c
commit e9a094b620
6 changed files with 253 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -75,20 +75,39 @@ class _TorchHTTPServer(http.server.HTTPServer):
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.total_requests: int = 0
self.failed_requests: int = 0
self.queue_depth: int = 0
self._stats_lock = threading.Lock()
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
@@ -100,7 +119,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self._handle_forward()
elif self.path == "/v1/infer":
self._handle_infer()
elif self.path == "/v1/chat/completions" and server.tracker_mode:
elif self.path == "/v1/chat/completions" and server.chat_enabled():
self._handle_chat_completions()
else:
self.send_response(404)
@@ -284,22 +303,26 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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 256)
temperature = float(body.get("temperature") or 1.0)
top_p = float(body.get("top_p") or 1.0)
# 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 server.backend.is_head and server.backend.is_tail:
if backend.is_head and backend.is_tail:
try:
if stream:
self._stream_openai_response(
server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
model_name,
)
else:
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
self._send_openai_response(text, model_name, False, messages)
text = backend.generate_text(messages, max_tokens, temperature, top_p)
self._send_openai_response(text, model_name, False, messages, backend=backend)
except Exception as exc:
self._record_failed_request()
self._send_json(500, {"error": f"generation failed: {exc}"})
@@ -309,7 +332,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# 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.
remaining_route = self._get_remaining_route(model_name)
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}",
@@ -318,11 +341,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if not remaining_route:
self._send_openai_response(
"error: no downstream route — check tracker connectivity",
model_name, False, messages,
model_name, False, messages, backend=backend,
)
return
backend = server.backend
# Format with chat template so the model knows it's in assistant mode.
try:
if hasattr(backend.tokenizer, "apply_chat_template"):
@@ -352,7 +374,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
except Exception as exc:
print(f" [node] distributed encode error: {exc}", flush=True)
break
token_str = self._run_downstream_pipeline(payload, remaining_route)
token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
if not token_str:
break
# Stop on error responses or EOS.
@@ -369,9 +391,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
if stream_emit is not None:
stream_emit(None)
return
self._send_openai_response(result_text, model_name, stream, messages)
self._send_openai_response(result_text, model_name, stream, messages, backend=backend)
def _get_remaining_route(self, model: str) -> list[dict]:
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.
@@ -404,9 +426,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
if server.tracker_url is None:
return []
route_model = getattr(server.backend, "model_id", None) or model
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:
@@ -433,18 +456,19 @@ 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]) -> str:
def _run_downstream_pipeline(self, payload: object, route: list[dict], *, backend: TorchModelShard | None = None) -> str:
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 server.backend.is_tail:
if active_backend.is_tail:
try:
tensor = server.backend.torch.frombuffer(
tensor = active_backend.torch.frombuffer(
bytearray(payload.body), # type: ignore[union-attr]
dtype=server.backend.torch.bfloat16,
).reshape(payload.shape).to(server.backend.device) # type: ignore[union-attr]
return server.backend.decode_tail(tensor)
dtype=active_backend.torch.bfloat16,
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
return active_backend.decode_tail(tensor)
except Exception as exc:
return f"decode error: {exc}"
return "no downstream route available for non-tail shard"
@@ -591,11 +615,13 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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(self.server.backend, messages or [], text) # type: ignore[attr-defined]
usage = _usage_for_response(active_backend, messages or [], text)
self._send_json(200, {
"id": chunk_id,
"object": "chat.completion",
@@ -706,9 +732,11 @@ class TorchNodeServer:
route_timeout: float = 30.0,
cache_dir: Path | None = None,
debug: bool = False,
max_loaded_shards: int = 1,
) -> 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,
@@ -716,6 +744,7 @@ class TorchNodeServer:
quantization,
cache_dir,
)
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
@@ -754,41 +783,64 @@ class TorchNodeServer:
def queue_depth(self) -> int:
return self._server.queue_depth if self._server is not None else 0
@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 LOAD_SHARD directives by hot-swapping the loaded backend."""
"""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,
)
if load_directive is None:
directive = add_directive or load_directive
if directive is None:
return None
shard_start = int(load_directive["shard_start"])
shard_end = int(load_directive["shard_end"])
quantization = str(load_directive.get("quantization") or self._backend.quantization)
model_id = str(load_directive.get("model") or self._backend.model_id)
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 reassigned shard: {model_id} layers {shard_start}-{shard_end}",
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._backend = new_backend
self._tracker_mode = shard_start == 0
if self._server is not None:
self._server.backend = new_backend
self._server.tracker_mode = self._tracker_mode
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 reassigned shard: {model_id} layers {shard_start}-{shard_end}",
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": self._tracker_mode,
"tracker_mode": shard_start == 0,
}
def start(self) -> int:
@@ -802,7 +854,9 @@ class TorchNodeServer:
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()