diff --git a/billing.sqlite b/billing.sqlite index 2a00b87..3452dcc 100644 Binary files a/billing.sqlite and b/billing.sqlite differ diff --git a/meshnet_registry.sqlite3 b/meshnet_registry.sqlite3 index 8060b82..4cfeef0 100644 Binary files a/meshnet_registry.sqlite3 and b/meshnet_registry.sqlite3 differ diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index da11d72..6e61b25 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/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 diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index dbc7f9b..ab3fa0a 100644 --- a/tests/test_tracker_routing.py +++ b/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", {}) - assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"] + 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 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") - widths = { - node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 - for node in route_resp["nodes"] - } - assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"] + net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map") + heads = [node for node in net["nodes"] if node["shard_start"] == 0] + assert len(heads) == 1 + assert heads[0]["endpoint"] == "http://127.0.0.1:9002" finally: tracker.stop()