You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
import subprocess
|
|
import time
|
|
import json
|
|
import os
|
|
from datetime import datetime
|
|
|
|
# Configuration
|
|
TEST_MODEL = "unsloth/gemma-4-E2B-it-GGUF"
|
|
PROMPTS_FILE = "../test/prompts.json"
|
|
RESULTS_FILE = "benchmark_results.md"
|
|
NUM_PROMPTS = 100 # Change this to test fewer prompts if needed
|
|
|
|
def load_prompts(file_path, limit=100):
|
|
if not os.path.exists(file_path):
|
|
print(f"Error: {file_path} not found.")
|
|
return ["Write a short greeting."]
|
|
with open(file_path, "r") as f:
|
|
prompts = json.load(f)
|
|
return prompts[:limit]
|
|
|
|
def benchmark_service(service_name, port, model_name, prompts):
|
|
url = f"http://localhost:{port}/v1"
|
|
|
|
print(f"\n--- Checking {service_name} on port {port} ---")
|
|
try:
|
|
res = subprocess.run(f"curl -s {url}/models", shell=True, capture_output=True, timeout=5)
|
|
if res.returncode != 0:
|
|
print(f"Skipping {service_name}: Service not found or not ready.")
|
|
return None
|
|
except Exception as e:
|
|
print(f"Skipping {service_name}: Connection failed ({e})")
|
|
return None
|
|
|
|
print(f"Benchmarking {service_name} with {len(prompts)} prompts...")
|
|
|
|
total_duration = 0
|
|
total_tokens = 0
|
|
results = []
|
|
|
|
for i, prompt in enumerate(prompts):
|
|
payload = {
|
|
"model": model_name,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0
|
|
}
|
|
|
|
start_request = time.time()
|
|
res = subprocess.run([
|
|
"curl", "-s", f"{url}/chat/completions",
|
|
"-H", "Content-Type: application/json",
|
|
"-d", json.dumps(payload)
|
|
], capture_output=True, text=True)
|
|
end_request = time.time()
|
|
|
|
duration = end_request - start_request
|
|
|
|
tokens = 0
|
|
try:
|
|
data = json.loads(res.stdout)
|
|
tokens = data.get("usage", {}).get("completion_tokens", 0)
|
|
if tokens == 0:
|
|
tokens = data.get("usage", {}).get("total_tokens", 0)
|
|
except:
|
|
print(f"Warning: Could not parse response for {service_name} on prompt {i+1}")
|
|
continue
|
|
|
|
total_duration += duration
|
|
total_tokens += tokens
|
|
tps = tokens / duration if duration > 0 else 0
|
|
|
|
print(f"[{i+1}/{len(prompts)}] {duration:.2f}s, {tokens} tokens, {tps:.2f} tps")
|
|
|
|
avg_duration = total_duration / len(prompts) if prompts else 0
|
|
avg_tps = total_tokens / total_duration if total_duration > 0 else 0
|
|
|
|
return {
|
|
"service": service_name,
|
|
"avg_duration": round(avg_duration, 2),
|
|
"total_tokens": total_tokens,
|
|
"avg_tps": round(avg_tps, 2),
|
|
"num_prompts": len(prompts)
|
|
}
|
|
|
|
def save_to_markdown(results):
|
|
if not results or all(r is None for r in results):
|
|
print("No results to save.")
|
|
return
|
|
|
|
header = f"\n## Benchmark Results (100 Prompts) - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
|
|
header += "| Service | Avg Latency (s) | Total Tokens | Avg Tokens/Sec | Prompts |\n"
|
|
header += "| :--- | :--- | :--- | :--- | :--- |\n"
|
|
|
|
rows = ""
|
|
for r in results:
|
|
if r:
|
|
rows += f"| {r['service']} | {r['avg_duration']}s | {r['total_tokens']} | {r['avg_tps']} | {r['num_prompts']} |\n"
|
|
|
|
with open(RESULTS_FILE, "a") as f:
|
|
f.write(header + rows + "\n")
|
|
print(f"\nResults saved to {RESULTS_FILE}")
|
|
|
|
if __name__ == "__main__":
|
|
prompts = load_prompts(PROMPTS_FILE, NUM_PROMPTS)
|
|
all_results = []
|
|
|
|
# Test vLLM (Port 8021)
|
|
vllm_res = benchmark_service("vllm-gemma", 8021, "google/gemma-4-E2B", prompts)
|
|
if vllm_res: all_results.append(vllm_res)
|
|
|
|
# Test llama.cpp (Port 8088)
|
|
llamacpp_res = benchmark_service("llamacpp-gemma", 8088, "unsloth/gemma-4-E2B-it-GGUF", prompts)
|
|
if llamacpp_res: all_results.append(llamacpp_res)
|
|
|
|
save_to_markdown(all_results)
|