refactor(bench): polish soup bench from PR #25

- Add -> None return type annotation (project convention)
- Replace broad 'except Exception' with specific exceptions
  (OSError, ImportError, RuntimeError, ValueError) + `raise ... from exc`
- Use `_` for unused response variable in tuple unpacking
- Add CPU warning: TPS on CPU is 10-100x slower, misleading users
- Add warmup run (discarded) to avoid CUDA JIT skewing averages
- Document that VRAM scope includes model load (deployment planning)
- Apply ruff style (trailing commas, en-dash -> ASCII, etc.)
This commit is contained in:
Alpamys 2026-04-15 22:06:45 +05:00
parent 3c339481d1
commit 8ea99d459f
1 changed files with 25 additions and 9 deletions

View File

@ -1,4 +1,4 @@
"""soup bench simple measuring tool for model speed and memory."""
"""soup bench -- simple measuring tool for model speed and memory."""
import time
from pathlib import Path
@ -34,7 +34,7 @@ def bench(
"-n",
help="Number of prompts to run for averaging",
),
):
) -> None:
"""Run an inference benchmark (speed and memory) on a loaded model."""
import torch
@ -48,6 +48,12 @@ def bench(
device, _ = detect_device()
if device == "cpu":
console.print(
"[yellow]Warning:[/] Running on CPU. Inference speed is typically "
"10-100x slower than GPU -- results will not reflect production TPS."
)
console.print(
Panel(
f"Model: [bold]{model_path}[/]\n"
@ -60,15 +66,17 @@ def bench(
console.print("[dim]Loading model to measure resource usage...[/]")
# Reset peak stats before load -- "Max VRAM" reflects total footprint
# (model load + inference), i.e. what users need for deployment planning.
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
start_load = time.time()
try:
model_obj, tokenizer = _load_model(str(model_path), base, device)
except Exception as exc:
except (OSError, ImportError, RuntimeError, ValueError) as exc:
console.print(f"[red]Failed to load model:[/] {exc}")
raise typer.Exit(1)
raise typer.Exit(1) from exc
load_time = time.time() - start_load
console.print(f"[green]Model loaded in {load_time:.2f}s.[/]\n")
@ -78,11 +86,19 @@ def bench(
"Write a short Python function to calculate fibonacci numbers.",
"What are the main consequences of the Industrial Revolution?",
"Compose a poem about a wandering space traveler.",
"Describe how a database index works under the hood."
"Describe how a database index works under the hood.",
]
# Loop over prompts
test_prompts = (prompts * (num_prompts // len(prompts) + 1))[:num_prompts]
# Warmup run: first inference includes CUDA kernel JIT compilation,
# which would skew the average. Discarded from timing.
console.print("[dim]Warmup run (discarded from timing)...[/]")
warmup_messages = [{"role": "user", "content": test_prompts[0]}]
_generate(
model_obj, tokenizer, warmup_messages,
max_tokens=min(max_tokens, 32), temperature=0.0,
)
total_tokens = 0
total_latency = 0.0
@ -92,7 +108,7 @@ def bench(
messages = [{"role": "user", "content": prompt_text}]
start_time = time.time()
response, token_count = _generate(
_, token_count = _generate(
model_obj, tokenizer, messages,
max_tokens=max_tokens, temperature=0.0,
)
@ -100,7 +116,7 @@ def bench(
latency = time.time() - start_time
total_tokens += token_count
total_latency += latency
console.print(f" [dim]Prompt {i+1}: {token_count} tokens in {latency:.2f}s[/]")
console.print(f" [dim]Prompt {i + 1}: {token_count} tokens in {latency:.2f}s[/]")
avg_tps = total_tokens / total_latency if total_latency > 0 else 0
@ -119,7 +135,7 @@ def bench(
"Transformers",
f"{avg_tps:.2f}",
f"{total_latency:.2f}s",
vram_str
vram_str,
)
console.print()