benchmark.py

#!/usr/bin/env python3
"""
Nap - Benchmark Script
Tests wake-on-demand latency through the AI gateway.
"""
import time
import requests
import subprocess
import yaml
import os

def load_config():
    config_path = os.path.join(os.path.dirname(__file__), "config.yaml")
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def check_pc_alive(ip):
    try:
        subprocess.run(["ping", "-c", "1", "-W", "500", ip], capture_output=True, check=True)
        return True
    except:
        return False

def run_benchmark():
    config = load_config()
    pc_ip = config['pc']['ip']
    gateway_port = config['proxy'].get('listen_port', 11435)
    gateway_url = f"http://localhost:{gateway_port}/api/tags"
    
    print("--- Nap AI Gateway Benchmark ---")
    
    if check_pc_alive(pc_ip):
        print("⚠️  PC is already ONLINE. Put it to sleep first for accurate results.")

    print("\n[1/3] Sending request to gateway...")
    start_time = time.time()
    
    try:
        print("      (Waiting for PC to wake and respond...)")
        response = requests.get(gateway_url, timeout=60)
        end_time = time.time()
        
        if response.status_code == 200:
            total_time = end_time - start_time
            print(f"\n✅ SUCCESS!")
            print(f"   Total Response Time: {total_time:.2f} seconds")
            models = response.json().get('models', [])
            print(f"   Model Count: {len(models)}")
        else:
            print(f"\n❌ Gateway returned error: {response.status_code}")
            print(response.text)
            
    except requests.exceptions.Timeout:
        print("\n❌ TIMEOUT: PC took too long to wake (>60s)")
    except Exception as e:
        print(f"\n❌ ERROR: {e}")

if __name__ == "__main__":
    run_benchmark()