baseline working model service and perf tooling

worklogs/remote_smoke_bench.py

@@ -0,0 +1,135 @@
+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()