Spaces:
Sleeping
Sleeping
File size: 5,192 Bytes
72c62ed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """Direct SentenceTransformers baseline benchmark.
This benchmark intentionally avoids HTTP, FastAPI, and DynamicBatcher overhead.
It measures one direct model.encode([text]) call per input text.
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import Any
import torch
from sentence_transformers import SentenceTransformer
from utils import base_metadata, build_texts, format_ms, percentile, seconds_to_ms, write_json
DEFAULT_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
DEFAULT_REQUESTS = 40
DEFAULT_TEXTS_PER_REQUEST = 2
DEFAULT_WARMUP = 0
def encode_one(model: SentenceTransformer, text: str) -> float:
start = time.perf_counter()
model.encode([text])
return time.perf_counter() - start
def run_naive(
model: SentenceTransformer,
requests: int,
texts_per_request: int,
) -> tuple[list[float], float]:
texts = build_texts(texts_per_request)
latencies: list[float] = []
start = time.perf_counter()
for _ in range(requests):
for text in texts:
latencies.append(encode_one(model, text))
wall_time = time.perf_counter() - start
return latencies, wall_time
def summarize(
latencies: list[float],
wall_time_seconds: float,
metadata: dict[str, Any],
) -> dict[str, Any]:
total_sequences = len(latencies)
return {
"metadata": metadata,
"summary": {
"total_requests": metadata["requests"],
"total_sequences": total_sequences,
"wall_time_seconds": wall_time_seconds,
"throughput_sequences_per_sec": total_sequences / wall_time_seconds if wall_time_seconds > 0 else 0.0,
"avg_latency_ms": seconds_to_ms(sum(latencies) / total_sequences) if total_sequences else None,
"p50_latency_ms": seconds_to_ms(percentile(latencies, 50)),
"p95_latency_ms": seconds_to_ms(percentile(latencies, 95)),
"p99_latency_ms": seconds_to_ms(percentile(latencies, 99)),
},
}
def print_summary(report: dict[str, Any]) -> None:
metadata = report["metadata"]
summary = report["summary"]
print("\nNaive Benchmark")
print(f"Label: {metadata['label']}")
print(f"Model: {metadata['model_name']}")
print(f"Device: {metadata['device']}")
print(f"Requests: {metadata['requests']}")
print(f"Texts/request: {metadata['texts_per_request']}")
print(f"Warmup requests: {metadata['warmup']}")
print("\nResults")
print(f"Total sequences: {summary['total_sequences']}")
print(f"Wall time: {summary['wall_time_seconds']:.3f}s")
print(f"Sequence throughput: {summary['throughput_sequences_per_sec']:.2f} seq/s")
print(f"Avg latency: {format_ms(summary['avg_latency_ms'])}")
print(f"p50 latency: {format_ms(summary['p50_latency_ms'])}")
print(f"p95 latency: {format_ms(summary['p95_latency_ms'])}")
print(f"p99 latency: {format_ms(summary['p99_latency_ms'])}")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Benchmark naive direct SentenceTransformers encoding.")
parser.add_argument("--label", default="naive-direct", help="Human-readable run label stored in output metadata.")
parser.add_argument("--model-name", default=DEFAULT_MODEL_NAME, help=f"SentenceTransformer model. Default: {DEFAULT_MODEL_NAME}")
parser.add_argument("--device", default=DEFAULT_DEVICE, help=f"Torch device passed to SentenceTransformer. Default: {DEFAULT_DEVICE}")
parser.add_argument("--requests", type=int, default=DEFAULT_REQUESTS, help="Request-equivalent loop count.")
parser.add_argument("--texts-per-request", type=int, default=DEFAULT_TEXTS_PER_REQUEST, help="Texts per request-equivalent loop.")
parser.add_argument("--warmup", type=int, default=DEFAULT_WARMUP, help="Warmup request-equivalent loops excluded from results.")
parser.add_argument("--output", type=Path, help="Optional JSON output path.")
return parser.parse_args()
def validate_args(args: argparse.Namespace) -> None:
if args.requests < 1:
raise SystemExit("--requests must be >= 1")
if args.texts_per_request < 1:
raise SystemExit("--texts-per-request must be >= 1")
if args.warmup < 0:
raise SystemExit("--warmup must be >= 0")
def main() -> None:
args = parse_args()
validate_args(args)
model = SentenceTransformer(args.model_name, device=args.device)
if args.warmup:
run_naive(model, args.warmup, args.texts_per_request)
latencies, wall_time = run_naive(model, args.requests, args.texts_per_request)
metadata = {
**base_metadata(),
"label": args.label,
"model_name": args.model_name,
"device": args.device,
"requests": args.requests,
"texts_per_request": args.texts_per_request,
"warmup": args.warmup,
}
report = summarize(latencies, wall_time, metadata)
print_summary(report)
if args.output:
write_json(report, args.output)
print(f"\nWrote JSON results to {args.output}")
if __name__ == "__main__":
main()
|