scripts/benchmark-batch.py - master - blanchefort/embedder_service

scripts/benchmark-batch.py
129 строк5 KB
Igor Shatalin
подготовка к проду
29 июл 2026, 20:41
29 июл 2026, 20:417875e8c
#!/usr/bin/env python3
"""Send repeatable batches of exactly 20 texts to the embedding endpoint."""

import argparse
import json
import math
import os
from pathlib import Path
import sys
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


SAMPLE_COUNT = 20
DEFAULT_URL = "http://127.0.0.1:8200/api/v1/embeddings/text/batch"


def positive_int(value: str) -> int:
    parsed = int(value)
    if parsed < 1:
        raise argparse.ArgumentTypeError("must be at least 1")
    return parsed


def non_negative_int(value: str) -> int:
    parsed = int(value)
    if parsed < 0:
        raise argparse.ArgumentTypeError("must be at least 0")
    return parsed


def positive_float(value: str) -> float:
    parsed = float(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("must be greater than 0")
    return parsed


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Benchmark embedding batches containing exactly 20 samples.")
    parser.add_argument("--url", default=DEFAULT_URL, help="Full batch endpoint URL.")
    parser.add_argument(
        "--token",
        default=os.environ.get("SECRET_AUTH_TOKEN"),
        help="Bearer token. Defaults to SECRET_AUTH_TOKEN; never printed.",
    )
    parser.add_argument("--requests", type=positive_int, default=10, help="Measured request count (default: 10).")
    parser.add_argument("--warmup", type=non_negative_int, default=1, help="Unmeasured warmup requests (default: 1).")
    parser.add_argument("--timeout", type=positive_float, default=120.0, help="Per-request timeout in seconds.")
    parser.add_argument(
        "--samples-file",
        type=Path,
        help="Optional UTF-8 JSON file containing exactly 20 non-empty strings.",
    )
    parser.add_argument(
        "--emb-type",
        choices=("classification", "clustering", "search_query", "search_document"),
        default="clustering",
    )
    args = parser.parse_args()
    if not args.token:
        parser.error("--token or SECRET_AUTH_TOKEN is required")
    return args


def load_samples(path: Path | None) -> list[str]:
    if path is None:
        return [f"Тестовый текст для CPU batch benchmark, образец {index + 1}" for index in range(SAMPLE_COUNT)]
    data: Any = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, list) or len(data) != SAMPLE_COUNT:
        raise ValueError(f"samples file must contain a JSON array of exactly {SAMPLE_COUNT} strings")
    if not all(isinstance(item, str) and item.strip() for item in data):
        raise ValueError("every benchmark sample must be a non-empty string")
    return data


def send_batch(url: str, token: str, samples: list[str], emb_type: str, timeout: float) -> float:
    payload = json.dumps({"texts": samples, "emb_type": emb_type}, ensure_ascii=False).encode()
    request = Request(
        url,
        data=payload,
        method="POST",
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        },
    )
    started = time.perf_counter()
    with urlopen(request, timeout=timeout) as response:
        body: Any = json.load(response)
    elapsed_ms = (time.perf_counter() - started) * 1000
    embeddings = body.get("embeddings") if isinstance(body, dict) else None
    received = len(embeddings) if isinstance(embeddings, list) else 0
    if received != SAMPLE_COUNT:
        raise ValueError(f"expected {SAMPLE_COUNT} embeddings, received {received}")
    return elapsed_ms


def percentile(values: list[float], rank: float) -> float:
    ordered = sorted(values)
    index = max(0, math.ceil((rank / 100) * len(ordered)) - 1)
    return ordered[index]


def main() -> int:
    args = parse_args()
    try:
        samples = load_samples(args.samples_file)
        for _ in range(args.warmup):
            send_batch(args.url, args.token, samples, args.emb_type, args.timeout)
        latencies = [
            send_batch(args.url, args.token, samples, args.emb_type, args.timeout) for _ in range(args.requests)
        ]
    except (HTTPError, URLError, TimeoutError, json.JSONDecodeError, OSError, ValueError) as exc:
        print(f"benchmark failed: {exc}", file=sys.stderr)
        return 1

    print(f"requests={len(latencies)}")
    print(f"samples_per_request={SAMPLE_COUNT}")
    print(f"success_rate={len(latencies) / args.requests:.2%}")
    print(f"p50_ms={percentile(latencies, 50):.2f}")
    print(f"p95_ms={percentile(latencies, 95):.2f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())