tracker rebalance events

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