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