464 lines
18 KiB
Python
464 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
压测每个 OCR 端点(IP)的实际 QPS，并把结果写成权重文件 config/ocr_weights.json，
|
||
供运行时按 QPS 加权分发请求(见 tool/Ocr/tencentcloud_ocr.py 的 OcrDispatcher)。
|
||
|
||
做法：对每个端点用固定并发、固定请求数发同一张样例图，
|
||
QPS = 成功请求数 / 墙钟耗时；同时报告平均时延和错误率。
|
||
|
||
本脚本完全自包含：单文件，不读 config、不 import 项目任何模块，拷过去就能跑。
|
||
|
||
用法:
|
||
python bench_ocr_endpoints.py [样例图片] [--requests N] [--concurrency C] [--no-write]
|
||
python bench_ocr_endpoints.py [样例图片] --url http://IP:PORT [--url http://IP2:PORT ...]
|
||
|
||
样例图片不给时，自动在 ./composite、./split 下找一张 .jpg。
|
||
端点：默认内置 Polaris 地址；要测别的 IP 用 --url(可多次)。
|
||
--requests 每个端点发的总请求数(默认 60)
|
||
--concurrency 每个端点的并发线程数(默认 10)
|
||
--no-write 只打印，不写 ocr_weights.json
|
||
"""
|
||
import argparse
|
||
import base64
|
||
import glob
|
||
import itertools
|
||
import json
|
||
import os
|
||
import random
|
||
import threading
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
import requests
|
||
|
||
OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"
|
||
OCR_PORT = 18081
|
||
|
||
# 内置的多台 OCR 服务 IP（自动负载均衡测试用）
|
||
DEFAULT_IPS = [
|
||
"30.154.78.83"
|
||
]
|
||
DEFAULT_ENDPOINTS = [f"http://{ip}:{OCR_PORT}" for ip in DEFAULT_IPS]
|
||
WEIGHTS_FILE = "./ocr_weights.json"
|
||
|
||
|
||
def get_endpoints(cli_urls):
|
||
"""命令行 --url 优先，否则用内置默认端点。"""
|
||
if cli_urls:
|
||
return [u.rstrip("/") for u in cli_urls]
|
||
return list(DEFAULT_ENDPOINTS)
|
||
|
||
|
||
def find_sample_image(given):
|
||
if given:
|
||
if os.path.exists(given):
|
||
return given
|
||
print(f"指定的图片不存在: {given}")
|
||
return None
|
||
for root in ("./composite", "./split"):
|
||
hits = glob.glob(os.path.join(root, "**", "*.jpg"), recursive=True)
|
||
if hits:
|
||
return hits[0]
|
||
return None
|
||
|
||
|
||
def img2base64(img_path):
|
||
with open(img_path, "rb") as f:
|
||
return base64.b64encode(f.read()).decode()
|
||
|
||
|
||
def ocr_request_once(base_url, image_b64, timeout=60):
|
||
"""向指定端点发一次 OCR 请求，返回原始响应文本；失败抛异常。"""
|
||
url = base_url.rstrip("/") + OCR_PATH
|
||
app_id = f"{int(time.time() * 1000)}{random.randint(100000, 999999)}"
|
||
payload = {
|
||
"app_id": app_id,
|
||
"image": image_b64,
|
||
"ocr_template": "OCR",
|
||
"options": {"pdf_page_index": 0, "detect_text": True},
|
||
}
|
||
resp = requests.post(url, headers={"Content-Type": "application/json"},
|
||
json=payload, timeout=timeout)
|
||
resp.raise_for_status()
|
||
return resp.text
|
||
|
||
|
||
def _ok(raw):
|
||
"""判断一次响应是否算成功：有内容且不是 error。"""
|
||
if not raw:
|
||
return False
|
||
try:
|
||
data = json.loads(raw)
|
||
except (json.JSONDecodeError, TypeError):
|
||
return False
|
||
if "TextDetections" in data:
|
||
return True
|
||
return data.get("error_code", 0) == 0
|
||
|
||
|
||
def bench_endpoint(base_url, image_b64, requests_n, concurrency):
|
||
"""对单个端点压测，返回统计 dict（含一个失败原因采样 sample_err）。"""
|
||
latencies = []
|
||
ok = 0
|
||
err = 0
|
||
sample_err = [None] # 采样第一个失败原因
|
||
|
||
def one_call():
|
||
t0 = time.time()
|
||
try:
|
||
raw = ocr_request_once(base_url, image_b64, timeout=60)
|
||
dt = time.time() - t0
|
||
if _ok(raw):
|
||
return (True, dt)
|
||
# HTTP 200 但内容是错误：采样响应体片段
|
||
if sample_err[0] is None:
|
||
sample_err[0] = f"响应非成功: {raw[:200]}"
|
||
return (False, dt)
|
||
except Exception as e:
|
||
if sample_err[0] is None:
|
||
sample_err[0] = f"{type(e).__name__}: {e}"
|
||
return (False, time.time() - t0)
|
||
|
||
t_start = time.time()
|
||
with ThreadPoolExecutor(max_workers=concurrency) as ex:
|
||
futures = [ex.submit(one_call) for _ in range(requests_n)]
|
||
for fut in as_completed(futures):
|
||
success, dt = fut.result()
|
||
latencies.append(dt)
|
||
if success:
|
||
ok += 1
|
||
else:
|
||
err += 1
|
||
wall = time.time() - t_start
|
||
|
||
return {
|
||
"qps": ok / wall if wall > 0 else 0.0,
|
||
"ok": ok,
|
||
"err": err,
|
||
"wall": wall,
|
||
"avg_latency": sum(latencies) / len(latencies) if latencies else 0.0,
|
||
"err_rate": err / requests_n if requests_n else 0.0,
|
||
"sample_err": sample_err[0],
|
||
}
|
||
|
||
|
||
def soak_endpoint(base_url, image_pool, concurrency, report_interval=5):
|
||
"""对单个端点持续满速压测：固定 concurrency 个线程不停发请求，
|
||
每次从 image_pool 随机抽一张图(避免被服务端缓存)，每 report_interval 秒打印滚动 QPS，
|
||
直到 Ctrl+C。image_pool 可以是单张或几千张。"""
|
||
stop = threading.Event()
|
||
lock = threading.Lock()
|
||
st = {"ok": 0, "err": 0, "lat": 0.0, "n": 0}
|
||
sample_err = [None]
|
||
|
||
def worker():
|
||
while not stop.is_set():
|
||
img = random.choice(image_pool) # 每次随机抽一张
|
||
t0 = time.time()
|
||
try:
|
||
raw = ocr_request_once(base_url, img, timeout=60)
|
||
success = _ok(raw)
|
||
if not success and sample_err[0] is None:
|
||
sample_err[0] = f"响应非成功: {raw[:160]}"
|
||
except Exception as e:
|
||
success = False
|
||
if sample_err[0] is None:
|
||
sample_err[0] = f"{type(e).__name__}: {e}"
|
||
dt = time.time() - t0
|
||
with lock:
|
||
st["n"] += 1
|
||
st["lat"] += dt
|
||
if success:
|
||
st["ok"] += 1
|
||
else:
|
||
st["err"] += 1
|
||
|
||
threads = [threading.Thread(target=worker, daemon=True) for _ in range(concurrency)]
|
||
for t in threads:
|
||
t.start()
|
||
|
||
print(f"持续满速压测 {base_url} 并发={concurrency} 每{report_interval}s 报一次 (Ctrl+C 停止)")
|
||
print(f"{'用时s':>6} {'累计OK':>9} {'瞬时QPS':>9} {'平均QPS':>9} {'时延ms':>8} {'错误率':>7}")
|
||
print("-" * 56)
|
||
t_start = time.time()
|
||
last_ok, last_t = 0, t_start
|
||
try:
|
||
while True:
|
||
time.sleep(report_interval)
|
||
now = time.time()
|
||
with lock:
|
||
ok, err, n, lat = st["ok"], st["err"], st["n"], st["lat"]
|
||
inst = (ok - last_ok) / (now - last_t) if now > last_t else 0
|
||
avg = ok / (now - t_start) if now > t_start else 0
|
||
avg_lat = (lat / n * 1000) if n else 0
|
||
er = err / n * 100 if n else 0
|
||
print(f"{now-t_start:6.0f} {ok:>9} {inst:>9.1f} {avg:>9.1f} {avg_lat:>8.0f} {er:>6.0f}%", flush=True)
|
||
last_ok, last_t = ok, now
|
||
except KeyboardInterrupt:
|
||
stop.set()
|
||
el = time.time() - t_start
|
||
with lock:
|
||
ok, n = st["ok"], st["n"]
|
||
print("-" * 56)
|
||
print(f"停止。{el:.0f}s 内成功 {ok} / 共发 {n}，平均 {ok/el if el else 0:.1f} QPS")
|
||
if sample_err[0]:
|
||
print(f"失败采样: {sample_err[0]}")
|
||
|
||
|
||
def health_check(endpoints, image_b64):
|
||
"""给每个端点发一次请求，返回 (存活列表, 死亡列表)。"""
|
||
alive, dead = [], []
|
||
|
||
def _check(ep):
|
||
try:
|
||
return ep, _ok(ocr_request_once(ep, image_b64, timeout=15))
|
||
except Exception:
|
||
return ep, False
|
||
|
||
with ThreadPoolExecutor(max_workers=min(32, len(endpoints))) as ex:
|
||
for ep, ok in ex.map(_check, endpoints):
|
||
(alive if ok else dead).append(ep)
|
||
return alive, dead
|
||
|
||
|
||
def bench_aggregate(endpoints, image_b64, requests_n, concurrency):
|
||
"""把请求轮询分发到所有端点（模拟运行时自动负载均衡），测合并 QPS。"""
|
||
rr = itertools.cycle(endpoints)
|
||
rr_lock = threading.Lock()
|
||
latencies = []
|
||
ok = 0
|
||
err = 0
|
||
stat_lock = threading.Lock()
|
||
|
||
def next_ep():
|
||
with rr_lock:
|
||
return next(rr)
|
||
|
||
def one_call(_):
|
||
ep = next_ep()
|
||
t0 = time.time()
|
||
try:
|
||
raw = ocr_request_once(ep, image_b64, timeout=60)
|
||
success = _ok(raw)
|
||
except Exception:
|
||
success = False
|
||
dt = time.time() - t0
|
||
with stat_lock:
|
||
latencies.append(dt)
|
||
return success
|
||
|
||
t_start = time.time()
|
||
with ThreadPoolExecutor(max_workers=concurrency) as ex:
|
||
for success in ex.map(one_call, range(requests_n)):
|
||
if success:
|
||
ok += 1
|
||
else:
|
||
err += 1
|
||
wall = time.time() - t_start
|
||
|
||
return {
|
||
"qps": ok / wall if wall > 0 else 0.0,
|
||
"ok": ok,
|
||
"err": err,
|
||
"wall": wall,
|
||
"avg_latency": sum(latencies) / len(latencies) if latencies else 0.0,
|
||
"err_rate": err / requests_n if requests_n else 0.0,
|
||
}
|
||
|
||
|
||
def run_aggregate(endpoints, image_b64, requests_n, concurrency):
|
||
"""先健康检查筛掉死端点，再在存活端点上做聚合负载测试。"""
|
||
print(f"健康检查 {len(endpoints)} 个端点 ...", flush=True)
|
||
alive, dead = health_check(endpoints, image_b64)
|
||
print(f" 存活: {len(alive)} 死亡: {len(dead)}")
|
||
if dead:
|
||
for ep in dead:
|
||
print(f" ✗ {ep}")
|
||
if not alive:
|
||
print("没有存活端点，无法测试。")
|
||
return
|
||
print("=" * 64)
|
||
print(f"聚合负载测试: {len(alive)} 个存活端点轮询分发，"
|
||
f"总请求 {requests_n}，总并发 {concurrency}")
|
||
s = bench_aggregate(alive, image_b64, requests_n, concurrency)
|
||
print("-" * 64)
|
||
print(f" 合并 QPS = {s['qps']:.2f} 成功={s['ok']} 失败={s['err']}"
|
||
f" 平均时延={s['avg_latency']*1000:.0f}ms 错误率={s['err_rate']*100:.0f}%")
|
||
if alive:
|
||
print(f" 单端点平均 ≈ {s['qps']/len(alive):.2f} QPS"
|
||
f" （{len(alive)} 台合计 {s['qps']:.0f} QPS）")
|
||
print("=" * 64)
|
||
|
||
|
||
def sweep_endpoint(base_url, image_b64, start, max_c, per_worker):
|
||
"""自动加压找极限 QPS：并发从 start 逐级翻倍到 max_c，
|
||
QPS 不再随并发上升(或错误率升高)即判定到极限，停止。"""
|
||
print(f"自动加压 {base_url} (并发 {start}→{max_c} 逐级翻倍)")
|
||
print(f"{'并发':>6} {'QPS':>9} {'时延ms':>9} {'错误率':>8}")
|
||
print("-" * 36)
|
||
best_qps, best_c = 0.0, 0
|
||
prev_qps = 0.0
|
||
c = start
|
||
while c <= max_c:
|
||
reqs = max(c * per_worker, 30) # 每个并发至少发 per_worker 次，样本够稳
|
||
s = bench_endpoint(base_url, image_b64, reqs, c)
|
||
print(f"{c:>6} {s['qps']:>9.2f} {s['avg_latency']*1000:>9.0f} {s['err_rate']*100:>7.0f}%")
|
||
if s["qps"] > best_qps:
|
||
best_qps, best_c = s["qps"], c
|
||
if s["err_rate"] > 0.1:
|
||
print(" → 错误率升高，已超过服务承载，停止。")
|
||
break
|
||
if prev_qps > 0 and s["qps"] < prev_qps * 1.08:
|
||
print(" → QPS 不再随并发上升，已到极限，停止。")
|
||
break
|
||
prev_qps = s["qps"]
|
||
c *= 2
|
||
print("-" * 36)
|
||
print(f"极限 QPS ≈ {best_qps:.2f} (最佳并发 {best_c})")
|
||
return best_qps, best_c
|
||
|
||
|
||
def gen_random_images(n, size=(1280, 200)):
|
||
"""生成 n 张带随机文字的图片(让 OCR 真干活)，返回 base64 列表。需要 Pillow。"""
|
||
try:
|
||
import io
|
||
import string
|
||
from PIL import Image, ImageDraw
|
||
except ImportError:
|
||
print("生成随机图需要 Pillow(pip install pillow)；改用 --pool 指向真实图目录")
|
||
raise
|
||
chars = string.ascii_letters + string.digits + " ,.!?-"
|
||
pool = []
|
||
print(f"生成 {n} 张随机文字图...", flush=True)
|
||
for i in range(n):
|
||
bg = random.randint(200, 255)
|
||
img = Image.new("RGB", size, (bg, bg, bg))
|
||
d = ImageDraw.Draw(img)
|
||
for row in range(random.randint(1, 3)):
|
||
text = "".join(random.choice(chars) for _ in range(random.randint(10, 40)))
|
||
d.text((20, 20 + row * 50), text, fill=(0, 0, 0))
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="JPEG", quality=85)
|
||
pool.append(base64.b64encode(buf.getvalue()).decode())
|
||
if (i + 1) % 1000 == 0:
|
||
print(f" 已生成 {i+1}/{n}", flush=True)
|
||
return pool
|
||
|
||
|
||
def build_image_pool(args):
|
||
"""构建测试图池(base64 列表)：
|
||
--pool DIR 用真实图目录(最贴近实际负载) > --gen N 随机生成 > 单张样例图。"""
|
||
if getattr(args, "pool", None):
|
||
files = (glob.glob(os.path.join(args.pool, "**", "*.jpg"), recursive=True) +
|
||
glob.glob(os.path.join(args.pool, "**", "*.png"), recursive=True))
|
||
if not files:
|
||
print(f"目录 {args.pool} 下没有图片")
|
||
return []
|
||
random.shuffle(files)
|
||
cap = args.gen if args.gen else 5000
|
||
files = files[:cap]
|
||
print(f"从 {args.pool} 取 {len(files)} 张真实图作为图池")
|
||
return [img2base64(f) for f in files]
|
||
if getattr(args, "gen", 0):
|
||
return gen_random_images(args.gen)
|
||
sample = find_sample_image(args.image)
|
||
if not sample:
|
||
return []
|
||
print(f"单张样例图: {sample}")
|
||
return [img2base64(sample)]
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("image", nargs="?", default=None)
|
||
ap.add_argument("--url", action="append", default=None, help="OCR 端点(可多次)")
|
||
ap.add_argument("--requests", type=int, default=60)
|
||
ap.add_argument("--concurrency", type=int, default=10)
|
||
ap.add_argument("--aggregate", action="store_true",
|
||
help="聚合负载测试:请求轮询分发到所有存活端点，测合并总 QPS")
|
||
ap.add_argument("--sweep", action="store_true",
|
||
help="自动加压找极限 QPS(并发逐级翻倍直到 QPS 不再上升)")
|
||
ap.add_argument("--soak", action="store_true",
|
||
help="单端点持续满速压测(固定并发不停发，实时滚动QPS，Ctrl+C 停)")
|
||
ap.add_argument("--interval", type=int, default=5, help="soak 模式打印间隔秒")
|
||
ap.add_argument("--gen", type=int, default=0,
|
||
help="随机生成 N 张带文字的图作为图池(每次请求随机抽，避免被缓存)")
|
||
ap.add_argument("--pool", default=None,
|
||
help="用该目录下的真实图作为图池(最贴近实际负载，如 ./split/某视频)")
|
||
ap.add_argument("--start-concurrency", type=int, default=10, help="sweep 起始并发")
|
||
ap.add_argument("--max-concurrency", type=int, default=320, help="sweep 最大并发")
|
||
ap.add_argument("--per-worker", type=int, default=6, help="sweep 每档每并发发多少次请求")
|
||
ap.add_argument("--no-write", action="store_true")
|
||
args = ap.parse_args()
|
||
|
||
endpoints = get_endpoints(args.url)
|
||
if not endpoints:
|
||
print("没有可用的 OCR 端点")
|
||
return
|
||
|
||
# 构建图池：--pool 真实图目录 / --gen 随机生成 / 默认单张样例
|
||
image_pool = build_image_pool(args)
|
||
if not image_pool:
|
||
print("没有可用图片：给一张 jpg 路径、或用 --gen N / --pool 目录。")
|
||
return
|
||
print(f"图池大小: {len(image_pool)} 张 | 端点数: {len(endpoints)} 并发: {args.concurrency}")
|
||
print("=" * 64)
|
||
|
||
image_b64 = image_pool[0] # 非 soak 模式用单张(测纯吞吐)
|
||
|
||
# ---- soak 模式：单端点持续满速压测，每次随机抽图，直到 Ctrl+C ----
|
||
if args.soak:
|
||
soak_endpoint(endpoints[0], image_pool, args.concurrency, args.interval)
|
||
return
|
||
|
||
# ---- aggregate 模式：所有端点轮询分发，测合并总 QPS ----
|
||
if args.aggregate:
|
||
run_aggregate(endpoints, image_b64, args.requests, args.concurrency)
|
||
return
|
||
|
||
# ---- sweep 模式：自动加压找极限，不写权重文件 ----
|
||
if args.sweep:
|
||
for i, ep in enumerate(endpoints, 1):
|
||
print(f"\n[{i}/{len(endpoints)}]", end=" ")
|
||
sweep_endpoint(ep, image_b64, args.start_concurrency,
|
||
args.max_concurrency, args.per_worker)
|
||
return
|
||
|
||
results = {}
|
||
for i, ep in enumerate(endpoints, 1):
|
||
print(f"[{i}/{len(endpoints)}] 压测 {ep} ...", flush=True)
|
||
stat = bench_endpoint(ep, image_b64, args.requests, args.concurrency)
|
||
results[ep] = stat
|
||
print(f" QPS={stat['qps']:.2f} 成功={stat['ok']} 失败={stat['err']}"
|
||
f" 平均时延={stat['avg_latency']*1000:.0f}ms"
|
||
f" 错误率={stat['err_rate']*100:.0f}%")
|
||
if stat["err"] and stat.get("sample_err"):
|
||
print(f" 失败原因(采样): {stat['sample_err']}")
|
||
|
||
print("=" * 64)
|
||
print("汇总(按 QPS 降序):")
|
||
for ep, s in sorted(results.items(), key=lambda kv: kv[1]["qps"], reverse=True):
|
||
print(f" {ep:48s} QPS={s['qps']:6.2f} 时延={s['avg_latency']*1000:5.0f}ms"
|
||
f" 错误率={s['err_rate']*100:.0f}%")
|
||
|
||
total_qps = sum(s["qps"] for s in results.values())
|
||
print(f"\n合计可用 QPS ≈ {total_qps:.2f}")
|
||
|
||
# 写权重文件：QPS 即权重；错误率过高(>50%)的端点权重清零(不参与分发)
|
||
weights = {ep: (round(s["qps"], 3) if s["err_rate"] <= 0.5 else 0.0)
|
||
for ep, s in results.items()}
|
||
|
||
if args.no_write:
|
||
print("\n--no-write：不写权重文件。")
|
||
return
|
||
|
||
os.makedirs(os.path.dirname(WEIGHTS_FILE), exist_ok=True)
|
||
with open(WEIGHTS_FILE, "w", encoding="utf-8") as f:
|
||
json.dump(weights, f, ensure_ascii=False, indent=2)
|
||
print(f"\n已写入权重文件: {WEIGHTS_FILE}")
|
||
print("运行时会按此权重把 OCR 请求加权分发到各端点（重启进程或新批次生效）。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|