feat(us-027/us-028): throughput-optimized routing + correctness tests
_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 <noreply@anthropic.com>
This commit is contained in:
@@ -628,6 +628,10 @@ class TorchNodeServer:
|
|||||||
self._thread: threading.Thread | None = None
|
self._thread: threading.Thread | None = None
|
||||||
self.port: int | None = None
|
self.port: int | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def route_timeout(self) -> float:
|
||||||
|
return self._route_timeout
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def backend(self) -> TorchModelShard:
|
def backend(self) -> TorchModelShard:
|
||||||
return self._backend
|
return self._backend
|
||||||
|
|||||||
@@ -292,12 +292,22 @@ class _NodeEntry:
|
|||||||
self.pending_new_assignment: dict | None = None
|
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(
|
def _select_route(
|
||||||
nodes: list[_NodeEntry],
|
nodes: list[_NodeEntry],
|
||||||
required_start: int,
|
required_start: int,
|
||||||
required_end: int,
|
required_end: int,
|
||||||
) -> tuple[list[_NodeEntry], str]:
|
) -> 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(
|
candidates = sorted(
|
||||||
[node for node in nodes if node.shard_start is not None and node.shard_end is not None],
|
[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]
|
key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator]
|
||||||
@@ -309,7 +319,11 @@ def _select_route(
|
|||||||
best: _NodeEntry | None = None
|
best: _NodeEntry | None = None
|
||||||
for node in candidates:
|
for node in candidates:
|
||||||
if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to:
|
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
|
best = node
|
||||||
if best is None:
|
if best is None:
|
||||||
missing = covered_up_to + 1
|
missing = covered_up_to + 1
|
||||||
|
|||||||
@@ -1185,3 +1185,133 @@ def test_stats_gossip_endpoint_merges_peer_slice():
|
|||||||
assert combined["models"]["remote-model"]["rpm_last_hour"] == 2.5
|
assert combined["models"]["remote-model"]["rpm_last_hour"] == 2.5
|
||||||
finally:
|
finally:
|
||||||
tracker.stop()
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user