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)
+    base_coverage = _coverage_map(base_nodes, required_start, required_end)
-    gaps = _coverage_gaps(coverage)
+    gaps = _coverage_gaps(base_coverage)
-    if not gaps:
-        gaps = [(required_start, required_end)]
 
     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]
+        _emit_shard_change_directives(
-        current_range = (node.shard_start, node.shard_end, node.quantization)
+            node,
-        if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
+            model,
-            continue
+            previous_ranges[node.node_id],
-        if previous_start is not None and previous_end is not None:
+            preset,
-            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),
-            )
         )
 
 
@@ -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)
+    base_coverage = _coverage_map(base_nodes, required_start, required_end)
-    gaps = _coverage_gaps(coverage)
+    gaps = _coverage_gaps(base_coverage)
-    if not gaps:
-        gaps = [(required_start, required_end)]
 
     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]
+        _emit_shard_change_directives(
-        current_range = (node.shard_start, node.shard_end, node.quantization)
+            node,
-        if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
+            hf_repo,
-            continue
+            previous_ranges[node.node_id],
-        if previous_start is not None and previous_end is not None:
+            preset,
-            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),
-            )
         )
 
 
@@ -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

tests/test_tracker_routing.py

@@ -1114,12 +1114,14 @@ def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap():
              "benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
         )
 
-        route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
+        net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
         widths = {
             node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
-            for node in route_resp["nodes"]
+            for node in net["nodes"]
         }
-        assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"]
+        assert widths["http://127.0.0.1:9011"] == 12
+        assert widths["http://127.0.0.1:9012"] == 12
+        assert net["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
     finally:
         tracker.stop()
 
@@ -1147,7 +1149,8 @@ def test_tracker_registration_directive_is_not_replayed_on_heartbeat():
         tracker.stop()
 
 
-def test_tracker_reassignment_emits_drop_before_load():
+def test_tracker_pool_join_adds_redundant_copy_without_splitting_incumbent():
+    """A second managed node with capacity for the full model keeps the first copy intact."""
     tracker = TrackerServer(model_presets={
         "tiny-model": {
             "total_layers": 4,
@@ -1156,21 +1159,69 @@ def test_tracker_reassignment_emits_drop_before_load():
     })
     tracker_port = tracker.start()
     try:
-        slow = _post_json(
+        first = _post_json(
             f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
             {"endpoint": "http://127.0.0.1:9015", "model": "tiny-model",
              "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
              "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
         )
-        _post_json(
+        second = _post_json(
             f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
             {"endpoint": "http://127.0.0.1:9016", "model": "tiny-model",
              "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
              "benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
         )
 
-        hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{slow['node_id']}/heartbeat", {})
+        net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
-        assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"]
+        widths = {
+            node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
+            for node in net["nodes"]
+        }
+        assert widths["http://127.0.0.1:9015"] == 4
+        assert widths["http://127.0.0.1:9016"] == 4
+        assert net["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
+
+        hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{first['node_id']}/heartbeat", {})
+        assert hb.get("directives", []) == []
+    finally:
+        tracker.stop()
+
+
+def test_tracker_explicit_full_copy_join_preserves_existing_serving_node():
+    """--model style joins with explicit shards add redundancy instead of reshuffling."""
+    tracker = TrackerServer(heartbeat_timeout=10.0)
+    tracker_port = tracker.start()
+    try:
+        base_reg = {
+            "model": "Qwen2.5-0.5B-Instruct",
+            "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
+            "num_layers": 24,
+            "shard_start": 0,
+            "shard_end": 23,
+            "managed_assignment": True,
+            "vram_bytes": 2_000_000_000,
+            "hardware_profile": {},
+            "score": 1.0,
+        }
+        first = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {**base_reg, "endpoint": "http://127.0.0.1:9201"},
+        )
+        second = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {**base_reg, "endpoint": "http://127.0.0.1:9202"},
+        )
+
+        coverage = _get_json(
+            f"http://127.0.0.1:{tracker_port}/v1/network/map"
+        )
+        assert coverage["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
+
+        hb = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/{first['node_id']}/heartbeat", {}
+        )
+        assert hb.get("directives", []) == []
+        assert second["node_id"] in tracker._registry
     finally:
         tracker.stop()
 
@@ -1208,22 +1259,20 @@ def test_tracker_faster_node_receives_wider_range_when_capacity_ties():
         _post_json(
             f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
             {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
-             "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
+             "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
              "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
         )
         _post_json(
             f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
             {"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
-             "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
+             "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
              "benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0},
         )
 
-        route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
+        net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
-        widths = {
+        heads = [node for node in net["nodes"] if node["shard_start"] == 0]
-            node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
+        assert len(heads) == 1
-            for node in route_resp["nodes"]
+        assert heads[0]["endpoint"] == "http://127.0.0.1:9002"
-        }
-        assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"]
     finally:
         tracker.stop()