relay preparing for public internet

_e2e_kv_check.py

@@ -0,0 +1,77 @@
+"""E2E check: two real TorchNodeServers over HTTP, user's live topology.
+
+Head owns layers 0-20, tail owns 10-23 (overlapping shards, clamped to @21),
+exactly like the reported 2-node run. Compares cached vs stateless tps and
+verifies identical output text.
+"""
+import json
+import sys
+import time
+import urllib.request
+
+sys.path.insert(0, "packages/node")
+
+from meshnet_node.model_backend import TorchModelShard
+from meshnet_node.torch_server import TorchNodeServer
+
+MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+MAX_TOKENS = 60
+
+print("loading shards...", flush=True)
+head_backend = TorchModelShard(MODEL, 0, 20)
+tail_backend = TorchModelShard(MODEL, 10, 23)
+
+head = TorchNodeServer(backend=head_backend, tracker_mode=True)
+tail = TorchNodeServer(backend=tail_backend)
+head_port = head.start()
+tail_port = tail.start()
+print(f"head :{head_port} layers 0-20 | tail :{tail_port} layers 10-23", flush=True)
+
+
+def chat(label):
+    body = json.dumps({
+        "model": MODEL,
+        "messages": [{"role": "user", "content": "Count from 1 to 10, then say done."}],
+        "max_tokens": MAX_TOKENS,
+        "stream": False,
+    }).encode()
+    req = urllib.request.Request(
+        f"http://127.0.0.1:{head_port}/v1/chat/completions",
+        data=body,
+        headers={
+            "Content-Type": "application/json",
+            # tracker-injected route: tail continues at layer 21 (clamped)
+            "X-Meshnet-Route": json.dumps(
+                [{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 21}]
+            ),
+        },
+        method="POST",
+    )
+    started = time.monotonic()
+    with urllib.request.urlopen(req, timeout=600) as r:
+        resp = json.loads(r.read())
+    elapsed = time.monotonic() - started
+    text = resp["choices"][0]["message"]["content"]
+    tokens = resp.get("usage", {}).get("completion_tokens", 0)
+    print(f"\n[{label}] tokens={tokens} elapsed={elapsed:.2f}s "
+          f"tps={tokens / max(elapsed, 1e-6):.2f}", flush=True)
+    print(f"[{label}] text={text!r}", flush=True)
+    return text, tokens, elapsed
+
+
+try:
+    cached_text, ct, ce = chat("cached-warmup")  # first run pays model warmup
+    cached_text, ct, ce = chat("cached")
+    forwards_cached = tail.forward_chunk_count
+
+    head_backend.supports_kv_cache = False  # force legacy stateless path
+    stateless_text, st, se = chat("stateless")
+    forwards_stateless = tail.forward_chunk_count - forwards_cached
+
+    print(f"\ntail /forward calls: cached-run total={forwards_cached} "
+          f"stateless-run={forwards_stateless}", flush=True)
+    print(f"speedup: {(ct / ce) / max(st / se, 1e-9):.2f}x", flush=True)
+    print(f"outputs identical: {cached_text == stateless_text}", flush=True)
+finally:
+    head.stop()
+    tail.stop()