tracker rebalancing tweaks

packages/tracker/meshnet_tracker/server.py

@@ -1099,6 +1099,76 @@ def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
     ]
 
 
+def _unassigned_managed_nodes(nodes: list["_NodeEntry"]) -> list["_NodeEntry"]:
+    return [
+        node for node in nodes
+        if node.managed_assignment
+        and (node.shard_start is None or node.shard_end is None)
+    ]
+
+
+def _emit_shard_change_directives(
+    node: "_NodeEntry",
+    model: str,
+    previous_range: tuple[int | None, int | None, str | None],
+    preset: dict,
+) -> None:
+    """Queue DROP/LOAD directives when a managed node's shard assignment changes."""
+    previous_start, previous_end, previous_quantization = previous_range
+    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_range:
+        return
+    if previous_start is not None and previous_end is not None:
+        node.pending_directives.append(
+            _drop_directive(
+                node,
+                model,
+                previous_start,
+                previous_end,
+                previous_quantization or _node_quantization(node, preset),
+            )
+        )
+    node.pending_directives.append(
+        _load_directive(
+            node,
+            model,
+            node.shard_start,
+            node.shard_end,
+            node.quantization or _node_quantization(node, preset),
+        )
+    )
+
+
+def _assign_redundant_managed_nodes(
+    managed_nodes: list["_NodeEntry"],
+    model: str,
+    preset: dict,
+    required_start: int,
+    required_end: int,
+) -> None:
+    """Give newly joined managed nodes their own copy without reshuffling incumbents."""
+    total_layers = required_end - required_start + 1
+    for node in sorted(
+        managed_nodes,
+        key=lambda entry: (
+            -entry.benchmark_tokens_per_sec,
+            -_node_layer_capacity(entry, preset),
+            entry.node_id,
+        ),
+    ):
+        if node.shard_start is not None and node.shard_end is not None:
+            continue
+        capacity = min(_node_layer_capacity(node, preset), total_layers)
+        if capacity <= 0:
+            continue
+        previous_range = (node.shard_start, node.shard_end, node.quantization)
+        quantization = _node_quantization(node, preset)
+        node.quantization = quantization
+        node.shard_start = required_start
+        node.shard_end = min(required_end, required_start + capacity - 1)
+        _emit_shard_change_directives(node, model, previous_range, preset)
+
+
 def _relay_http_request_frames(
     relay_addr: str,
     path: str,
@@ -1439,6 +1509,17 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
     if not managed_nodes:
         return
 
+    coverage = _coverage_map(model_nodes, required_start, required_end)
+    gaps = _coverage_gaps(coverage)
+    unassigned = _unassigned_managed_nodes(managed_nodes)
+    if not gaps and not unassigned:
+        return
+    if not gaps and unassigned:
+        _assign_redundant_managed_nodes(
+            unassigned, model, preset, required_start, required_end,
+        )
+        return
+
     previous_ranges = {
         node.node_id: (node.shard_start, node.shard_end, node.quantization)
         for node in managed_nodes
@@ -1455,10 +1536,8 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
         )
     )
     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)]
+    base_coverage = _coverage_map(base_nodes, required_start, required_end)
+    gaps = _coverage_gaps(base_coverage)
 
     eligible_nodes = [
         node for node in managed_nodes
@@ -1487,28 +1566,11 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
             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,
-                    model,
-                    previous_start,
-                    previous_end,
-                    previous_quantization or _node_quantization(node, preset),
-                )
-            )
-        node.pending_directives.append(
-            _load_directive(
-                node,
-                model,
-                node.shard_start,
-                node.shard_end,
-                node.quantization or _node_quantization(node, preset),
-            )
+        _emit_shard_change_directives(
+            node,
+            model,
+            previous_ranges[node.node_id],
+            preset,
         )
 
 
@@ -1537,6 +1599,17 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
     if total_layers <= 0:
         return
 
+    coverage = _coverage_map(model_nodes, required_start, required_end)
+    gaps = _coverage_gaps(coverage)
+    unassigned = _unassigned_managed_nodes(managed_nodes)
+    if not gaps and not unassigned:
+        return
+    if not gaps and unassigned:
+        _assign_redundant_managed_nodes(
+            unassigned, hf_repo, preset, required_start, required_end,
+        )
+        return
+
     previous_ranges = {
         node.node_id: (node.shard_start, node.shard_end, node.quantization)
         for node in managed_nodes
@@ -1553,10 +1626,8 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
         )
     )
     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)]
+    base_coverage = _coverage_map(base_nodes, required_start, required_end)
+    gaps = _coverage_gaps(base_coverage)
 
     eligible_nodes = [
         node for node in managed_nodes
@@ -1585,28 +1656,11 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
             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),
-            )
+        _emit_shard_change_directives(
+            node,
+            hf_repo,
+            previous_ranges[node.node_id],
+            preset,
         )
 
 
@@ -3257,7 +3311,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 old = server.registry.pop(eid)
                 stale_entries.append((eid, old))
             server.registry[node_id] = entry
-            if entry.managed_assignment:
+            if entry.managed_assignment and not explicit_shard:
                 if entry.hf_repo:
                     _rebalance_hf_model_locked(server, entry.hf_repo)
                 else:
@@ -3358,10 +3412,6 @@ 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"]
-            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