Frank
214bad946f
The script runs ON Dankclaw but was calling ssh() to SSH back to dankclaw for service management. Nested SSH fails (no key auth for localhost). Replaced ssh() with local() using subprocess.run(shell=True) for systemctl and sed commands.
230 lines
7.4 KiB
Python
Executable File
230 lines
7.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Benchmark decode throughput across speculative decoding configs.
|
|
|
|
Sweeps 4 configs: none, ngram-mod, draft-mtp, draft-mtp+ngram-mod.
|
|
Each config generates 16k tokens from a minimal prompt at temp 0.
|
|
|
|
Runs ON the target host (Dankclaw). Service management uses local sudo.
|
|
Usage:
|
|
python3 bench-spec-decode.py [--max-tokens N] [--host H] [--port P]
|
|
"""
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
SERVICE = "frank-default-large"
|
|
ENV_FILE = "/home/dank/serverAppConfigs/scripts/frank-default-large.env"
|
|
ALIAS = "frank-default-large"
|
|
DEFAULT_HOST = "localhost"
|
|
DEFAULT_PORT = 7770
|
|
DEFAULT_MAX_TOKENS = 16384
|
|
HEALTH_TIMEOUT = 1200 # 20 min for cold load
|
|
HEALTH_INTERVAL = 5
|
|
|
|
PROMPT = (
|
|
"Write a comprehensive, detailed essay about the complete history of "
|
|
"computing, from the earliest mechanical calculators through modern AI. "
|
|
"Cover every major era in detail, including key figures, machines, and "
|
|
"breakthroughs. Be as thorough as possible."
|
|
)
|
|
|
|
CONFIGS = [
|
|
("none", "none"),
|
|
("ngram-mod", "ngram-mod"),
|
|
("draft-mtp", "draft-mtp"),
|
|
("draft-mtp+ngram-mod", "draft-mtp,ngram-mod"),
|
|
]
|
|
|
|
|
|
def local(cmd):
|
|
"""Run a command locally (the script runs on the target host)."""
|
|
r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=120)
|
|
if r.returncode != 0:
|
|
print(f" WARNING: {cmd} -> exit {r.returncode}: {r.stderr.strip()}",
|
|
file=sys.stderr)
|
|
return r
|
|
|
|
|
|
def sed_spec_type(value):
|
|
local(f"sudo sed -i 's/^SPEC_TYPE=.*/SPEC_TYPE={value}/' {ENV_FILE}")
|
|
|
|
|
|
def wait_health(host, port, timeout=HEALTH_TIMEOUT):
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
resp = urllib.request.urlopen(
|
|
f"http://{host}:{port}/health", timeout=10
|
|
)
|
|
if resp.status == 200:
|
|
return True
|
|
except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError,
|
|
OSError):
|
|
pass
|
|
time.sleep(HEALTH_INTERVAL)
|
|
return False
|
|
|
|
|
|
def benchmark(host, port, alias, max_tokens):
|
|
body = json.dumps({
|
|
"model": alias,
|
|
"messages": [{"role": "user", "content": PROMPT}],
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0,
|
|
"stream": False,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://{host}:{port}/v1/chat/completions",
|
|
data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
t0 = time.time()
|
|
resp = urllib.request.urlopen(req, timeout=1200)
|
|
wall = time.time() - t0
|
|
data = json.loads(resp.read())
|
|
return data, wall
|
|
|
|
|
|
def main():
|
|
import argparse
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument("--host", default=DEFAULT_HOST)
|
|
ap.add_argument("--port", type=int, default=DEFAULT_PORT)
|
|
ap.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS)
|
|
args = ap.parse_args()
|
|
|
|
# Read original SPEC_TYPE from the env file
|
|
r = local(f"grep '^SPEC_TYPE=' {ENV_FILE}")
|
|
original_spec = r.stdout.strip()
|
|
orig_val = (original_spec.split("=", 1)[1]
|
|
if "=" in original_spec else "draft-mtp,ngram-mod")
|
|
print(f"Original: {original_spec}")
|
|
print(f"Max tokens: {args.max_tokens}")
|
|
print(f"Configs: {[c[0] for c in CONFIGS]}")
|
|
print()
|
|
|
|
# Stop + mask the production service
|
|
print("Stopping + masking production service...")
|
|
local(f"sudo systemctl stop {SERVICE}")
|
|
local(f"sudo systemctl mask {SERVICE}")
|
|
time.sleep(2)
|
|
|
|
results = []
|
|
for name, spec_type in CONFIGS:
|
|
print(f"\n{'='*60}")
|
|
print(f"Config: {name} (SPEC_TYPE={spec_type})")
|
|
print(f"{'='*60}")
|
|
|
|
# Set SPEC_TYPE
|
|
sed_spec_type(spec_type)
|
|
print(f"Set SPEC_TYPE={spec_type}")
|
|
|
|
# Start the service
|
|
local(f"sudo systemctl unmask {SERVICE}")
|
|
local(f"sudo systemctl start {SERVICE}")
|
|
print("Starting service...")
|
|
|
|
# Wait for health
|
|
print(f"Waiting for health (up to {HEALTH_TIMEOUT}s)...")
|
|
if not wait_health(args.host, args.port):
|
|
print(f"FAILED: service did not become healthy for {name}")
|
|
results.append({
|
|
"config": name, "spec_type": spec_type,
|
|
"error": "health timeout",
|
|
})
|
|
local(f"sudo systemctl stop {SERVICE}")
|
|
local(f"sudo systemctl mask {SERVICE}")
|
|
continue
|
|
|
|
print("Healthy. Running benchmark...")
|
|
|
|
# Run benchmark
|
|
try:
|
|
data, wall = benchmark(
|
|
args.host, args.port, ALIAS, args.max_tokens
|
|
)
|
|
except Exception as e:
|
|
print(f"FAILED: {e}")
|
|
results.append({
|
|
"config": name, "spec_type": spec_type,
|
|
"error": str(e),
|
|
})
|
|
local(f"sudo systemctl stop {SERVICE}")
|
|
local(f"sudo systemctl mask {SERVICE}")
|
|
continue
|
|
|
|
timings = data.get("timings", {})
|
|
usage = data.get("usage", {})
|
|
pred_n = timings.get("predicted_n", usage.get("completion_tokens", 0))
|
|
pred_ms = timings.get("predicted_ms", 0)
|
|
pred_ps = timings.get("predicted_per_second", 0)
|
|
prompt_n = timings.get("prompt_n", usage.get("prompt_tokens", 0))
|
|
prompt_ms = timings.get("prompt_ms", 0)
|
|
|
|
print(f" Generated: {pred_n} tokens")
|
|
print(f" Wall time: {wall:.1f}s")
|
|
print(f" Decode: {pred_ps:.2f} tok/s")
|
|
print(f" Prefill: {prompt_n} tokens in {prompt_ms:.0f}ms"
|
|
f" ({prompt_n / max(prompt_ms, 1) * 1000:.1f} tok/s)")
|
|
|
|
results.append({
|
|
"config": name,
|
|
"spec_type": spec_type,
|
|
"predicted_tokens": pred_n,
|
|
"predicted_ms": pred_ms,
|
|
"predicted_per_second": pred_ps,
|
|
"prompt_tokens": prompt_n,
|
|
"prompt_ms": prompt_ms,
|
|
"wall_time": wall,
|
|
})
|
|
|
|
# Stop + mask between configs
|
|
local(f"sudo systemctl stop {SERVICE}")
|
|
local(f"sudo systemctl mask {SERVICE}")
|
|
time.sleep(2)
|
|
|
|
# Restore original SPEC_TYPE
|
|
sed_spec_type(orig_val)
|
|
print(f"\nRestoring SPEC_TYPE={orig_val}")
|
|
print("Unmasking + starting production service...")
|
|
local(f"sudo systemctl unmask {SERVICE}")
|
|
local(f"sudo systemctl start {SERVICE}")
|
|
|
|
# Health check
|
|
if wait_health(args.host, args.port):
|
|
print("Production service healthy.")
|
|
else:
|
|
print("WARNING: production service not healthy after restore!")
|
|
|
|
# Print results table
|
|
print(f"\n{'='*70}")
|
|
print("RESULTS")
|
|
print(f"{'='*70}")
|
|
print(f"{'Config':<30} {'Tokens':>8} {'Wall(s)':>10} {'Decode tok/s':>14} "
|
|
f"{'Error':>20}")
|
|
print(f"{'-'*30} {'-'*8} {'-'*10} {'-'*14} {'-'*20}")
|
|
for r in results:
|
|
if "error" in r:
|
|
print(f"{r['config']:<30} {'':>8} {'':>10} {'':>14} "
|
|
f"{r['error']:>20}")
|
|
else:
|
|
print(f"{r['config']:<30} {r['predicted_tokens']:>8} "
|
|
f"{r['wall_time']:>10.1f} {r['predicted_per_second']:>14.2f} "
|
|
f"{'':>20}")
|
|
print(f"{'='*70}")
|
|
|
|
# Save JSON results
|
|
ts = time.strftime("%Y%m%d-%H%M%S")
|
|
out = f"/tmp/bench-spec-decode-{ts}.json"
|
|
with open(out, "w") as f:
|
|
json.dump({"results": results, "max_tokens": args.max_tokens,
|
|
"timestamp": ts}, f, indent=2)
|
|
print(f"Results saved: {out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|