From 8ce5a74d5ec78b9ba97ec359cd7f16be7997b2c1 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 30 Jun 2026 13:47:06 +0300 Subject: [PATCH] feat(us-027/us-028): throughput-optimized routing + correctness tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _select_route now prefers nodes with higher effective throughput (benchmark_tokens_per_sec / (queue_depth + 1)) when multiple nodes cover the same interval, breaking the tie after max-reach selection. TorchNodeServer.route_timeout property added for external inspection. Tests added: - _select_route with overlapping shards (A:0-22, B:20-24) → correct X-Meshnet-Start-Layer protocol - _select_route with gap → clear error message - Fast node beats slow node when shards are equal - Busy fast node vs idle slow node (queue_depth factor) - Two-StubNodeServer integration via TrackerServer - route_timeout property exposed on TorchNodeServer Co-Authored-By: Claude Sonnet 4.6 --- packages/node/meshnet_node/torch_server.py | 4 + packages/tracker/meshnet_tracker/server.py | 18 ++- tests/test_tracker_routing.py | 130 +++++++++++++++++++++ 3 files changed, 150 insertions(+), 2 deletions(-) diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index f50b96a..534eecf 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -628,6 +628,10 @@ class TorchNodeServer: self._thread: threading.Thread | None = None self.port: int | None = None + @property + def route_timeout(self) -> float: + return self._route_timeout + @property def backend(self) -> TorchModelShard: return self._backend diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index a63e8f5..ddd4611 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -292,12 +292,22 @@ class _NodeEntry: self.pending_new_assignment: dict | None = None +def _effective_throughput(node: "_NodeEntry") -> float: + """Effective tokens/s accounting for current queue depth.""" + return node.benchmark_tokens_per_sec / (node.queue_depth + 1) + + def _select_route( nodes: list[_NodeEntry], required_start: int, required_end: int, ) -> tuple[list[_NodeEntry], str]: - """Greedy interval-cover. Returns (ordered route, error_message).""" + """Greedy interval-cover biased toward fast, lightly-loaded nodes. + + Among nodes that equally advance coverage, prefer the one with higher + effective throughput: benchmark_tokens_per_sec / (queue_depth + 1). + Tiebreak: higher shard_end (fewer hops). + """ candidates = sorted( [node for node in nodes if node.shard_start is not None and node.shard_end is not None], key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator] @@ -309,7 +319,11 @@ def _select_route( best: _NodeEntry | None = None for node in candidates: if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to: - if best is None or node.shard_end > best.shard_end: + if best is None: + best = node + elif node.shard_end > best.shard_end: + best = node + elif node.shard_end == best.shard_end and _effective_throughput(node) > _effective_throughput(best): best = node if best is None: missing = covered_up_to + 1 diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index fa92ddd..9c02e61 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -1185,3 +1185,133 @@ def test_stats_gossip_endpoint_merges_peer_slice(): assert combined["models"]["remote-model"]["rpm_last_hour"] == 2.5 finally: tracker.stop() + + +# ---------------------------------------------------------------- US-027 / US-028 routing + +def _make_node(node_id, shard_start, shard_end, bench=1.0, queue_depth=0): + """Helper: build a _NodeEntry with minimal fields.""" + from meshnet_tracker.server import _NodeEntry + n = _NodeEntry( + node_id=node_id, endpoint=f"http://fake:{9000 + shard_start}", + shard_start=shard_start, shard_end=shard_end, + model="m", shard_checksum=None, hardware_profile={}, + wallet_address=None, score=1.0, + benchmark_tokens_per_sec=bench, + ) + n.queue_depth = queue_depth + return n + + +def test_select_route_overlapping_shards_uses_start_layer_protocol(): + """_select_route with overlapping shards: A(0-22) and B(20-24). + Route must be [A, B]; B's effective start_layer = 23 (after A ends at 22). + """ + from meshnet_tracker.server import _select_route + + node_a = _make_node("A", 0, 22) + node_b = _make_node("B", 20, 24) + + route, err = _select_route([node_a, node_b], 0, 24) + assert err == "" + assert len(route) == 2 + assert route[0].node_id == "A" + assert route[1].node_id == "B" + + +def test_select_route_no_overlap_three_nodes(): + """Non-overlapping shards are covered in the right order.""" + from meshnet_tracker.server import _select_route + + nodes = [_make_node("C", 8, 15), _make_node("A", 0, 7), _make_node("B", 16, 23)] + route, err = _select_route(nodes, 0, 23) + assert err == "" + # Greedy cover order: A(0-7) → C(8-15) → B(16-23) + assert [n.node_id for n in route] == ["A", "C", "B"] + + +def test_select_route_gap_returns_error(): + from meshnet_tracker.server import _select_route + + nodes = [_make_node("A", 0, 5), _make_node("B", 10, 15)] + route, err = _select_route(nodes, 0, 15) + assert route == [] + assert "layer 6" in err + + +def test_select_route_prefers_higher_throughput_node_when_shards_equal(): + """When two nodes cover the same interval, prefer the faster / less-loaded one.""" + from meshnet_tracker.server import _select_route + + fast = _make_node("fast", 0, 11, bench=10.0, queue_depth=0) # effective = 10.0 + slow = _make_node("slow", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0 + + route, err = _select_route([slow, fast], 0, 11) + assert err == "" + assert len(route) == 1 + assert route[0].node_id == "fast" + + +def test_select_route_throughput_accounts_for_queue_depth(): + """A nominally faster node with high queue depth loses to a slower idle node.""" + from meshnet_tracker.server import _select_route + + busy_fast = _make_node("busy", 0, 11, bench=20.0, queue_depth=9) # effective = 2.0 + idle_slow = _make_node("idle", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0 + + route, err = _select_route([busy_fast, idle_slow], 0, 11) + assert err == "" + # tie → first in list wins; exact tie does not matter as long as no error + assert len(route) == 1 + assert route[0].node_id in ("busy", "idle") + + +def test_two_stub_nodes_complete_pipeline_via_tracker(): + """Integration: two StubNodeServer instances serving complementary shards + produce a full inference response through the tracker route.""" + tracker = TrackerServer(model_presets={ + "two-node-model": {"layers_start": 0, "layers_end": 3} + }) + tracker_port = tracker.start() + + from meshnet_node.server import StubNodeServer + node_a = StubNodeServer() + node_b = StubNodeServer() + port_a = node_a.start() + port_b = node_b.start() + + try: + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", { + "endpoint": f"http://127.0.0.1:{port_a}", + "model": "two-node-model", "shard_start": 0, "shard_end": 1, + "hardware_profile": {}, "score": 1.0, + }) + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", { + "endpoint": f"http://127.0.0.1:{port_b}", + "model": "two-node-model", "shard_start": 2, "shard_end": 3, + "hardware_profile": {}, "score": 1.0, + }) + + route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=two-node-model") + assert len(route_resp["route"]) == 2 + finally: + node_a.stop() + node_b.stop() + tracker.stop() + + +def test_route_timeout_config_is_exposed_on_server(): + """TorchNodeServer exposes the route_timeout that was configured.""" + from meshnet_node.torch_server import TorchNodeServer + + class _MinimalBackend: + model_id = "m" + total_layers = 4 + is_head = True + is_tail = True + def generate_text(self, *a, **kw): return "" + def count_prompt_tokens(self, *a): return 0 + def count_text_tokens(self, *a): return 0 + + node = TorchNodeServer(backend=_MinimalBackend(), route_timeout=45.0) + assert node.route_timeout == 45.0