136 lines
5.1 KiB
Python
136 lines
5.1 KiB
Python
import json
|
|
import statistics
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
|
|
BASE_URL = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:1111"
|
|
OUT_DIR = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/root/work/logs")
|
|
MODEL = sys.argv[3] if len(sys.argv) > 3 else "llm"
|
|
|
|
|
|
def post_json(path, payload, timeout=300):
|
|
req = urllib.request.Request(
|
|
BASE_URL + path,
|
|
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
return resp.status, resp.read().decode("utf-8")
|
|
|
|
|
|
def health():
|
|
with urllib.request.urlopen(BASE_URL + "/health", timeout=10) as resp:
|
|
return resp.status
|
|
|
|
|
|
def chat_once(prompt, max_tokens=32):
|
|
payload = {
|
|
"model": MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0,
|
|
}
|
|
start = time.perf_counter()
|
|
status, body = post_json("/v1/chat/completions", payload)
|
|
elapsed = time.perf_counter() - start
|
|
obj = json.loads(body)
|
|
usage = obj.get("usage") or {}
|
|
completion_tokens = usage.get("completion_tokens") or 0
|
|
return {
|
|
"status": status,
|
|
"elapsed_sec": elapsed,
|
|
"usage": usage,
|
|
"output_tps": completion_tokens / elapsed if elapsed and completion_tokens else None,
|
|
"content_preview": (obj["choices"][0]["message"].get("content") or "")[:200],
|
|
}
|
|
|
|
|
|
def chat_stream(prompt, max_tokens=64):
|
|
payload = {
|
|
"model": MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
}
|
|
req = urllib.request.Request(
|
|
BASE_URL + "/v1/chat/completions",
|
|
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
start = time.perf_counter()
|
|
first_token_at = None
|
|
usage = None
|
|
pieces = []
|
|
with urllib.request.urlopen(req, timeout=300) as resp:
|
|
for raw in resp:
|
|
line = raw.decode("utf-8", errors="replace").strip()
|
|
if not line or not line.startswith("data:"):
|
|
continue
|
|
data = line[5:].strip()
|
|
if data == "[DONE]":
|
|
break
|
|
obj = json.loads(data)
|
|
if obj.get("usage"):
|
|
usage = obj["usage"]
|
|
choices = obj.get("choices") or []
|
|
if choices:
|
|
delta = choices[0].get("delta") or {}
|
|
text = delta.get("content") or delta.get("reasoning_content") or ""
|
|
if text:
|
|
if first_token_at is None:
|
|
first_token_at = time.perf_counter()
|
|
pieces.append(text)
|
|
elapsed = time.perf_counter() - start
|
|
completion_tokens = (usage or {}).get("completion_tokens") or 0
|
|
gen_time = elapsed - (first_token_at - start) if first_token_at else elapsed
|
|
return {
|
|
"elapsed_sec": elapsed,
|
|
"ttft_sec": (first_token_at - start) if first_token_at else None,
|
|
"usage": usage,
|
|
"output_tps_after_ttft": completion_tokens / gen_time if gen_time and completion_tokens else None,
|
|
"content_preview": "".join(pieces)[:200],
|
|
}
|
|
|
|
|
|
def main():
|
|
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
results = {
|
|
"base_url": BASE_URL,
|
|
"model": MODEL,
|
|
"created_at": datetime.now().isoformat(timespec="seconds"),
|
|
}
|
|
results["health_status"] = health()
|
|
results["smoke_nonstream"] = chat_once("你好，请用一句话介绍你自己。", max_tokens=16)
|
|
|
|
short_prompt = "请用中文简要说明什么是模型推理服务。"
|
|
stream_runs = [chat_stream(short_prompt, max_tokens=64) for _ in range(3)]
|
|
results["stream_runs"] = stream_runs
|
|
ttfts = [x["ttft_sec"] for x in stream_runs if x["ttft_sec"] is not None]
|
|
tps = [x["output_tps_after_ttft"] for x in stream_runs if x["output_tps_after_ttft"] is not None]
|
|
results["stream_summary"] = {
|
|
"ttft_avg_sec": statistics.mean(ttfts) if ttfts else None,
|
|
"ttft_p90_sec": sorted(ttfts)[int(0.9 * (len(ttfts) - 1))] if ttfts else None,
|
|
"output_tps_avg": statistics.mean(tps) if tps else None,
|
|
}
|
|
|
|
repeated = "以下是一段用于测试前缀缓存的公共上下文：" + ("模型部署竞赛关注吞吐、延迟、缓存命中和稳定性。" * 80)
|
|
results["prefix_cache_probe_1"] = chat_once(repeated + "\n请总结一句话。", max_tokens=16)
|
|
results["prefix_cache_probe_2"] = chat_once(repeated + "\n请换一种说法总结一句话。", max_tokens=16)
|
|
|
|
out = OUT_DIR / f"baseline_smoke_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
|
out.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
print(json.dumps(results, ensure_ascii=False, indent=2), flush=True)
|
|
print("RESULT_FILE", out, flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|