78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
"""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()
|