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

BIN
billing.sqlite
View File

Binary file not shown.

BIN
meshnet_registry.sqlite3
View File

Binary file not shown.

35
packages/node/meshnet_node/startup.py
View File

@@ -372,6 +372,27 @@ def _start_heartbeat(
except Exception:
return False
def _register_additional_assignment(applied: dict) -> None:
model_id = str(applied.get("model") or register_payload.get("hf_repo") or register_payload.get("model"))
extra_payload = {
**register_payload,
"model": model_id.split("/")[-1],
"hf_repo": model_id if "/" in model_id else register_payload.get("hf_repo"),
"shard_start": applied["shard_start"],
"shard_end": applied["shard_end"],
"quantization": applied.get("quantization", register_payload.get("quantization")),
"tracker_mode": bool(applied.get("tracker_mode", False)),
"managed_assignment": True,
}
try:
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", extra_payload)
print(
f" [node] registered additional model — node ID: {reg_resp.get('node_id')}",
flush=True,
)
except Exception as exc:
print(f" [node] WARNING: additional model registration failed: {exc}", flush=True)
def _apply_directives(directives: list[dict]) -> None:
if not directives:
return
@@ -384,6 +405,9 @@ def _start_heartbeat(
print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
return
if applied:
if applied.get("action") == "ADD_SHARD":
_register_additional_assignment(applied)
return
model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
register_payload["model"] = str(model_id).split("/")[-1]
register_payload["hf_repo"] = model_id
@@ -423,11 +447,12 @@ def _start_heartbeat(
new_asgn = resp.get("new_assignment")
if new_asgn:
print(
f" [node] tracker reassignment received: "
f"model={new_asgn.get('model')!r} "
f" [node] tracker assignment received: "
f"action={new_asgn.get('action')!r} model={new_asgn.get('model')!r} "
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
flush=True,
)
_apply_directives([new_asgn])
except urllib.error.HTTPError as exc:
if exc.code == 404:
# Node was purged (e.g. long gap before restart noticed) — re-register now.
@@ -668,6 +693,7 @@ def run_startup(
route_timeout=route_timeout,
cache_dir=cache_dir,
debug=debug,
max_loaded_shards=max_loaded_shards,
)
_node_start_time = time.monotonic()
actual_port = node.start()
@@ -810,6 +836,7 @@ def run_startup(
route_timeout=route_timeout,
cache_dir=cache_dir,
debug=debug,
max_loaded_shards=max_loaded_shards,
)
_node_start_time = time.monotonic()
actual_port = node.start()
@@ -883,13 +910,13 @@ def run_startup(
)
return node
if model != "stub-model":
if not assigned_hf_repo and model is None:
raise RuntimeError(
"Tracker did not assign a model. Join a network that already serves one, "
"or start with --model <HF_REPO>."
)
# 3b. Stub preset path (tests / explicit stub-model only)
# 3b. Stub preset path (tests / explicit stub-model) or named preset models.
print("Querying tracker for shard assignment...", flush=True)
assign_qs = urllib.parse.urlencode({
"model": model or "stub-model",

116
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] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
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
if self._server is not None:
self._server.backend = new_backend
self._server.tracker_mode = self._tracker_mode
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()

41
packages/tracker/meshnet_tracker/server.py
View File

@@ -944,6 +944,20 @@ def _pool_summary(nodes: list[_NodeEntry], preset: dict | None = None) -> dict:
}
def _repo_demand_rpm(
server: "_TrackerHTTPServer",
repo: str,
local_rpms: dict[str, dict],
) -> float:
"""Look up recent request rate by hf_repo or preset alias."""
if repo in local_rpms:
return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0)
resolved_name, _preset = _resolve_model_preset(server.model_presets, repo)
if resolved_name and resolved_name in local_rpms:
return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0)
return 0.0
def _preset_for_node(server: "_TrackerHTTPServer", node: _NodeEntry) -> dict | None:
if node.hf_repo:
_resolved, preset = _resolve_model_preset(server.model_presets, node.hf_repo)
@@ -1392,20 +1406,6 @@ def _assign_redundant_managed_nodes(
_emit_shard_change_directives(node, model, previous_range, preset)
def _repo_demand_rpm(
server: "_TrackerHTTPServer",
repo: str,
local_rpms: dict[str, dict],
) -> float:
"""Look up recent request rate by hf_repo or preset alias."""
if repo in local_rpms:
return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0)
resolved_name, _preset = _resolve_model_preset(server.model_presets, repo)
if resolved_name and resolved_name in local_rpms:
return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0)
return 0.0
def _relay_http_request_frames(
relay_addr: str,
path: str,
@@ -2587,6 +2587,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
with server.lock:
self._purge_expired_nodes()
nodes = list(server.registry.values())
memory_pool = _memory_pool_map(server)
def capacity_for(node: _NodeEntry) -> dict:
preset = None
@@ -2611,6 +2612,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._send_json(200, {
"relay_url": server.relay_url,
"pool": _pool_summary(nodes),
"memory_pool": memory_pool,
"recommended_models": [
{
"id": name,
@@ -3538,10 +3540,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
)
with server.lock:
self._purge_expired_nodes()
# Dedup: if this endpoint is already registered, remove the old entry first.
# Dedup: replace the same node id or the same endpoint+model assignment.
endpoint_key = entry.endpoint.rstrip("/")
model_key = entry.hf_repo or entry.model
stale_ids = [
eid for eid, e in server.registry.items()
if e.endpoint == entry.endpoint.rstrip("/")
if eid == node_id
or (
e.endpoint.rstrip("/") == endpoint_key
and (e.hf_repo or e.model) == model_key
)
]
stale_entries: list[tuple[str, _NodeEntry]] = []
for eid in stale_ids:
@@ -5436,6 +5444,7 @@ class TrackerServer:
with self._lock:
_purge_expired_nodes_locked(server)
_rebalance_all_locked(server)
_scale_demanded_models_locked(server)
def _push_to_peers(self, path: str, body: bytes) -> bool:
"""POST a hive-signed payload to every cluster peer; True if all succeeded."""

111
tests/test_tracker_routing.py
View File

@@ -13,7 +13,13 @@ from meshnet_gateway.server import GatewayServer, _banned_route_wallet
from meshnet_node.server import StubNodeServer
from meshnet_contracts import LocalSolanaContracts
from meshnet_tracker.auth import sign_hive_request
from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error
from meshnet_tracker.server import (
TrackerServer,
_NodeEntry,
_memory_pool_map,
_registration_ban_error,
_scale_demanded_models_locked,
)
_TEST_HIVE_SECRET = "test-hive-secret"
@@ -2340,6 +2346,7 @@ def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
assert applied == {
"action": "LOAD_SHARD",
"model": "Qwen/Qwen2.5-0.5B-Instruct",
"shard_start": 0,
"shard_end": 23,
@@ -2456,3 +2463,105 @@ def test_shard_heal_cycle_surviving_node_covers_dead_peers_gap(monkeypatch):
)
finally:
tracker.stop()
def test_network_map_exposes_memory_pool():
tracker = TrackerServer(model_presets={
"tiny-model": {
"total_layers": 8,
"bytes_per_layer": {"bfloat16": 1_000},
"hf_repo": "org/TinyModel",
},
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9050", "model": "tiny-model",
"hf_repo": "org/TinyModel", "num_layers": 8,
"shard_start": 0, "shard_end": 7, "max_loaded_shards": 2,
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
pool = net["memory_pool"]
assert pool["total_spare_slots"] == 1
assert pool["hosts"][0]["loaded_slots"] == 1
assert pool["hosts"][0]["max_loaded_shards"] == 2
assert pool["hosts"][0]["memory_spare_bytes"] > 0
finally:
tracker.stop()
def test_same_endpoint_can_register_multiple_models():
tracker = TrackerServer()
tracker_port = tracker.start()
try:
base = {
"endpoint": "http://127.0.0.1:9055",
"num_layers": 24,
"shard_start": 0,
"shard_end": 23,
"hardware_profile": {},
"score": 1.0,
"max_loaded_shards": 2,
"vram_bytes": 50_000_000,
"ram_bytes": 50_000_000,
}
first = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base, "model": "Qwen2.5-0.5B-Instruct", "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct"},
)
second = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base, "model": "OtherModel", "hf_repo": "org/OtherModel"},
)
assert first["node_id"] != second["node_id"]
assert len(tracker._registry) == 2
finally:
tracker.stop()
def test_scale_demanded_models_queues_add_shard_on_spare_host():
tracker = TrackerServer(model_presets={
"model-a": {
"total_layers": 4,
"bytes_per_layer": {"bfloat16": 1_000},
"hf_repo": "org/ModelA",
},
"model-b": {
"total_layers": 4,
"bytes_per_layer": {"bfloat16": 1_000},
"hf_repo": "org/ModelB",
},
})
tracker_port = tracker.start()
try:
reg_b = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9061", "model": "model-b",
"hf_repo": "org/ModelB", "num_layers": 4,
"shard_start": 0, "shard_end": 3, "max_loaded_shards": 2,
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9060", "model": "model-a",
"hf_repo": "org/ModelA", "num_layers": 4,
"shard_start": 0, "shard_end": 3, "max_loaded_shards": 1,
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
assert tracker._stats is not None
for _ in range(400):
tracker._stats.record_request("org/ModelA")
with tracker._lock:
_scale_demanded_models_locked(tracker._server) # type: ignore[arg-type]
node_b = tracker._registry[reg_b["node_id"]]
assignment = node_b.pending_new_assignment
assert assignment is not None
assert assignment["action"] == "ADD_SHARD"
assert assignment["model"] == "org/ModelA"
finally:
tracker.stop()