120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import json
|
|
import math
|
|
import statistics
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
|
|
def percentile(sorted_values, pct):
|
|
if not sorted_values:
|
|
return 0.0
|
|
if len(sorted_values) == 1:
|
|
return float(sorted_values[0])
|
|
rank = (len(sorted_values) - 1) * pct
|
|
low = math.floor(rank)
|
|
high = math.ceil(rank)
|
|
if low == high:
|
|
return float(sorted_values[low])
|
|
weight = rank - low
|
|
return float(sorted_values[low] * (1.0 - weight) + sorted_values[high] * weight)
|
|
|
|
|
|
def wait_for_health(base_url, timeout_s, server_log):
|
|
deadline = time.time() + timeout_s
|
|
last_error = "server did not become healthy"
|
|
while time.time() < deadline:
|
|
try:
|
|
with urllib.request.urlopen(base_url + "/health", timeout=2) as response:
|
|
if response.status == 200:
|
|
return
|
|
except Exception as exc: # noqa: BLE001
|
|
last_error = str(exc)
|
|
time.sleep(0.5)
|
|
|
|
print("[ERROR] health check failed: {}".format(last_error), file=sys.stderr)
|
|
if server_log:
|
|
try:
|
|
with open(server_log, "r", encoding="utf-8", errors="replace") as fh:
|
|
lines = fh.readlines()[-40:]
|
|
print("[ERROR] tail of server log:", file=sys.stderr)
|
|
for line in lines:
|
|
print(line.rstrip(), file=sys.stderr)
|
|
except OSError as exc: # noqa: BLE001
|
|
print("[ERROR] failed to read server log: {}".format(exc), file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
def build_payload(model_id):
|
|
return {
|
|
"model": model_id,
|
|
"stream": False,
|
|
"temperature": 0.0,
|
|
"top_p": 1.0,
|
|
"max_completion_tokens": 48,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Write a compact Rust function that parses a comma-separated list of u32 values and returns a Vec<u32> while ignoring empty segments.",
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def run_once(base_url, model_id):
|
|
body = json.dumps(build_payload(model_id)).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
base_url + "/v1/chat/completions",
|
|
data=body,
|
|
headers={"content-type": "application/json"},
|
|
method="POST",
|
|
)
|
|
started = time.perf_counter()
|
|
with urllib.request.urlopen(req, timeout=180) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
|
usage = payload.get("usage") or {}
|
|
return elapsed_ms, float(usage.get("tokens_per_second") or 0.0)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-url", required=True)
|
|
parser.add_argument("--model-id", required=True)
|
|
parser.add_argument("--server-log", default="")
|
|
parser.add_argument("--warmups", type=int, default=1)
|
|
parser.add_argument("--samples", type=int, default=7)
|
|
parser.add_argument("--health-timeout", type=float, default=120.0)
|
|
args = parser.parse_args()
|
|
|
|
wait_for_health(args.base_url, args.health_timeout, args.server_log)
|
|
|
|
for _ in range(args.warmups):
|
|
run_once(args.base_url, args.model_id)
|
|
|
|
latencies = []
|
|
tps_values = []
|
|
for _ in range(args.samples):
|
|
latency_ms, tps = run_once(args.base_url, args.model_id)
|
|
latencies.append(latency_ms)
|
|
tps_values.append(tps)
|
|
|
|
latencies_sorted = sorted(latencies)
|
|
p50_ms = percentile(latencies_sorted, 0.50)
|
|
p95_ms = percentile(latencies_sorted, 0.95)
|
|
mean_ms = statistics.fmean(latencies)
|
|
completion_tps_mean = statistics.fmean(tps_values) if tps_values else 0.0
|
|
|
|
print("METRIC p95_ms={:.3f}".format(p95_ms))
|
|
print("METRIC p50_ms={:.3f}".format(p50_ms))
|
|
print("METRIC mean_ms={:.3f}".format(mean_ms))
|
|
print("METRIC completion_tps_mean={:.3f}".format(completion_tps_mean))
|
|
print("METRIC success_count={}".format(len(latencies)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|