290 lines
10 KiB
Python
290 lines
10 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=["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, "你好，请只回复：OK", {"max_tokens": 16, "avoid_cache": False})
|
|
return jsonify({"ok": True, "total_ms": m["total_ms"], "metrics": m})
|
|
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():
|
|
return jsonify(db.list_tests())
|
|
|
|
|
|
@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)
|
|
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})
|
|
|
|
|
|
# ───────────────────────── 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("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(14, 1).font = title_font
|
|
r0 = len(info) + 2
|
|
overall = [
|
|
["平均首字延迟(ms)", s.get("avg_ttft_ms")],
|
|
["最佳首字延迟(ms)", s.get("best_ttft_ms")],
|
|
["平均预填充速度(tok/s)", s.get("avg_prefill_speed")],
|
|
["平均解码速度(tok/s)", s.get("avg_decode_speed")],
|
|
["平均提示词(tok)", s.get("avg_prompt_tokens")],
|
|
["平均输出(tok)", s.get("avg_output_tokens")],
|
|
["平均总耗时(ms)", s.get("avg_total_ms")],
|
|
]
|
|
ws.cell(r0, 1, "整体平均指标").font = title_font
|
|
for i, row in enumerate(overall, start=r0 + 1):
|
|
ws.append([])
|
|
for j, v in enumerate(row, start=1):
|
|
ws.cell(row=i, column=j, value=v)
|
|
|
|
# 按上下文长度分组
|
|
r1 = r0 + len(overall) + 2
|
|
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)
|