mirror of https://github.com/razor-ai/soup.git
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:
parent
3c339481d1
commit
8ea99d459f
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue