331 lines
12 KiB
Python
331 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""LLM 速度测试台 - Flask 主应用"""
|
||
import io
|
||
import json
|
||
|
||
from flask import Flask, jsonify, request, send_file, send_from_directory
|
||
|
||
import config
|
||
import database as db
|
||
from llm_providers import DEFAULT_URLS, ProviderError, call_stream
|
||
from tester import TestRunner
|
||
|
||
app = Flask(__name__, static_folder="static", static_url_path="")
|
||
app.json.ensure_ascii = False
|
||
|
||
db.init_db()
|
||
|
||
RUNNERS = {} # test_id -> TestRunner
|
||
|
||
|
||
@app.route("/")
|
||
def index():
|
||
return send_from_directory(app.static_folder, "index.html")
|
||
|
||
|
||
@app.route("/api/health")
|
||
def health():
|
||
running = [tid for tid, r in RUNNERS.items() if r.is_alive()]
|
||
return jsonify({"ok": True, "port": config.PORT, "running_tests": running})
|
||
|
||
|
||
def _fill_defaults(cfg):
|
||
p = cfg.get("provider", "openai")
|
||
if not cfg.get("base_url"):
|
||
cfg["base_url"] = DEFAULT_URLS.get(p, "")
|
||
if cfg.get("temperature") is None:
|
||
cfg["temperature"] = 0.7
|
||
return cfg
|
||
|
||
|
||
# ───────────────────────── 提供商配置 ─────────────────────────
|
||
|
||
@app.route("/api/configs", methods=["GET"])
|
||
def list_configs():
|
||
return jsonify(db.list_configs())
|
||
|
||
|
||
@app.route("/api/configs", methods=["POST"])
|
||
def add_config():
|
||
cfg = request.get_json(force=True) or {}
|
||
if not cfg.get("name"):
|
||
return jsonify({"ok": False, "error": "请填写配置名称"}), 400
|
||
cid = db.add_config(cfg)
|
||
return jsonify({"ok": True, "id": cid})
|
||
|
||
|
||
@app.route("/api/configs/<int:cid>", methods=["GET"])
|
||
def get_one_config(cid):
|
||
c = db.get_config(cid)
|
||
if not c:
|
||
return jsonify({"ok": False, "error": "配置不存在"}), 404
|
||
return jsonify(c)
|
||
|
||
|
||
@app.route("/api/configs/<int:cid>", methods=["PUT"])
|
||
def update_config(cid):
|
||
cfg = request.get_json(force=True) or {}
|
||
old = db.get_config(cid)
|
||
if not old:
|
||
return jsonify({"ok": False, "error": "配置不存在"}), 404
|
||
db.update_config(cid, cfg)
|
||
return jsonify({"ok": True, "id": cid})
|
||
|
||
|
||
@app.route("/api/configs/<int:cid>", methods=["DELETE"])
|
||
def del_config(cid):
|
||
db.delete_config(cid)
|
||
return jsonify({"ok": True})
|
||
|
||
|
||
@app.route("/api/configs/test", methods=["POST"])
|
||
def test_config():
|
||
cfg = _fill_defaults(request.get_json(force=True) or {})
|
||
if not cfg.get("api_key"):
|
||
return jsonify({"ok": False, "error": "请填写 API Key"}), 400
|
||
try:
|
||
# 连接测试：只要流式请求成功返回（哪怕正文为空/只有思维链）都算连通
|
||
m = call_stream(cfg, "你好，请简要回答：1+1=？",
|
||
{"max_tokens": 32, "avoid_cache": False})
|
||
note = ""
|
||
if not (m.get("output_tokens") or m.get("output_chars")):
|
||
note = "（连接正常，但本次未返回正文内容，可能为推理型模型）"
|
||
return jsonify({"ok": True, "total_ms": m["total_ms"], "metrics": m, "note": note})
|
||
except ProviderError as e:
|
||
return jsonify({"ok": False, "error": str(e)})
|
||
except Exception as e:
|
||
return jsonify({"ok": False, "error": str(e)})
|
||
|
||
|
||
# ───────────────────────── 测试 ─────────────────────────
|
||
|
||
@app.route("/api/tests", methods=["POST"])
|
||
def start_test():
|
||
body = request.get_json(force=True) or {}
|
||
cfg = _fill_defaults(body.get("config") or {})
|
||
gen = body.get("gen") or {}
|
||
if not cfg.get("api_key"):
|
||
return jsonify({"ok": False, "error": "请填写 API Key"}), 400
|
||
if not cfg.get("model"):
|
||
return jsonify({"ok": False, "error": "请填写模型名称"}), 400
|
||
tid = db.create_test(cfg, gen)
|
||
runner = TestRunner(tid, cfg, gen)
|
||
RUNNERS[tid] = runner
|
||
runner.start()
|
||
return jsonify({"ok": True, "id": tid})
|
||
|
||
|
||
@app.route("/api/tests", methods=["GET"])
|
||
def list_tests():
|
||
try:
|
||
limit = int(request.args.get("limit", 100))
|
||
except ValueError:
|
||
limit = 100
|
||
return jsonify(db.list_tests(max(1, min(limit, 1000))))
|
||
|
||
|
||
@app.route("/api/tests/<int:tid>", methods=["GET"])
|
||
def get_test(tid):
|
||
t = db.get_test(tid)
|
||
if not t:
|
||
return jsonify({"ok": False, "error": "测试不存在"}), 404
|
||
t["runs"] = db.get_runs(tid)
|
||
t["logs"] = db.get_logs(tid)
|
||
_mask_cfg(t.get("config"))
|
||
return jsonify(t)
|
||
|
||
|
||
@app.route("/api/tests/<int:tid>/logs", methods=["GET"])
|
||
def get_logs(tid):
|
||
after = int(request.args.get("after", 0))
|
||
data = db.get_logs_after(tid, after)
|
||
if data is None:
|
||
return jsonify({"ok": False, "error": "测试不存在"}), 404
|
||
return jsonify(data)
|
||
|
||
|
||
@app.route("/api/tests/<int:tid>/cancel", methods=["POST"])
|
||
def cancel_test(tid):
|
||
r = RUNNERS.get(tid)
|
||
if r and r.is_alive():
|
||
r.request_cancel()
|
||
return jsonify({"ok": True, "msg": "正在停止..."})
|
||
return jsonify({"ok": False, "msg": "测试未在运行"})
|
||
|
||
|
||
@app.route("/api/tests/<int:tid>", methods=["DELETE"])
|
||
def del_test(tid):
|
||
db.delete_test(tid)
|
||
RUNNERS.pop(tid, None)
|
||
return jsonify({"ok": True})
|
||
|
||
|
||
@app.route("/api/tests/<int:tid>/export.json")
|
||
def export_json(tid):
|
||
t = db.get_test(tid)
|
||
if not t:
|
||
return jsonify({"ok": False, "error": "测试不存在"}), 404
|
||
t["runs"] = db.get_runs(tid)
|
||
t["logs"] = db.get_logs(tid)
|
||
_mask_cfg(t.get("config"))
|
||
return jsonify(t)
|
||
|
||
|
||
def _mask_cfg(cfg):
|
||
"""对外隐藏 API Key，仅保留前几位便于识别"""
|
||
if cfg and cfg.get("api_key"):
|
||
k = cfg["api_key"]
|
||
cfg["api_key"] = k[:4] + "****" if len(k) > 6 else "****"
|
||
return cfg
|
||
|
||
|
||
# ───────────────────────── Excel 导出 ─────────────────────────
|
||
|
||
@app.route("/api/tests/<int:tid>/export.xlsx")
|
||
def export_xlsx(tid):
|
||
t = db.get_test(tid)
|
||
if not t:
|
||
return jsonify({"ok": False, "error": "测试不存在"}), 404
|
||
t["runs"] = db.get_runs(tid)
|
||
t["logs"] = db.get_logs(tid)
|
||
try:
|
||
data = _build_xlsx(t)
|
||
except Exception as e:
|
||
return jsonify({"ok": False, "error": "导出失败: %s" % e}), 500
|
||
return send_file(data, as_attachment=True,
|
||
download_name="llm_speed_test_%d.xlsx" % tid,
|
||
mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
|
||
|
||
|
||
def _build_xlsx(t):
|
||
from openpyxl import Workbook
|
||
from openpyxl.styles import Alignment, Font, PatternFill
|
||
|
||
s = t.get("summary") or {}
|
||
g = t.get("gen") or {}
|
||
cfg = t.get("config") or {}
|
||
by_length = s.get("by_length") or {}
|
||
runs = t.get("runs") or []
|
||
logs = t.get("logs") or []
|
||
|
||
wb = Workbook()
|
||
head_fill = PatternFill("solid", fgColor="2A3550")
|
||
head_font = Font(color="FFFFFF", bold=True)
|
||
title_font = Font(bold=True, size=12)
|
||
|
||
def style_header(ws, row, ncol):
|
||
for c in range(1, ncol + 1):
|
||
cell = ws.cell(row=row, column=c)
|
||
cell.fill = head_fill
|
||
cell.font = head_font
|
||
cell.alignment = Alignment(horizontal="center", vertical="center")
|
||
|
||
# ── Sheet1 汇总 ──
|
||
ws = wb.active
|
||
ws.title = "汇总"
|
||
ws.append(["LLM 速度测试报告"])
|
||
ws.cell(1, 1).font = Font(bold=True, size=14)
|
||
info = [
|
||
["测试编号", "#%d" % t["id"]],
|
||
["测试名称", t.get("name") or "（未命名）"],
|
||
["创建时间", t.get("created_at", "")],
|
||
["状态", t.get("status", "")],
|
||
["提供商", t.get("provider", "")],
|
||
["模型", t.get("model", "")],
|
||
["Base URL", cfg.get("base_url") or "（默认）"],
|
||
["上下文长度列表", " / ".join(str(x) for x in (g.get("context_lengths") or []))],
|
||
["生成长度(max tokens)", g.get("max_tokens", 128)],
|
||
["每个长度采样次数", g.get("samples", 2)],
|
||
["预热(空转)", "开" if g.get("warmup", True) else "关"],
|
||
["避免缓存", "开" if g.get("avoid_cache") else "关"],
|
||
["采样(成功/总数)", "%s / %s" % (s.get("samples_ok"), s.get("samples_total"))],
|
||
["校准 字符/token", s.get("calibration_chars_per_token") or "—"],
|
||
["错误信息", t.get("error") or ""],
|
||
]
|
||
for row in info:
|
||
ws.append(row)
|
||
ws.cell(15, 1).font = title_font
|
||
r0 = len(info) + 2
|
||
overall = [
|
||
["首字延迟(ms)", s.get("avg_ttft_ms"), s.get("max_ttft_ms"), s.get("min_ttft_ms")],
|
||
["预填充速度(tok/s)", s.get("avg_prefill_speed"), s.get("max_prefill_speed"), s.get("min_prefill_speed")],
|
||
["解码速度(tok/s)", s.get("avg_decode_speed"), s.get("max_decode_speed"), s.get("min_decode_speed")],
|
||
["提示词(tok)", s.get("avg_prompt_tokens"), None, None],
|
||
["输出(tok)", s.get("avg_output_tokens"), None, None],
|
||
["总耗时(ms)", s.get("avg_total_ms"), s.get("max_total_ms"), s.get("min_total_ms")],
|
||
]
|
||
ws.cell(r0, 1, "整体统计指标（平均 / 最大 / 最小）").font = title_font
|
||
for j, c in enumerate(["指标", "平均", "最大", "最小"], start=1):
|
||
ws.cell(row=r0 + 1, column=j, value=c)
|
||
style_header(ws, r0 + 1, 4)
|
||
for i, row in enumerate(overall, start=r0 + 2):
|
||
for j, v in enumerate(row, start=1):
|
||
ws.cell(row=i, column=j, value=v)
|
||
|
||
# 按上下文长度分组
|
||
r1 = r0 + len(overall) + 3
|
||
ws.cell(r1, 1, "按上下文长度分组").font = title_font
|
||
cols = ["上下文长度(tok)", "采样(成功/总数)", "首字ms", "预填充tok/s", "解码tok/s", "提示词tok", "输出tok", "总耗时ms"]
|
||
ws.append([])
|
||
for j, c in enumerate(cols, start=1):
|
||
ws.cell(row=r1 + 1, column=j, value=c)
|
||
style_header(ws, r1 + 1, len(cols))
|
||
if by_length:
|
||
rr = r1 + 2
|
||
for L in sorted(int(k) for k in by_length):
|
||
bl = by_length[str(L)] if str(L) in by_length else by_length[L]
|
||
ws.cell(row=rr, column=1, value=L)
|
||
ws.cell(row=rr, column=2, value="%s / %s" % (bl.get("samples_ok"), bl.get("samples_total")))
|
||
ws.cell(row=rr, column=3, value=bl.get("avg_ttft_ms"))
|
||
ws.cell(row=rr, column=4, value=bl.get("avg_prefill_speed"))
|
||
ws.cell(row=rr, column=5, value=bl.get("avg_decode_speed"))
|
||
ws.cell(row=rr, column=6, value=bl.get("avg_prompt_tokens"))
|
||
ws.cell(row=rr, column=7, value=bl.get("avg_output_tokens"))
|
||
ws.cell(row=rr, column=8, value=bl.get("avg_total_ms"))
|
||
rr += 1
|
||
else:
|
||
ws.cell(row=r1 + 2, column=1, value="（无成功采样数据）")
|
||
for col, w in zip("ABCDEFGH", [22, 20, 12, 14, 14, 12, 12, 14]):
|
||
ws.column_dimensions[col].width = w
|
||
|
||
# ── Sheet2 采样明细 ──
|
||
ws2 = wb.create_sheet("采样明细")
|
||
h2 = ["序号", "上下文长度tok", "提示词tok", "缓存tok", "首字ms", "预填充tok/s",
|
||
"输出tok", "解码tok/s", "总耗时ms", "备注"]
|
||
ws2.append(h2)
|
||
style_header(ws2, 1, len(h2))
|
||
for i, r in enumerate(runs, start=1):
|
||
m = r.get("metrics") or {}
|
||
ws2.append([
|
||
i,
|
||
r.get("context_length") or m.get("context_length") or "",
|
||
m.get("prompt_tokens") or "",
|
||
m.get("cached_tokens") if m.get("cached_tokens") else "",
|
||
m.get("ttft_ms"),
|
||
m.get("prefill_speed"),
|
||
m.get("output_tokens"),
|
||
m.get("decode_speed"),
|
||
m.get("total_ms"),
|
||
r.get("error") or "OK",
|
||
])
|
||
for col, w in zip("ABCDEFGHIJ", [8, 14, 12, 10, 12, 14, 12, 14, 12, 30]):
|
||
ws2.column_dimensions[col].width = w
|
||
|
||
# ── Sheet3 日志 ──
|
||
ws3 = wb.create_sheet("日志")
|
||
ws3.append(["相对时间(s)", "级别", "内容"])
|
||
style_header(ws3, 1, 3)
|
||
for l in logs:
|
||
ws3.append([l.get("rel", 0), l.get("level", ""), l.get("msg", "")])
|
||
for col, w in zip("ABC", [14, 10, 90]):
|
||
ws3.column_dimensions[col].width = w
|
||
|
||
bio = io.BytesIO()
|
||
wb.save(bio)
|
||
bio.seek(0)
|
||
return bio
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app.run(host=config.HOST, port=config.PORT, threaded=True, debug=False)
|