168 lines
7.2 KiB
Python
168 lines
7.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""速度测试执行器：校准 -> 采样 -> 汇总，全程写日志与指标入库"""
|
|
import json
|
|
import statistics
|
|
import threading
|
|
import time
|
|
import uuid
|
|
|
|
import database as db
|
|
import llm_providers as lp
|
|
from llm_providers import ProviderError, StopRequested
|
|
|
|
|
|
class TestRunner(threading.Thread):
|
|
def __init__(self, test_id, cfg, gen):
|
|
super().__init__(daemon=True)
|
|
self.test_id = test_id
|
|
self.cfg = cfg
|
|
self.gen = gen
|
|
self.cancel_flag = False
|
|
self.start_wall = time.time()
|
|
self.ratio = None
|
|
self.samples = []
|
|
|
|
def request_cancel(self):
|
|
self.cancel_flag = True
|
|
|
|
def should_stop(self):
|
|
return self.cancel_flag
|
|
|
|
def log(self, level, msg):
|
|
db.add_log(self.test_id, level, msg)
|
|
|
|
# ───────────────────────── 主流程 ─────────────────────────
|
|
|
|
def run(self):
|
|
try:
|
|
self._run()
|
|
except StopRequested:
|
|
self.log("WARN", "用户请求停止测试")
|
|
db.update_status(self.test_id, "canceled",
|
|
summary=self._make_summary(), error="用户取消")
|
|
except Exception as e:
|
|
self.log("ERROR", "测试异常终止: %s" % e)
|
|
db.update_status(self.test_id, "error",
|
|
summary=self._make_summary(), error=str(e))
|
|
|
|
def _run(self):
|
|
provider = self.cfg.get("provider", "openai")
|
|
model = self.cfg.get("model", "")
|
|
n = max(1, int(self.gen.get("samples", 3)))
|
|
target_tokens = max(16, int(self.gen.get("prompt_tokens", 2048)))
|
|
max_tokens = max(1, int(self.gen.get("max_tokens", 256)))
|
|
avoid_cache = bool(self.gen.get("avoid_cache"))
|
|
|
|
self.log("INFO", "═══ 开始速度测试 ═══")
|
|
self.log("INFO", "提供商: %s | 模型: %s" % (lp.PROVIDER_LABELS.get(provider, provider), model))
|
|
self.log("INFO", "目标上文: %d tokens | 生成长度: %d tokens | 采样: %d 次 | 避免缓存: %s"
|
|
% (target_tokens, max_tokens, n, "开" if avoid_cache else "关"))
|
|
|
|
ratio = self._calibrate()
|
|
self.ratio = ratio
|
|
base_prompt = self._build_prompt(target_tokens, ratio)
|
|
self.log("INFO", "构造基准提示词完成，目标约 %d tokens" % target_tokens)
|
|
|
|
for i in range(1, n + 1):
|
|
if self.should_stop():
|
|
raise StopRequested()
|
|
prompt = self._finalize_prompt(base_prompt)
|
|
self.log("INFO", "── 采样 %d/%d 开始 ──" % (i, n))
|
|
try:
|
|
m = lp.call_stream(
|
|
self.cfg, prompt,
|
|
{"max_tokens": max_tokens, "avoid_cache": avoid_cache},
|
|
log=lambda lv, msg: self.log(lv, msg),
|
|
should_stop=self.should_stop)
|
|
m["run_index"] = i
|
|
self.samples.append({"run_index": i, "ok": True, "metrics": m})
|
|
db.add_run(self.test_id, i, m)
|
|
self.log("METRIC", self._fmt_metric(i, n, m))
|
|
except StopRequested:
|
|
raise
|
|
except ProviderError as e:
|
|
self.log("ERROR", "采样 %d/%d 失败: %s" % (i, n, e))
|
|
self.samples.append({"run_index": i, "ok": False, "error": str(e)})
|
|
db.add_run(self.test_id, i, {}, str(e))
|
|
raise e
|
|
|
|
summary = self._make_summary()
|
|
db.update_status(self.test_id, "done", summary=summary)
|
|
self.log("INFO", "═══ 测试完成 ═══")
|
|
self.log("INFO", "汇总: 平均首字 %.1f ms | 平均预填充 %.1f tok/s | 平均解码 %.1f tok/s"
|
|
% (summary.get("avg_ttft_ms") or 0,
|
|
summary.get("avg_prefill_speed") or 0,
|
|
summary.get("avg_decode_speed") or 0))
|
|
|
|
# ───────────────────────── 工具方法 ─────────────────────────
|
|
|
|
def _calibrate(self):
|
|
probe = ("The quick brown fox jumps over the lazy dog. 人工智能大模型推理速度基准语料，"
|
|
"用于测量提示词预填充与流式解码性能。\n") * 40
|
|
self.log("INFO", "正在校准 token/字符 比例（发送小探测请求）...")
|
|
try:
|
|
m = lp.call_stream(self.cfg, probe,
|
|
{"max_tokens": 8, "avoid_cache": False},
|
|
log=lambda lv, msg: self.log(lv, msg),
|
|
should_stop=self.should_stop)
|
|
pt = m.get("prompt_tokens") or 0
|
|
if pt and len(probe):
|
|
ratio = pt / len(probe)
|
|
self.log("INFO", "探测提示词 %d tokens / %d 字符 = %.3f tok/字符"
|
|
% (pt, len(probe), ratio))
|
|
return max(ratio, 0.001)
|
|
except StopRequested:
|
|
raise
|
|
except Exception as e:
|
|
self.log("WARN", "校准失败（%s），使用默认估算 0.55 tok/字符" % e)
|
|
return 0.55
|
|
|
|
def _build_prompt(self, target_tokens, ratio):
|
|
seg = ("基准语料：The quick brown fox jumps over the lazy dog. "
|
|
"人工智能大模型推理性能测试文本，用于测量提示词预填充速度、首字延迟与流式解码吞吐。\n")
|
|
target_chars = max(64, int(target_tokens / ratio))
|
|
repeats = max(1, target_chars // len(seg))
|
|
return seg * repeats
|
|
|
|
def _finalize_prompt(self, base):
|
|
if self.gen.get("avoid_cache"):
|
|
return "[cache-bust %s]\n%s" % (uuid.uuid4().hex, base)
|
|
return base
|
|
|
|
def _fmt_metric(self, i, n, m):
|
|
return ("采样 %d/%d 完成 | 提示词 %d tok | 缓存 %d tok | 首字 %s ms | 预填充 %s tok/s"
|
|
" | 输出 %d tok | 解码 %s tok/s | 总耗时 %s ms"
|
|
% (i, n, m.get("prompt_tokens") or 0, m.get("cached_tokens") or 0,
|
|
m.get("ttft_ms"), m.get("prefill_speed"), m.get("output_tokens") or 0,
|
|
m.get("decode_speed"), m.get("total_ms")))
|
|
|
|
def _make_summary(self):
|
|
ok = [s["metrics"] for s in self.samples if s.get("ok")]
|
|
base = {
|
|
"provider": self.cfg.get("provider"),
|
|
"model": self.cfg.get("model"),
|
|
"gen": self.gen,
|
|
"samples_total": len(self.samples),
|
|
"samples_ok": len(ok),
|
|
"calibration_chars_per_token": round(1 / self.ratio, 2) if self.ratio else None,
|
|
}
|
|
if not ok:
|
|
return base
|
|
|
|
def avg(k):
|
|
vals = [m[k] for m in ok if m.get(k) is not None]
|
|
return round(statistics.mean(vals), 1) if vals else None
|
|
|
|
summary = dict(base)
|
|
summary.update({
|
|
"avg_ttft_ms": avg("ttft_ms"),
|
|
"avg_prefill_speed": avg("prefill_speed"),
|
|
"avg_decode_speed": avg("decode_speed"),
|
|
"avg_prompt_tokens": avg("prompt_tokens"),
|
|
"avg_output_tokens": avg("output_tokens"),
|
|
"avg_cached_tokens": avg("cached_tokens"),
|
|
"avg_total_ms": avg("total_ms"),
|
|
"best_ttft_ms": min([m["ttft_ms"] for m in ok if m.get("ttft_ms") is not None], default=None),
|
|
})
|
|
return summary
|