neuron-tai/tests/test_model_speed_latency.py

"""Tracker-backed latency experiments and model-speed drill-down."""

import http.server
import json
import threading
import time
import urllib.request

import pytest

from meshnet_tracker.server import TrackerServer


MODELS = {
    "qwen2.5-0.5b-instruct": (24, "Qwen/Qwen2.5-0.5B-Instruct"),
    "qwen3.6-35b-a3b": (40, "unsloth/Qwen3.6-35B-A3B"),
}


def _post_json(url: str, payload: dict) -> dict:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=10.0) as response:
        return json.loads(response.read())


def _get_json(url: str) -> dict:
    with urllib.request.urlopen(url, timeout=10.0) as response:
        return json.loads(response.read())


class _LatencyNode(http.server.BaseHTTPRequestHandler):
    """Synthetic node: every downstream Activation Seam adds deterministic delay."""

    base_delay_seconds = 0.004
    seam_delay_seconds = 0.006

    def log_message(self, *_args):
        pass

    def do_POST(self):
        self.rfile.read(int(self.headers.get("Content-Length", 0)))
        downstream = json.loads(self.headers.get("X-Meshnet-Route", "[]"))
        time.sleep(self.base_delay_seconds + self.seam_delay_seconds * len(downstream))
        body = json.dumps({
            "choices": [{"message": {"role": "assistant", "content": "ok " * 40}}],
            "usage": {"prompt_tokens": 10, "completion_tokens": 40},
        }).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)


def _start_latency_nodes(count: int):
    nodes = []
    threads = []
    for _ in range(count):
        node = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _LatencyNode)
        thread = threading.Thread(target=node.serve_forever, daemon=True)
        thread.start()
        nodes.append(node)
        threads.append(thread)
    return nodes, threads


@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("hardware", ["gpu", "gpu-cpu", "cpu"])
def test_tracker_records_increasing_hop_latency_for_model_and_hardware(model, hardware):
    """One through five hops must preserve a measurable seam penalty in tracker stats."""
    layer_count, hf_repo = MODELS[model]
    nodes, threads = _start_latency_nodes(5)
    tracker = TrackerServer(model_presets={
        model: {
            "layers_start": 0,
            "layers_end": layer_count - 1,
            "hf_repo": hf_repo,
            "aliases": [model],
        }
    })
    tracker_port = tracker.start()
    try:
        registered_ids = []
        for index, node in enumerate(nodes):
            start = (layer_count * index) // 5
            end = (layer_count * (index + 1)) // 5 - 1
            device = "cuda" if hardware == "gpu" or (hardware == "gpu-cpu" and index == 0) else "cpu"
            data = _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": model,
                    "hf_repo": hf_repo,
                    "num_layers": layer_count,
                    "shard_start": start,
                    "shard_end": end,
                    "tracker_mode": index == 0,
                    "hardware_profile": {"device": device},
                    "vram_bytes": 8_000_000_000 if device == "cuda" else 0,
                    "ram_bytes": 32_000_000_000,
                    "benchmark_tokens_per_sec": 100.0,
                },
            )
            registered_ids.append(data["node_id"])

        for hops in range(1, 6):
            route = registered_ids[:hops]
            # Widen the final shard to make each pinned prefix a complete route.
            with tracker._server.lock:
                tracker._server.registry[route[-1]].shard_end = layer_count - 1
            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
                {
                    "model": model,
                    "messages": [{"role": "user", "content": "measure"}],
                    "route": route,
                },
            )

        report = _get_json(f"http://127.0.0.1:{tracker_port}/v1/model-speed?model={model}")
    finally:
        tracker.stop()
        for node, thread in zip(nodes, threads):
            node.shutdown()
            node.server_close()
            thread.join(timeout=1.0)

    routes = {entry["hop_count"]: entry for entry in report["routes"]}
    assert set(routes) == {1, 2, 3, 4, 5}
    assert routes[5]["latency_ms"] > routes[1]["latency_ms"]
    assert routes[5]["latency_penalty_ms"] > 0
    assert routes[5]["device_mix"] == hardware
    assert report["model"] == model
    assert report["nodes"]


def test_model_speed_dashboard_includes_visualization_and_route_drilldown():
    tracker = TrackerServer()
    port = tracker.start()
    try:
        html = urllib.request.urlopen(f"http://127.0.0.1:{port}/dashboard").read().decode()
    finally:
        tracker.stop()

    assert "Model inference speed" in html
    assert "model-speed-chart" in html
    assert "renderModelSpeed" in html
    assert "/v1/model-speed" in html