chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Benchmark: Add typed DatasetArgs/Loader registry and CPU dataset unit tests (#19147)

Browse Source 
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
This commit is contained in:
Ratish P
2026-02-25 01:52:01 +05:30
committed by GitHub
parent 65de904130
commit ae6f6e1495
19 changed files with 807 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/bench_offline_throughput.py
View File

@@ -23,7 +23,8 @@ from typing import Dict, List, Optional
import numpy as np
from sglang.benchmark.datasets import DatasetRow, get_dataset, sample_random_requests
from sglang.benchmark.datasets import DatasetRow, get_dataset
from sglang.benchmark.datasets.random import sample_random_requests
from sglang.benchmark.utils import get_tokenizer, set_ulimit
from sglang.lang.backend.runtime_endpoint import Runtime
from sglang.srt.entrypoints.engine import Engine

7
python/sglang/bench_serving.py
View File

@@ -36,11 +36,8 @@ import requests
from tqdm.asyncio import tqdm
from transformers import AutoTokenizer, PreTrainedTokenizerBase
from sglang.benchmark.datasets import (
DatasetRow,
get_dataset,
get_mooncake_request_over_time,
)
from sglang.benchmark.datasets import DatasetRow, get_dataset
from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time
from sglang.benchmark.utils import (
get_tokenizer,
parse_custom_headers,

175
python/sglang/benchmark/datasets/__init__.py
View File

@@ -1,156 +1,47 @@
import json
import os
from typing import Dict, Type
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.common import BaseDataset, DatasetRow
from sglang.benchmark.datasets.custom import CustomDataset
from sglang.benchmark.datasets.generated_shared_prefix import (
get_gen_prefix_cache_path,
sample_generated_shared_prefix_requests,
GeneratedSharedPrefixDataset,
)
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
from sglang.benchmark.datasets.image import ImageDataset
from sglang.benchmark.datasets.mmmu import MMMUDataset
from sglang.benchmark.datasets.mooncake import MooncakeDataset
from sglang.benchmark.datasets.openai_dataset import OpenAIDataset
from sglang.benchmark.datasets.random import RandomDataset
from sglang.benchmark.datasets.sharegpt import ShareGPTDataset
DATASET_MAPPING: Dict[str, Type[BaseDataset]] = {
"sharegpt": ShareGPTDataset,
"custom": CustomDataset,
"openai": OpenAIDataset,
# TODO: "random" vs "random-ids" should be a flag (e.g. --random-source=sharegpt|integers),
# not two separate dataset names sharing the same class.
"random": RandomDataset,
"random-ids": RandomDataset,
"generated-shared-prefix": GeneratedSharedPrefixDataset,
"mmmu": MMMUDataset,
"image": ImageDataset,
"mooncake": MooncakeDataset,
}
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
dataset_name = args.dataset_name
if dataset_name.startswith("random") and dataset_name not in DATASET_MAPPING:
dataset_name = "random-ids"
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:
if dataset_name not in DATASET_MAPPING:
raise ValueError(f"Unknown dataset: {args.dataset_name}")
return input_requests
dataset_cls = DATASET_MAPPING[dataset_name]
dataset = dataset_cls.from_args(args)
return dataset.load(tokenizer=tokenizer, model_id=model_id)
__all__ = [
"ASSISTANT_SUFFIX",
"MOONCAKE_DATASET_URL",
"SHAREGPT_FILENAME",
"SHAREGPT_REPO_ID",
"DATASET_MAPPING",
"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",
]

16
python/sglang/benchmark/datasets/common.py
View File

@@ -1,4 +1,6 @@
import random
from abc import ABC, abstractmethod
from argparse import Namespace
from dataclasses import dataclass
from functools import lru_cache
from typing import Any, Dict, List, Optional
@@ -37,6 +39,20 @@ class DatasetRow:
self.extra_request_body = {}
@dataclass
class BaseDataset(ABC):
@classmethod
@abstractmethod
def from_args(cls, args: Namespace) -> "BaseDataset": ...
@abstractmethod
def load(
self,
tokenizer: Any,
model_id: Optional[str] = None,
) -> List[DatasetRow]: ...
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),

43
python/sglang/benchmark/datasets/custom.py
View File

@@ -1,15 +1,56 @@
import json
import os
import random
from argparse import Namespace
from dataclasses import dataclass
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.datasets.common import (
ASSISTANT_SUFFIX,
BaseDataset,
DatasetRow,
)
from sglang.benchmark.utils import remove_suffix
@dataclass
class CustomDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
context_len: Optional[int]
prompt_suffix: str
apply_chat_template: bool
@classmethod
def from_args(cls, args: Namespace) -> "CustomDataset":
assert not getattr(args, "tokenize_prompt", False)
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
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,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_custom_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
context_len=self.context_len,
prompt_suffix=self.prompt_suffix,
apply_chat_template=self.apply_chat_template,
)
def sample_custom_requests(
dataset_path: str,
num_requests: int,

55
python/sglang/benchmark/datasets/generated_shared_prefix.py
View File

@@ -2,6 +2,8 @@ import argparse
import pickle
import random
import uuid
from argparse import Namespace
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import List
@@ -10,7 +12,58 @@ 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
from sglang.benchmark.datasets.common import (
BaseDataset,
DatasetRow,
compute_random_lens,
gen_prompt,
)
@dataclass
class GeneratedSharedPrefixDataset(BaseDataset):
num_groups: int
prompts_per_group: int
system_prompt_len: int
question_len: int
output_len: int
range_ratio: float
seed: int
gsp_fast_prepare: bool
gsp_send_routing_key: bool
gsp_num_turns: int
gsp_ordered: bool
@classmethod
def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset":
assert not getattr(args, "tokenize_prompt", False)
return cls(
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),
seed=args.seed,
gsp_fast_prepare=getattr(args, "gsp_fast_prepare", False),
gsp_send_routing_key=getattr(args, "gsp_send_routing_key", False),
gsp_num_turns=getattr(args, "gsp_num_turns", 1),
gsp_ordered=getattr(args, "gsp_ordered", False),
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_generated_shared_prefix_requests(
num_groups=self.num_groups,
prompts_per_group=self.prompts_per_group,
system_prompt_len=self.system_prompt_len,
question_len=self.question_len,
output_len=self.output_len,
range_ratio=self.range_ratio,
tokenizer=tokenizer,
args=self,
)
def get_gen_prefix_cache_path(args, tokenizer):

49
python/sglang/benchmark/datasets/image.py
View File

@@ -1,5 +1,7 @@
import io
import warnings
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Tuple
import numpy as np
@@ -8,10 +10,57 @@ from PIL import Image
from transformers import AutoProcessor
from sglang.benchmark.datasets.common import (
BaseDataset,
DatasetRow,
compute_random_lens,
gen_mm_prompt,
)
from sglang.benchmark.utils import get_processor
@dataclass
class ImageDataset(BaseDataset):
num_requests: int
image_count: int
input_len: int
output_len: int
range_ratio: float
image_content: str
image_format: str
image_resolution: str
backend: str
random_image_count: bool
@classmethod
def from_args(cls, args: Namespace) -> "ImageDataset":
return cls(
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,
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,
)
def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]:
processor = get_processor(model_id)
return sample_image_requests(
num_requests=self.num_requests,
image_count=self.image_count,
input_len=self.input_len,
output_len=self.output_len,
range_ratio=self.range_ratio,
processor=processor,
image_content=self.image_content,
image_format=self.image_format,
image_resolution=self.image_resolution,
backend=self.backend,
random_image_count=self.random_image_count,
)
def parse_image_resolution(image_resolution: str) -> Tuple[int, int]:

