tracker rebalance events

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

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")

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

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