feat: add random serving benchmark workload

benchmark/README.md

@@ -14,7 +14,7 @@ This tool will sample prompts from dataset, and run benchmark with serialized re
 ```bash
 python benchmark_latency.py --dataset /path/to/ShareGPT_V3_unfiltered_cleaned_split.json \
                             --tokenizer /path/to/tokenizer \
-                            --num-prompt 100 \
+                            --num-prompts 100 \
                             --model-uid ${model_uid}
 ```
 
@@ -26,7 +26,20 @@ This tool will sample prompts from dataset, and run benchmark with parallel requ
 python benchmark_serving.py --dataset /path/to/ShareGPT_V3_unfiltered_cleaned_split.json \
                             --tokenizer /path/to/tokenizer \
                             --model-uid ${model_uid} \
-                            --num-prompt 100 --concurrency 50
+                            --num-prompts 100 --concurrency 50
+```
+
+It can also generate synthetic prompts without a dataset, similar to vLLM's
+random benchmark mode:
+
+```bash
+python benchmark_serving.py --dataset-name random \
+                            --tokenizer /path/to/tokenizer \
+                            --model-uid ${model_uid} \
+                            --input-len 1024 --output-len 128 \
+                            --random-range-ratio 0.2 \
+                            --num-prompts 1000 --concurrency 50 \
+                            --stream --ignore-eos
 ```
 
 ## Benchmarking long context serving

benchmark/benchmark_runner.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import aiohttp
 import json
+import logging
 import sys
+import time
 import traceback
 import warnings
-import logging
 from dataclasses import dataclass, field
-import time
 from typing import List, Optional, Tuple
 
+import aiohttp
 import numpy as np
 
 logging.basicConfig(level=logging.INFO)
@@ -56,6 +56,7 @@ def __init__(
         stream: bool,
         api_key: Optional[str] = None,
         print_error: bool = False,
+        ignore_eos: bool = False,
     ):
         self.api_url = api_url
         self.model_uid = model_uid
@@ -65,6 +66,7 @@ def __init__(
         self.stream = stream
         self.api_key = api_key
         self.print_error = print_error
+        self.ignore_eos = ignore_eos
 
     async def run(self):
         await self.warm_up()
@@ -107,6 +109,8 @@ async def send_request(self, request: tuple, warming_up: bool = False):
                 "stream": False,
                 "messages": [{"role": "user", "content": prompt}],
             }
+        if self.ignore_eos:
+            pload["ignore_eos"] = True
 
         headers = {"User-Agent": "Benchmark Client"}
         if self.api_key:
@@ -124,6 +128,7 @@ async def send_request(self, request: tuple, warming_up: bool = False):
                 ) as response:
                     if response.status == 200:
                         if self.stream:
+                            usage = None
                             async for chunk_bytes in response.content:
                                 # {
                                 #     "id": "chataec79465-dfea-46af-81b9-c28124063fc0",
@@ -142,13 +147,16 @@ async def send_request(self, request: tuple, warming_up: bool = False):
                                 if not chunk_bytes:
                                     continue
 
-                                chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:")
+                                chunk = remove_prefix(
+                                    chunk_bytes.decode("utf-8"), "data:"
+                                )
 
                                 if chunk == "[DONE]":
                                     latency = time.perf_counter() - st
                                 else:
                                     timestamp = time.perf_counter()
                                     data = json.loads(chunk)
+                                    usage = data.get("usage") or usage
 
                                     # First token
                                     if ttft == 0.0:
@@ -157,18 +165,36 @@ async def send_request(self, request: tuple, warming_up: bool = False):
 
                                     # Decoding phase
                                     else:
-                                        output.itl.append(timestamp - most_recent_timestamp)
+                                        output.itl.append(
+                                            timestamp - most_recent_timestamp
+                                        )
 
                                     most_recent_timestamp = timestamp
 
                             output.latency = latency
                             output.success = True
-                            output.completion_tokens = data["usage"]["completion_tokens"]
+                            output.completion_tokens = (
+                                usage["completion_tokens"]
+                                if isinstance(usage, dict)
+                                and isinstance(usage.get("completion_tokens"), int)
+                                else len(output.itl) + 1
+                            )
                         else:
                             resp = await response.json()
                             output.latency = time.perf_counter() - st
                             output.success = True
-                            output.completion_tokens = resp["usage"]["completion_tokens"]
+                            usage = resp.get("usage")
+                            if isinstance(usage, dict) and isinstance(
+                                usage.get("completion_tokens"), int
+                            ):
+                                output.completion_tokens = usage["completion_tokens"]
+                            else:
+                                logger.warning(
+                                    "usage missing or malformed in non-streaming "
+                                    "response; completion_tokens set to 0 and "
+                                    "throughput stats will be inaccurate."
+                                )
+                                output.completion_tokens = 0
             except Exception:
                 output.success = False
                 exc_info = sys.exc_info()
@@ -368,7 +394,9 @@ def print_stats(self):
                 logger.info("Errors encountered during benchmark:")
                 for output in self.outputs:
                     if not output.success:
-                        print(f"Error for prompt with length {output.prompt_len}: {output.error}")
+                        print(
+                            f"Error for prompt with length {output.prompt_len}: {output.error}"
+                        )
             else:
                 logger.info(
                     "Errors were encountered during the benchmark. Run with --print-error to see detailed error messages."
@@ -385,6 +413,7 @@ def __init__(
         concurrency: int,
         api_key: Optional[str] = None,
         print_error: bool = False,
+        ignore_eos: bool = False,
     ):
         super().__init__(
             api_url,
@@ -393,6 +422,7 @@ def __init__(
             stream,
             api_key,
             print_error,
+            ignore_eos,
         )
         self.concurrency = concurrency
         self.left = len(input_requests)

benchmark/benchmark_serving.py

@@ -14,20 +14,61 @@
 
 import argparse
 import asyncio
+import json
 import logging
 import random
-from typing import List, Tuple, Optional
+from typing import List, Optional, Tuple
 
 import numpy as np
-
-from utils import sample_requests, get_tokenizer
 from benchmark_runner import ConcurrentBenchmarkRunner
-
+from utils import get_tokenizer, sample_random_requests, sample_requests
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
+def _validate_random_range_ratio(value: float, name: str) -> float:
+    if not (0.0 <= value < 1.0):
+        raise argparse.ArgumentTypeError(f"{name} must be in [0, 1), got {value}.")
+    return value
+
+
+def parse_random_range_ratio(value):
+    try:
+        return _validate_random_range_ratio(float(value), name="--random-range-ratio")
+    except ValueError:
+        pass
+
+    try:
+        parsed = json.loads(value)
+    except json.JSONDecodeError as exc:
+        raise argparse.ArgumentTypeError(
+            "--random-range-ratio must be a float or a JSON object with "
+            "'input' and 'output' keys."
+        ) from exc
+    if not isinstance(parsed, dict):
+        raise argparse.ArgumentTypeError(
+            "--random-range-ratio JSON value must be an object."
+        )
+    try:
+        return {
+            "input": _validate_random_range_ratio(
+                float(parsed["input"]), name="--random-range-ratio input"
+            ),
+            "output": _validate_random_range_ratio(
+                float(parsed["output"]), name="--random-range-ratio output"
+            ),
+        }
+    except KeyError as exc:
+        raise argparse.ArgumentTypeError(
+            "--random-range-ratio JSON object must contain 'input' and 'output'."
+        ) from exc
+    except (TypeError, ValueError) as exc:
+        raise argparse.ArgumentTypeError(
+            "--random-range-ratio input and output values must be numbers."
+        ) from exc
+
+
 class ServingBenchmarkRunner(ConcurrentBenchmarkRunner):
     def __init__(
         self,
@@ -39,6 +80,7 @@ def __init__(
         request_rate: float,
         api_key: Optional[str] = None,
         print_error: bool = False,
+        ignore_eos: bool = False,
     ):
         super().__init__(
             api_url,
@@ -48,6 +90,7 @@ def __init__(
             concurrency,
             api_key,
             print_error,
+            ignore_eos,
         )
         self.request_rate = request_rate
         self.queue = None  # delay the creation of the queue
@@ -102,13 +145,31 @@ def main(args: argparse.Namespace):
     model_uid = args.model_uid
 
     logger.info("Preparing for benchmark.")
+    if args.request_rate <= 0.0:
+        raise ValueError("--request-rate must be positive.")
+
     tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code)
-    input_requests = sample_requests(
-        args.dataset,
-        args.num_prompts,
-        tokenizer,
-        prompt_len_limit=args.prompt_len_limit,
-    )
+    if args.dataset_name == "sharegpt":
+        if not args.dataset:
+            raise ValueError("--dataset is required when --dataset-name=sharegpt.")
+        input_requests = sample_requests(
+            args.dataset,
+            args.num_prompts,
+            tokenizer,
+            prompt_len_limit=args.prompt_len_limit,
+        )
+    elif args.dataset_name == "random":
+        input_requests = sample_random_requests(
+            args.num_prompts,
+            tokenizer,
+            input_len=args.input_len,
+            output_len=args.output_len,
+            range_ratio=args.random_range_ratio,
+            prefix_len=args.random_prefix_len,
+            seed=args.seed,
+        )
+    else:
+        raise ValueError(f"Unsupported dataset name: {args.dataset_name}")
 
     logger.info("Benchmark starts.")
 
@@ -121,6 +182,7 @@ def main(args: argparse.Namespace):
         concurrency=args.concurrency,
         api_key=args.api_key,
         print_error=args.print_error,
+        ignore_eos=args.ignore_eos,
     )
     asyncio.run(benchmark.run())
 
@@ -134,7 +196,17 @@ def main(args: argparse.Namespace):
     parser.add_argument("--host", type=str, default="localhost")
     parser.add_argument("--port", type=int, default=9997)
     parser.add_argument(
-        "--dataset", type=str, required=True, help="Path to the dataset."
+        "--dataset-name",
+        type=str,
+        choices=["sharegpt", "random"],
+        default="sharegpt",
+        help="Dataset source to use for benchmark requests.",
+    )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help="Path to the dataset. Required when --dataset-name=sharegpt.",
     )
     parser.add_argument(
         "--tokenizer", type=str, required=True, help="Name or path of the tokenizer."
@@ -145,6 +217,43 @@ def main(args: argparse.Namespace):
     parser.add_argument(
         "--prompt-len-limit", type=int, default=1024, help="Prompt length limitation."
     )
+    parser.add_argument(
+        "--input-len",
+        type=int,
+        default=1024,
+        help="Input token length for random dataset requests.",
+    )
+    parser.add_argument(
+        "--output-len",
+        type=int,
+        default=128,
+        help="Output token length for random dataset requests.",
+    )
+    parser.add_argument(
+        "--random-input-len",
+        dest="input_len",
+        type=int,
+        help="Alias for --input-len.",
+    )
+    parser.add_argument(
+        "--random-output-len",
+        dest="output_len",
+        type=int,
+        help="Alias for --output-len.",
+    )
+    parser.add_argument(
+        "--random-range-ratio",
+        type=parse_random_range_ratio,
+        default=0.0,
+        help="Sample random input/output lengths from length * (1 +/- ratio). "
+        'Can also be a JSON object like \'{"input": 0.3, "output": 0.5}\'.',
+    )
+    parser.add_argument(
+        "--random-prefix-len",
+        type=int,
+        default=0,
+        help="Fixed prefix token length prepended to random dataset prompts.",
+    )
     parser.add_argument(
         "--api-key",
         type=str,
@@ -177,10 +286,15 @@ def main(args: argparse.Namespace):
     parser.add_argument(
         "--stream", action="store_true", help="Enable streaming responses."
     )
+    parser.add_argument(
+        "--ignore-eos",
+        action="store_true",
+        help="Ask the backend to ignore EOS so random outputs approach max_tokens.",
+    )
     parser.add_argument(
         "--print-error",
         action="store_true",
-        help="Print detailed error messages if any errors encountered."
+        help="Print detailed error messages if any errors encountered.",
     )
     args = parser.parse_args()
     main(args)