diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 0c4cedb..e2c5c0c 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -123,6 +123,26 @@ def _start_heartbeat( except Exception: return False + def _apply_directives(directives: list[dict]) -> None: + if not directives: + return + if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"): + print(f" [node] tracker directives received: {directives}", flush=True) + return + try: + applied = node_ref.apply_tracker_directives(directives) + except Exception as exc: + print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True) + return + if applied: + 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 + register_payload["shard_start"] = applied["shard_start"] + register_payload["shard_end"] = applied["shard_end"] + register_payload["quantization"] = applied.get("quantization", register_payload.get("quantization")) + register_payload["tracker_mode"] = bool(applied.get("tracker_mode", False)) + def _loop() -> None: nonlocal node_id hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat" @@ -150,6 +170,7 @@ def _start_heartbeat( try: resp = _post_json(hb_url, _get_stats()) + _apply_directives(resp.get("directives", [])) new_asgn = resp.get("new_assignment") if new_asgn: print( @@ -282,6 +303,7 @@ def run_startup( print(f" {probationary_line}", flush=True) if model_id: # treat "" the same as None — no explicit model given + user_pinned_shard = shard_start is not None or shard_end is not None # Auto-detect shard range from model config if not explicitly provided if shard_start is None or shard_end is None: detected = _detect_num_layers(model_id) @@ -355,6 +377,7 @@ def run_startup( "quantization": quantization, "score": 1.0, "tracker_mode": (shard_start == 0), + "managed_assignment": not user_pinned_shard, **relay_fields, } tracker_node_id: str | None = None @@ -442,6 +465,7 @@ def run_startup( "quantization": quantization, "score": 1.0, "tracker_mode": (assigned_shard_start == 0), + "managed_assignment": True, **relay_fields, } tracker_node_id = None diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 16e7ffd..1f64c2c 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -729,6 +729,40 @@ class TorchNodeServer: def queue_depth(self) -> int: return self._server.queue_depth if self._server is not None else 0 + def apply_tracker_directives(self, directives: list[dict]) -> dict | None: + """Apply tracker LOAD_SHARD directives by hot-swapping the loaded backend.""" + load_directive = next( + (directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"), + None, + ) + if load_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) + print( + f" [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}", + flush=True, + ) + 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 + print( + f" [node] loaded reassigned shard: {model_id} layers {shard_start}-{shard_end}", + flush=True, + ) + return { + "model": model_id, + "shard_start": shard_start, + "shard_end": shard_end, + "quantization": quantization, + "tracker_mode": self._tracker_mode, + } + def start(self) -> int: if self._server is not None: raise RuntimeError("TorchNodeServer is already running") diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 6d3a81b..be02866 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -685,9 +685,109 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: ) +def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict: + total_layers = max(node.num_layers or 0 for node in nodes) + return { + "layers_start": 0, + "layers_end": total_layers - 1, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, + } + + +def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None: + model_nodes = [ + node for node in server.registry.values() + if node.hf_repo == hf_repo + and node.num_layers is not None + ] + managed_nodes = [node for node in model_nodes if node.managed_assignment] + if not model_nodes or not managed_nodes: + return + + preset = _hf_rebalance_preset(model_nodes) + required_start, required_end = _preset_layer_bounds(preset) + total_layers = required_end - required_start + 1 + if total_layers <= 0: + return + + previous_ranges = { + node.node_id: (node.shard_start, node.shard_end, node.quantization) + for node in managed_nodes + } + for node in managed_nodes: + node.shard_start = None + node.shard_end = None + + managed_nodes.sort( + key=lambda node: ( + -node.benchmark_tokens_per_sec, + -_node_layer_capacity(node, preset), + node.node_id, + ) + ) + base_nodes = [node for node in model_nodes if not node.managed_assignment] + coverage = _coverage_map(base_nodes, required_start, required_end) + gaps = _coverage_gaps(coverage) + if not gaps: + gaps = [(required_start, required_end)] + + eligible_nodes = [ + node for node in managed_nodes + if _node_layer_capacity(node, preset) > 0 + ] + node_index = 0 + for gap_start, gap_end in gaps: + cursor = gap_start + while cursor <= gap_end and node_index < len(eligible_nodes): + node = eligible_nodes[node_index] + remaining_layers = gap_end - cursor + 1 + remaining_nodes_after = len(eligible_nodes) - node_index - 1 + capacity = min( + _node_layer_capacity(node, preset), + total_layers, + max(1, remaining_layers - remaining_nodes_after), + ) + if capacity <= 0: + node_index += 1 + continue + quantization = _node_quantization(node, preset) + node.quantization = quantization + node.shard_start = cursor + node.shard_end = min(gap_end, cursor + capacity - 1) + cursor = node.shard_end + 1 + node_index += 1 + + for node in managed_nodes: + previous_start, previous_end, previous_quantization = previous_ranges[node.node_id] + current_range = (node.shard_start, node.shard_end, node.quantization) + if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]: + continue + if previous_start is not None and previous_end is not None: + node.pending_directives.append( + _drop_directive( + node, + hf_repo, + previous_start, + previous_end, + previous_quantization or _node_quantization(node, preset), + ) + ) + node.pending_directives.append( + _load_directive( + node, + hf_repo, + node.shard_start, + node.shard_end, + node.quantization or _node_quantization(node, preset), + ) + ) + + def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: for model in list(server.model_presets): _rebalance_model_locked(server, model) + for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}): + _rebalance_hf_model_locked(server, hf_repo) def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: @@ -1314,6 +1414,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): return tracker_mode = bool(body.get("tracker_mode", False)) + managed_assignment = bool(body.get("managed_assignment", False)) hf_repo = body.get("hf_repo") if hf_repo is not None and not isinstance(hf_repo, str): self._send_json(400, {"error": "hf_repo must be a string"}) @@ -1353,7 +1454,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): quantizations=quantizations, benchmark_tokens_per_sec=benchmark_tokens_per_sec, quantization=quantization, - managed_assignment=not explicit_shard, + managed_assignment=managed_assignment or not explicit_shard, tracker_mode=tracker_mode, hf_repo=hf_repo, num_layers=num_layers, @@ -1377,7 +1478,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ) server.registry[node_id] = entry if entry.managed_assignment: - _rebalance_model_locked(server, model) + if entry.hf_repo: + _rebalance_hf_model_locked(server, entry.hf_repo) + else: + _rebalance_model_locked(server, model) assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None if assignment_directive is not None: entry.pending_directives.clear() @@ -1441,7 +1545,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): entry.uptime_seconds = float(body["uptime_seconds"]) if "status" in body and body["status"] in ("ready", "loading"): entry.status = body["status"] - _rebalance_model_locked(server, entry.model or "stub-model") + if entry.hf_repo: + _rebalance_hf_model_locked(server, entry.hf_repo) + else: + _rebalance_model_locked(server, entry.model or "stub-model") directives = list(entry.pending_directives) entry.pending_directives.clear() new_assignment = entry.pending_new_assignment diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index 9c02e61..c04030f 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -587,6 +587,42 @@ def test_tracker_rebalances_after_middle_range_node_timeout(): tracker.stop() +def test_tracker_rebalances_managed_hf_node_after_peer_timeout(): + """HF nodes auto-assigned by the tracker receive LOAD_SHARD after a peer dies.""" + tracker = TrackerServer(heartbeat_timeout=0.15, rebalance_interval=10.0) + tracker_port = tracker.start() + try: + managed = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 21, "managed_assignment": True, + "vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0}, + ) + expired = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 22, "shard_end": 23, + "vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0}, + ) + + time.sleep(0.10) + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {}) + time.sleep(0.10) + + hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {}) + + assert expired["node_id"] not in tracker._registry + load_directives = [d for d in hb.get("directives", []) if d["action"] == "LOAD_SHARD"] + assert load_directives + assert load_directives[-1]["model"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert load_directives[-1]["start_layer"] == 0 + assert load_directives[-1]["end_layer"] == 23 + finally: + tracker.stop() + + def test_tracker_route_error_no_nodes(): """Tracker returns 503 with clear error when the registry is empty.""" tracker = TrackerServer() @@ -1315,3 +1351,47 @@ def test_route_timeout_config_is_exposed_on_server(): node = TorchNodeServer(backend=_MinimalBackend(), route_timeout=45.0) assert node.route_timeout == 45.0 + + +def test_torch_node_applies_tracker_load_shard_directive(monkeypatch): + from meshnet_node import torch_server + from meshnet_node.torch_server import TorchNodeServer + + class _MinimalBackend: + def __init__(self, model_id="Qwen/Qwen2.5-0.5B-Instruct", shard_start=0, shard_end=21, quantization="bfloat16"): + self.model_id = model_id + self.shard_start = shard_start + self.shard_end = shard_end + self.quantization = quantization + self.total_layers = 24 + self.is_head = shard_start == 0 + self.is_tail = shard_end == 23 + + def generate_text(self, *a, **kw): return "" + def count_prompt_tokens(self, *a): return 0 + def count_text_tokens(self, *a): return 0 + + loaded = [] + + def fake_load_backend(model_id, shard_start, shard_end, quantization): + loaded.append((model_id, shard_start, shard_end, quantization)) + return _MinimalBackend(model_id, shard_start, shard_end, quantization) + + monkeypatch.setattr(torch_server, "_load_backend", fake_load_backend) + node = TorchNodeServer(backend=_MinimalBackend()) + + applied = node.apply_tracker_directives([ + {"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 21}, + {"action": "LOAD_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 23, + "quantization": "bfloat16"}, + ]) + + assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")] + assert applied == { + "model": "Qwen/Qwen2.5-0.5B-Instruct", + "shard_start": 0, + "shard_end": 23, + "quantization": "bfloat16", + "tracker_mode": True, + } + assert node.backend.shard_end == 23