[Bugfix] Correct prompt lengths for timed_traces benchmark (#45423)

Signed-off-by: Stanislaw Wozniak <stw@zurich.ibm.com>
Co-authored-by: Thomas Parnell <tpa@zurich.ibm.com>

Author: s3woz

vllm/benchmarks/datasets/datasets.py

@@ -79,7 +79,7 @@ class SampleRequest:
     Represents a single inference request for benchmarking.
     """
 
-    prompt: str | list[str] | list[dict]
+    prompt: str | list[str] | list[int] | list[dict]
     prompt_len: int
     expected_output_len: int = 0
     multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
@@ -1572,7 +1572,6 @@ def sample(
             prompt_ids = self._expand_prompt(
                 entry.get(self.label_hash_ids, []), input_length, tokenizer
             )
-            prompt = tokenizer.decode(prompt_ids)
 
             # Get timestamp with proper error handling
             ts_value = entry.get(self.label_ts)
@@ -1588,7 +1587,7 @@ def sample(
 
             samples.append(
                 SampleRequest(
-                    prompt=prompt,
+                    prompt=prompt_ids,
                     prompt_len=prompt_len,
                     expected_output_len=new_output_len,
                     lora_request=None,

vllm/benchmarks/lib/endpoint_request_func.py

@@ -66,7 +66,7 @@ def add_chunk(self, chunk_bytes: bytes) -> list[str]:
 class RequestFuncInput:
     """The input for the request function."""
 
-    prompt: str | list[str] | list[dict[str, Any]]
+    prompt: str | list[str] | list[int] | list[dict[str, Any]]
     api_url: str
     prompt_len: int
     output_len: int

vllm/benchmarks/serve.py

@@ -2097,6 +2097,11 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
         args.ignore_eos = True
 
     if args.dataset_name == "timed_trace":
+        if args.backend not in ("vllm", "openai"):
+            raise ValueError(
+                "timed_trace dataset passes pre-tokenized prompts (list[int])"
+                " and requires a completions backend ('vllm' or 'openai')."
+            )
         # timed_trace carries per-request timestamps;
         # ignore EOS so generation runs to the trace's specified output length,
         # and default to using those timestamps for scheduling unless the user