From 8ea99d459fcc8ae712a9f0149ef178528bb4c7af Mon Sep 17 00:00:00 2001
From: Alpamys <vpn.alpamys@gmail.com>
Date: Wed, 15 Apr 2026 22:06:45 +0500
Subject: [PATCH] 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.)
---
 soup_cli/commands/bench.py | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/soup_cli/commands/bench.py b/soup_cli/commands/bench.py
index 6da0d9a..57c83bd 100644
--- a/soup_cli/commands/bench.py
+++ b/soup_cli/commands/bench.py
@@ -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()