@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 31436e8b4f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Misc] Add request_id into benchmark_serve.py (#23065)

Signed-off-by: yangxia <yangxiast@gmail.com>

hustxiayang committed Aug 19, 2025 at 04:32 UTC 31436e8b4f0a210cf5c6f9ec9b47fc1b5d7cc47a

6 files changed +243 -46

benchmarks/backend_request_func.py

+21 -2

@@ -34,6 +34,7 @@ class RequestFuncInput:

34 multi_modal_content: Optional[dict | list[dict]] = None

35 ignore_eos: bool = False

36 language: Optional[str] = None

37 + request_id: Optional[str] = None

40 @dataclass

@@ -71,6 +72,9 @@ async def async_request_tgi(

72 "inputs": request_func_input.prompt,

73 "parameters": params,

74 }

75 + headers = None

76 + if request_func_input.request_id:

77 + headers = {"x-request-id": request_func_input.request_id}

78 output = RequestFuncOutput()

79 output.prompt_len = request_func_input.prompt_len

80 if request_func_input.ignore_eos:

@@ -82,7 +86,9 @@ async def async_request_tgi(

86 st = time.perf_counter()

87 most_recent_timestamp = st

88 try:

85 - async with session.post(url=api_url, json=payload) as response:

89 + async with session.post(

90 + url=api_url, json=payload, headers=headers

91 + ) as response:

92 if response.status == 200:

93 async for chunk_bytes in response.content:

94 chunk_bytes = chunk_bytes.strip()

@@ -145,6 +151,9 @@ async def async_request_trt_llm(

151 }

152 if request_func_input.ignore_eos:

153 payload["min_length"] = request_func_input.output_len

154 + headers = None

155 + if request_func_input.request_id:

156 + headers = {"x-request-id": request_func_input.request_id}

157 output = RequestFuncOutput()

158 output.prompt_len = request_func_input.prompt_len

159

@@ -152,7 +161,9 @@ async def async_request_trt_llm(

161 st = time.perf_counter()

162 most_recent_timestamp = st

163 try:

155 - async with session.post(url=api_url, json=payload) as response:

164 + async with session.post(

165 + url=api_url, json=payload, headers=headers

166 + ) as response:

167 if response.status == 200:

168 async for chunk_bytes in response.content:

169 chunk_bytes = chunk_bytes.strip()

@@ -211,6 +222,8 @@ async def async_request_deepspeed_mii(

222 "top_p": 1.0,

223 }

224 headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}

225 + if request_func_input.request_id:

226 + headers["x-request-id"] = request_func_input.request_id

227

228 output = RequestFuncOutput()

229 output.prompt_len = request_func_input.prompt_len

@@ -283,6 +296,8 @@ async def async_request_openai_completions(

296 if request_func_input.extra_body:

297 payload.update(request_func_input.extra_body)

298 headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}

299 + if request_func_input.request_id:

300 + headers["x-request-id"] = request_func_input.request_id

301

302 output = RequestFuncOutput()

303 output.prompt_len = request_func_input.prompt_len

@@ -395,6 +410,8 @@ async def async_request_openai_chat_completions(

410 "Content-Type": "application/json",

411 "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",

412 }

413 + if request_func_input.request_id:

414 + headers["x-request-id"] = request_func_input.request_id

415

416 output = RequestFuncOutput()

417 output.prompt_len = request_func_input.prompt_len

@@ -491,6 +508,8 @@ async def async_request_openai_audio(

508 headers = {

509 "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",

510 }

511 + if request_func_input.request_id:

512 + headers["x-request-id"] = request_func_input.request_id

513

514 # Send audio file

515 def to_bytes(y, sr):

benchmarks/benchmark_dataset.py

+90 -19

@@ -19,6 +19,7 @@ import logging

19 import random

20 from abc import ABC, abstractmethod

21 from collections.abc import Mapping

22 +from copy import deepcopy

23 from dataclasses import dataclass

24 from functools import cache

25 from io import BytesIO

@@ -54,6 +55,7 @@ class SampleRequest:

55 expected_output_len: int

56 multi_modal_data: Optional[Union[MultiModalDataDict, dict, list[dict]]] = None

57 lora_request: Optional[LoRARequest] = None

58 + request_id: Optional[str] = None

61 # -----------------------------------------------------------------------------

@@ -155,7 +157,10 @@ class BenchmarkDataset(ABC):

157

158 @abstractmethod

159 def sample(

158 - self, tokenizer: PreTrainedTokenizerBase, num_requests: int

160 + self,

161 + tokenizer: PreTrainedTokenizerBase,

162 + num_requests: int,

163 + request_id_prefix: str = "",

164 ) -> list[SampleRequest]:

165 """

166 Abstract method to generate sample requests from the dataset.

@@ -167,6 +172,7 @@ class BenchmarkDataset(ABC):

172 tokenizer (PreTrainedTokenizerBase): The tokenizer to be used

173 for processing the dataset's text.

174 num_requests (int): The number of sample requests to generate.

175 + request_id_prefix (str) The prefix of request_id.

176

177 Returns:

178 list[SampleRequest]: A list of sample requests generated from the

@@ -175,7 +181,10 @@ class BenchmarkDataset(ABC):

181 raise NotImplementedError("sample must be implemented in subclasses.")

182

183 def maybe_oversample_requests(

178 - self, requests: list[SampleRequest], num_requests: int

184 + self,

185 + requests: list[SampleRequest],

186 + num_requests: int,

187 + request_id_prefix: str = "",

188 ) -> None:

189 """

190 Oversamples the list of requests if its size is less than the desired

@@ -183,11 +192,18 @@ class BenchmarkDataset(ABC):

192

193 Args:

194 requests (List[SampleRequest]): The current list of sampled

186 - requests. num_requests (int): The target number of requests.

195 + requests.

196 + num_requests (int): The target number of requests.

197 + request_id_prefix (str) The prefix of the request ids.

198 """

199 if len(requests) < num_requests:

200 random.seed(self.random_seed)

190 - additional = random.choices(requests, k=num_requests - len(requests))

201 + additional = deepcopy(

202 + random.choices(requests, k=num_requests - len(requests))

203 + )

204 + for i in range(len(additional)):

205 + req = additional[i]

206 + req.request_id = request_id_prefix + str(len(requests) + i)

207 requests.extend(additional)

208 logger.info("Oversampled requests to reach %d total samples.", num_requests)

209

@@ -303,6 +319,7 @@ class RandomDataset(BenchmarkDataset):

319 range_ratio: float = DEFAULT_RANGE_RATIO,

320 input_len: int = DEFAULT_INPUT_LEN,

321 output_len: int = DEFAULT_OUTPUT_LEN,

322 + request_id_prefix: str = "",

323 **kwargs,

324 ) -> list[SampleRequest]:

325 # Enforce range_ratio < 1

@@ -363,8 +380,10 @@ class RandomDataset(BenchmarkDataset):

380 prompt=prompt,

381 prompt_len=total_input_len,

382 expected_output_len=int(output_lens[i]),

383 + request_id=request_id_prefix + str(i),

384 )

385 )

386 +

387 return requests

388

389

@@ -406,9 +425,11 @@ class ShareGPTDataset(BenchmarkDataset):

425 max_loras: Optional[int] = None,

426 output_len: Optional[int] = None,

427 enable_multimodal_chat: bool = False,

428 + request_id_prefix: str = "",

429 **kwargs,

430 ) -> list:

431 samples: list = []

432 + ind = 0

433 for entry in self.data:

434 if len(samples) >= num_requests:

435 break

@@ -444,9 +465,11 @@ class ShareGPTDataset(BenchmarkDataset):

465 expected_output_len=new_output_len,

466 lora_request=lora_request,

467 multi_modal_data=mm_content,

468 + request_id=request_id_prefix + str(ind),

469 )

470 )

449 - self.maybe_oversample_requests(samples, num_requests)

471 + ind += 1

472 + self.maybe_oversample_requests(samples, num_requests, request_id_prefix)

473 return samples

474

475

@@ -512,10 +535,11 @@ class CustomDataset(BenchmarkDataset):

535 output_len: Optional[int] = None,

536 enable_multimodal_chat: bool = False,

537 skip_chat_template: bool = False,

538 + request_id_prefix: str = "",

539 **kwargs,

540 ) -> list:

541 sampled_requests = []

518 - for item in self.data:

542 + for i, item in enumerate(self.data):

543 if len(sampled_requests) >= num_requests:

544 break

545 prompt = item["prompt"]

@@ -534,9 +558,12 @@ class CustomDataset(BenchmarkDataset):

558 prompt=prompt,

559 prompt_len=prompt_len,

560 expected_output_len=output_len,

561 + request_id=request_id_prefix + str(i),

562 )

563 )

539 - self.maybe_oversample_requests(sampled_requests, num_requests)

564 + self.maybe_oversample_requests(

565 + sampled_requests, num_requests, request_id_prefix

566 + )

567

568 return sampled_requests

569

@@ -578,6 +605,7 @@ class SonnetDataset(BenchmarkDataset):

605 input_len: int = DEFAULT_INPUT_LEN,

606 output_len: int = DEFAULT_OUTPUT_LEN,

607 return_prompt_formatted: bool = False,

608 + request_id_prefix: str = "",

609 **kwargs,

610 ) -> list:

611 # Calculate average token length for a poem line.

@@ -603,6 +631,7 @@ class SonnetDataset(BenchmarkDataset):

631 prefix_lines = self.data[:num_prefix_lines]

632

633 samples = []

634 + ind = 0

635 while len(samples) < num_requests:

636 extra_lines = random.choices(

637 self.data, k=num_input_lines - num_prefix_lines

@@ -613,14 +642,17 @@ class SonnetDataset(BenchmarkDataset):

642 msg, add_generation_prompt=True, tokenize=False

643 )

644 prompt_len = len(tokenizer(prompt_formatted).input_ids)

645 +

646 if prompt_len <= input_len:

647 samples.append(

648 SampleRequest(

649 prompt=prompt_formatted if return_prompt_formatted else prompt,

650 prompt_len=prompt_len,

651 expected_output_len=output_len,

652 + request_id=request_id_prefix + str(ind),

653 )

654 )

655 + ind += 1

656 return samples

657

658

@@ -672,6 +704,7 @@ class BurstGPTDataset(BenchmarkDataset):

704 num_requests: int,

705 max_loras: Optional[int] = None,

706 lora_path: Optional[str] = None,

707 + request_id_prefix: str = "",

708 **kwargs,

709 ) -> list[SampleRequest]:

710 samples = []

@@ -693,6 +726,7 @@ class BurstGPTDataset(BenchmarkDataset):

726 prompt_len=input_len,

727 expected_output_len=output_len,

728 lora_request=lora_req,

729 + request_id=request_id_prefix + str(i),

730 )

731 )

732 return samples

@@ -752,12 +786,14 @@ class ConversationDataset(HuggingFaceDataset):

786 num_requests: int,

787 output_len: Optional[int] = None,

788 enable_multimodal_chat: bool = False,

789 + request_id_prefix: str = "",

790 **kwargs,

791 ) -> list:

792 # Filter examples with at least 2 conversations

793 filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)

794 sampled_requests = []

795 dynamic_output = output_len is None

796 + ind = 0

797

798 for item in filtered_data:

799 if len(sampled_requests) >= num_requests:

@@ -785,9 +821,13 @@ class ConversationDataset(HuggingFaceDataset):

821 prompt_len=prompt_len,

822 expected_output_len=output_len,

823 multi_modal_data=mm_content,

824 + request_id=request_id_prefix + str(ind),

825 )

826 )

790 - self.maybe_oversample_requests(sampled_requests, num_requests)

827 + ind += 1

828 + self.maybe_oversample_requests(

829 + sampled_requests, num_requests, request_id_prefix

830 + )

831 return sampled_requests

832

833

@@ -814,11 +854,12 @@ class VisionArenaDataset(HuggingFaceDataset):

854 num_requests: int,

855 output_len: Optional[int] = None,

856 enable_multimodal_chat: bool = False,

857 + request_id_prefix: str = "",

858 **kwargs,

859 ) -> list:

860 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

861 sampled_requests = []

821 - for item in self.data:

862 + for i, item in enumerate(self.data):

863 if len(sampled_requests) >= num_requests:

864 break

865 parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.dataset_path)

@@ -838,9 +879,12 @@ class VisionArenaDataset(HuggingFaceDataset):

879 prompt_len=prompt_len,

880 expected_output_len=output_len,

881 multi_modal_data=mm_content,

882 + request_id=request_id_prefix + str(i),

883 )

884 )

843 - self.maybe_oversample_requests(sampled_requests, num_requests)

885 + self.maybe_oversample_requests(

886 + sampled_requests, num_requests, request_id_prefix

887 + )

888 return sampled_requests

889

890

@@ -870,11 +914,12 @@ class InstructCoderDataset(HuggingFaceDataset):

914 num_requests: int,

915 output_len: Optional[int] = None,

916 enable_multimodal_chat: bool = False,

917 + request_id_prefix: str = "",

918 **kwargs,

919 ) -> list:

920 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

921 sampled_requests = []

877 - for item in self.data:

922 + for i, item in enumerate(self.data):

923 if len(sampled_requests) >= num_requests:

924 break

925 prompt = f"{item['input']}\n\n{item['instruction']} Just output \

@@ -892,9 +937,12 @@ class InstructCoderDataset(HuggingFaceDataset):

937 prompt=prompt,

938 prompt_len=prompt_len,

939 expected_output_len=output_len,

940 + request_id=request_id_prefix + str(i),

941 )

942 )

897 - self.maybe_oversample_requests(sampled_requests, num_requests)

943 + self.maybe_oversample_requests(

944 + sampled_requests, num_requests, request_id_prefix

945 + )

946 return sampled_requests

947

948

@@ -924,12 +972,13 @@ class MTBenchDataset(HuggingFaceDataset):

972 num_requests: int,

973 output_len: Optional[int] = None,

974 enable_multimodal_chat: bool = False,

975 + request_id_prefix: str = "",

976 **kwargs,

977 ) -> list:

978 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN

979 sampled_requests = []

980

932 - for item in self.data:

981 + for i, item in enumerate(self.data):

982 if len(sampled_requests) >= num_requests:

983 break

984 prompt = item["turns"][0]

@@ -947,9 +996,12 @@ class MTBenchDataset(HuggingFaceDataset):

996 prompt=prompt,

997 prompt_len=prompt_len,

998 expected_output_len=output_len,

999 + request_id=request_id_prefix + str(i),

1000 )

1001 )

952 - self.maybe_oversample_requests(sampled_requests, num_requests)

1002 + self.maybe_oversample_requests(

1003 + sampled_requests, num_requests, request_id_prefix

1004 + )

1005 return sampled_requests

1006

1007

@@ -974,10 +1026,12 @@ class AIMODataset(HuggingFaceDataset):

1026 tokenizer: PreTrainedTokenizerBase,

1027 num_requests: int,

1028 output_len: Optional[int] = None,

1029 + request_id_prefix: str = "",

1030 **kwargs,

1031 ) -> list:

1032 sampled_requests = []

1033 dynamic_output = output_len is None

1034 + ind = 0

1035

1036 for item in self.data:

1037 if len(sampled_requests) >= num_requests:

@@ -1000,9 +1054,13 @@ class AIMODataset(HuggingFaceDataset):

1054 prompt_len=prompt_len,

1055 expected_output_len=output_len,

1056 multi_modal_data=None,

1057 + request_id=request_id_prefix + str(ind),

1058 )

1059 )

1005 - self.maybe_oversample_requests(sampled_requests, num_requests)

1060 + ind += 1

1061 + self.maybe_oversample_requests(

1062 + sampled_requests, num_requests, request_id_prefix

1063 + )

1064 return sampled_requests

1065

1066

@@ -1072,12 +1130,18 @@ class NextEditPredictionDataset(HuggingFaceDataset):

1130 "zed-industries/zeta": _format_zeta_prompt,

1131 }

1132

1075 - def sample(self, tokenizer: PreTrainedTokenizerBase, num_requests: int, **kwargs):

1133 + def sample(

1134 + self,

1135 + tokenizer: PreTrainedTokenizerBase,

1136 + num_requests: int,

1137 + request_id_prefix: str = "",

1138 + **kwargs,

1139 + ):

1140 formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(self.dataset_path)

1141 if formatting_prompt_func is None:

1142 raise ValueError(f"Unsupported dataset path: {self.dataset_path}")

1143 samples = []

1080 - for sample in self.data:

1144 + for i, sample in enumerate(self.data):

1145 sample = formatting_prompt_func(sample)

1146 samples.append(

1147 SampleRequest(

@@ -1086,11 +1150,12 @@ class NextEditPredictionDataset(HuggingFaceDataset):

1150 expected_output_len=len(

1151 tokenizer(sample["expected_output"]).input_ids

1152 ),

1153 + request_id=request_id_prefix + str(i),

1154 )

1155 )

1156 if len(samples) >= num_requests:

1157 break

1093 - self.maybe_oversample_requests(samples, num_requests)

1158 + self.maybe_oversample_requests(samples, num_requests, request_id_prefix)

1159 return samples

1160

1161

@@ -1139,6 +1204,7 @@ class ASRDataset(HuggingFaceDataset):

1204 tokenizer: PreTrainedTokenizerBase,

1205 num_requests: int,

1206 output_len: Optional[int] = None,

1207 + request_id_prefix: str = "",

1208 **kwargs,

1209 ) -> list:

1210 import librosa

@@ -1148,6 +1214,7 @@ class ASRDataset(HuggingFaceDataset):

1214 prompt_len = len(tokenizer(prompt).input_ids)

1215 sampled_requests = []

1216 skipped = 0

1217 + ind = 0

1218 for item in self.data:

1219 if len(sampled_requests) >= num_requests:

1220 break

@@ -1166,8 +1233,10 @@ class ASRDataset(HuggingFaceDataset):

1233 prompt_len=prompt_len,

1234 expected_output_len=output_len,

1235 multi_modal_data=mm_content,

1236 + request_id=request_id_prefix + str(ind),

1237 )

1238 )

1239 + ind += 1

1240 if skipped:

1241 logger.warning(

1242 "%d samples discarded from dataset due to"

@@ -1175,5 +1244,7 @@ class ASRDataset(HuggingFaceDataset):

1244 " what Whisper supports.",

1245 skipped,

1246 )

1178 - self.maybe_oversample_requests(sampled_requests, num_requests)

1247 + self.maybe_oversample_requests(

1248 + sampled_requests, num_requests, request_id_prefix

1249 + )

1250 return sampled_requests

benchmarks/benchmark_serving.py

+21 -2

@@ -375,11 +375,12 @@ async def benchmark(

375 rps_change_events.append({"rps": rps_val, "timestamp": timestamp})

376 last_int_rps = current_int_rps

377

378 - prompt, prompt_len, output_len, mm_content = (

378 + prompt, prompt_len, output_len, mm_content, request_id = (

379 request.prompt,

380 request.prompt_len,

381 request.expected_output_len,

382 request.multi_modal_data,

383 + request.request_id,

384 )

385 req_model_id, req_model_name = model_id, model_name

386 if lora_modules:

@@ -397,6 +398,7 @@ async def benchmark(

398 multi_modal_content=mm_content,

399 ignore_eos=ignore_eos,

400 extra_body=extra_body,

401 + request_id=request_id,

402 )

403 task = limited_request_func(request_func_input=request_func_input, pbar=pbar)

404 tasks.append(asyncio.create_task(task))

@@ -665,6 +667,7 @@ def main(args: argparse.Namespace):

667 tokenizer=tokenizer,

668 output_len=args.custom_output_len,

669 skip_chat_template=args.custom_skip_chat_template,

670 + request_id_prefix=args.request_id_prefix,

671 )

672

673 elif args.dataset_name == "sonnet":

@@ -678,6 +681,7 @@ def main(args: argparse.Namespace):

681 prefix_len=args.sonnet_prefix_len,

682 tokenizer=tokenizer,

683 return_prompt_formatted=False,

684 + request_id_prefix=args.request_id_prefix,

685 )

686 else:

687 assert tokenizer.chat_template or tokenizer.default_chat_template, (

@@ -690,6 +694,7 @@ def main(args: argparse.Namespace):

694 prefix_len=args.sonnet_prefix_len,

695 tokenizer=tokenizer,

696 return_prompt_formatted=True,

697 + request_id_prefix=args.request_id_prefix,

698 )

699

700 elif args.dataset_name == "hf":

@@ -751,6 +756,7 @@ def main(args: argparse.Namespace):

756 num_requests=args.num_prompts,

757 tokenizer=tokenizer,

758 output_len=args.hf_output_len,

759 + request_id_prefix=args.request_id_prefix,

760 )

761

762 else:

@@ -762,10 +768,15 @@ def main(args: argparse.Namespace):

768 tokenizer=tokenizer,

769 num_requests=args.num_prompts,

770 output_len=args.sharegpt_output_len,

771 + request_id_prefix=args.request_id_prefix,

772 ),

773 "burstgpt": lambda: BurstGPTDataset(

774 random_seed=args.seed, dataset_path=args.dataset_path

768 - ).sample(tokenizer=tokenizer, num_requests=args.num_prompts),

775 + ).sample(

776 + tokenizer=tokenizer,

777 + num_requests=args.num_prompts,

778 + request_id_prefix=args.request_id_prefix,

779 + ),

780 "random": lambda: RandomDataset(dataset_path=args.dataset_path).sample(

781 tokenizer=tokenizer,

782 num_requests=args.num_prompts,

@@ -773,6 +784,7 @@ def main(args: argparse.Namespace):

784 input_len=args.random_input_len,

785 output_len=args.random_output_len,

786 range_ratio=args.random_range_ratio,

787 + request_id_prefix=args.request_id_prefix,

788 ),

789 }

790

@@ -1118,6 +1130,13 @@ def create_argument_parser():

1130 "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "

1131 "and the blog: https://hao-ai-lab.github.io/blogs/distserve",

1132 )

1133 + parser.add_argument(

1134 + "--request-id-prefix",

1135 + type=str,

1136 + required=False,

1137 + default="benchmark-serving",

1138 + help="Specify the prefix of request id.",

1139 + )

1140

1141 # group for dataset specific arguments

1142 custom_group = parser.add_argument_group("custom dataset options")

vllm/benchmarks/datasets.py

+92 -21

@@ -18,6 +18,7 @@ import logging

18 import random

19 from abc import ABC, abstractmethod

20 from collections.abc import Mapping

21 +from copy import deepcopy

22 from dataclasses import dataclass

23 from functools import cache

24 from io import BytesIO

@@ -76,6 +77,7 @@ class SampleRequest:

77 Union[MultiModalDataDict, dict, list[dict]]

78 ] = None

79 lora_request: Optional[LoRARequest] = None

80 + request_id: Optional[str] = None

83 # -----------------------------------------------------------------------------

@@ -183,7 +185,8 @@ class BenchmarkDataset(ABC):

185

186 @abstractmethod

187 def sample(self, tokenizer: PreTrainedTokenizerBase,

186 - num_requests: int) -> list[SampleRequest]:

188 + num_requests: int,

189 + request_id_prefix: str = "") -> list[SampleRequest]:

190 """

191 Abstract method to generate sample requests from the dataset.

192

@@ -194,6 +197,8 @@ class BenchmarkDataset(ABC):

197 tokenizer (PreTrainedTokenizerBase): The tokenizer to be used

198 for processing the dataset's text.

199 num_requests (int): The number of sample requests to generate.

200 + request_id_prefix (str) The prefix of request_id.

201 +

202

203 Returns:

204 list[SampleRequest]: A list of sample requests generated from the

@@ -201,8 +206,12 @@ class BenchmarkDataset(ABC):

206 """

207 raise NotImplementedError("sample must be implemented in subclasses.")

208

204 - def maybe_oversample_requests(self, requests: list[SampleRequest],

205 - num_requests: int) -> None:

209 + def maybe_oversample_requests(

210 + self,

211 + requests: list[SampleRequest],

212 + num_requests: int,

213 + request_id_prefix: str = "",

214 + ) -> None:

215 """

216 Oversamples the list of requests if its size is less than the desired

217 number.

@@ -211,11 +220,17 @@ class BenchmarkDataset(ABC):

220 requests (List[SampleRequest]): The current list of sampled

221 requests.

222 num_requests (int): The target number of requests.

223 + request_id_prefix (str) The prefix of the request ids.

224 +

225 """

226 if len(requests) < num_requests:

227 random.seed(self.random_seed)

217 - additional = random.choices(requests,

218 - k=num_requests - len(requests))

228 + additional = deepcopy(

229 + random.choices(requests, k=num_requests - len(requests))

230 + )

231 + for i in range(len(additional)):

232 + req = additional[i]

233 + req.request_id = request_id_prefix + str(len(requests) + i)

234 requests.extend(additional)

235 logger.info("Oversampled requests to reach %d total samples.",

236 num_requests)

@@ -334,6 +349,7 @@ class RandomDataset(BenchmarkDataset):

349 range_ratio: float = DEFAULT_RANGE_RATIO,

350 input_len: int = DEFAULT_INPUT_LEN,

351 output_len: int = DEFAULT_OUTPUT_LEN,

352 + request_id_prefix: str = "",

353 **kwargs,

354 ) -> list[SampleRequest]:

355 # Enforce range_ratio < 1

@@ -391,6 +407,7 @@ class RandomDataset(BenchmarkDataset):

407 prompt=prompt,

408 prompt_len=total_input_len,

409 expected_output_len=int(output_lens[i]),

410 + request_id=request_id_prefix + str(i),

411 ))

412 return requests

413

@@ -432,9 +449,11 @@ class ShareGPTDataset(BenchmarkDataset):

449 max_loras: Optional[int] = None,

450 output_len: Optional[int] = None,

451 enable_multimodal_chat: bool = False,

452 + request_id_prefix: str = "",

453 **kwargs,

454 ) -> list:

455 samples: list = []

456 + ind = 0

457 for entry in self.data:

458 if len(samples) >= num_requests:

459 break

@@ -470,8 +489,10 @@ class ShareGPTDataset(BenchmarkDataset):

489 expected_output_len=new_output_len,

490 lora_request=lora_request,

491 multi_modal_data=mm_content,

492 + request_id=request_id_prefix + str(ind),

493 ))

474 - self.maybe_oversample_requests(samples, num_requests)

494 + ind += 1

495 + self.maybe_oversample_requests(samples, num_requests, request_id_prefix)

496 return samples

497

498

@@ -647,6 +668,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

668 tokenizer=tokenizer,

669 output_len=args.custom_output_len,

670 skip_chat_template=args.custom_skip_chat_template,

671 + request_id_prefix=args.request_id_prefix,

672 )

673

674 elif args.dataset_name == "sonnet":

@@ -660,6 +682,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

682 prefix_len=args.sonnet_prefix_len,

683 tokenizer=tokenizer,

684 return_prompt_formatted=False,

685 + request_id_prefix=args.request_id_prefix,

686 )

687 else:

688 assert tokenizer.chat_template or tokenizer.default_chat_template, (

@@ -671,6 +694,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

694 prefix_len=args.sonnet_prefix_len,

695 tokenizer=tokenizer,

696 return_prompt_formatted=True,

697 + request_id_prefix=args.request_id_prefix,

698 )

699

700 elif args.dataset_name == "hf":

@@ -730,6 +754,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

754 num_requests=args.num_prompts,

755 tokenizer=tokenizer,

756 output_len=args.hf_output_len,

757 + request_id_prefix=args.request_id_prefix,

758 )

759

760 else:

@@ -741,11 +766,13 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

766 tokenizer=tokenizer,

767 num_requests=args.num_prompts,

768 output_len=args.sharegpt_output_len,

769 + request_id_prefix=args.request_id_prefix,

770 ),

771 "burstgpt":

772 lambda: BurstGPTDataset(random_seed=args.seed,

773 dataset_path=args.dataset_path).

748 - sample(tokenizer=tokenizer, num_requests=args.num_prompts),

774 + sample(tokenizer=tokenizer, num_requests=args.num_prompts,

775 + request_id_prefix=args.request_id_prefix,),

776 "random":

777 lambda: RandomDataset(random_seed=args.seed,

778 dataset_path=args.dataset_path).sample(

@@ -755,6 +782,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

782 input_len=args.random_input_len,

783 output_len=args.random_output_len,

784 range_ratio=args.random_range_ratio,

785 + request_id_prefix=args.request_id_prefix,

786 ),

787 "prefix_repetition":

788 lambda: PrefixRepetitionRandomDataset(

@@ -766,6 +794,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

794 suffix_len=args.prefix_repetition_suffix_len,

795 num_prefixes=args.prefix_repetition_num_prefixes,

796 output_len=args.prefix_repetition_output_len,

797 + request_id_prefix=args.request_id_prefix,

798 ),

799 }

800

@@ -839,10 +868,11 @@ class CustomDataset(BenchmarkDataset):

868 output_len: Optional[int] = None,

869 enable_multimodal_chat: bool = False,

870 skip_chat_template: bool = False,

871 + request_id_prefix: str = "",

872 **kwargs,

873 ) -> list:

874 sampled_requests = []

845 - for item in self.data:

875 + for i, item in enumerate(self.data):

876 if len(sampled_requests) >= num_requests:

877 break

878 prompt = item["prompt"]

@@ -864,8 +894,10 @@ class CustomDataset(BenchmarkDataset):

894 prompt=prompt,

895 prompt_len=prompt_len,

896 expected_output_len=output_len,

897 + request_id=request_id_prefix + str(i),

898 ))

868 - self.maybe_oversample_requests(sampled_requests, num_requests)

899 + self.maybe_oversample_requests(sampled_requests, num_requests,

900 + request_id_prefix)

901

902 return sampled_requests

903

@@ -909,6 +941,7 @@ class SonnetDataset(BenchmarkDataset):

941 input_len: int = DEFAULT_INPUT_LEN,

942 output_len: int = DEFAULT_OUTPUT_LEN,

943 return_prompt_formatted: bool = False,

944 + request_id_prefix: str = "",

945 **kwargs,

946 ) -> list:

947 # Calculate average token length for a poem line.

@@ -934,6 +967,7 @@ class SonnetDataset(BenchmarkDataset):

967 prefix_lines = self.data[:num_prefix_lines]

968

969 samples = []

970 + ind = 0

971 while len(samples) < num_requests:

972 extra_lines = random.choices(self.data,

973 k=num_input_lines - num_prefix_lines)

@@ -949,7 +983,9 @@ class SonnetDataset(BenchmarkDataset):

983 if return_prompt_formatted else prompt,

984 prompt_len=prompt_len,

985 expected_output_len=output_len,

986 + request_id=request_id_prefix + str(ind),

987 ))

988 + ind += 1

989 return samples

990

991

@@ -1000,6 +1036,7 @@ class BurstGPTDataset(BenchmarkDataset):

1036 num_requests: int,

1037 max_loras: Optional[int] = None,

1038 lora_path: Optional[str] = None,

1039 + request_id_prefix: str = "",

1040 **kwargs,

1041 ) -> list[SampleRequest]:

1042 samples = []

@@ -1020,6 +1057,7 @@ class BurstGPTDataset(BenchmarkDataset):

1057 prompt_len=input_len,

1058 expected_output_len=output_len,

1059 lora_request=lora_req,

1060 + request_id=request_id_prefix + str(i),

1061 ))

1062 return samples

1063

@@ -1075,11 +1113,13 @@ class ConversationDataset(HuggingFaceDataset):

1113 num_requests: int,

1114 output_len: Optional[int] = None,

1115 enable_multimodal_chat: bool = False,

1116 + request_id_prefix: str = "",

1117 **kwargs) -> list:

1118 # Filter examples with at least 2 conversations

1119 filtered_data = self.data.filter(

1120 lambda x: len(x["conversations"]) >= 2)

1121 sampled_requests = []

1122 + ind = 0

1123 dynamic_output = output_len is None

1124

1125 for item in filtered_data:

@@ -1111,8 +1151,11 @@ class ConversationDataset(HuggingFaceDataset):

1151 prompt_len=prompt_len,

1152 expected_output_len=output_len,

1153 multi_modal_data=mm_content,

1154 + request_id=request_id_prefix + str(ind),

1155 ))

1115 - self.maybe_oversample_requests(sampled_requests, num_requests)

1156 + ind += 1

1157 + self.maybe_oversample_requests(sampled_requests, num_requests,

1158 + request_id_prefix)

1159 return sampled_requests

1160

1161

@@ -1141,12 +1184,13 @@ class VisionArenaDataset(HuggingFaceDataset):

1184 num_requests: int,

1185 output_len: Optional[int] = None,

1186 enable_multimodal_chat: bool = False,

1187 + request_id_prefix: str = "",

1188 **kwargs,

1189 ) -> list:

1190 output_len = (output_len

1191 if output_len is not None else self.DEFAULT_OUTPUT_LEN)

1192 sampled_requests = []

1149 - for item in self.data:

1193 + for i, item in enumerate(self.data):

1194 if len(sampled_requests) >= num_requests:

1195 break

1196 parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.dataset_path)

@@ -1168,8 +1212,10 @@ class VisionArenaDataset(HuggingFaceDataset):

1212 prompt_len=prompt_len,

1213 expected_output_len=output_len,

1214 multi_modal_data=mm_content,

1215 + request_id=request_id_prefix + str(i),

1216 ))

1172 - self.maybe_oversample_requests(sampled_requests, num_requests)

1217 + self.maybe_oversample_requests(sampled_requests, num_requests,

1218 + request_id_prefix)

1219 return sampled_requests

1220

1221

@@ -1198,11 +1244,12 @@ class InstructCoderDataset(HuggingFaceDataset):

1244 num_requests: int,

1245 output_len: Optional[int] = None,

1246 enable_multimodal_chat: bool = False,

1247 + request_id_prefix: str = "",

1248 **kwargs) -> list:

1249 output_len = (output_len

1250 if output_len is not None else self.DEFAULT_OUTPUT_LEN)

1251 sampled_requests = []

1205 - for item in self.data:

1252 + for i, item in enumerate(self.data):

1253 if len(sampled_requests) >= num_requests:

1254 break

1255 prompt = f"{item['input']}\n\n{item['instruction']} Just output \

@@ -1224,8 +1271,10 @@ class InstructCoderDataset(HuggingFaceDataset):

1271 prompt=prompt,

1272 prompt_len=prompt_len,

1273 expected_output_len=output_len,

1274 + request_id=request_id_prefix + str(i),

1275 ))

1228 - self.maybe_oversample_requests(sampled_requests, num_requests)

1276 + self.maybe_oversample_requests(sampled_requests, num_requests,

1277 + request_id_prefix)

1278 return sampled_requests

1279

1280

@@ -1255,13 +1304,14 @@ class MTBenchDataset(HuggingFaceDataset):

1304 num_requests: int,

1305 output_len: Optional[int] = None,

1306 enable_multimodal_chat: bool = False,

1307 + request_id_prefix: str = "",

1308 **kwargs,

1309 ) -> list:

1310 output_len = (output_len

1311 if output_len is not None else self.DEFAULT_OUTPUT_LEN)

1312 sampled_requests = []

1313

1264 - for item in self.data:

1314 + for i, item in enumerate(self.data):

1315 if len(sampled_requests) >= num_requests:

1316 break

1317 prompt = item["turns"][0]

@@ -1282,8 +1332,10 @@ class MTBenchDataset(HuggingFaceDataset):

1332 prompt=prompt,

1333 prompt_len=prompt_len,

1334 expected_output_len=output_len,

1335 + request_id=request_id_prefix + str(i),

1336 ))

1286 - self.maybe_oversample_requests(sampled_requests, num_requests)

1337 + self.maybe_oversample_requests(sampled_requests, num_requests,

1338 + request_id_prefix)

1339 return sampled_requests

1340

1341

@@ -1305,8 +1357,10 @@ class AIMODataset(HuggingFaceDataset):

1357 tokenizer: PreTrainedTokenizerBase,

1358 num_requests: int,

1359 output_len: Optional[int] = None,

1360 + request_id_prefix: str = "",

1361 **kwargs) -> list:

1362 sampled_requests = []

1363 + ind = 0

1364 dynamic_output = output_len is None

1365

1366 for item in self.data:

@@ -1331,8 +1385,12 @@ class AIMODataset(HuggingFaceDataset):

1385 prompt_len=prompt_len,

1386 expected_output_len=output_len,

1387 multi_modal_data=None,

1388 + request_id=request_id_prefix + str(ind),

1389 +

1390 ))

1335 - self.maybe_oversample_requests(sampled_requests, num_requests)

1391 + ind += 1

1392 + self.maybe_oversample_requests(sampled_requests, num_requests,

1393 + request_id_prefix)

1394 return sampled_requests

1395

1396

@@ -1403,13 +1461,14 @@ class NextEditPredictionDataset(HuggingFaceDataset):

1461 }

1462

1463 def sample(self, tokenizer: PreTrainedTokenizerBase, num_requests: int,

1464 + request_id_prefix: str = "",

1465 **kwargs):

1466 formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(

1467 self.dataset_path)

1468 if formatting_prompt_func is None:

1469 raise ValueError(f"Unsupported dataset path: {self.dataset_path}")

1470 samples = []

1412 - for sample in self.data:

1471 + for i, sample in enumerate(self.data):

1472 sample = formatting_prompt_func(sample)

1473 samples.append(

1474 SampleRequest(

@@ -1417,10 +1476,11 @@ class NextEditPredictionDataset(HuggingFaceDataset):

1476 prompt_len=len(tokenizer(sample["prompt"]).input_ids),

1477 expected_output_len=len(

1478 tokenizer(sample["expected_output"]).input_ids),

1479 + request_id=request_id_prefix + str(i),

1480 ))

1481 if len(samples) >= num_requests:

1482 break

1423 - self.maybe_oversample_requests(samples, num_requests)

1483 + self.maybe_oversample_requests(samples, num_requests, request_id_prefix)

1484 return samples

1485

1486

@@ -1470,6 +1530,7 @@ class ASRDataset(HuggingFaceDataset):

1530 tokenizer: PreTrainedTokenizerBase,

1531 num_requests: int,

1532 output_len: Optional[int] = None,

1533 + request_id_prefix: str = "",

1534 **kwargs,

1535 ) -> list:

1536 output_len = (output_len

@@ -1477,6 +1538,7 @@ class ASRDataset(HuggingFaceDataset):

1538 prompt = ASRDataset.TRANSCRIPTION_PREAMBLE

1539 prompt_len = len(tokenizer(prompt).input_ids)

1540 sampled_requests = []

1541 + ind = 0

1542 skipped = 0

1543 for item in self.data:

1544 if len(sampled_requests) >= num_requests:

@@ -1496,7 +1558,9 @@ class ASRDataset(HuggingFaceDataset):

1558 prompt_len=prompt_len,

1559 expected_output_len=output_len,

1560 multi_modal_data=mm_content,

1561 + request_id=request_id_prefix + str(ind),

1562 ))

1563 + ind += 1

1564 if skipped:

1565 logger.warning(

1566 "%d samples discarded from dataset due to"

@@ -1504,7 +1568,8 @@ class ASRDataset(HuggingFaceDataset):

1568 " what Whisper supports.",

1569 skipped,

1570 )

1507 - self.maybe_oversample_requests(sampled_requests, num_requests)

1571 + self.maybe_oversample_requests(sampled_requests, num_requests,

1572 + request_id_prefix)

1573 return sampled_requests

1574

1575

@@ -1541,11 +1606,13 @@ class MLPerfDataset(HuggingFaceDataset):

1606 tokenizer: PreTrainedTokenizerBase,

1607 num_requests: int,

1608 output_len: Optional[int] = None,

1609 + request_id_prefix: str = "",

1610 **kwargs,

1611 ) -> list[SampleRequest]:

1612 # Force dynamic output length based on reference completion.

1613 dynamic_output = output_len is None

1614 sampled_requests: list[SampleRequest] = []

1615 + ind = 0

1616

1617 for item in self.data:

1618 if len(sampled_requests) >= num_requests:

@@ -1580,10 +1647,13 @@ class MLPerfDataset(HuggingFaceDataset):

1647 prompt=prompt_formatted,

1648 prompt_len=prompt_len,

1649 expected_output_len=expected_output_len,

1650 + request_id=request_id_prefix + str(ind),

1651 )

1652 )

1653 + ind += 1

1654

1586 - self.maybe_oversample_requests(sampled_requests, num_requests)

1655 + self.maybe_oversample_requests(sampled_requests, num_requests,

1656 + request_id_prefix)

1657 return sampled_requests

1658

1659

@@ -1616,6 +1686,7 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):

1686 suffix_len: int = DEFAULT_SUFFIX_LEN,

1687 num_prefixes: int = DEFAULT_NUM_PREFIXES,

1688 output_len: int = DEFAULT_OUTPUT_LEN,

1689 + request_id_prefix: str = "",

1690 **kwargs,

1691 ) -> list[SampleRequest]:

1692 vocab_size = tokenizer.vocab_size

vllm/benchmarks/lib/endpoint_request_func.py

@@ -31,6 +31,7 @@ class RequestFuncInput:

31 multi_modal_content: Optional[dict | list[dict]] = None

32 ignore_eos: bool = False

33 language: Optional[str] = None

34 + request_id: Optional[str] = None

37 @dataclass

@@ -87,6 +88,8 @@ async def async_request_openai_completions(

88 headers = {

89 "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"

90 }

91 + if request_func_input.request_id:

92 + headers["x-request-id"] = request_func_input.request_id

94 output = RequestFuncOutput()

95 output.prompt_len = request_func_input.prompt_len

@@ -210,6 +213,8 @@ async def async_request_openai_chat_completions(

213 "Content-Type": "application/json",

214 "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",

215 }

216 + if request_func_input.request_id:

217 + headers["x-request-id"] = request_func_input.request_id

218

219 output = RequestFuncOutput()

220 output.prompt_len = request_func_input.prompt_len

@@ -311,6 +316,8 @@ async def async_request_openai_audio(

316 headers = {

317 "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",

318 }

319 + if request_func_input.request_id:

320 + headers["x-request-id"] = request_func_input.request_id

321

322 # Send audio file

323 def to_bytes(y, sr):

vllm/benchmarks/serve.py

+12 -2

@@ -478,11 +478,12 @@ async def benchmark(

478 "timestamp": timestamp

479 })

480 last_int_rps = current_int_rps

481 - prompt, prompt_len, output_len, mm_content = (

481 + prompt, prompt_len, output_len, mm_content, request_id = (

482 request.prompt,

483 request.prompt_len,

484 request.expected_output_len,

485 request.multi_modal_data,

486 + request.request_id,

487 )

488 req_model_id, req_model_name = model_id, model_name

489 if lora_modules:

@@ -498,7 +499,8 @@ async def benchmark(

499 logprobs=logprobs,

500 multi_modal_content=mm_content,

501 ignore_eos=ignore_eos,

501 - extra_body=extra_body)

502 + extra_body=extra_body,

503 + request_id=request_id,)

504 tasks.append(

505 asyncio.create_task(

506 limited_request_func(request_func_input=request_func_input,

@@ -865,6 +867,14 @@ def add_cli_args(parser: argparse.ArgumentParser):

867 "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "

868 "and the blog: https://hao-ai-lab.github.io/blogs/distserve",

869 )

870 + parser.add_argument(

871 + "--request-id-prefix",

872 + type=str,

873 + required=False,

874 + default="benchmark-serving",

875 + help="Specify the prefix of request id.",

876 + )

877 +

878

879 sampling_group = parser.add_argument_group("sampling parameters")

880 sampling_group.add_argument(