From f158869c2cc3e4915669bc1e2a94049b06264723 Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Sun, 22 Feb 2026 03:20:11 +0530 Subject: [PATCH] [Refactor] Benchmark Phase 1: extract utils and datasets from bench_serving (#19077) Co-authored-by: Xuchun Shang <107600043+xucsh@users.noreply.github.com> --- python/sglang/bench_serving.py | 1279 +---------------- python/sglang/benchmark/__init__.py | 0 python/sglang/benchmark/datasets/__init__.py | 156 ++ python/sglang/benchmark/datasets/common.py | 67 + python/sglang/benchmark/datasets/custom.py | 106 ++ .../datasets/generated_shared_prefix.py | 163 +++ python/sglang/benchmark/datasets/image.py | 239 +++ python/sglang/benchmark/datasets/mmmu.py | 97 ++ python/sglang/benchmark/datasets/mooncake.py | 83 ++ .../benchmark/datasets/openai_dataset.py | 86 ++ python/sglang/benchmark/datasets/random.py | 127 ++ python/sglang/benchmark/datasets/sharegpt.py | 113 ++ python/sglang/benchmark/utils.py | 159 ++ 13 files changed, 1432 insertions(+), 1243 deletions(-) create mode 100644 python/sglang/benchmark/__init__.py create mode 100644 python/sglang/benchmark/datasets/__init__.py create mode 100644 python/sglang/benchmark/datasets/common.py create mode 100644 python/sglang/benchmark/datasets/custom.py create mode 100644 python/sglang/benchmark/datasets/generated_shared_prefix.py create mode 100644 python/sglang/benchmark/datasets/image.py create mode 100644 python/sglang/benchmark/datasets/mmmu.py create mode 100644 python/sglang/benchmark/datasets/mooncake.py create mode 100644 python/sglang/benchmark/datasets/openai_dataset.py create mode 100644 python/sglang/benchmark/datasets/random.py create mode 100644 python/sglang/benchmark/datasets/sharegpt.py create mode 100644 python/sglang/benchmark/utils.py diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index f2e6621b4..ff62c6cec 100644 --- a/python/sglang/bench_serving.py +++ b/python/sglang/bench_serving.py @@ -14,12 +14,9 @@ import argparse import asyncio import copy import importlib.util -import io import json import os -import pickle import random -import resource import shutil import sys import time @@ -30,27 +27,51 @@ from argparse import ArgumentParser from copy import deepcopy from dataclasses import dataclass, field, replace from datetime import datetime -from functools import lru_cache -from json import JSONDecodeError from pathlib import Path from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, Union import aiohttp import numpy as np -import pybase64 import requests -from datasets import load_dataset -from PIL import Image from tqdm.asyncio import tqdm -from transformers import ( - AutoProcessor, - AutoTokenizer, - PreTrainedTokenizer, - PreTrainedTokenizerBase, - PreTrainedTokenizerFast, +from transformers import AutoTokenizer, PreTrainedTokenizerBase + +from sglang.benchmark.datasets import ( # noqa: F401 + ASSISTANT_SUFFIX, + MOONCAKE_DATASET_URL, + SHAREGPT_FILENAME, + SHAREGPT_REPO_ID, + DatasetRow, + compute_random_lens, + create_mm_data_row, + gen_mm_prompt, + gen_prompt, + get_available_tokens, + get_dataset, + get_gen_prefix_cache_path, + get_mooncake_request_over_time, + parse_image_resolution, + sample_custom_requests, + sample_generated_shared_prefix_requests, + sample_image_requests, + sample_mmmu_requests, + sample_openai_requests, + sample_random_requests, + sample_sharegpt_requests, +) +from sglang.benchmark.utils import ( # noqa: F401 + download_and_cache_file, + download_and_cache_hf_file, + get_model, + get_processor, + get_tokenizer, + is_file_valid_json, + parse_custom_headers, + remove_prefix, + remove_suffix, + set_ulimit, ) -ASSISTANT_SUFFIX = "Assistant:" _ROUTING_KEY_HEADER = "X-SMG-Routing-Key" TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and ( @@ -113,14 +134,6 @@ class RequestFuncOutput: return output -def remove_prefix(text: str, prefix: str) -> str: - return text[len(prefix) :] if text.startswith(prefix) else text - - -def remove_suffix(text: str, suffix: str) -> str: - return text[: -len(suffix)] if text.endswith(suffix) else text - - def get_auth_headers() -> Dict[str, str]: openai_api_key = os.environ.get("OPENAI_API_KEY") if openai_api_key: @@ -132,10 +145,6 @@ def get_auth_headers() -> Dict[str, str]: return {} -def parse_custom_headers(header_list: List[str]) -> Dict[str, str]: - return {k: v for h in header_list for k, _, v in [h.partition("=")] if k and v} - - def get_request_headers() -> Dict[str, str]: headers = get_auth_headers() if h := getattr(args, "header", None): @@ -785,168 +794,6 @@ async def _call_profile_pd(profile_urls: List[Tuple[str, str]], mode: str) -> No ) -def get_model(pretrained_model_name_or_path: str) -> str: - if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true": - import huggingface_hub.constants - from modelscope import snapshot_download - - model_path = snapshot_download( - model_id=pretrained_model_name_or_path, - local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, - ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], - ) - - return model_path - return pretrained_model_name_or_path - - -def get_tokenizer( - pretrained_model_name_or_path: str, -) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: - assert ( - pretrained_model_name_or_path is not None - and pretrained_model_name_or_path != "" - ) - if pretrained_model_name_or_path.endswith( - ".json" - ) or pretrained_model_name_or_path.endswith(".model"): - from sglang.srt.utils.hf_transformers_utils import get_tokenizer - - return get_tokenizer(pretrained_model_name_or_path) - - if pretrained_model_name_or_path is not None and not os.path.exists( - pretrained_model_name_or_path - ): - pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) - return AutoTokenizer.from_pretrained( - pretrained_model_name_or_path, trust_remote_code=True - ) - - -def get_processor( - pretrained_model_name_or_path: str, -) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: - assert ( - pretrained_model_name_or_path is not None - and pretrained_model_name_or_path != "" - ) - if pretrained_model_name_or_path.endswith( - ".json" - ) or pretrained_model_name_or_path.endswith(".model"): - from sglang.srt.utils.hf_transformers_utils import get_processor - - return get_processor(pretrained_model_name_or_path) - - if pretrained_model_name_or_path is not None and not os.path.exists( - pretrained_model_name_or_path - ): - pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) - return AutoProcessor.from_pretrained( - pretrained_model_name_or_path, trust_remote_code=True - ) - - -def get_dataset(args, tokenizer, model_id=None): - tokenize_prompt = getattr(args, "tokenize_prompt", False) - if args.dataset_name == "sharegpt": - assert not tokenize_prompt - input_requests = sample_sharegpt_requests( - dataset_path=args.dataset_path, - num_requests=args.num_prompts, - tokenizer=tokenizer, - fixed_output_len=args.sharegpt_output_len, - context_len=args.sharegpt_context_len, - prompt_suffix=args.prompt_suffix, - apply_chat_template=args.apply_chat_template, - ) - elif args.dataset_name.startswith("random"): - input_requests = sample_random_requests( - input_len=args.random_input_len, - output_len=args.random_output_len, - num_prompts=args.num_prompts, - range_ratio=args.random_range_ratio, - tokenizer=tokenizer, - dataset_path=args.dataset_path, - random_sample=args.dataset_name == "random", - return_text=not tokenize_prompt, - ) - elif args.dataset_name == "image": - processor = get_processor(model_id) - input_requests = sample_image_requests( - num_requests=args.num_prompts, - image_count=args.image_count, - input_len=args.random_input_len, - output_len=args.random_output_len, - range_ratio=args.random_range_ratio, - processor=processor, - image_content=args.image_content, - image_format=args.image_format, - image_resolution=args.image_resolution, - backend=args.backend, - random_image_count=args.random_image_count, - ) - elif args.dataset_name == "generated-shared-prefix": - assert not tokenize_prompt - input_requests = sample_generated_shared_prefix_requests( - num_groups=args.gsp_num_groups, - prompts_per_group=args.gsp_prompts_per_group, - system_prompt_len=args.gsp_system_prompt_len, - question_len=args.gsp_question_len, - output_len=args.gsp_output_len, - range_ratio=getattr(args, "gsp_range_ratio", 1.0), - tokenizer=tokenizer, - args=args, - ) - elif args.dataset_name == "mmmu": - processor = get_processor(model_id) - input_requests = sample_mmmu_requests( - num_requests=args.num_prompts, - processor=processor, - backend=args.backend, - fixed_output_len=args.random_output_len, - random_sample=True, - ) - elif args.dataset_name == "mooncake": - # For mooncake, we don't generate the prompts here. - # We just load the raw trace data. The async generator will handle the rest. - if not args.dataset_path: - local_path = os.path.join("/tmp", args.mooncake_workload + "_trace.jsonl") - else: - local_path = args.dataset_path - - if not os.path.exists(local_path): - download_and_cache_file( - MOONCAKE_DATASET_URL[args.mooncake_workload], local_path - ) - - with open(local_path, "r") as f: - all_requests_data = [json.loads(line) for line in f if line.strip()] - - # Limit the number of requests based on --num-prompts - input_requests = all_requests_data[: args.num_prompts] - elif args.dataset_name == "custom": - assert not tokenize_prompt - input_requests = sample_custom_requests( - dataset_path=args.dataset_path, - num_requests=args.num_prompts, - tokenizer=tokenizer, - fixed_output_len=args.sharegpt_output_len, - context_len=args.sharegpt_context_len, - prompt_suffix=args.prompt_suffix, - apply_chat_template=args.apply_chat_template, - ) - elif args.dataset_name == "openai": - input_requests = sample_openai_requests( - dataset_path=args.dataset_path, - num_requests=args.num_prompts, - tokenizer=tokenizer, - fixed_output_len=args.sharegpt_output_len, - ) - else: - raise ValueError(f"Unknown dataset: {args.dataset_name}") - return input_requests - - ASYNC_REQUEST_FUNCS = { "sglang": async_request_sglang_generate, "sglang-native": async_request_sglang_generate, @@ -1000,1049 +847,6 @@ class BenchmarkMetrics: max_concurrent_requests: int = 0 -SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered" -SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json" -MOONCAKE_DATASET_URL = { - "mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl", - "conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl", - "synthetic": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/synthetic_trace.jsonl", - "toolagent": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/toolagent_trace.jsonl", -} - - -def download_and_cache_hf_file( - repo_id: str, - filename: str, - repo_type: str = "dataset", -): - """Download a file from Hugging Face and cache it locally.""" - from huggingface_hub import hf_hub_download - - return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type) - - -def download_and_cache_file(url: str, filename: Optional[str] = None): - """Read and cache a file from a url.""" - if filename is None: - filename = os.path.join("/tmp", url.split("/")[-1]) - - # Check if the cache file already exists - if is_file_valid_json(filename): - return filename - - print(f"Downloading from {url} to {filename}") - - # Stream the response to show the progress bar - response = requests.get(url, stream=True) - response.raise_for_status() # Check for request errors - - # Total size of the file in bytes - total_size = int(response.headers.get("content-length", 0)) - chunk_size = 1024 # Download in chunks of 1KB - - # Use tqdm to display the progress bar - with open(filename, "wb") as f, tqdm( - desc=filename, - total=total_size, - unit="B", - unit_scale=True, - unit_divisor=1024, - ) as bar: - for chunk in response.iter_content(chunk_size=chunk_size): - f.write(chunk) - bar.update(len(chunk)) - - return filename - - -def is_file_valid_json(path): - if not os.path.isfile(path): - return False - - # TODO can fuse into the real file open later - try: - with open(path) as f: - json.load(f) - return True - except JSONDecodeError as e: - print( - f"{path} exists but json loading fails ({e=}), thus treat as invalid file" - ) - return False - - -@dataclass -class DatasetRow: - prompt: str - prompt_len: int - output_len: int - text_prompt_len: Optional[int] = None - vision_prompt_len: Optional[int] = None - image_data: Optional[List[str]] = None - timestamp: Optional[float] = None - routing_key: Optional[str] = None - extra_request_body: Optional[Dict[str, Any]] = None # Per-request API parameters - - def __post_init__(self): - if self.text_prompt_len is None: - self.text_prompt_len = self.prompt_len - if self.vision_prompt_len is None: - self.vision_prompt_len = 0 - if self.extra_request_body is None: - self.extra_request_body = {} - - -async def get_mooncake_request_over_time( - input_requests: List[Dict], - tokenizer: PreTrainedTokenizerBase, - slowdown_factor: float, - num_rounds: int, -) -> AsyncGenerator[DatasetRow, None]: - """ - An async generator that yields requests based on the timestamps in the Mooncake trace file, - with support for multi-round sessions. - """ - if not input_requests: - return - - input_requests.sort(key=lambda r: r["timestamp"]) - - start_time = time.perf_counter() - trace_start_time_ms = input_requests[0]["timestamp"] - - for record in input_requests: - # Calculate when this entire session should start - relative_arrival_time_s = (record["timestamp"] - trace_start_time_ms) / 1000.0 - target_arrival_time_s = relative_arrival_time_s * slowdown_factor - - current_elapsed_time_s = time.perf_counter() - start_time - sleep_duration_s = target_arrival_time_s - current_elapsed_time_s - if sleep_duration_s > 0: - await asyncio.sleep(sleep_duration_s) - - # Once the session starts, generate all rounds for it as a burst - # This simulates a user engaging in a multi-turn conversation - - # Base user query constructed from hash_ids - user_query_base = "" - hash_ids = record.get("hash_ids", []) - for hash_id in hash_ids: - user_query_base += f"{hash_id}" + " ".join( - ["hi"] * 128 - ) # Shorter for multi-round - user_query_base += "Tell me a story based on this context." - - output_len_per_round = record.get("output_length", 256) - chat_history = [] - - for i in range(num_rounds): - # Add user query for the current round - chat_history.append( - {"role": "user", "content": f"Round {i + 1}: {user_query_base}"} - ) - - # Form the full prompt from history - try: - full_prompt_text = tokenizer.apply_chat_template( - chat_history, - tokenize=False, - add_generation_prompt=True, - return_dict=False, - ) - except Exception: - full_prompt_text = "\n".join( - [f"{msg['role']}: {msg['content']}" for msg in chat_history] - ) - - prompt_len = len(tokenizer.encode(full_prompt_text)) - - yield DatasetRow( - prompt=full_prompt_text, - prompt_len=prompt_len, - output_len=output_len_per_round, - ) - - # Add a placeholder assistant response for the next round's context - # We use a placeholder because we don't know the real response - placeholder_response = " ".join(["story"] * output_len_per_round) - chat_history.append({"role": "assistant", "content": placeholder_response}) - - -def sample_mmmu_requests( - num_requests: int, - processor: AutoProcessor | AutoTokenizer, - backend: str = "sglang", - fixed_output_len: Optional[int] = None, - random_sample: bool = True, -) -> List[DatasetRow]: - """ - Sample requests from the MMMU dataset using HuggingFace datasets. - - Args: - num_requests: Number of requests to sample. - fixed_output_len: If provided, use this fixed output length for all requests. - random_sample: Whether to randomly sample or take the first N. - - Returns: - List of tuples (prompt, prompt_token_len, output_token_len). - """ - print("Loading MMMU dataset from HuggingFace...") - - try: - print("Attempting to load MMMU Math dataset...") - mmmu_dataset = load_dataset("MMMU/MMMU", "Math", split="test") - print( - f"Successfully loaded MMMU Math dataset from HuggingFace with {len(mmmu_dataset)} examples" - ) - except Exception as e: - print(f"Failed to load MMMU Math dataset: {e}") - raise ValueError(f"Failed to load MMMU dataset: {e}") - - # Sample from the dataset - if len(mmmu_dataset) > num_requests: - if random_sample: - # Random sample - indices = random.sample(range(len(mmmu_dataset)), num_requests) - sample_dataset = mmmu_dataset.select(indices) - else: - # Take first N - sample_dataset = mmmu_dataset.select( - range(min(num_requests, len(mmmu_dataset))) - ) - else: - print(f"Dataset has less than {num_requests} examples, using all examples") - sample_dataset = mmmu_dataset - - print(f"Selected {len(sample_dataset)} examples for benchmarking") - - # Create prompts - filtered_dataset = [] - - for i, example in enumerate(sample_dataset): - try: - # Extract image_1 - image = example.get("image_1") - - if image is not None: - if hasattr(image, "save"): - # Convert RGBA images to RGB before encoding - if image.mode == "RGBA": - image = image.convert("RGB") - - # Encode image to base64 (save as PNG to support palette/alpha modes) - buffered = io.BytesIO() - image.save(buffered, format="PNG") - img_str = pybase64.b64encode(buffered.getvalue()).decode("utf-8") - image_data = f"data:image/png;base64,{img_str}" - else: - continue - - # Extract the question - question = example.get("question") - - # Construct the prompt - text_prompt = f"Question: {question}\n\nAnswer: " - output_len = fixed_output_len if fixed_output_len is not None else 256 - data_row = create_mm_data_row( - text_prompt, [image], [image_data], output_len, processor, backend - ) - filtered_dataset.append(data_row) - - except Exception as e: - print(f"Error processing example {i}: {e}") - - print(f"\nCreated {len(filtered_dataset)} MMMU prompts") - return filtered_dataset - - -def sample_sharegpt_requests( - dataset_path: str, - num_requests: int, - tokenizer: PreTrainedTokenizerBase, - fixed_output_len: Optional[int] = None, - context_len: Optional[int] = None, - prompt_suffix: Optional[str] = "", - apply_chat_template=False, -) -> List[DatasetRow]: - if fixed_output_len is not None and fixed_output_len < 4: - raise ValueError("output_len too small") - - # Download sharegpt if necessary - if not is_file_valid_json(dataset_path) and dataset_path == "": - dataset_path = download_and_cache_hf_file( - repo_id=SHAREGPT_REPO_ID, - filename=SHAREGPT_FILENAME, - ) - - # Load the dataset. - with open(dataset_path) as f: - dataset = json.load(f) - - # Filter out the conversations with less than 2 turns. - dataset = [ - data - for data in dataset - if len(data.get("conversations", data.get("conversation", []))) >= 2 - ] - # Only keep the first two turns of each conversation. - dataset = [ - ( - data.get("conversations", data.get("conversation", []))[0]["value"], - data.get("conversations", data.get("conversation", []))[1]["value"], - ) - for data in dataset - ] - - # Shuffle the dataset. - random.shuffle(dataset) - - # Filter out sequences that are too long or too short - filtered_dataset: List[DatasetRow] = [] - for i in range(len(dataset)): - if len(filtered_dataset) == num_requests: - break - - # Tokenize the prompts and completions. - prompt = dataset[i][0] - if prompt_suffix: - prompt = ( - remove_suffix(prompt, ASSISTANT_SUFFIX) - + prompt_suffix - + ASSISTANT_SUFFIX - ) - - if apply_chat_template: - prompt = tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], - add_generation_prompt=True, - tokenize=False, - return_dict=False, - ) - if tokenizer.bos_token: - prompt = prompt.replace(tokenizer.bos_token, "") - - prompt_token_ids = tokenizer.encode(prompt) - completion = dataset[i][1] - completion_token_ids = tokenizer.encode(completion) - prompt_len = len(prompt_token_ids) - output_len = ( - len(completion_token_ids) if fixed_output_len is None else fixed_output_len - ) - - if prompt_len < 2 or output_len < 2: - # Prune too short sequences. - continue - - if context_len and prompt_len + output_len > context_len: - # Prune too long sequences. - continue - - filtered_dataset.append( - DatasetRow( - prompt=prompt, - prompt_len=prompt_len, - output_len=output_len, - ) - ) - - print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") - print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") - return filtered_dataset - - -def sample_openai_requests( - dataset_path: str, - num_requests: int, - tokenizer: PreTrainedTokenizerBase, - fixed_output_len: Optional[int] = None, -) -> List[DatasetRow]: - """ - Load OpenAI-compatible chat completion requests from a JSONL file. - - Each line should be a JSON object with: - - "messages": list of {"role": str, "content": str} - - "max_tokens": int (used as output_len if fixed_output_len not set) - - "tools": optional list of tool definitions - - "temperature": optional temperature value - - "top_p": optional top_p value - - Other OpenAI API parameters are also extracted and passed through - """ - dataset = [] - with open(dataset_path, "r") as f: - for line in f: - if num_requests > 0 and len(dataset) >= num_requests: - break - if line.strip(): - try: - dataset.append(json.loads(line)) - except json.JSONDecodeError: - # Skip invalid JSON lines - continue - - # Fields that should NOT be passed through extra_request_body - # These are either handled separately or are metadata - # max_tokens is excluded because it's handled via output_len -> max_completion_tokens - # max_completion_tokens is also excluded to avoid conflicts - EXCLUDED_FIELDS = {"messages", "max_tokens", "max_completion_tokens", "model"} - - filtered_dataset: List[DatasetRow] = [] - for data in dataset: - messages = data.get("messages", []) - if not messages: - continue - - # Use max_tokens from the request, or fall back to fixed_output_len - output_len = fixed_output_len or data.get("max_tokens", 256) - - # Extract extra request body parameters (tools, temperature, top_p, etc.) - extra_body = {k: v for k, v in data.items() if k not in EXCLUDED_FIELDS} - - # Calculate prompt length by applying chat template - # This includes the messages but not the tools - prompt_len = len( - tokenizer.apply_chat_template( - messages, tokenize=True, add_generation_prompt=True - ) - ) - - # If tools are present, we need to add their token count - # Tools are sent as part of the request and count toward input tokens - if "tools" in extra_body: - # Encode tools as JSON string to estimate token count - tools_str = json.dumps(extra_body["tools"]) - tools_tokens = len(tokenizer.encode(tools_str)) - prompt_len += tools_tokens - - # Pass messages list directly - bench_serving handles List[Dict] prompts - filtered_dataset.append( - DatasetRow( - prompt=messages, - prompt_len=prompt_len, - output_len=output_len, - extra_request_body=extra_body, # Store per-request parameters - ) - ) - - print(f"Loaded {len(filtered_dataset)} OpenAI-format requests") - print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") - print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") - return filtered_dataset - - -def sample_custom_requests( - dataset_path: str, - num_requests: int, - tokenizer: PreTrainedTokenizerBase, - fixed_output_len: Optional[int] = None, - context_len: Optional[int] = None, - prompt_suffix: Optional[str] = "", - apply_chat_template=False, -) -> List[DatasetRow]: - """ - Sample requests from a custom JSONL dataset: supports 'content'/'value' as conversation keys. - """ - if fixed_output_len is not None and fixed_output_len < 4: - raise ValueError("output_len too small") - - # Load the dataset - dataset = [] - if not os.path.isfile(dataset_path): - raise FileNotFoundError(f"Dataset not found at {dataset_path}") - - with open(dataset_path, "r", encoding="utf-8") as f: - for line in f: - line = line.strip() - if line: # skip empty lines - try: - dataset.append(json.loads(line)) - except json.JSONDecodeError: - continue # skip lines with JSON errors - - # Filter out the conversations with less than 2 turns. - processed_dataset = [] - for data in dataset: - convs = data.get("conversations", data.get("conversation", [])) - if len(convs) >= 2: - user_turn = convs[0].get("content", convs[0].get("value", "")) - assist_turn = convs[1].get("content", convs[1].get("value", "")) - processed_dataset.append((user_turn, assist_turn)) - dataset = processed_dataset - random.shuffle(dataset) - - # Filter out sequences that are too long or too short - filtered_dataset: List[DatasetRow] = [] - - for i in range(len(dataset)): - if len(filtered_dataset) == num_requests: - break - - # Tokenize the prompts and completions. - prompt = dataset[i][0] - - if prompt_suffix: - prompt = ( - remove_suffix(prompt, ASSISTANT_SUFFIX) - + prompt_suffix - + ASSISTANT_SUFFIX - ) - - if apply_chat_template: - prompt = tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], - add_generation_prompt=True, - tokenize=False, - return_dict=False, - ) - if tokenizer.bos_token: - prompt = prompt.replace(tokenizer.bos_token, "") - - prompt_token_ids = tokenizer.encode(prompt) - completion = dataset[i][1] - completion_token_ids = tokenizer.encode(completion) - prompt_len = len(prompt_token_ids) - output_len = ( - len(completion_token_ids) if fixed_output_len is None else fixed_output_len - ) - - if prompt_len < 2 or output_len < 2: - # Prune too short sequences. - continue - - if context_len and prompt_len + output_len > context_len: - # Prune too long sequences. - continue - - filtered_dataset.append( - DatasetRow( - prompt=prompt, - prompt_len=prompt_len, - output_len=output_len, - ) - ) - - print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") - print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") - return filtered_dataset - - -def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]: - return np.random.randint( - max(int(full_len * range_ratio), 1), - full_len + 1, - size=num, - ).tolist() - - -def sample_random_requests( - input_len: int, - output_len: int, - num_prompts: int, - range_ratio: float, - tokenizer: PreTrainedTokenizerBase, - dataset_path: str, - random_sample: bool = True, - return_text: bool = True, -) -> List[DatasetRow]: - input_lens = compute_random_lens( - full_len=input_len, - range_ratio=range_ratio, - num=num_prompts, - ) - output_lens = compute_random_lens( - full_len=output_len, - range_ratio=range_ratio, - num=num_prompts, - ) - - if return_text: - # Need to truncate input_len as server encode will add special token. - num_special_tokens = int(tokenizer.num_special_tokens_to_add()) - for i in range(num_prompts): - input_lens[i] = max(0, input_lens[i] - num_special_tokens) - - if random_sample: - # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens - - # Download sharegpt if necessary - if not is_file_valid_json(dataset_path): - dataset_path = download_and_cache_hf_file( - repo_id=SHAREGPT_REPO_ID, - filename=SHAREGPT_FILENAME, - ) - - # Load the dataset. - with open(dataset_path) as f: - dataset = json.load(f) - # Filter out the conversations with less than 2 turns. - dataset = [ - data - for data in dataset - if len(data.get("conversations", data.get("conversation", []))) >= 2 - ] - # Only keep the first two turns of each conversation. - dataset = [ - ( - data.get("conversations", data.get("conversation", []))[0]["value"], - data.get("conversations", data.get("conversation", []))[1]["value"], - ) - for data in dataset - ] - # Shuffle the dataset. - random.shuffle(dataset) - - # Filter out sequences that are too long or too short - input_requests: List[DatasetRow] = [] - for data in dataset: - i = len(input_requests) - if i == num_prompts: - break - - # Tokenize the prompts and completions. - prompt = data[0] - prompt_token_ids = tokenizer.encode(prompt) - prompt_len = len(prompt_token_ids) - - # Skip empty prompt - if prompt_len == 0: - continue - - if prompt_len > input_lens[i]: - input_ids = prompt_token_ids[: input_lens[i]] - else: - ratio = (input_lens[i] + prompt_len - 1) // prompt_len - input_ids = (prompt_token_ids * ratio)[: input_lens[i]] - input_content = input_ids - if return_text: - input_content = tokenizer.decode(input_content) - input_requests.append( - DatasetRow( - prompt=input_content, - prompt_len=input_lens[i], - output_len=output_lens[i], - ) - ) - else: - # Sample token ids from random integers. This can cause some NaN issues. - offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) - input_requests = [] - for i in range(num_prompts): - # Use int() to convert numpy.int64 to native Python int for JSON serialization - input_content = [ - int((offsets[i] + i + j) % tokenizer.vocab_size) - for j in range(input_lens[i]) - ] - if return_text: - input_content = tokenizer.decode(input_content) - input_requests.append( - DatasetRow( - prompt=input_content, - prompt_len=input_lens[i], - output_len=output_lens[i], - ) - ) - - print(f"#Input tokens: {np.sum(input_lens)}") - print(f"#Output tokens: {np.sum(output_lens)}") - return input_requests - - -def parse_image_resolution(image_resolution: str) -> Tuple[int, int]: - """Parse image resolution into (width, height). - - Supports presets '1080p', '720p', '360p' and custom 'heightxwidth' format - (e.g., '1080x1920' means height=1080, width=1920). - """ - resolution_to_size = { - "4k": (3840, 2160), - "1080p": (1920, 1080), - "720p": (1280, 720), - "360p": (640, 360), - } - if image_resolution in resolution_to_size: - return resolution_to_size[image_resolution] - - res = image_resolution.strip().lower() - if "x" in res: - parts = res.split("x") - if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit(): - height = int(parts[0]) - width = int(parts[1]) - if height > 0 and width > 0: - return (width, height) - - raise ValueError( - f"Unsupported image resolution: {image_resolution}. " - "Choose from 4k, 1080p, 720p, 360p, or provide custom 'heightxwidth' (e.g., 1080x1920)." - ) - - -def create_mm_data_row( - text_prompt, images: list, images_base64, output_len, processor, backend -): - try: - if type(processor).__name__ == "Phi4MMProcessor": - # <|endoftext10|> is the image token used in the phi-4-multimodal model. - content_items = text_prompt.replace("image 1", "|endoftext10|") - else: - content_items = [ - {"type": "image", "image": {"url": image_base64}} - for image_base64 in images_base64 - ] - content_items.append({"type": "text", "text": text_prompt}) - prompt_str = processor.apply_chat_template( - [{"role": "user", "content": content_items}], - add_generation_prompt=True, - tokenize=False, - ) - except Exception as e: - # Note (Xinyuan): This is a workaround for an issue where some tokenizers do not support content as a list. (e.g. InternVL) - print(f"Error applying chat template: {e}, fallback to tag") - # Some tokenizers do not support list content; fall back to a placeholder in the text - prompt_str = f"{text_prompt}" - - # Calculate total tokens (text + vision) - prompt_len = processor( - text=[prompt_str], - images=images, - padding=False, - return_tensors="pt", - )["input_ids"].numel() - - # Calculate text-only tokens - try: - # Create text-only version of the prompt - text_only_prompt = processor.apply_chat_template( - [{"role": "user", "content": text_prompt}], - add_generation_prompt=True, - tokenize=False, - ) - text_prompt_len = processor( - text=[text_only_prompt], - padding=False, - return_tensors="pt", - )["input_ids"].numel() - except Exception: - # Fallback: just tokenize the text prompt directly - tokenizer_to_use = ( - processor.tokenizer if hasattr(processor, "tokenizer") else processor - ) - text_prompt_len = len(tokenizer_to_use.encode(text_prompt)) - - # Vision tokens = total tokens - text tokens - vision_prompt_len = prompt_len - text_prompt_len - - use_raw_prompt = backend in [ - "sglang", - "sglang-oai", - "sglang-oai-chat", - "vllm", - "vllm-chat", - "lmdeploy", - "lmdeploy-chat", - ] - return DatasetRow( - prompt=text_prompt if use_raw_prompt else prompt_str, - prompt_len=prompt_len, - output_len=output_len, - text_prompt_len=text_prompt_len, - vision_prompt_len=vision_prompt_len, - image_data=images_base64, - ) - - -def sample_image_requests( - num_requests: int, - image_count: int, - input_len: int, - output_len: int, - range_ratio: float, - processor: AutoProcessor, - image_content: str, - image_format: str, - image_resolution: str, - backend: str, - random_image_count: bool = False, -) -> List[DatasetRow]: - """Generate requests with images. - - - If ``random_image_count`` is True, each request includes a random number of images between 1 and ``image_count``. - - If ``random_image_count`` is False, each request includes exactly ``image_count`` images. - - Supported resolutions: 4k (3840x2160), 1080p (1920x1080), 720p (1280x720), 360p (640x360), - or custom 'heightxwidth' (e.g., 1080x1920). - - Text lengths follow the 'random' dataset sampling rule. ``prompt_len`` - only counts text tokens and excludes image data. - """ - - # Parse resolution (supports presets and 'heightxwidth') - width, height = parse_image_resolution(image_resolution) - - # Determine image counts for each request - if random_image_count: - # Random number of images per request - image_counts = np.random.randint(1, image_count + 1, size=num_requests) - total_images = np.sum(image_counts) - else: - # Fixed number of images per request - image_counts = np.full(num_requests, image_count) - total_images = image_count * num_requests - - # Check for potentially problematic combinations and warn user - if width * height >= 1920 * 1080 and total_images >= 100: - warnings.warn( - f"High resolution ({width}x{height}) with {total_images} total images " - f"may take a long time. Consider reducing resolution or image count.", - UserWarning, - stacklevel=2, - ) - - # Sample text lengths - input_lens = compute_random_lens( - full_len=input_len, - range_ratio=range_ratio, - num=num_requests, - ) - output_lens = compute_random_lens( - full_len=output_len, - range_ratio=range_ratio, - num=num_requests, - ) - - def _gen_random_image_data_uri( - width: int = width, height: int = height - ) -> (Image, str, int): - if image_content == "blank": - # Generate blank white image - arr = np.full((height, width, 3), 255, dtype=np.uint8) - else: - # Generate random colored image - arr = (np.random.rand(height, width, 3) * 255).astype(np.uint8) - img = Image.fromarray(arr) - buf = io.BytesIO() - img.save(buf, format=image_format, quality=85) - encoded = pybase64.b64encode(buf.getvalue()).decode("utf-8") - image_data = f"data:image/{image_format};base64,{encoded}" - image_bytes = len(image_data.encode("utf-8")) - return img, image_data, image_bytes - - dataset: List[DatasetRow] = [] - total_image_bytes = 0 - for i in range(num_requests): - # Get the number of images for this request - request_image_count = int(image_counts[i]) - - # Generate text prompt - text_prompt = gen_mm_prompt( - processor.tokenizer, - processor.image_token_id if hasattr(processor, "image_token_id") else None, - int(input_lens[i]), - ) - - # Generate image list - images, images_base64, images_bytes = zip( - *[_gen_random_image_data_uri() for _ in range(request_image_count)] - ) - total_image_bytes += sum(list(images_bytes)) - - data_row = create_mm_data_row( - text_prompt, - list(images), - list(images_base64), - int(output_lens[i]), - processor, - backend, - ) - dataset.append(data_row) - - # Print statistics - print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}") - print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}") - print(f"#Total images: {total_images}") - - if random_image_count: - print( - f"#Images per request: min={np.min(image_counts)}, max={np.max(image_counts)}, mean={np.mean(image_counts):.2f}" - ) - else: - print(f"#Images per request: {image_count} (fixed)") - - print( - f"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request" - ) - return dataset - - -@lru_cache(maxsize=1) -def get_available_tokens(tokenizer): - """Get all available token ids from the tokenizer vocabulary.""" - return list(tokenizer.get_vocab().values()) - - -def gen_prompt(tokenizer, token_num): - """Generate a random prompt of specified token length using tokenizer vocabulary.""" - all_available_tokens = get_available_tokens(tokenizer) - selected_tokens = random.choices(all_available_tokens, k=token_num) - return tokenizer.decode(selected_tokens) - - -def gen_mm_prompt(tokenizer, image_pad_id, token_num): - """Generate a random prompt of specified token length using tokenizer vocabulary.""" - all_available_tokens = list(tokenizer.get_vocab().values()) - if image_pad_id: - all_available_tokens.remove(image_pad_id) - selected_tokens = random.choices(all_available_tokens, k=token_num) - return tokenizer.decode(selected_tokens) - - -def get_gen_prefix_cache_path(args, tokenizer): - """Create cache directory under ~/.cache/sglang/benchmark""" - cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" - - # Create a unique cache filename based on the generation parameters - cache_key = ( - f"gen_shared_prefix_{args.seed}_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" - f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_" - f"{tokenizer.__class__.__name__}.pkl" - ) - return cache_dir / cache_key - - -def sample_generated_shared_prefix_requests( - num_groups: int, - prompts_per_group: int, - system_prompt_len: int, - question_len: int, - output_len: int, - range_ratio: float, - tokenizer: PreTrainedTokenizerBase, - args: argparse.Namespace, -) -> List[DatasetRow]: - """Generate benchmark requests with shared system prompts using random tokens and caching.""" - send_routing_key = getattr(args, "gsp_send_routing_key", False) - num_turns = getattr(args, "gsp_num_turns", 1) - - cache_path = get_gen_prefix_cache_path(args, tokenizer) - should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1 - - # Try to load from cache first - if cache_path.exists() and should_cache: - print(f"\nLoading cached generated input data from {cache_path}") - with open(cache_path, "rb") as f: - return pickle.load(f) - - print( - f"\nGenerating new input data... " - f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})" - ) - - run_random_str = uuid.uuid4().hex[:8] - run_start_timestamp = datetime.now().strftime("%Y%m%d%H%M%S") - - system_prompt_lens = compute_random_lens( - full_len=system_prompt_len, - range_ratio=range_ratio, - num=num_groups, - ) - question_lens = np.array( - compute_random_lens( - full_len=question_len, - range_ratio=range_ratio, - num=num_groups * prompts_per_group * num_turns, - ) - ).reshape(num_groups, prompts_per_group, num_turns) - output_lens = np.array( - compute_random_lens( - full_len=output_len, - range_ratio=range_ratio, - num=num_groups * prompts_per_group, - ) - ).reshape(num_groups, prompts_per_group) - del system_prompt_len, question_len, output_len - - # Generate system prompts for each group - system_prompts = [ - gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups) - ] - - # Generate questions: shape (num_groups, prompts_per_group, num_turns) - questions = [ - [ - [ - gen_prompt(tokenizer, int(question_lens[g, p, t])) - for t in range(num_turns) - ] - for p in range(prompts_per_group) - ] - for g in range(num_groups) - ] - - # Combine system prompts with questions - input_requests = [] - total_input_tokens = 0 - total_output_tokens = 0 - - for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): - system_prompt = system_prompts[group_idx] - routing_key = ( - f"{run_random_str}_{run_start_timestamp}_{group_idx}" - if send_routing_key - else None - ) - for prompt_idx in tqdm( - range(prompts_per_group), desc="Generating questions", leave=False - ): - turn_questions = questions[group_idx][prompt_idx] - turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[ - 1: - ] - full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts - prompt_len = ( - 1 - if getattr(args, "gsp_fast_prepare", False) - else len(tokenizer.encode(turn_prompts[0])) - ) - output_len_val = int(output_lens[group_idx, prompt_idx]) - - input_requests.append( - DatasetRow( - prompt=full_prompt, - prompt_len=prompt_len, - output_len=output_len_val, - routing_key=routing_key, - ) - ) - total_input_tokens += prompt_len - total_output_tokens += output_len_val - - if not getattr(args, "gsp_ordered", False): - random.shuffle(input_requests) - - # Print statistics - print(f"\nGenerated shared prefix dataset statistics:") - print(f"Number of groups: {num_groups}") - print(f"Prompts per group: {prompts_per_group}") - print(f"Number of turns: {num_turns}") - print(f"Total prompts: {len(input_requests)}") - if not getattr(args, "gsp_fast_prepare", False): - print(f"Total input tokens: {total_input_tokens}") - print(f"Total output tokens: {total_output_tokens}") - print( - f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" - ) - all_questions = [q for group in questions for conv in group for q in conv] - print( - f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n" - ) - - # Save to cache - if should_cache: - cache_path.parent.mkdir(parents=True, exist_ok=True) - print(f"Caching generated input data to {cache_path}") - with open(cache_path, "wb") as f: - pickle.dump(input_requests, f) - - return input_requests - - async def get_request( input_requests: List[DatasetRow], request_rate: float, @@ -2981,17 +1785,6 @@ def run_benchmark(args_: argparse.Namespace): ) -def set_ulimit(target_soft_limit=65535): - resource_type = resource.RLIMIT_NOFILE - current_soft, current_hard = resource.getrlimit(resource_type) - - if current_soft < target_soft_limit: - try: - resource.setrlimit(resource_type, (target_soft_limit, current_hard)) - except ValueError as e: - print(f"Fail to set RLIMIT_NOFILE: {e}") - - class LoRAPathAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, []) diff --git a/python/sglang/benchmark/__init__.py b/python/sglang/benchmark/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/benchmark/datasets/__init__.py b/python/sglang/benchmark/datasets/__init__.py new file mode 100644 index 000000000..5b771c1bc --- /dev/null +++ b/python/sglang/benchmark/datasets/__init__.py @@ -0,0 +1,156 @@ +import json +import os + +from sglang.benchmark.datasets.common import ( + ASSISTANT_SUFFIX, + MOONCAKE_DATASET_URL, + SHAREGPT_FILENAME, + SHAREGPT_REPO_ID, + DatasetRow, + compute_random_lens, + gen_mm_prompt, + gen_prompt, + get_available_tokens, +) +from sglang.benchmark.datasets.custom import sample_custom_requests +from sglang.benchmark.datasets.generated_shared_prefix import ( + get_gen_prefix_cache_path, + sample_generated_shared_prefix_requests, +) +from sglang.benchmark.datasets.image import ( + create_mm_data_row, + parse_image_resolution, + sample_image_requests, +) +from sglang.benchmark.datasets.mmmu import sample_mmmu_requests +from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time +from sglang.benchmark.datasets.openai_dataset import sample_openai_requests +from sglang.benchmark.datasets.random import sample_random_requests +from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests +from sglang.benchmark.utils import download_and_cache_file, get_processor + + +def get_dataset(args, tokenizer, model_id=None): + tokenize_prompt = getattr(args, "tokenize_prompt", False) + if args.dataset_name == "sharegpt": + assert not tokenize_prompt + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + context_len=args.sharegpt_context_len, + prompt_suffix=args.prompt_suffix, + apply_chat_template=args.apply_chat_template, + ) + elif args.dataset_name.startswith("random"): + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path=args.dataset_path, + random_sample=args.dataset_name == "random", + return_text=not tokenize_prompt, + ) + elif args.dataset_name == "image": + processor = get_processor(model_id) + input_requests = sample_image_requests( + num_requests=args.num_prompts, + image_count=args.image_count, + input_len=args.random_input_len, + output_len=args.random_output_len, + range_ratio=args.random_range_ratio, + processor=processor, + image_content=args.image_content, + image_format=args.image_format, + image_resolution=args.image_resolution, + backend=args.backend, + random_image_count=args.random_image_count, + ) + elif args.dataset_name == "generated-shared-prefix": + assert not tokenize_prompt + input_requests = sample_generated_shared_prefix_requests( + num_groups=args.gsp_num_groups, + prompts_per_group=args.gsp_prompts_per_group, + system_prompt_len=args.gsp_system_prompt_len, + question_len=args.gsp_question_len, + output_len=args.gsp_output_len, + range_ratio=getattr(args, "gsp_range_ratio", 1.0), + tokenizer=tokenizer, + args=args, + ) + elif args.dataset_name == "mmmu": + processor = get_processor(model_id) + input_requests = sample_mmmu_requests( + num_requests=args.num_prompts, + processor=processor, + backend=args.backend, + fixed_output_len=args.random_output_len, + random_sample=True, + ) + elif args.dataset_name == "mooncake": + # For mooncake, we don't generate the prompts here. + # We just load the raw trace data. The async generator will handle the rest. + if not args.dataset_path: + local_path = os.path.join("/tmp", args.mooncake_workload + "_trace.jsonl") + else: + local_path = args.dataset_path + + if not os.path.exists(local_path): + download_and_cache_file( + MOONCAKE_DATASET_URL[args.mooncake_workload], local_path + ) + + with open(local_path, "r") as f: + all_requests_data = [json.loads(line) for line in f if line.strip()] + + # Limit the number of requests based on --num-prompts + input_requests = all_requests_data[: args.num_prompts] + elif args.dataset_name == "custom": + assert not tokenize_prompt + input_requests = sample_custom_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + context_len=args.sharegpt_context_len, + prompt_suffix=args.prompt_suffix, + apply_chat_template=args.apply_chat_template, + ) + elif args.dataset_name == "openai": + input_requests = sample_openai_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + ) + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests + + +__all__ = [ + "ASSISTANT_SUFFIX", + "MOONCAKE_DATASET_URL", + "SHAREGPT_FILENAME", + "SHAREGPT_REPO_ID", + "DatasetRow", + "compute_random_lens", + "create_mm_data_row", + "gen_mm_prompt", + "gen_prompt", + "get_available_tokens", + "get_dataset", + "get_gen_prefix_cache_path", + "get_mooncake_request_over_time", + "parse_image_resolution", + "sample_custom_requests", + "sample_generated_shared_prefix_requests", + "sample_image_requests", + "sample_mmmu_requests", + "sample_openai_requests", + "sample_random_requests", + "sample_sharegpt_requests", +] diff --git a/python/sglang/benchmark/datasets/common.py b/python/sglang/benchmark/datasets/common.py new file mode 100644 index 000000000..900c9b9e4 --- /dev/null +++ b/python/sglang/benchmark/datasets/common.py @@ -0,0 +1,67 @@ +import random +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, Dict, List, Optional + +import numpy as np + +ASSISTANT_SUFFIX = "Assistant:" +SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered" +SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json" +MOONCAKE_DATASET_URL = { + "mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl", + "conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl", + "synthetic": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/synthetic_trace.jsonl", + "toolagent": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/toolagent_trace.jsonl", +} + + +@dataclass +class DatasetRow: + prompt: Any + prompt_len: int + output_len: int + text_prompt_len: Optional[int] = None + vision_prompt_len: Optional[int] = None + image_data: Optional[List[str]] = None + timestamp: Optional[float] = None + routing_key: Optional[str] = None + extra_request_body: Optional[Dict[str, Any]] = None # Per-request API parameters + + def __post_init__(self): + if self.text_prompt_len is None: + self.text_prompt_len = self.prompt_len + if self.vision_prompt_len is None: + self.vision_prompt_len = 0 + if self.extra_request_body is None: + self.extra_request_body = {} + + +def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]: + return np.random.randint( + max(int(full_len * range_ratio), 1), + full_len + 1, + size=num, + ).tolist() + + +@lru_cache(maxsize=1) +def get_available_tokens(tokenizer): + """Get all available token ids from the tokenizer vocabulary.""" + return list(tokenizer.get_vocab().values()) + + +def gen_prompt(tokenizer, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = get_available_tokens(tokenizer) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) + + +def gen_mm_prompt(tokenizer, image_pad_id, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = list(tokenizer.get_vocab().values()) + if image_pad_id: + all_available_tokens.remove(image_pad_id) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) diff --git a/python/sglang/benchmark/datasets/custom.py b/python/sglang/benchmark/datasets/custom.py new file mode 100644 index 000000000..5b1ec2353 --- /dev/null +++ b/python/sglang/benchmark/datasets/custom.py @@ -0,0 +1,106 @@ +import json +import os +import random +from typing import List, Optional + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import ASSISTANT_SUFFIX, DatasetRow +from sglang.benchmark.utils import remove_suffix + + +def sample_custom_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, + context_len: Optional[int] = None, + prompt_suffix: Optional[str] = "", + apply_chat_template=False, +) -> List[DatasetRow]: + """ + Sample requests from a custom JSONL dataset: supports 'content'/'value' as conversation keys. + """ + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Load the dataset + dataset = [] + if not os.path.isfile(dataset_path): + raise FileNotFoundError(f"Dataset not found at {dataset_path}") + + with open(dataset_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: # skip empty lines + try: + dataset.append(json.loads(line)) + except json.JSONDecodeError: + continue # skip lines with JSON errors + + # Filter out the conversations with less than 2 turns. + processed_dataset = [] + for data in dataset: + convs = data.get("conversations", data.get("conversation", [])) + if len(convs) >= 2: + user_turn = convs[0].get("content", convs[0].get("value", "")) + assist_turn = convs[1].get("content", convs[1].get("value", "")) + processed_dataset.append((user_turn, assist_turn)) + dataset = processed_dataset + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: List[DatasetRow] = [] + + for i in range(len(dataset)): + if len(filtered_dataset) == num_requests: + break + + # Tokenize the prompts and completions. + prompt = dataset[i][0] + + if prompt_suffix: + prompt = ( + remove_suffix(prompt, ASSISTANT_SUFFIX) + + prompt_suffix + + ASSISTANT_SUFFIX + ) + + if apply_chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + return_dict=False, + ) + if tokenizer.bos_token: + prompt = prompt.replace(tokenizer.bos_token, "") + + prompt_token_ids = tokenizer.encode(prompt) + completion = dataset[i][1] + completion_token_ids = tokenizer.encode(completion) + prompt_len = len(prompt_token_ids) + output_len = ( + len(completion_token_ids) if fixed_output_len is None else fixed_output_len + ) + + if prompt_len < 2 or output_len < 2: + # Prune too short sequences. + continue + + if context_len and prompt_len + output_len > context_len: + # Prune too long sequences. + continue + + filtered_dataset.append( + DatasetRow( + prompt=prompt, + prompt_len=prompt_len, + output_len=output_len, + ) + ) + + print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") + return filtered_dataset diff --git a/python/sglang/benchmark/datasets/generated_shared_prefix.py b/python/sglang/benchmark/datasets/generated_shared_prefix.py new file mode 100644 index 000000000..d1f99d5e4 --- /dev/null +++ b/python/sglang/benchmark/datasets/generated_shared_prefix.py @@ -0,0 +1,163 @@ +import argparse +import pickle +import random +import uuid +from datetime import datetime +from pathlib import Path +from typing import List + +import numpy as np +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import DatasetRow, compute_random_lens, gen_prompt + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gen_shared_prefix_{args.seed}_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" + f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_" + f"{tokenizer.__class__.__name__}.pkl" + ) + return cache_dir / cache_key + + +def sample_generated_shared_prefix_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + args: argparse.Namespace, +) -> List[DatasetRow]: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + send_routing_key = getattr(args, "gsp_send_routing_key", False) + num_turns = getattr(args, "gsp_num_turns", 1) + + cache_path = get_gen_prefix_cache_path(args, tokenizer) + should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1 + + # Try to load from cache first + if cache_path.exists() and should_cache: + print(f"\nLoading cached generated input data from {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + print( + f"\nGenerating new input data... " + f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})" + ) + + run_random_str = uuid.uuid4().hex[:8] + run_start_timestamp = datetime.now().strftime("%Y%m%d%H%M%S") + + system_prompt_lens = compute_random_lens( + full_len=system_prompt_len, + range_ratio=range_ratio, + num=num_groups, + ) + question_lens = np.array( + compute_random_lens( + full_len=question_len, + range_ratio=range_ratio, + num=num_groups * prompts_per_group * num_turns, + ) + ).reshape(num_groups, prompts_per_group, num_turns) + output_lens = np.array( + compute_random_lens( + full_len=output_len, + range_ratio=range_ratio, + num=num_groups * prompts_per_group, + ) + ).reshape(num_groups, prompts_per_group) + del system_prompt_len, question_len, output_len + + # Generate system prompts for each group + system_prompts = [ + gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups) + ] + + # Generate questions: shape (num_groups, prompts_per_group, num_turns) + questions = [ + [ + [ + gen_prompt(tokenizer, int(question_lens[g, p, t])) + for t in range(num_turns) + ] + for p in range(prompts_per_group) + ] + for g in range(num_groups) + ] + + # Combine system prompts with questions + input_requests = [] + total_input_tokens = 0 + total_output_tokens = 0 + + for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): + system_prompt = system_prompts[group_idx] + routing_key = ( + f"{run_random_str}_{run_start_timestamp}_{group_idx}" + if send_routing_key + else None + ) + for prompt_idx in tqdm( + range(prompts_per_group), desc="Generating questions", leave=False + ): + turn_questions = questions[group_idx][prompt_idx] + turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[ + 1: + ] + full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts + prompt_len = ( + 1 + if getattr(args, "gsp_fast_prepare", False) + else len(tokenizer.encode(turn_prompts[0])) + ) + output_len_val = int(output_lens[group_idx, prompt_idx]) + + input_requests.append( + DatasetRow( + prompt=full_prompt, + prompt_len=prompt_len, + output_len=output_len_val, + routing_key=routing_key, + ) + ) + total_input_tokens += prompt_len + total_output_tokens += output_len_val + + if not getattr(args, "gsp_ordered", False): + random.shuffle(input_requests) + + # Print statistics + print(f"\nGenerated shared prefix dataset statistics:") + print(f"Number of groups: {num_groups}") + print(f"Prompts per group: {prompts_per_group}") + print(f"Number of turns: {num_turns}") + print(f"Total prompts: {len(input_requests)}") + if not getattr(args, "gsp_fast_prepare", False): + print(f"Total input tokens: {total_input_tokens}") + print(f"Total output tokens: {total_output_tokens}") + print( + f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" + ) + all_questions = [q for group in questions for conv in group for q in conv] + print( + f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n" + ) + + # Save to cache + if should_cache: + cache_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Caching generated input data to {cache_path}") + with open(cache_path, "wb") as f: + pickle.dump(input_requests, f) + + return input_requests diff --git a/python/sglang/benchmark/datasets/image.py b/python/sglang/benchmark/datasets/image.py new file mode 100644 index 000000000..0c2ab98f3 --- /dev/null +++ b/python/sglang/benchmark/datasets/image.py @@ -0,0 +1,239 @@ +import io +import warnings +from typing import List, Tuple + +import numpy as np +import pybase64 +from PIL import Image +from transformers import AutoProcessor + +from sglang.benchmark.datasets.common import ( + DatasetRow, + compute_random_lens, + gen_mm_prompt, +) + + +def parse_image_resolution(image_resolution: str) -> Tuple[int, int]: + """Parse image resolution into (width, height). + + Supports presets '1080p', '720p', '360p' and custom 'heightxwidth' format + (e.g., '1080x1920' means height=1080, width=1920). + """ + resolution_to_size = { + "4k": (3840, 2160), + "1080p": (1920, 1080), + "720p": (1280, 720), + "360p": (640, 360), + } + if image_resolution in resolution_to_size: + return resolution_to_size[image_resolution] + + res = image_resolution.strip().lower() + if "x" in res: + parts = res.split("x") + if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit(): + height = int(parts[0]) + width = int(parts[1]) + if height > 0 and width > 0: + return (width, height) + + raise ValueError( + f"Unsupported image resolution: {image_resolution}. " + "Choose from 4k, 1080p, 720p, 360p, or provide custom 'heightxwidth' (e.g., 1080x1920)." + ) + + +def create_mm_data_row( + text_prompt, images: list, images_base64, output_len, processor, backend +): + try: + if type(processor).__name__ == "Phi4MMProcessor": + # <|endoftext10|> is the image token used in the phi-4-multimodal model. + content_items = text_prompt.replace("image 1", "|endoftext10|") + else: + content_items = [ + {"type": "image", "image": {"url": image_base64}} + for image_base64 in images_base64 + ] + content_items.append({"type": "text", "text": text_prompt}) + prompt_str = processor.apply_chat_template( + [{"role": "user", "content": content_items}], + add_generation_prompt=True, + tokenize=False, + ) + except Exception as e: + # Note (Xinyuan): This is a workaround for an issue where some tokenizers do not support content as a list. (e.g. InternVL) + print(f"Error applying chat template: {e}, fallback to tag") + # Some tokenizers do not support list content; fall back to a placeholder in the text + prompt_str = f"{text_prompt}" + + # Calculate total tokens (text + vision) + prompt_len = processor( + text=[prompt_str], + images=images, + padding=False, + return_tensors="pt", + )["input_ids"].numel() + + # Calculate text-only tokens + try: + # Create text-only version of the prompt + text_only_prompt = processor.apply_chat_template( + [{"role": "user", "content": text_prompt}], + add_generation_prompt=True, + tokenize=False, + ) + text_prompt_len = processor( + text=[text_only_prompt], + padding=False, + return_tensors="pt", + )["input_ids"].numel() + except Exception: + # Fallback: just tokenize the text prompt directly + tokenizer_to_use = ( + processor.tokenizer if hasattr(processor, "tokenizer") else processor + ) + text_prompt_len = len(tokenizer_to_use.encode(text_prompt)) + + # Vision tokens = total tokens - text tokens + vision_prompt_len = prompt_len - text_prompt_len + + use_raw_prompt = backend in [ + "sglang", + "sglang-oai", + "sglang-oai-chat", + "vllm", + "vllm-chat", + "lmdeploy", + "lmdeploy-chat", + ] + return DatasetRow( + prompt=text_prompt if use_raw_prompt else prompt_str, + prompt_len=prompt_len, + output_len=output_len, + text_prompt_len=text_prompt_len, + vision_prompt_len=vision_prompt_len, + image_data=images_base64, + ) + + +def sample_image_requests( + num_requests: int, + image_count: int, + input_len: int, + output_len: int, + range_ratio: float, + processor: AutoProcessor, + image_content: str, + image_format: str, + image_resolution: str, + backend: str, + random_image_count: bool = False, +) -> List[DatasetRow]: + """Generate requests with images. + + - If ``random_image_count`` is True, each request includes a random number of images between 1 and ``image_count``. + - If ``random_image_count`` is False, each request includes exactly ``image_count`` images. + - Supported resolutions: 4k (3840x2160), 1080p (1920x1080), 720p (1280x720), 360p (640x360), + or custom 'heightxwidth' (e.g., 1080x1920). + - Text lengths follow the 'random' dataset sampling rule. ``prompt_len`` + only counts text tokens and excludes image data. + """ + + # Parse resolution (supports presets and 'heightxwidth') + width, height = parse_image_resolution(image_resolution) + + # Determine image counts for each request + if random_image_count: + # Random number of images per request + image_counts = np.random.randint(1, image_count + 1, size=num_requests) + total_images = np.sum(image_counts) + else: + # Fixed number of images per request + image_counts = np.full(num_requests, image_count) + total_images = image_count * num_requests + + # Check for potentially problematic combinations and warn user + if width * height >= 1920 * 1080 and total_images >= 100: + warnings.warn( + f"High resolution ({width}x{height}) with {total_images} total images " + f"may take a long time. Consider reducing resolution or image count.", + UserWarning, + stacklevel=2, + ) + + # Sample text lengths + input_lens = compute_random_lens( + full_len=input_len, + range_ratio=range_ratio, + num=num_requests, + ) + output_lens = compute_random_lens( + full_len=output_len, + range_ratio=range_ratio, + num=num_requests, + ) + + def _gen_random_image_data_uri( + width: int = width, height: int = height + ) -> Tuple[Image.Image, str, int]: + if image_content == "blank": + # Generate blank white image + arr = np.full((height, width, 3), 255, dtype=np.uint8) + else: + # Generate random colored image + arr = (np.random.rand(height, width, 3) * 255).astype(np.uint8) + img = Image.fromarray(arr) + buf = io.BytesIO() + img.save(buf, format=image_format, quality=85) + encoded = pybase64.b64encode(buf.getvalue()).decode("utf-8") + image_data = f"data:image/{image_format};base64,{encoded}" + image_bytes = len(image_data.encode("utf-8")) + return img, image_data, image_bytes + + dataset: List[DatasetRow] = [] + total_image_bytes = 0 + for i in range(num_requests): + # Get the number of images for this request + request_image_count = int(image_counts[i]) + + # Generate text prompt + text_prompt = gen_mm_prompt( + processor.tokenizer, + processor.image_token_id if hasattr(processor, "image_token_id") else None, + int(input_lens[i]), + ) + + # Generate image list + images, images_base64, images_bytes = zip( + *[_gen_random_image_data_uri() for _ in range(request_image_count)] + ) + total_image_bytes += sum(images_bytes) + + data_row = create_mm_data_row( + text_prompt, + list(images), + list(images_base64), + int(output_lens[i]), + processor, + backend, + ) + dataset.append(data_row) + + # Print statistics + print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}") + print(f"#Total images: {total_images}") + + if random_image_count: + print( + f"#Images per request: min={np.min(image_counts)}, max={np.max(image_counts)}, mean={np.mean(image_counts):.2f}" + ) + else: + print(f"#Images per request: {image_count} (fixed)") + + print( + f"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request" + ) + return dataset diff --git a/python/sglang/benchmark/datasets/mmmu.py b/python/sglang/benchmark/datasets/mmmu.py new file mode 100644 index 000000000..762365548 --- /dev/null +++ b/python/sglang/benchmark/datasets/mmmu.py @@ -0,0 +1,97 @@ +import io +import random +from typing import List, Optional + +import pybase64 +from datasets import load_dataset +from transformers import AutoProcessor, AutoTokenizer + +from sglang.benchmark.datasets.common import DatasetRow +from sglang.benchmark.datasets.image import create_mm_data_row + + +def sample_mmmu_requests( + num_requests: int, + processor: AutoProcessor | AutoTokenizer, + backend: str = "sglang", + fixed_output_len: Optional[int] = None, + random_sample: bool = True, +) -> List[DatasetRow]: + """ + Sample requests from the MMMU dataset using HuggingFace datasets. + + Args: + num_requests: Number of requests to sample. + fixed_output_len: If provided, use this fixed output length for all requests. + random_sample: Whether to randomly sample or take the first N. + + Returns: + List of tuples (prompt, prompt_token_len, output_token_len). + """ + print("Loading MMMU dataset from HuggingFace...") + + try: + print("Attempting to load MMMU Math dataset...") + mmmu_dataset = load_dataset("MMMU/MMMU", "Math", split="test") + print( + f"Successfully loaded MMMU Math dataset from HuggingFace with {len(mmmu_dataset)} examples" + ) + except Exception as e: + print(f"Failed to load MMMU Math dataset: {e}") + raise ValueError(f"Failed to load MMMU dataset: {e}") + + # Sample from the dataset + if len(mmmu_dataset) > num_requests: + if random_sample: + # Random sample + indices = random.sample(range(len(mmmu_dataset)), num_requests) + sample_dataset = mmmu_dataset.select(indices) + else: + # Take first N + sample_dataset = mmmu_dataset.select( + range(min(num_requests, len(mmmu_dataset))) + ) + else: + print(f"Dataset has less than {num_requests} examples, using all examples") + sample_dataset = mmmu_dataset + + print(f"Selected {len(sample_dataset)} examples for benchmarking") + + # Create prompts + filtered_dataset = [] + + for i, example in enumerate(sample_dataset): + try: + # Extract image_1 + image = example.get("image_1") + + if image is not None: + if hasattr(image, "save"): + # Convert RGBA images to RGB before encoding + if image.mode == "RGBA": + image = image.convert("RGB") + + # Encode image to base64 (save as PNG to support palette/alpha modes) + buffered = io.BytesIO() + image.save(buffered, format="PNG") + img_str = pybase64.b64encode(buffered.getvalue()).decode("utf-8") + image_data = f"data:image/png;base64,{img_str}" + else: + continue + + # Extract the question + question = example.get("question") + + # Construct the prompt + text_prompt = f"Question: {question}\n\nAnswer: " + output_len = fixed_output_len if fixed_output_len is not None else 256 + data_row = create_mm_data_row( + text_prompt, [image], [image_data], output_len, processor, backend + ) + filtered_dataset.append(data_row) + + except Exception as e: + print(f"Error processing example {i}: {e}") + + print(f"\nCreated {len(filtered_dataset)} MMMU prompts") + return filtered_dataset diff --git a/python/sglang/benchmark/datasets/mooncake.py b/python/sglang/benchmark/datasets/mooncake.py new file mode 100644 index 000000000..bc3a7f1ce --- /dev/null +++ b/python/sglang/benchmark/datasets/mooncake.py @@ -0,0 +1,83 @@ +import asyncio +import time +from typing import AsyncGenerator, Dict, List + +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import DatasetRow + + +async def get_mooncake_request_over_time( + input_requests: List[Dict], + tokenizer: PreTrainedTokenizerBase, + slowdown_factor: float, + num_rounds: int, +) -> AsyncGenerator[DatasetRow, None]: + """ + An async generator that yields requests based on the timestamps in the Mooncake trace file, + with support for multi-round sessions. + """ + if not input_requests: + return + + input_requests.sort(key=lambda r: r["timestamp"]) + + start_time = time.perf_counter() + trace_start_time_ms = input_requests[0]["timestamp"] + + for record in input_requests: + # Calculate when this entire session should start + relative_arrival_time_s = (record["timestamp"] - trace_start_time_ms) / 1000.0 + target_arrival_time_s = relative_arrival_time_s * slowdown_factor + + current_elapsed_time_s = time.perf_counter() - start_time + sleep_duration_s = target_arrival_time_s - current_elapsed_time_s + if sleep_duration_s > 0: + await asyncio.sleep(sleep_duration_s) + + # Once the session starts, generate all rounds for it as a burst + # This simulates a user engaging in a multi-turn conversation + + # Base user query constructed from hash_ids + user_query_base = "" + hash_ids = record.get("hash_ids", []) + for hash_id in hash_ids: + user_query_base += f"{hash_id}" + " ".join( + ["hi"] * 128 + ) # Shorter for multi-round + user_query_base += "Tell me a story based on this context." + + output_len_per_round = record.get("output_length", 256) + chat_history = [] + + for i in range(num_rounds): + # Add user query for the current round + chat_history.append( + {"role": "user", "content": f"Round {i + 1}: {user_query_base}"} + ) + + # Form the full prompt from history + try: + full_prompt_text = tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=True, + return_dict=False, + ) + except Exception: + full_prompt_text = "\n".join( + [f"{msg['role']}: {msg['content']}" for msg in chat_history] + ) + + prompt_len = len(tokenizer.encode(full_prompt_text)) + + yield DatasetRow( + prompt=full_prompt_text, + prompt_len=prompt_len, + output_len=output_len_per_round, + ) + + # Add a placeholder assistant response for the next round's context + # We use a placeholder because we don't know the real response + placeholder_response = " ".join(["story"] * output_len_per_round) + chat_history.append({"role": "assistant", "content": placeholder_response}) diff --git a/python/sglang/benchmark/datasets/openai_dataset.py b/python/sglang/benchmark/datasets/openai_dataset.py new file mode 100644 index 000000000..c4d6c6c0b --- /dev/null +++ b/python/sglang/benchmark/datasets/openai_dataset.py @@ -0,0 +1,86 @@ +import json +from typing import List, Optional + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import DatasetRow + + +def sample_openai_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, +) -> List[DatasetRow]: + """ + Load OpenAI-compatible chat completion requests from a JSONL file. + + Each line should be a JSON object with: + - "messages": list of {"role": str, "content": str} + - "max_tokens": int (used as output_len if fixed_output_len not set) + - "tools": optional list of tool definitions + - "temperature": optional temperature value + - "top_p": optional top_p value + - Other OpenAI API parameters are also extracted and passed through + """ + dataset = [] + with open(dataset_path, "r") as f: + for line in f: + if num_requests > 0 and len(dataset) >= num_requests: + break + if line.strip(): + try: + dataset.append(json.loads(line)) + except json.JSONDecodeError: + # Skip invalid JSON lines + continue + + # Fields that should NOT be passed through extra_request_body + # These are either handled separately or are metadata + # max_tokens is excluded because it's handled via output_len -> max_completion_tokens + # max_completion_tokens is also excluded to avoid conflicts + EXCLUDED_FIELDS = {"messages", "max_tokens", "max_completion_tokens", "model"} + + filtered_dataset: List[DatasetRow] = [] + for data in dataset: + messages = data.get("messages", []) + if not messages: + continue + + # Use max_tokens from the request, or fall back to fixed_output_len + output_len = fixed_output_len or data.get("max_tokens", 256) + + # Extract extra request body parameters (tools, temperature, top_p, etc.) + extra_body = {k: v for k, v in data.items() if k not in EXCLUDED_FIELDS} + + # Calculate prompt length by applying chat template + # This includes the messages but not the tools + prompt_len = len( + tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + ) + ) + + # If tools are present, we need to add their token count + # Tools are sent as part of the request and count toward input tokens + if "tools" in extra_body: + # Encode tools as JSON string to estimate token count + tools_str = json.dumps(extra_body["tools"]) + tools_tokens = len(tokenizer.encode(tools_str)) + prompt_len += tools_tokens + + # Pass messages list directly - bench_serving handles List[Dict] prompts + filtered_dataset.append( + DatasetRow( + prompt=messages, + prompt_len=prompt_len, + output_len=output_len, + extra_request_body=extra_body, # Store per-request parameters + ) + ) + + print(f"Loaded {len(filtered_dataset)} OpenAI-format requests") + print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") + return filtered_dataset diff --git a/python/sglang/benchmark/datasets/random.py b/python/sglang/benchmark/datasets/random.py new file mode 100644 index 000000000..b7d38e96f --- /dev/null +++ b/python/sglang/benchmark/datasets/random.py @@ -0,0 +1,127 @@ +import json +import random +from typing import List + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import ( + SHAREGPT_FILENAME, + SHAREGPT_REPO_ID, + DatasetRow, + compute_random_lens, +) +from sglang.benchmark.utils import download_and_cache_hf_file, is_file_valid_json + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, + random_sample: bool = True, + return_text: bool = True, +) -> List[DatasetRow]: + input_lens = compute_random_lens( + full_len=input_len, + range_ratio=range_ratio, + num=num_prompts, + ) + output_lens = compute_random_lens( + full_len=output_len, + range_ratio=range_ratio, + num=num_prompts, + ) + + if return_text: + # Need to truncate input_len as server encode will add special token. + num_special_tokens = int(tokenizer.num_special_tokens_to_add()) + for i in range(num_prompts): + input_lens[i] = max(0, input_lens[i] - num_special_tokens) + + if random_sample: + # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens + + # Download sharegpt if necessary + if not is_file_valid_json(dataset_path): + dataset_path = download_and_cache_hf_file( + repo_id=SHAREGPT_REPO_ID, + filename=SHAREGPT_FILENAME, + ) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + input_requests: List[DatasetRow] = [] + for data in dataset: + i = len(input_requests) + if i == num_prompts: + break + + # Tokenize the prompts and completions. + prompt = data[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + # Skip empty prompt + if prompt_len == 0: + continue + + if prompt_len > input_lens[i]: + input_ids = prompt_token_ids[: input_lens[i]] + else: + ratio = (input_lens[i] + prompt_len - 1) // prompt_len + input_ids = (prompt_token_ids * ratio)[: input_lens[i]] + input_content = input_ids + if return_text: + input_content = tokenizer.decode(input_content) + input_requests.append( + DatasetRow( + prompt=input_content, + prompt_len=input_lens[i], + output_len=output_lens[i], + ) + ) + else: + # Sample token ids from random integers. This can cause some NaN issues. + offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) + input_requests = [] + for i in range(num_prompts): + # Use int() to convert numpy.int64 to native Python int for JSON serialization + input_content = [ + int((offsets[i] + i + j) % tokenizer.vocab_size) + for j in range(input_lens[i]) + ] + if return_text: + input_content = tokenizer.decode(input_content) + input_requests.append( + DatasetRow( + prompt=input_content, + prompt_len=input_lens[i], + output_len=output_lens[i], + ) + ) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests diff --git a/python/sglang/benchmark/datasets/sharegpt.py b/python/sglang/benchmark/datasets/sharegpt.py new file mode 100644 index 000000000..c7255fd59 --- /dev/null +++ b/python/sglang/benchmark/datasets/sharegpt.py @@ -0,0 +1,113 @@ +import json +import random +from typing import List, Optional + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import ( + ASSISTANT_SUFFIX, + SHAREGPT_FILENAME, + SHAREGPT_REPO_ID, + DatasetRow, +) +from sglang.benchmark.utils import ( + download_and_cache_hf_file, + is_file_valid_json, + remove_suffix, +) + + +def sample_sharegpt_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, + context_len: Optional[int] = None, + prompt_suffix: Optional[str] = "", + apply_chat_template=False, +) -> List[DatasetRow]: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Download sharegpt if necessary + if not is_file_valid_json(dataset_path) and dataset_path == "": + dataset_path = download_and_cache_hf_file( + repo_id=SHAREGPT_REPO_ID, + filename=SHAREGPT_FILENAME, + ) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: List[DatasetRow] = [] + for i in range(len(dataset)): + if len(filtered_dataset) == num_requests: + break + + # Tokenize the prompts and completions. + prompt = dataset[i][0] + if prompt_suffix: + prompt = ( + remove_suffix(prompt, ASSISTANT_SUFFIX) + + prompt_suffix + + ASSISTANT_SUFFIX + ) + + if apply_chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + return_dict=False, + ) + if tokenizer.bos_token: + prompt = prompt.replace(tokenizer.bos_token, "") + + prompt_token_ids = tokenizer.encode(prompt) + completion = dataset[i][1] + completion_token_ids = tokenizer.encode(completion) + prompt_len = len(prompt_token_ids) + output_len = ( + len(completion_token_ids) if fixed_output_len is None else fixed_output_len + ) + + if prompt_len < 2 or output_len < 2: + # Prune too short sequences. + continue + + if context_len and prompt_len + output_len > context_len: + # Prune too long sequences. + continue + + filtered_dataset.append( + DatasetRow( + prompt=prompt, + prompt_len=prompt_len, + output_len=output_len, + ) + ) + + print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") + return filtered_dataset diff --git a/python/sglang/benchmark/utils.py b/python/sglang/benchmark/utils.py new file mode 100644 index 000000000..7bf6494b5 --- /dev/null +++ b/python/sglang/benchmark/utils.py @@ -0,0 +1,159 @@ +import json +import os +import resource +from json import JSONDecodeError +from typing import Dict, List, Optional, Union + +import requests +from tqdm.asyncio import tqdm +from transformers import ( + AutoProcessor, + AutoTokenizer, + PreTrainedTokenizer, + PreTrainedTokenizerFast, +) + + +def remove_prefix(text: str, prefix: str) -> str: + return text[len(prefix) :] if text.startswith(prefix) else text + + +def remove_suffix(text: str, suffix: str) -> str: + return text[: -len(suffix)] if text.endswith(suffix) else text + + +def parse_custom_headers(header_list: List[str]) -> Dict[str, str]: + return {k: v for h in header_list for k, _, v in [h.partition("=")] if k and v} + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true": + import huggingface_hub.constants + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], + ) + + return model_path + return pretrained_model_name_or_path + + +def get_tokenizer( + pretrained_model_name_or_path: str, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + assert ( + pretrained_model_name_or_path is not None + and pretrained_model_name_or_path != "" + ) + if pretrained_model_name_or_path.endswith( + ".json" + ) or pretrained_model_name_or_path.endswith(".model"): + from sglang.srt.utils.hf_transformers_utils import get_tokenizer + + return get_tokenizer(pretrained_model_name_or_path) + + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path + ): + pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) + return AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=True + ) + + +def get_processor( + pretrained_model_name_or_path: str, +) -> AutoProcessor: + assert ( + pretrained_model_name_or_path is not None + and pretrained_model_name_or_path != "" + ) + if pretrained_model_name_or_path.endswith( + ".json" + ) or pretrained_model_name_or_path.endswith(".model"): + from sglang.srt.utils.hf_transformers_utils import get_processor + + return get_processor(pretrained_model_name_or_path) + + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path + ): + pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) + return AutoProcessor.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=True + ) + + +def download_and_cache_hf_file( + repo_id: str, + filename: str, + repo_type: str = "dataset", +): + """Download a file from Hugging Face and cache it locally.""" + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type) + + +def download_and_cache_file(url: str, filename: Optional[str] = None): + """Read and cache a file from a url.""" + if filename is None: + filename = os.path.join("/tmp", url.split("/")[-1]) + + # Check if the cache file already exists + if is_file_valid_json(filename): + return filename + + print(f"Downloading from {url} to {filename}") + + # Stream the response to show the progress bar + response = requests.get(url, stream=True) + response.raise_for_status() # Check for request errors + + # Total size of the file in bytes + total_size = int(response.headers.get("content-length", 0)) + chunk_size = 1024 # Download in chunks of 1KB + + # Use tqdm to display the progress bar + with open(filename, "wb") as f, tqdm( + desc=filename, + total=total_size, + unit="B", + unit_scale=True, + unit_divisor=1024, + ) as bar: + for chunk in response.iter_content(chunk_size=chunk_size): + f.write(chunk) + bar.update(len(chunk)) + + return filename + + +def is_file_valid_json(path): + if not os.path.isfile(path): + return False + + # TODO can fuse into the real file open later + try: + with open(path) as f: + json.load(f) + return True + except JSONDecodeError as e: + print( + f"{path} exists but json loading fails ({e=}), thus treat as invalid file" + ) + return False + + +def set_ulimit(target_soft_limit=65535): + resource_type = resource.RLIMIT_NOFILE + current_soft, current_hard = resource.getrlimit(resource_type) + + if current_soft < target_soft_limit: + try: + resource.setrlimit(resource_type, (target_soft_limit, current_hard)) + except ValueError as e: + print(f"Fail to set RLIMIT_NOFILE: {e}")