diff --git a/scripts/bench-spec-decode.py b/scripts/bench-spec-decode.py index 15f9485..1c540ea 100755 --- a/scripts/bench-spec-decode.py +++ b/scripts/bench-spec-decode.py @@ -4,14 +4,11 @@ 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] - -Requires SSH access to the target host (for service management) and -HTTP access to the llama.cpp server endpoint. """ import json -import os import subprocess import sys import time @@ -42,18 +39,17 @@ CONFIGS = [ ] -def ssh(cmd): - r = subprocess.run( - ["ssh", "dankclaw", cmd], - capture_output=True, text=True, timeout=30 - ) +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" ssh stderr: {r.stderr.strip()}", file=sys.stderr) + print(f" WARNING: {cmd} -> exit {r.returncode}: {r.stderr.strip()}", + file=sys.stderr) return r def sed_spec_type(value): - ssh(f"sudo sed -i 's/^SPEC_TYPE=.*/SPEC_TYPE={value}/' {ENV_FILE}") + local(f"sudo sed -i 's/^SPEC_TYPE=.*/SPEC_TYPE={value}/' {ENV_FILE}") def wait_health(host, port, timeout=HEALTH_TIMEOUT): @@ -65,7 +61,8 @@ def wait_health(host, port, timeout=HEALTH_TIMEOUT): ) if resp.status == 200: return True - except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError): + except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError, + OSError): pass time.sleep(HEALTH_INTERVAL) return False @@ -85,7 +82,7 @@ def benchmark(host, port, alias, max_tokens): headers={"Content-Type": "application/json"}, ) t0 = time.time() - resp = urllib.request.urlopen(req, timeout=600) + resp = urllib.request.urlopen(req, timeout=1200) wall = time.time() - t0 data = json.loads(resp.read()) return data, wall @@ -97,12 +94,13 @@ def main(): 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) - ap.add_argument("--ssh-host", default="dankclaw") args = ap.parse_args() # Read original SPEC_TYPE from the env file - r = ssh(f"grep '^SPEC_TYPE=' {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]}") @@ -110,8 +108,8 @@ def main(): # Stop + mask the production service print("Stopping + masking production service...") - ssh(f"sudo systemctl stop {SERVICE}") - ssh(f"sudo systemctl mask {SERVICE}") + local(f"sudo systemctl stop {SERVICE}") + local(f"sudo systemctl mask {SERVICE}") time.sleep(2) results = [] @@ -125,8 +123,8 @@ def main(): print(f"Set SPEC_TYPE={spec_type}") # Start the service - ssh(f"sudo systemctl unmask {SERVICE}") - ssh(f"sudo systemctl start {SERVICE}") + local(f"sudo systemctl unmask {SERVICE}") + local(f"sudo systemctl start {SERVICE}") print("Starting service...") # Wait for health @@ -137,11 +135,10 @@ def main(): "config": name, "spec_type": spec_type, "error": "health timeout", }) - ssh(f"sudo systemctl stop {SERVICE}") - ssh(f"sudo systemctl mask {SERVICE}") + local(f"sudo systemctl stop {SERVICE}") + local(f"sudo systemctl mask {SERVICE}") continue - load_time = 0 # not tracking separately print("Healthy. Running benchmark...") # Run benchmark @@ -155,8 +152,8 @@ def main(): "config": name, "spec_type": spec_type, "error": str(e), }) - ssh(f"sudo systemctl stop {SERVICE}") - ssh(f"sudo systemctl mask {SERVICE}") + local(f"sudo systemctl stop {SERVICE}") + local(f"sudo systemctl mask {SERVICE}") continue timings = data.get("timings", {}) @@ -185,19 +182,16 @@ def main(): }) # Stop + mask between configs - ssh(f"sudo systemctl stop {SERVICE}") - ssh(f"sudo systemctl mask {SERVICE}") + local(f"sudo systemctl stop {SERVICE}") + local(f"sudo systemctl mask {SERVICE}") time.sleep(2) # Restore original SPEC_TYPE - orig_val = original_spec.split("=", 1)[1] if "=" in original_spec else "none" sed_spec_type(orig_val) - - # Restart production service print(f"\nRestoring SPEC_TYPE={orig_val}") print("Unmasking + starting production service...") - ssh(f"sudo systemctl unmask {SERVICE}") - ssh(f"sudo systemctl start {SERVICE}") + local(f"sudo systemctl unmask {SERVICE}") + local(f"sudo systemctl start {SERVICE}") # Health check if wait_health(args.host, args.port): @@ -209,14 +203,17 @@ def main(): print(f"\n{'='*70}") print("RESULTS") print(f"{'='*70}") - print(f"{'Config':<30} {'Tokens':>8} {'Wall(s)':>10} {'Decode tok/s':>14} {'Error':>20}") + 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} {r['error']:>20}") + 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} {'':>20}") + f"{r['wall_time']:>10.1f} {r['predicted_per_second']:>14.2f} " + f"{'':>20}") print(f"{'='*70}") # Save JSON results