29
python/sglang/benchmark/datasets/mmmu.py
View File

@@ -1,13 +1,40 @@
import io
import random
from argparse import Namespace
from dataclasses import dataclass
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.common import BaseDataset, DatasetRow
from sglang.benchmark.datasets.image import create_mm_data_row
from sglang.benchmark.utils import get_processor
@dataclass
class MMMUDataset(BaseDataset):
num_requests: int
backend: str
fixed_output_len: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "MMMUDataset":
return cls(
num_requests=args.num_prompts,
backend=args.backend,
fixed_output_len=args.random_output_len,
)
def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]:
processor = get_processor(model_id)
return sample_mmmu_requests(
num_requests=self.num_requests,
processor=processor,
backend=self.backend,
fixed_output_len=self.fixed_output_len,
)
def sample_mmmu_requests(

42
python/sglang/benchmark/datasets/mooncake.py
View File

@@ -1,10 +1,50 @@
import asyncio
import json
import os
import time
from argparse import Namespace
from dataclasses import dataclass
from typing import AsyncGenerator, Dict, List
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import DatasetRow
from sglang.benchmark.datasets.common import (
MOONCAKE_DATASET_URL,
BaseDataset,
DatasetRow,
)
from sglang.benchmark.utils import download_and_cache_file
@dataclass
class MooncakeDataset(BaseDataset):
dataset_path: str
mooncake_workload: str
num_requests: int
@classmethod
def from_args(cls, args: Namespace) -> "MooncakeDataset":
return cls(
dataset_path=args.dataset_path,
mooncake_workload=args.mooncake_workload,
num_requests=args.num_prompts,
)
def load(self, tokenizer=None, model_id=None) -> List[Dict]:
if not self.dataset_path:
local_path = os.path.join("/tmp", self.mooncake_workload + "_trace.jsonl")
else:
local_path = self.dataset_path
if not os.path.exists(local_path):
download_and_cache_file(
MOONCAKE_DATASET_URL[self.mooncake_workload], local_path
)
with open(local_path, "r") as f:
all_requests_data = [json.loads(line) for line in f if line.strip()]
return all_requests_data[: self.num_requests]
async def get_mooncake_request_over_time(

29
python/sglang/benchmark/datasets/openai_dataset.py
View File

@@ -1,10 +1,37 @@
import json
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
import numpy as np
from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import DatasetRow
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
@dataclass
class OpenAIDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
@classmethod
def from_args(cls, args: Namespace) -> "OpenAIDataset":
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
fixed_output_len=args.sharegpt_output_len,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_openai_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
)
def sample_openai_requests(

40
python/sglang/benchmark/datasets/random.py
View File

@@ -1,5 +1,7 @@
import json
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List
import numpy as np
@@ -8,12 +10,50 @@ from transformers import PreTrainedTokenizerBase
from sglang.benchmark.datasets.common import (
SHAREGPT_FILENAME,
SHAREGPT_REPO_ID,
BaseDataset,
DatasetRow,
compute_random_lens,
)
from sglang.benchmark.utils import download_and_cache_hf_file, is_file_valid_json
@dataclass
class RandomDataset(BaseDataset):
input_len: int
output_len: int
num_requests: int
range_ratio: float
dataset_path: str
return_text: bool
random_sample: bool
@classmethod
def from_args(cls, args: Namespace) -> "RandomDataset":
return cls(
input_len=args.random_input_len,
output_len=args.random_output_len,
num_requests=args.num_prompts,
range_ratio=args.random_range_ratio,
dataset_path=args.dataset_path,
return_text=not getattr(args, "tokenize_prompt", False),
random_sample=(args.dataset_name == "random"),
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_random_requests(
input_len=self.input_len,
output_len=self.output_len,
num_prompts=self.num_requests,
range_ratio=self.range_ratio,
tokenizer=tokenizer,
dataset_path=self.dataset_path,
random_sample=self.random_sample,
return_text=self.return_text,
)
def sample_random_requests(
input_len: int,
output_len: int,

38
python/sglang/benchmark/datasets/sharegpt.py
View File

@@ -1,5 +1,7 @@
import json
import random
from argparse import Namespace
from dataclasses import dataclass
from typing import List, Optional
import numpy as np
@@ -9,6 +11,7 @@ from sglang.benchmark.datasets.common import (
ASSISTANT_SUFFIX,
SHAREGPT_FILENAME,
SHAREGPT_REPO_ID,
BaseDataset,
DatasetRow,
)
from sglang.benchmark.utils import (
@@ -18,6 +21,41 @@ from sglang.benchmark.utils import (
)
@dataclass
class ShareGPTDataset(BaseDataset):
dataset_path: str
num_requests: int
fixed_output_len: Optional[int]
context_len: Optional[int]
prompt_suffix: str
apply_chat_template: bool
@classmethod
def from_args(cls, args: Namespace) -> "ShareGPTDataset":
assert not getattr(args, "tokenize_prompt", False)
return cls(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
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,
)
def load(
self, tokenizer: PreTrainedTokenizerBase, model_id=None
) -> List[DatasetRow]:
return sample_sharegpt_requests(
dataset_path=self.dataset_path,
num_requests=self.num_requests,
tokenizer=tokenizer,
fixed_output_len=self.fixed_output_len,
context_len=self.context_len,
prompt_suffix=self.prompt_suffix,
apply_chat_template=self.apply_chat_template,
)
def sample_sharegpt_requests(
dataset_path: str,
num_requests: int,

2
python/sglang/test/kits/cache_hit_kit.py
View File

@@ -6,7 +6,7 @@ import aiohttp
import requests
from sglang.bench_serving import RequestFuncOutput
from sglang.benchmark.datasets import sample_random_requests
from sglang.benchmark.datasets.random import sample_random_requests
from sglang.benchmark.utils import get_tokenizer, remove_prefix
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60)