wip -more responsive UI, better routing

tests/test_real_model_backend.py

@@ -376,6 +376,92 @@ def test_split_shard_chat_streams_each_generated_token_incrementally():
     assert "data: [DONE]" in rest
 
 
+def test_current_requests_snapshot_while_generating():
+    release_second = threading.Event()
+    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
+    tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
+    head_port = head.start()
+    tail_port = tail.start()
+    response = None
+    try:
+        payload = json.dumps({
+            "model": "fake-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "stream": True,
+            "max_tokens": 2,
+        }).encode()
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{head_port}/v1/chat/completions",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "X-Meshnet-Request-Id": "req-live-1",
+                "X-Meshnet-Route": json.dumps([
+                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
+                ]),
+            },
+            method="POST",
+        )
+        response = urllib.request.urlopen(req, timeout=5)
+        deadline = time.time() + 2.0
+        while time.time() < deadline:
+            live = head.current_requests
+            if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1:
+                break
+            time.sleep(0.02)
+        assert head.current_requests
+        snap = head.current_requests[0]
+        assert snap["request_id"] == "req-live-1"
+        assert snap["tokens"] >= 1
+        assert snap["tokens_per_sec"] >= 0
+        assert snap["routing_complete"] is True
+        release_second.set()
+        response.read()
+    finally:
+        release_second.set()
+        if response is not None:
+            response.close()
+        head.stop()
+        tail.stop()
+
+    assert head.current_requests == []
+
+
+def test_distributed_generating_log_includes_tps(capsys):
+    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
+    tail = TorchNodeServer(backend=_FakePipelineTailBackend())
+    head_port = head.start()
+    tail_port = tail.start()
+    try:
+        payload = json.dumps({
+            "model": "fake-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "max_tokens": 1,
+        }).encode()
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{head_port}/v1/chat/completions",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "X-Meshnet-Route": json.dumps([
+                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
+                ]),
+            },
+            method="POST",
+        )
+        with urllib.request.urlopen(req, timeout=5) as resp:
+            json.loads(resp.read())
+    finally:
+        head.stop()
+        tail.stop()
+
+    out = capsys.readouterr().out
+    assert "generating step=1/1" in out
+    assert " tps=" in out
+    assert "generation complete tokens=1" in out
+    assert out.count("generating step=1/1") == 1
+
+
 def test_int_tensor_header_serializes_torch_tensors():
     torch = pytest.importorskip("torch")