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

better dash and inference api QOL

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 15:51:27 +02:00
parent 938a0a721b
commit 640ef78711
4 changed files with 404 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

266
tests/test_tracker_routing.py
View File

@@ -324,6 +324,177 @@ def test_tracker_serves_health_while_proxy_request_is_in_flight():
slow_thread.join(timeout=1.0)
def test_tracker_route_log_counts_proxy_inflight_requests():
entered = threading.Event()
release = threading.Event()
class SlowChatHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args):
pass
def do_POST(self):
if self.path != "/v1/chat/completions":
self.send_response(404)
self.end_headers()
return
length = int(self.headers.get("Content-Length", 0))
self.rfile.read(length)
entered.set()
release.wait(timeout=3.0)
body = json.dumps({"choices": [{"message": {"content": "ok"}}]}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
slow_node = http.server.HTTPServer(("127.0.0.1", 0), SlowChatHandler)
slow_thread = threading.Thread(target=slow_node.serve_forever, daemon=True)
slow_thread.start()
tracker = TrackerServer(heartbeat_timeout=60.0)
tracker_port = tracker.start()
errors = []
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{slow_node.server_address[1]}",
"model": "burst-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0,
"hardware_profile": {}, "score": 1.0},
)
def call_proxy():
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "burst-model", "messages": [{"role": "user", "content": "hi"}]},
)
except Exception as exc:
errors.append(exc)
first = threading.Thread(target=call_proxy)
second = threading.Thread(target=call_proxy)
first.start()
assert entered.wait(timeout=2.0)
second.start()
selected_events = []
deadline = time.time() + 2.0
while time.time() < deadline:
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
selected_events = [
event for event in console["events"]
if event["message"] == "proxy route selected"
]
if len(selected_events) >= 2:
break
time.sleep(0.05)
assert len(selected_events) >= 2
second_nodes = selected_events[-1]["fields"]["nodes"]
assert second_nodes[0]["queue_depth"] == 2
assert second_nodes[0]["proxy_inflight"] == 2
finally:
release.set()
first.join(timeout=3.0)
second.join(timeout=3.0)
tracker.stop()
slow_node.shutdown()
slow_node.server_close()
slow_thread.join(timeout=1.0)
assert not first.is_alive()
assert not second.is_alive()
assert not errors
def test_tracker_logs_stream_progress_before_request_completes():
chunk_sent = threading.Event()
release = threading.Event()
class StreamingChatHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args):
pass
def do_POST(self):
if self.path != "/v1/chat/completions":
self.send_response(404)
self.end_headers()
return
self.rfile.read(int(self.headers.get("Content-Length", 0)))
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.end_headers()
payload = json.dumps({
"choices": [{"delta": {"content": "hello world"}}],
}).encode()
self.wfile.write(b"data: " + payload + b"\n\n")
self.wfile.flush()
chunk_sent.set()
release.wait(timeout=3.0)
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
node = http.server.HTTPServer(("127.0.0.1", 0), StreamingChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
tracker = TrackerServer(heartbeat_timeout=60.0)
tracker_port = tracker.start()
response = None
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "stream-progress-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0,
"hardware_profile": {}, "score": 1.0},
)
req = urllib.request.Request(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
data=json.dumps({
"model": "stream-progress-model",
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
}).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
response = urllib.request.urlopen(req, timeout=3.0)
first_line = response.readline()
assert first_line.startswith(b"data:")
assert chunk_sent.wait(timeout=1.0)
progress_events = []
deadline = time.time() + 2.0
while time.time() < deadline:
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
progress_events = [
event for event in console["events"]
if event["message"] == "proxy progress"
]
if progress_events:
break
time.sleep(0.05)
assert progress_events
fields = progress_events[-1]["fields"]
assert fields["tokens"] == 2
assert fields["tokens_per_sec"] > 0
active = [
event for event in console["events"]
if event["message"] == "proxy route selected"
]
assert active
finally:
release.set()
if response is not None:
response.close()
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
def test_tracker_routes_hf_model_alias_from_quickstart():
"""The documented qwen2.5-0.5b alias resolves a full HF repo registration."""
tracker = TrackerServer()
@@ -536,6 +707,101 @@ def test_tracker_route_endpoint_routes_split_preset_nodes_by_alias():
assert [node["start_layer"] for node in response["nodes"]] == [0, 22]
def test_tracker_route_endpoint_ignores_model_case_and_outer_whitespace():
tracker = TrackerServer(model_presets={
"qwen3.6-35b-a3b": {
"layers_start": 0,
"layers_end": 39,
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"aliases": ["Qwen3.6-35B-A3B"],
}
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9101",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/route?model=%20Qwen3.6-35B-A3B%20"
)
finally:
tracker.stop()
assert response["route"] == ["http://127.0.0.1:9101"]
def test_tracker_proxy_ignores_model_case_and_outer_whitespace():
class ChatHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args):
pass
def do_POST(self):
if self.path != "/v1/chat/completions":
self.send_response(404)
self.end_headers()
return
length = int(self.headers.get("Content-Length", 0))
request_body = json.loads(self.rfile.read(length) or b"{}")
body = json.dumps({
"model": request_body["model"],
"choices": [{"message": {"content": "ok"}}],
}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
node = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
tracker = TrackerServer(model_presets={
"qwen3.6-35b-a3b": {
"layers_start": 0,
"layers_end": 39,
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"aliases": ["Qwen3.6-35B-A3B"],
}
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "Qwen3.6-35B-A3B ",
"messages": [{"role": "user", "content": "hi"}]},
)
finally:
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
assert response["choices"][0]["message"]["content"] == "ok"
def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix():
tracker = TrackerServer()
tracker_port = tracker.start()