popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dash works !!! good data. billing seems to work

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-02 23:07:41 +02:00
parent 1e0aa6ea8f
commit a938c19a82
14 changed files with 617 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
tests/test_tracker_as_node.py
View File

@@ -1,8 +1,9 @@
"""US-014 integration test: tracker-as-first-layer-node inference entry point.
"""US-014 integration test: head-worker inference entry point.
Two stub tracker-nodes (shard 0-5, tracker_mode=True) + two mid-shard stub nodes
(shard 6-11) for openai-community/gpt2. Ten requests via gateway assert round-robin
load distribution across tracker-nodes and valid OpenAI-format responses.
Two stub head workers (shard 0-5, tracker_mode=True for legacy wire
compatibility) + two mid-shard stub nodes (shard 6-11) for openai-community/gpt2.
Ten requests via gateway assert round-robin load distribution across head
workers and valid OpenAI-format responses.
"""
import json
@@ -51,7 +52,7 @@ def _register_node(
@pytest.fixture
def tracker_node_setup():
"""Start tracker, two tracker-nodes (shard 0-5), two mid-shard nodes (shard 6-11),
"""Start tracker, two head workers (shard 0-5), two mid-shard nodes (shard 6-11),
and a gateway. Yields (gateway_url, tracker_node_a, tracker_node_b)."""
tracker = TrackerServer(
@@ -66,12 +67,12 @@ def tracker_node_setup():
tracker_port = tracker.start()
tracker_url = f"http://127.0.0.1:{tracker_port}"
# Two tracker-nodes: serve shard 0-5, expose /v1/chat/completions
# Two head workers: serve shard 0-5, expose /v1/chat/completions.
tracker_node_a = StubNodeServer(
shard_start=0,
shard_end=5,
is_last_shard=False,
response_prefix="tracker-node-A:",
response_prefix="head-worker-A:",
model=GPT2_MODEL,
tracker_mode=True,
)
@@ -81,7 +82,7 @@ def tracker_node_setup():
shard_start=0,
shard_end=5,
is_last_shard=False,
response_prefix="tracker-node-B:",
response_prefix="head-worker-B:",
model=GPT2_MODEL,
tracker_mode=True,
)
@@ -106,7 +107,7 @@ def tracker_node_setup():
)
mid_port_b = mid_node_b.start()
# Register all nodes with tracker (tracker-nodes get tracker_mode=True)
# Register all nodes with tracker (head workers keep legacy tracker_mode=True).
_register_node(tracker_url, f"http://127.0.0.1:{port_a}", 0, 5, tracker_mode=True)
_register_node(tracker_url, f"http://127.0.0.1:{port_b}", 0, 5, tracker_mode=True)
_register_node(tracker_url, f"http://127.0.0.1:{mid_port_a}", 6, 11)
@@ -155,22 +156,22 @@ def test_all_responses_valid_openai_format(tracker_node_setup):
def test_both_tracker_nodes_receive_load(tracker_node_setup):
"""Both tracker-nodes handle at least one request each out of ten."""
"""Both head workers handle at least one request each out of ten."""
gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup
for i in range(10):
_send_chat_request(gateway_url, f"message {i}")
assert tracker_node_a.chat_completion_count >= 1, (
f"tracker-node-A received no requests (count={tracker_node_a.chat_completion_count})"
f"head-worker-A received no requests (count={tracker_node_a.chat_completion_count})"
)
assert tracker_node_b.chat_completion_count >= 1, (
f"tracker-node-B received no requests (count={tracker_node_b.chat_completion_count})"
f"head-worker-B received no requests (count={tracker_node_b.chat_completion_count})"
)
total = tracker_node_a.chat_completion_count + tracker_node_b.chat_completion_count
assert total == 10, f"total requests handled by tracker-nodes: {total}, expected 10"
assert total == 10, f"total requests handled by head workers: {total}, expected 10"
def test_tracker_nodes_endpoint_returns_registered_nodes(tracker_node_setup):
"""GET /v1/tracker-nodes/<model> on the tracker returns both registered tracker-nodes."""
"""GET /v1/tracker-nodes/<model> remains as a legacy alias for head workers."""
_, tracker_node_a, tracker_node_b = tracker_node_setup
# Find the tracker URL by inspecting the fixture indirectly
# We need the tracker URL — use the gateway's tracker_url
@@ -181,13 +182,13 @@ def test_tracker_nodes_endpoint_returns_registered_nodes(tracker_node_setup):
def test_load_is_distributed_evenly(tracker_node_setup):
"""With 10 requests and round-robin, each tracker-node gets exactly 5."""
"""With 10 requests and round-robin, each head worker gets exactly 5."""
gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup
for i in range(10):
_send_chat_request(gateway_url, f"round-robin test {i}")
assert tracker_node_a.chat_completion_count == 5, (
f"expected 5 requests on tracker-node-A, got {tracker_node_a.chat_completion_count}"
f"expected 5 requests on head-worker-A, got {tracker_node_a.chat_completion_count}"
)
assert tracker_node_b.chat_completion_count == 5, (
f"expected 5 requests on tracker-node-B, got {tracker_node_b.chat_completion_count}"
f"expected 5 requests on head-worker-B, got {tracker_node_b.chat_completion_count}"
)