sa-bench: update sa bench serving to upstream InferenceMAX code

[cquil11]

• updated bench_serving.py to be consistent with InferenceMAX/InferenceMAX as of 6eecc0e660c7a695240732b35a2bc2c8f1210ccd (see here)
  • removes unused code and imports
  • adds backed in warmup code
• updates bench.sh to remove separate warmup

Note: I tried to remove superfluous diffs (whitespace, newline, ie any non functional changes) but couldn't get all of them

Summary by CodeRabbit

• New Features

  • Added --num-warmups parameter to specify the number of warmup prompts during benchmarking.

• Refactor

  • Warmup configuration is now integrated into the main benchmark invocation, eliminating the need for separate pre-run steps and streamlining the workflow.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

---

📝 Walkthrough

Walkthrough

Consolidates warmup handling by removing the separate pre-run warmup invocation in the shell script and adding a built-in warmup phase to the main benchmarking tool via a new --num-warmups CLI flag and corresponding num_warmups parameter in benchmark().

Changes

Cohort / File(s)	Summary
Bench script src/srtctl/benchmarks/scripts/sa-bench/bench.sh	Removed the separate Python pre-run that executed warmups; now passes --num-warmups "$num_warmup_prompts" to the primary benchmark_serving.py invocation.
Benchmark runtime + CLI src/srtctl/benchmarks/scripts/sa-bench/benchmark_serving.py	Added num_warmups: int parameter to benchmark() and a --num-warmups CLI option; implements a warmup phase that issues configured warmup requests under concurrency control (semaphore), optionally shows a tqdm progress bar, and prints completion before the main benchmark run. Public function signature and CLI wiring updated accordingly.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• Fix benchmark order #95: Edits the same bench.sh warmup flow to remove a separate warmup run and consolidate warmups into the main benchmark invocation.

Poem

🐰 Hopping prompts in morning light,
I warm my circuits, snug and bright.
One call now wakes the testing crew,
Seeded runs and progress too — woohoo! 🥕✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main objective of the pull request: updating sa-bench serving code to align with upstream InferenceMAX code, including warmup improvements.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

src/srtctl/benchmarks/scripts/sa-bench/benchmark_serving.py (2)

109-118: ⚠️ Potential issue | 🟡 Minor

Guard against negative input lengths when applying chat templates.

If input_len is smaller than the chat template overhead, this will yield negative lengths and crash np.random.randint.

🛡️ Suggested guard

         tokenized_chat_template_dummy = tokenizer.encode(chat_template_dummy, add_special_tokens=False)
         chat_template_len = len(tokenized_chat_template_dummy) - 1
-        input_len = input_len - chat_template_len
+        if input_len <= chat_template_len:
+            raise ValueError("random_input_len must exceed chat template length.")
+        input_len = input_len - chat_template_len

---

626-684: ⚠️ Potential issue | 🟡 Minor

Validate request_rate at startup to prevent divide-by-zero.

request_rate == 0 causes a ZeroDivisionError when calculating theta = 1.0 / (request_rate * burstiness) in get_request(). The code already safely handles max_concurrency == 0 via truthiness checks, so validation is only needed for request_rate.

✅ Suggested validation

 def main(args: argparse.Namespace):
     print(args)
     random.seed(args.seed)
     np.random.seed(args.seed)
+    if args.request_rate != float("inf") and args.request_rate <= 0:
+        raise ValueError("--request-rate must be positive or 'inf'.")

🧹 Nitpick comments (2)

src/srtctl/benchmarks/scripts/sa-bench/benchmark_serving.py (2)

98-160: Fix request tuple typing and add missing type hints to keep ty happy.

sample_random_requests now returns a 4‑tuple (extra None), but the return type and downstream annotations still imply a 3‑tuple. Also, sample_uniform lacks type hints. This will likely fail ty/Ruff typing and makes the API inconsistent.

✅ Suggested typing update

-def sample_random_requests(
+def sample_random_requests(
     prefix_len: int,
     input_len: int,
     output_len: int,
     num_prompts: int,
     range_ratio: float,
     tokenizer: PreTrainedTokenizerBase,
     use_chat_template: bool = False,
-) -> list[tuple[str, int, int]]:
+) -> list[tuple[str, int, int, Any | None]]:
@@
-    def sample_uniform(seq_len):
+    def sample_uniform(seq_len: int) -> list[int]:
         lower = int(seq_len * range_ratio)
         upper = seq_len
         seq_lens = np.random.randint(lower, upper + 1, size=num_prompts).tolist()
         return seq_lens

Please also align input_requests type hints in get_request(...) and benchmark(...) to match the 4‑tuple shape.

As per coding guidelines, include type hints everywhere and use ty for type checking.

Also applies to: 181-185, 322-330

---

1052-1052: Add help text for --num-warmups to match other CLI flags.

This improves discoverability and makes CLI usage consistent with surrounding options.

✍️ Suggested tweak

-    parser.add_argument("--num-warmups", type=int, default=0)
+    parser.add_argument(
+        "--num-warmups",
+        type=int,
+        default=0,
+        help="Number of warmup requests to run before the main benchmark.",
+    )