@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 275b43183c

[MyPy] Fix mypy for `vllm/benchmarks` (#39896)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Martin Hickey committed Jun 23, 2026 at 16:22 UTC 275b43183c7253a45bd0214d165b0d5347de1ec6

24 files changed +296 -150

tools/pre_commit/mypy.py

-2

@@ -31,8 +31,6 @@ SEPARATE_GROUPS = [

31 EXCLUDE = [

32 "vllm/model_executor/models",

33 "vllm/model_executor/layers/fla/ops",

34 - # TODO: Remove these entries after fixing mypy errors.

35 - "vllm/benchmarks",

34 ]

vllm/benchmarks/datasets/create_txt_slices_dataset.py

+2 -2

@@ -30,7 +30,6 @@ The resulting JSONL file can then be used with the serving benchmark::

31 from __future__ import annotations

33 -import argparse

33 import json

34 import logging

35 import random

@@ -40,6 +39,7 @@ import numpy as np

39 from transformers import AutoTokenizer

41 from vllm.benchmarks.datasets.utils import RangeRatio, get_sampling_params

42 +from vllm.utils.argparse_utils import FlexibleArgumentParser

44 logger = logging.getLogger(__name__)

@@ -121,7 +121,7 @@ def create_txt_slices_jsonl(

121

122

123 def main(argv: list[str] | None = None) -> None:

124 - parser = argparse.ArgumentParser(

124 + parser = FlexibleArgumentParser(

125 description="Convert a plain-text file into a JSONL dataset "

126 "for CustomDataset (--dataset-name custom).",

127 )

vllm/benchmarks/datasets/datasets.py

+140 -63

@@ -81,7 +81,7 @@ class SampleRequest:

82 prompt: str | list[str] | list[dict]

83 prompt_len: int

84 - expected_output_len: int | None

84 + expected_output_len: int = 0

85 multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None

86 lora_request: LoRARequest | None = None

87 request_id: str | None = None

@@ -634,7 +634,7 @@ class RandomDataset(BenchmarkDataset):

634 # Generate prefix once

635 prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)

636

637 - requests = []

637 + requests: list[SampleRequest] = []

638 token_mismatch_total = 0

639 for i in range(num_requests):

640 prompt, total_input_len, token_mismatch = self.generate_token_sequence( # noqa: E501

@@ -665,15 +665,14 @@ class RandomDataset(BenchmarkDataset):

665 )

666 # only used for embeddings benchmark.

667 if batchsize > 1:

668 - batch_requests = []

668 + batch_requests: list[SampleRequest] = []

669 # Create batched requests

670 for i in range(0, num_requests, batchsize):

671 batch = requests[i : i + batchsize]

672 batch_requests.append(

673 SampleRequest(

674 - prompt=[req.prompt for req in batch],

674 + prompt=[req.prompt for req in batch], # type: ignore[arg-type]

675 prompt_len=sum(req.prompt_len for req in batch),

676 - expected_output_len=0,

676 request_id=request_id_prefix + str(i // batchsize),

677 )

678 )

@@ -797,6 +796,9 @@ class RandomDatasetForReranking(RandomDataset):

796 input_len: int = RandomDataset.DEFAULT_INPUT_LEN,

797 output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,

798 batchsize: int = 1,

799 + max_loras: int | None = None,

800 + lora_path: str | None = None,

801 + lora_assignment: str = "random",

802 is_reranker: bool = True,

803 **kwargs,

804 ) -> list[SampleRequest]:

@@ -1207,6 +1209,9 @@ class RandomMultiModalDataset(RandomDataset):

1209 input_len: int = RandomDataset.DEFAULT_INPUT_LEN,

1210 output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,

1211 batchsize: int = 1,

1212 + max_loras: int | None = None,

1213 + lora_path: str | None = None,

1214 + lora_assignment: str = "random",

1215 limit_mm_per_prompt: dict[str, int] = DEFAULT_LIMIT_MM_PER_PROMPT,

1216 base_items_per_request: int = DEFAULT_BASE_ITEMS_PER_REQUEST,

1217 num_mm_items_range_ratio: float = DEFAULT_NUM_MM_ITEMS_RANGE_RATIO,

@@ -1251,6 +1256,7 @@ class RandomMultiModalDataset(RandomDataset):

1256 # We want to exclude placeholder tokens and all

1257 # tokens that indicate start/end of image as it

1258 # may break prompt replacement logic.

1259 + assert hasattr(tokenizer, "added_tokens_decoder")

1260 prohibited_tokens = list(

1261 tok_id

1262 for tok_id, token in tokenizer.added_tokens_decoder.items()

@@ -1376,6 +1382,7 @@ class ShareGPTDataset(BenchmarkDataset):

1382 lora_assignment: str = "random",

1383 **kwargs,

1384 ) -> list[SampleRequest]:

1385 + assert self.data is not None, "Dataset must be loaded before sampling"

1386 samples: list[SampleRequest] = []

1387 ind = 0

1388 for entry in self.data:

@@ -1454,7 +1461,7 @@ class TimedTrace(BenchmarkDataset):

1461 f'label_output_length: "{self.label_output_length}", '

1462 f'label_hash_ids: "{self.label_hash_ids}"'

1463 )

1457 - self._expanded_generated_prompts = {}

1464 + self._expanded_generated_prompts: dict[str, Any] = {}

1465 self.load_data()

1466

1467 def load_data(self) -> None:

@@ -1529,10 +1536,12 @@ class TimedTrace(BenchmarkDataset):

1536 tokenizer: TokenizerLike,

1537 num_requests: int,

1538 request_id_prefix: str = "",

1539 + no_oversample: bool = False,

1540 **kwargs,

1533 - ) -> list:

1534 - samples: list = []

1541 + ) -> list[SampleRequest]:

1542 + samples: list[SampleRequest] = []

1543 assert tokenizer is not None, "Tokenizer must be provided, now is Null"

1544 + assert self.data is not None, "Data must be loaded before sampling"

1545

1546 for ind, entry in enumerate(self.data):

1547 if len(samples) >= num_requests:

@@ -2139,12 +2148,12 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:

2148 )

2149

2150 elif args.dataset_name == "sonnet":

2142 - dataset = SonnetDataset(

2151 + sonnet_dataset = SonnetDataset(

2152 dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle

2153 )

2154 # For the "sonnet" dataset, formatting depends on the backend.

2155 if args.backend == "openai-chat":

2147 - input_requests = dataset.sample(

2156 + input_requests = sonnet_dataset.sample(

2157 num_requests=args.num_prompts,

2158 input_len=args.sonnet_input_len,

2159 output_len=args.sonnet_output_len,

@@ -2155,10 +2164,13 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:

2164 no_oversample=args.no_oversample,

2165 )

2166 else:

2158 - assert tokenizer.chat_template or tokenizer.default_chat_template, (

2159 - "Tokenizer/model must have chat template for sonnet dataset."

2160 - )

2161 - input_requests = dataset.sample(

2167 + assert (

2168 + hasattr(tokenizer, "chat_template") and tokenizer.chat_template

2169 + ) or (

2170 + hasattr(tokenizer, "default_chat_template")

2171 + and tokenizer.default_chat_template

2172 + ), "Tokenizer/model must have chat template for sonnet dataset."

2173 + input_requests = sonnet_dataset.sample(

2174 num_requests=args.num_prompts,

2175 input_len=args.sonnet_input_len,

2176 output_len=args.sonnet_output_len,

@@ -2173,6 +2185,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:

2185 # all following datasets are implemented from the

2186 # HuggingFaceDataset base class

2187 hf_kwargs = {}

2188 + dataset_class: type[BenchmarkDataset]

2189 if (

2190 args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS

2191 or args.hf_name in VisionArenaDataset.SUPPORTED_DATASET_PATHS

@@ -2548,23 +2561,23 @@ class CustomDataset(BenchmarkDataset):

2561

2562 if tokenizer is None:

2563 new_output_len = 1

2564 + elif output_len is None or output_len == -1:

2565 + # check that the request has an 'output_tokens' field

2566 + if "output_tokens" not in item:

2567 + raise ValueError(

2568 + "If no output length is provided the "

2569 + "custom dataset must contain an 'output_tokens' field."

2570 + )

2571 + # Use number of output tokens from the request data

2572 + try:

2573 + new_output_len = int(item["output_tokens"])

2574 + except (ValueError, TypeError) as e:

2575 + raise ValueError(

2576 + f"Invalid value for 'output_tokens' in custom dataset: "

2577 + f"'{item['output_tokens']}'. Must be an integer."

2578 + ) from e

2579 else:

2580 new_output_len = output_len

2553 - if output_len is None or output_len == -1:

2554 - # check that the request has an 'output_tokens' field

2555 - if "output_tokens" not in item:

2556 - raise ValueError(

2557 - "If no output length is provided the "

2558 - "custom dataset must contain an 'output_tokens' field."

2559 - )

2560 - # Use number of output tokens from the request data

2561 - try:

2562 - new_output_len = int(item["output_tokens"])

2563 - except (ValueError, TypeError) as e:

2564 - raise ValueError(

2565 - f"Invalid value for 'output_tokens' in custom dataset: "

2566 - f"'{item['output_tokens']}'. Must be an integer."

2567 - ) from e

2581

2582 if tokenizer is None:

2583 prompt_len = 1

@@ -2797,11 +2810,15 @@ class CustomImageDataset(CustomDataset):

2810 self,

2811 tokenizer: TokenizerLike,

2812 num_requests: int,

2813 + request_id_prefix: str = "",

2814 + no_oversample: bool = False,

2815 + lora_path: str | None = None,

2816 + max_loras: int | None = None,

2817 output_len: int | None = None,

2818 enable_multimodal_chat: bool = False,

2819 + skip_chat_template: bool = False,

2820 + chat_template_kwargs: dict | None = None,

2821 ensure_client_side_data: bool = False,

2803 - request_id_prefix: str = "",

2804 - no_oversample: bool = False,

2822 **kwargs,

2823 ) -> list[SampleRequest]:

2824 # load all data if needed

@@ -2814,7 +2831,7 @@ class CustomImageDataset(CustomDataset):

2831 num_requests,

2832 )

2833

2817 - sampled_requests = []

2834 + sampled_requests: list[SampleRequest] = []

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

2836 if len(sampled_requests) >= num_requests:

2837 break

@@ -2837,7 +2854,7 @@ class CustomImageDataset(CustomDataset):

2854 SampleRequest(

2855 prompt=prompt,

2856 prompt_len=prompt_len,

2840 - expected_output_len=output_len,

2857 + expected_output_len=output_len or 0,

2858 multi_modal_data=None,

2859 request_id=request_id_prefix + str(i),

2860 )

@@ -2863,7 +2880,7 @@ class CustomImageDataset(CustomDataset):

2880 SampleRequest(

2881 prompt=prompt,

2882 prompt_len=prompt_len,

2866 - expected_output_len=output_len,

2883 + expected_output_len=output_len or 0,

2884 multi_modal_data=mm_content,

2885 request_id=request_id_prefix + str(i),

2886 )

@@ -2891,17 +2908,20 @@ class CustomAudioDataset(CustomDataset):

2908 self,

2909 tokenizer: TokenizerLike,

2910 num_requests: int,

2894 - output_len: int | None = None,

2911 request_id_prefix: str = "",

2912 no_oversample: bool = False,

2897 - skip_chat_template: bool = False,

2913 + lora_path: str | None = None,

2914 + max_loras: int | None = None,

2915 + output_len: int | None = None,

2916 enable_multimodal_chat: bool = False,

2917 + skip_chat_template: bool = False,

2918 + chat_template_kwargs: dict | None = None,

2919 **kwargs,

2920 ) -> list[SampleRequest]:

2921 self.num_available_samples = len(self.data)

2922 if num_requests <= 0:

2923 num_requests = self.num_available_samples

2904 - sampled_requests = []

2924 + sampled_requests: list[SampleRequest] = []

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

2926 if len(sampled_requests) >= num_requests:

2927 break

@@ -2950,7 +2970,6 @@ class CustomAudioDataset(CustomDataset):

2970 prompt_len = (

2971 len(tokenizer(prompt).input_ids) if isinstance(prompt, str) else 1

2972 )

2953 - new_output_len = output_len

2973 if output_len is None or output_len == -1:

2974 if "output_tokens" not in item:

2975 raise ValueError(

@@ -2958,6 +2977,8 @@ class CustomAudioDataset(CustomDataset):

2977 "custom dataset must contain an 'output_tokens' field."

2978 )

2979 new_output_len = int(item["output_tokens"])

2980 + else:

2981 + new_output_len = output_len

2982 sampled_requests.append(

2983 SampleRequest(

2984 prompt=prompt,

@@ -3015,10 +3036,30 @@ class SpecBench(CustomDataset):

3036

3037 def sample(

3038 self,

3039 + tokenizer: TokenizerLike,

3040 + num_requests: int,

3041 + request_id_prefix: str = "",

3042 + no_oversample: bool = False,

3043 + lora_path: str | None = None,

3044 + max_loras: int | None = None,

3045 + output_len: int | None = None,

3046 + enable_multimodal_chat: bool = False,

3047 + skip_chat_template: bool = False,

3048 + chat_template_kwargs: dict | None = None,

3049 **kwargs,

3050 ) -> list[SampleRequest]:

3051 # leverage CustomDataset sample

3052 return super().sample(

3053 + tokenizer=tokenizer,

3054 + num_requests=num_requests,

3055 + request_id_prefix=request_id_prefix,

3056 + no_oversample=no_oversample,

3057 + lora_path=lora_path,

3058 + max_loras=max_loras,

3059 + output_len=output_len,

3060 + enable_multimodal_chat=enable_multimodal_chat,

3061 + skip_chat_template=skip_chat_template,

3062 + chat_template_kwargs=chat_template_kwargs,

3063 **kwargs,

3064 )

3065

@@ -3067,16 +3108,21 @@ class SonnetDataset(BenchmarkDataset):

3108 return_prompt_formatted: bool = False,

3109 **kwargs,

3110 ) -> list[SampleRequest]:

3111 + poem_lines = self.data

3112 + assert poem_lines is not None

3113 # Calculate average token length for a poem line.

3071 - tokenized_lines = [tokenizer(line).input_ids for line in self.data]

3114 + tokenized_lines = [tokenizer(line).input_ids for line in poem_lines]

3115 avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)

3116

3117 # Build the base prompt.

3118 base_prompt = "Pick as many lines as you can from these poem lines:\n"

3119 base_msg = [{"role": "user", "content": base_prompt}]

3120 base_fmt = tokenizer.apply_chat_template(

3078 - base_msg, add_generation_prompt=True, tokenize=False

3121 + base_msg, # type: ignore[arg-type]

3122 + add_generation_prompt=True,

3123 + tokenize=False,

3124 )

3125 + assert isinstance(base_fmt, str)

3126 base_offset = len(tokenizer(base_fmt).input_ids)

3127 if input_len <= base_offset:

3128 raise ValueError(

@@ -3085,26 +3131,31 @@ class SonnetDataset(BenchmarkDataset):

3131 )

3132

3133 # Determine how many poem lines to use.

3088 - num_input_lines = round((input_len - base_offset) / avg_len)

3134 + num_input_lines = max(round((input_len - base_offset) / avg_len), 1)

3135 num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0)

3090 - prefix_lines = self.data[:num_prefix_lines]

3136 + prefix_lines = poem_lines[:num_prefix_lines]

3137

3138 samples: list[SampleRequest] = []

3139 ind = 0

3140 while len(samples) < num_requests:

3141 extra_lines = random.choices(

3096 - self.data, k=num_input_lines - num_prefix_lines

3142 + poem_lines, k=num_input_lines - num_prefix_lines

3143 )

3144 prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"

3145 msg = [{"role": "user", "content": prompt}]

3146 prompt_formatted = tokenizer.apply_chat_template(

3101 - msg, add_generation_prompt=True, tokenize=False

3147 + msg, # type: ignore[arg-type]

3148 + add_generation_prompt=True,

3149 + tokenize=False,

3150 )

3151 + assert isinstance(prompt_formatted, str)

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

3153 if prompt_len <= input_len:

3154 samples.append(

3155 SampleRequest(

3107 - prompt=prompt_formatted if return_prompt_formatted else prompt,

3156 + prompt=(

3157 + prompt_formatted if return_prompt_formatted else prompt

3158 + ),

3159 prompt_len=prompt_len,

3160 expected_output_len=output_len,

3161 request_id=request_id_prefix + str(ind),

@@ -3145,6 +3196,7 @@ class BurstGPTDataset(BenchmarkDataset):

3196 self.data = gpt4_df

3197

3198 def _sample_loaded_data(self, num_requests: int) -> list:

3199 + assert self.data is not None, "Dataset must be loaded before sampling"

3200 if num_requests <= len(self.data):

3201 data = self.data.sample(n=num_requests, random_state=self.random_seed)

3202 else:

@@ -3258,6 +3310,7 @@ class ConversationDataset(HuggingFaceDataset):

3310 enable_multimodal_chat: bool = False,

3311 **kwargs,

3312 ) -> list[SampleRequest]:

3313 + assert self.data is not None, "Dataset must be loaded before sampling"

3314 # Filter examples with at least 2 conversations

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

3316 sampled_requests: list[SampleRequest] = []

@@ -3318,6 +3371,7 @@ class MultiModalConversationDataset(HuggingFaceDataset):

3371 enable_multimodal_chat: bool = False,

3372 **kwargs,

3373 ) -> list[SampleRequest]:

3374 + assert self.data is not None, "Dataset must be loaded before sampling"

3375 # Filter examples with at least 2 conversations

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

3377 sampled_requests: list[SampleRequest] = []

@@ -3393,7 +3447,8 @@ class VisionArenaDataset(HuggingFaceDataset):

3447

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

3449

3396 - sampled_requests = []

3450 + assert self.data is not None, "Dataset must be loaded before sampling"

3451 + sampled_requests: list[SampleRequest] = []

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

3453 if len(sampled_requests) >= num_requests:

3454 break

@@ -3465,7 +3520,8 @@ class MMVUDataset(HuggingFaceDataset):

3520

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

3522

3468 - sampled_requests = []

3523 + assert self.data is not None, "Dataset must be loaded before sampling"

3524 + sampled_requests: list[SampleRequest] = []

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

3526 if len(sampled_requests) >= num_requests:

3527 break

@@ -3531,18 +3587,21 @@ class InstructCoderDataset(HuggingFaceDataset):

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

3588 sampled_requests: list[SampleRequest] = []

3589 for i, prompt in enumerate(self.sample_prompts(n=num_requests)):

3590 + prompt_text = prompt

3591 # apply template

3592 if not skip_chat_template:

3536 - prompt = tokenizer.apply_chat_template(

3593 + prompt_text_result = tokenizer.apply_chat_template(

3594 [{"role": "user", "content": prompt}],

3595 add_generation_prompt=True,

3596 tokenize=False,

3597 )

3598 + assert isinstance(prompt_text_result, str)

3599 + prompt_text = prompt_text_result

3600

3542 - prompt_len = len(tokenizer(prompt).input_ids)

3601 + prompt_len = len(tokenizer(prompt_text).input_ids)

3602 sampled_requests.append(

3603 SampleRequest(

3545 - prompt=prompt,

3604 + prompt=prompt_text,

3605 prompt_len=prompt_len,

3606 expected_output_len=output_len,

3607 request_id=request_id_prefix + str(i),

@@ -3554,6 +3613,7 @@ class InstructCoderDataset(HuggingFaceDataset):

3613 return sampled_requests

3614

3615 def sample_prompts(self, n: int) -> Iterator[str]:

3616 + assert self.data is not None, "Dataset must be loaded before sampling"

3617 for item in self.data.take(n):

3618 prompt = (

3619 f"{item['input']}\n\n{item['instruction']} Just output "

@@ -3594,6 +3654,7 @@ class MTBenchDataset(HuggingFaceDataset):

3654 **kwargs,

3655 ) -> list[SampleRequest]:

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

3657 + assert self.data is not None, "Dataset must be loaded before sampling"

3658 sampled_requests: list[SampleRequest] = []

3659

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

@@ -3654,7 +3715,8 @@ class HumanEvalDataset(HuggingFaceDataset):

3715 **kwargs,

3716 ) -> list[SampleRequest]:

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

3657 - sampled_requests = []

3718 + sampled_requests: list[SampleRequest] = []

3719 + assert self.data is not None, "Data must be loaded before sampling"

3720

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

3722 if len(sampled_requests) >= num_requests:

@@ -3714,7 +3776,8 @@ class GSM8KDataset(HuggingFaceDataset):

3776 **kwargs,

3777 ) -> list[SampleRequest]:

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

3717 - sampled_requests = []

3779 + sampled_requests: list[SampleRequest] = []

3780 + assert self.data is not None, "Data must be loaded before sampling"

3781

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

3783 if len(sampled_requests) >= num_requests:

@@ -3772,17 +3835,18 @@ class BlazeditDataset(HuggingFaceDataset):

3835 self,

3836 tokenizer: TokenizerLike,

3837 num_requests: int,

3775 - output_len: int | None = None,

3776 - skip_chat_template: bool = False,

3838 request_id_prefix: str = "",

3839 no_oversample: bool = False,

3840 + output_len: int | None = None,

3841 + skip_chat_template: bool = False,

3842 min_distance: float = 0.0,

3843 max_distance: float = 1.0,

3844 **kwargs,

3845 ) -> list[SampleRequest]:

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

3784 - sampled_requests = []

3847 + sampled_requests: list[SampleRequest] = []

3848

3849 + assert self.data is not None, "Dataset must be loaded before sampling"

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

3851 if len(sampled_requests) >= num_requests:

3852 break

@@ -3810,11 +3874,13 @@ Please generate the new code file in the "New file" section below.""" # noqa: E

3874

3875 # apply template

3876 if not skip_chat_template:

3813 - prompt = tokenizer.apply_chat_template(

3877 + prompt_result = tokenizer.apply_chat_template(

3878 [{"role": "user", "content": prompt}],

3879 add_generation_prompt=True,

3880 tokenize=False,

3881 )

3882 + assert isinstance(prompt_result, str)

3883 + prompt = prompt_result

3884

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

3886

@@ -3858,6 +3924,7 @@ class AIMODataset(HuggingFaceDataset):

3924 output_len: int | None = None,

3925 **kwargs,

3926 ) -> list[SampleRequest]:

3927 + assert self.data is not None, "Dataset must be loaded before sampling"

3928 sampled_requests: list[SampleRequest] = []

3929 ind = 0

3930 dynamic_output = output_len is None

@@ -3970,6 +4037,7 @@ class NextEditPredictionDataset(HuggingFaceDataset):

4037 formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(self.hf_name)

4038 if formatting_prompt_func is None:

4039 raise ValueError(f"Unsupported dataset path: {self.hf_name}")

4040 + assert self.data is not None, "Dataset must be loaded before sampling"

4041 samples = []

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

4043 sample = formatting_prompt_func(sample)

@@ -4061,12 +4129,14 @@ class ASRDataset(HuggingFaceDataset):

4129 def _disable_audio_decode(self) -> None:

4130 from datasets import Audio

4131

4132 + assert self.data is not None, "Dataset must be loaded before sampling"

4133 self.data = self.data.cast_column("audio", Audio(decode=False))

4134

4135 def _materialize_local_audio_column(self) -> None:

4136 local_path_root = Path(

4137 hf_api().snapshot_download(self.hf_name, repo_type="dataset")

4138 )

4139 + assert self.data is not None, "Dataset must be loaded before sampling"

4140 self.data = self.data.map(

4141 lambda item: {

4142 "audio": str(local_path_root / item["url"]),

@@ -4090,15 +4160,19 @@ class ASRDataset(HuggingFaceDataset):

4160 else:

4161 prompt = ""

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

4163 + assert self.data is not None, "Dataset must be loaded before sampling"

4164 sampled_requests: list[SampleRequest] = []

4165 ind = 0

4166 skipped = 0

4096 - asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec")

4097 - asr_max_audio_len_sec = kwargs.get("asr_max_audio_len_sec")

4167 + asr_min_audio_len_sec: float = float(kwargs.get("asr_min_audio_len_sec") or 0.0)

4168 + asr_max_audio_len_sec: float = float(

4169 + kwargs.get("asr_max_audio_len_sec") or float("inf")

4170 + )

4171 durations = []

4172 for item in self.data:

4173 if len(sampled_requests) >= num_requests:

4174 break

4175 + mm_content: dict[str, Any]

4176 audio = item["audio"]

4177 if (

4178 isinstance(audio, dict)

@@ -4203,8 +4277,7 @@ class MLPerfDataset(HuggingFaceDataset):

4277 output_len: int | None = None,

4278 **kwargs,

4279 ) -> list[SampleRequest]:

4206 - # Force dynamic output length based on reference completion.

4207 - dynamic_output = output_len is None

4280 + assert self.data is not None, "Dataset must be loaded before sampling"

4281 sampled_requests: list[SampleRequest] = []

4282 ind = 0

4283

@@ -4222,15 +4295,18 @@ class MLPerfDataset(HuggingFaceDataset):

4295 {"role": "user", "content": question},

4296 ]

4297 prompt_formatted = tokenizer.apply_chat_template(

4225 - messages, add_generation_prompt=True, tokenize=False

4298 + messages, # type: ignore[arg-type]

4299 + add_generation_prompt=True,

4300 + tokenize=False,

4301 )

4302 + assert isinstance(prompt_formatted, str)

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

4304

4305 # Determine output length from reference answer tokens.

4306 ref_out_len = len(

4307 tokenizer(reference_answer, add_special_tokens=False).input_ids

4308 )

4233 - expected_output_len = ref_out_len if dynamic_output else output_len

4309 + expected_output_len = ref_out_len if output_len is None else output_len

4310

4311 # Validate sequence lengths.

4312 if not is_valid_sequence(prompt_len, expected_output_len):

@@ -4371,6 +4447,7 @@ class MMStarDataset(HuggingFaceDataset):

4447 ) -> list[SampleRequest]:

4448 # If --hf-output-len is not set, use the default output length.

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

4450 + assert self.data is not None, "Dataset must be loaded before sampling"

4451 sampled_requests: list[SampleRequest] = []

4452

4453 for ind, item in enumerate(self.data):

@@ -4603,7 +4680,7 @@ class BFCLDataset(HuggingFaceDataset):

4680 exc_info=True,

4681 )

4682 rendered = None

4606 - if rendered is not None:

4683 + if rendered is not None and isinstance(rendered, str):

4684 prompt_len = len(tokenizer(rendered).input_ids)

4685 else:

4686 text = "\n".join(m.get("content", "") for m in messages)

vllm/benchmarks/latency.py

+6 -4

@@ -13,8 +13,9 @@ from tqdm import tqdm

14 from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json

15 from vllm.engine.arg_utils import EngineArgs

16 -from vllm.inputs import PromptType

16 +from vllm.inputs import TextPrompt, TokensPrompt

17 from vllm.sampling_params import BeamSearchParams

18 +from vllm.utils.argparse_utils import FlexibleArgumentParser

21 def save_to_pytorch_benchmark_format(

@@ -30,7 +31,7 @@ def save_to_pytorch_benchmark_format(

31 write_to_json(pt_file, pt_records)

33 -def add_cli_args(parser: argparse.ArgumentParser):

34 +def add_cli_args(parser: FlexibleArgumentParser):

35 parser.add_argument("--input-len", type=int, default=32)

36 parser.add_argument("--output-len", type=int, default=128)

37 parser.add_argument("--batch-size", type=int, default=8)

@@ -103,8 +104,9 @@ def main(args: argparse.Namespace):

104 dummy_prompt_token_ids = np.random.randint(

105 10000, size=(args.batch_size, args.input_len)

106 )

106 - dummy_prompts: list[PromptType] = [

107 - {"prompt_token_ids": batch} for batch in dummy_prompt_token_ids.tolist()

107 + dummy_prompts: list[TokensPrompt | TextPrompt] = [

108 + TokensPrompt(prompt_token_ids=batch)

109 + for batch in dummy_prompt_token_ids.tolist()

110 ]

111

112 def llm_generate():

vllm/benchmarks/lib/endpoint_request_func.py

+24 -20

@@ -75,7 +75,7 @@ class RequestFuncInput:

75 logprobs: int | None = None

76 extra_headers: dict | None = None

77 extra_body: dict | None = None

78 - multi_modal_content: dict | list[dict] | None = None

78 + multi_modal_content: dict[str, Any] | list[dict[str, Any]] | None = None

79 ignore_eos: bool = False

80 language: str | None = None

81 request_id: str | None = None

@@ -272,13 +272,13 @@ def _get_chat_content(

272 request_func_input: RequestFuncInput,

273 mm_position: Literal["first", "last"] = "last",

274 ) -> list[dict[str, Any]]:

275 - mm_contents = []

275 + mm_contents: list[dict[str, Any]] = []

276 if request_func_input.multi_modal_content:

277 mm_content = request_func_input.multi_modal_content

278 if isinstance(mm_content, list):

279 - mm_contents.extend(request_func_input.multi_modal_content)

279 + mm_contents.extend(mm_content)

280 elif isinstance(mm_content, dict):

281 - mm_contents.append(request_func_input.multi_modal_content)

281 + mm_contents.append(mm_content)

282 else:

283 raise TypeError(

284 "multi_modal_content must be a dict or list[dict] for openai-chat"

@@ -293,10 +293,11 @@ def _get_chat_content(

293 for item in prompt

294 )

295 ):

296 + prompt_dicts: list[dict[str, Any]] = prompt # type: ignore[assignment]

297 if mm_position == "first":

297 - return mm_contents + prompt

298 + return mm_contents + prompt_dicts

299

299 - return prompt + mm_contents

300 + return prompt_dicts + mm_contents

301

302 text_contents = [{"type": "text", "text": prompt}]

303

@@ -307,15 +308,15 @@ def _get_chat_content(

308

309

310 def _is_chat_messages(prompt: Any) -> bool:

310 - return (

311 - isinstance(prompt, list)

312 - and prompt

313 - and all(

314 - isinstance(item, dict)

315 - and isinstance(item.get("role"), str)

316 - and isinstance(item.get("content"), (str, list))

317 - for item in prompt

318 - )

311 + if not isinstance(prompt, list):

312 + return False

313 + if not prompt:

314 + return False

315 + return all(

316 + isinstance(item, dict)

317 + and isinstance(item.get("role"), str)

318 + and isinstance(item.get("content"), (str, list))

319 + for item in prompt

320 )

321

322

@@ -325,7 +326,7 @@ def _get_chat_messages(

326 ) -> list[dict[str, Any]]:

327 prompt = request_func_input.prompt

328 if _is_chat_messages(prompt):

328 - return prompt

329 + return prompt # type: ignore[return-value]

330

331 return [

332 {

@@ -385,8 +386,8 @@ async def async_request_openai_chat_completions(

386 if not chunk_bytes:

387 continue

388

388 - messages = handler.add_chunk(chunk_bytes)

389 - for message in messages:

389 + message_strings = handler.add_chunk(chunk_bytes)

390 + for message in message_strings:

391 # NOTE: SSE comments (often used as pings) start with

392 # a colon. These are not JSON data payload and should

393 # be skipped.

@@ -790,7 +791,7 @@ async def async_request_infinity_embeddings(

791 api_url = request_func_input.api_url

792 _validate_api_url(api_url, "Infinity Embeddings API", "embeddings")

793

793 - payload = {

794 + payload: dict[str, Any] = {

795 "model": request_func_input.model_name

796 if request_func_input.model_name

797 else request_func_input.model,

@@ -849,7 +850,10 @@ async def async_request_vllm_pooling(

850 "truncate_prompt_tokens": -1,

851 }

852

852 - payload = payload | request_func_input.prompt

853 + if isinstance(request_func_input.prompt, dict):

854 + payload = payload | request_func_input.prompt

855 + else:

856 + payload["input"] = request_func_input.prompt

857

858 _update_payload_common(payload, request_func_input)

859

vllm/benchmarks/lib/utils.py

+2 -2

@@ -15,7 +15,7 @@ def extract_field(

15 if field_name in extra_info:

16 return extra_info[field_name]

18 - v = args

18 + v: Any = args

19 # For example, args.compilation_config.mode

20 for nested_field in field_name.split("."):

21 if not hasattr(v, nested_field):

@@ -43,7 +43,7 @@ def convert_to_pytorch_benchmark_format(

43 on metric per record

44 https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database

45 """

46 - records = []

46 + records: list[Any] = []

47 if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False):

48 return records

vllm/benchmarks/mm_processor.py

+3 -2

@@ -28,6 +28,7 @@ from vllm.benchmarks.datasets import (

28 )

29 from vllm.benchmarks.throughput import get_requests

30 from vllm.engine.arg_utils import EngineArgs

31 +from vllm.utils.argparse_utils import FlexibleArgumentParser

32 from vllm.utils.gc_utils import freeze_gc_heap

33 from vllm.utils.import_utils import PlaceholderModule

@@ -369,7 +370,7 @@ def benchmark_multimodal_processor(

370 return benchmark_result

371

372

372 -def add_cli_args(parser: argparse.ArgumentParser) -> None:

373 +def add_cli_args(parser: FlexibleArgumentParser) -> None:

374 """Add CLI arguments for the multimodal processor benchmark."""

375 from vllm.engine.arg_utils import EngineArgs

376

@@ -532,7 +533,7 @@ def main(args: argparse.Namespace) -> None:

533

534

535 if __name__ == "__main__":

535 - parser = argparse.ArgumentParser(description="Benchmark mm processor latency")

536 + parser = FlexibleArgumentParser(description="Benchmark mm processor latency")

537 add_cli_args(parser)

538 args = parser.parse_args()

539 main(args)

vllm/benchmarks/serve.py

+37 -14

@@ -30,7 +30,7 @@ import ssl

30 import time

31 import uuid

32 import warnings

33 -from collections.abc import AsyncGenerator, Iterable

33 +from collections.abc import AsyncGenerator, Iterable, Iterator

34 from dataclasses import dataclass, replace

35 from datetime import datetime

36 from enum import Enum

@@ -52,6 +52,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (

52 from vllm.benchmarks.lib.ready_checker import wait_for_endpoint

53 from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json

54 from vllm.tokenizers import TokenizerLike, get_tokenizer

55 +from vllm.utils.argparse_utils import FlexibleArgumentParser

56 from vllm.utils.gc_utils import freeze_gc_heap

57 from vllm.utils.network_utils import join_host_port

@@ -129,6 +130,7 @@ async def _align_prompts_to_server_tokenizer(

130 )

131

132 async def _fix_one(req: SampleRequest) -> SampleRequest:

133 + assert isinstance(req.prompt, str)

134 tokens = await _tokenize(req.prompt)

135 if len(tokens) <= req.prompt_len:

136 return req

@@ -360,7 +362,7 @@ class EmbedBenchmarkMetrics:

362 mean_e2el_ms: float

363 std_e2el_ms: float

364 median_e2el_ms: float

363 - percentiles_e2el_ms: float

365 + percentiles_e2el_ms: list[tuple[float, float]]

366

367

368 def _get_current_request_rate(

@@ -433,8 +435,8 @@ async def get_request(

435 assert total_requests > 0, "No requests provided."

436

437 # Precompute delays among requests to minimize request send laggings

436 - request_rates = []

437 - delay_ts = []

438 + request_rates: list[float] = []

439 + delay_ts: list[float] = []

440

441 # if the traces have timing info then:

442 if not self_timed:

@@ -484,7 +486,10 @@ async def get_request(

486 else:

487 for request_index, request in enumerate(input_requests):

488 # this is cumulative running ts, from which sleep is calculated later

487 - delay_ts.append(request.timestamp)

489 + if request.timestamp is not None:

490 + delay_ts.append(request.timestamp)

491 + else:

492 + delay_ts.append(0.0)

493 # TODO: there is no notion of RPS here, may be we can calculate

494 # from the trace.

495 request_rates.append(0.0)

@@ -599,7 +604,7 @@ def calculate_metrics(

604 )

605 actual_output_lens.append(output_len)

606 total_input += outputs[i].prompt_len

602 - tpot = 0

607 + tpot = 0.0

608 if output_len > 1:

609 latency_minus_ttft = outputs[i].latency - outputs[i].ttft

610 tpot = latency_minus_ttft / (output_len - 1)

@@ -894,11 +899,12 @@ async def benchmark(

899

900 print("Starting main benchmark run...")

901

902 + lora_modules_iter: Iterator[str] | None = None

903 if lora_modules:

904 lora_modules_list = list(lora_modules)

905 if lora_assignment == "round-robin":

906 # Deterministic round-robin assignment across requests.

901 - lora_modules = iter(

907 + lora_modules_iter = iter(

908 [

909 lora_modules_list[i % len(lora_modules_list)]

910 for i in range(len(input_requests))

@@ -906,7 +912,7 @@ async def benchmark(

912 )

913 else:

914 # For each input request, choose a LoRA module at random.

909 - lora_modules = iter(

915 + lora_modules_iter = iter(

916 [random.choice(lora_modules_list) for _ in range(len(input_requests))]

917 )

918

@@ -1004,19 +1010,23 @@ async def benchmark(

1010 )

1011 per_request_extra_body = _merge_overrides(extra_body, request.request_overrides)

1012 req_model_id, req_model_name = model_id, model_name

1007 - if lora_modules:

1008 - req_lora_module = next(lora_modules)

1013 + if lora_modules_iter:

1014 + req_lora_module = next(lora_modules_iter)

1015 req_model_id, req_model_name = req_lora_module, req_lora_module

1016

1017 + mm_content_typed: dict[str, Any] | list[dict[str, Any]] | None = None

1018 + if isinstance(mm_content, (dict, list)):

1019 + mm_content_typed = mm_content

1020 +

1021 request_func_input = RequestFuncInput(

1022 model=req_model_id,

1023 model_name=req_model_name,

1024 prompt=prompt,

1025 api_url=api_url,

1026 prompt_len=prompt_len,

1017 - output_len=output_len,

1027 + output_len=output_len or 0,

1028 logprobs=logprobs,

1019 - multi_modal_content=mm_content,

1029 + multi_modal_content=mm_content_typed,

1030 ignore_eos=ignore_eos,

1031 extra_headers=extra_headers,

1032 extra_body=per_request_extra_body,

@@ -1107,6 +1117,8 @@ async def benchmark(

1117 "committed_per_step": delta_committed / denoising_steps,

1118 }

1119

1120 + metrics: BenchmarkMetrics | EmbedBenchmarkMetrics

1121 + actual_output_lens: list[int] | int

1122 if task_type == TaskType.GENERATION:

1123 metrics, actual_output_lens = calculate_metrics(

1124 input_requests=input_requests,

@@ -1140,7 +1152,7 @@ async def benchmark(

1152 "Request throughput (req/s):", metrics.request_throughput

1153 )

1154 )

1143 - if goodput_config_dict:

1155 + if goodput_config_dict and isinstance(metrics, BenchmarkMetrics):

1156 print(

1157 "{:<40} {:<10.2f}".format(

1158 "Request goodput (req/s):", metrics.request_goodput

@@ -1177,6 +1189,7 @@ async def benchmark(

1189 )

1190 )

1191

1192 + result: dict[str, Any]

1193 if isinstance(metrics, BenchmarkMetrics):

1194 result = {

1195 "duration": benchmark_duration,

@@ -1462,7 +1475,7 @@ def compute_result_filename(

1475 return file_name

1476

1477

1465 -def add_cli_args(parser: argparse.ArgumentParser):

1478 +def add_cli_args(parser: FlexibleArgumentParser):

1479 add_dataset_parser(parser)

1480 parser.add_argument(

1481 "--label",

@@ -2023,6 +2036,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:

2036 args.self_timed = False

2037

2038 # Load the dataset.

2039 + assert tokenizer is not None, "Tokenizer must be initialized before loading dataset"

2040 input_requests = get_samples(args, tokenizer)

2041

2042 if args.dataset_name in ("random", "prefix_repetition"):

@@ -2154,6 +2168,9 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:

2168

2169 # Generate timeline plot if requested

2170 if args.plot_timeline:

2171 + assert file_name is not None, (

2172 + "file_name must be set when plot_timeline is enabled"

2173 + )

2174 try:

2175 from vllm.benchmarks.plot import generate_timeline_plot

2176

@@ -2200,6 +2217,9 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:

2217

2218 # Generate dataset statistics plot if requested

2219 if args.plot_dataset_stats:

2220 + assert file_name is not None, (

2221 + "file_name must be set when plot_dataset_stats is enabled"

2222 + )

2223 try:

2224 from vllm.benchmarks.plot import generate_dataset_stats_plot

2225

@@ -2249,6 +2269,9 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:

2269

2270 # Save to file

2271 if args.save_result or args.append_result:

2272 + assert file_name is not None, (

2273 + "file_name must be set when save_result or append_result is enabled"

2274 + )

2275 with open(

2276 file_name, mode="a+" if args.append_result else "w", encoding="utf-8"

2277 ) as outfile:

vllm/benchmarks/startup.py

+3 -2

@@ -26,6 +26,7 @@ from vllm.benchmarks.lib.utils import (

26 write_to_json,

27 )

28 from vllm.engine.arg_utils import EngineArgs

29 +from vllm.utils.argparse_utils import FlexibleArgumentParser

31 PERCENTAGES = [10, 25, 50, 75, 90, 99]

@@ -190,7 +191,7 @@ def save_to_pytorch_benchmark_format(

191 write_to_json(f"{base_name}.{m.key}.pytorch.json", records)

192

193

193 -def add_cli_args(parser: argparse.ArgumentParser):

194 +def add_cli_args(parser: FlexibleArgumentParser):

195 parser.add_argument(

196 "--num-iters-cold",

197 type=int,

@@ -234,7 +235,7 @@ def main(args: argparse.Namespace):

235 """

236

237 # Create a queue for inter-process communication

237 - result_queue = multiprocessing.Queue()

238 + result_queue: multiprocessing.Queue[Any] = multiprocessing.Queue()

239 process = multiprocessing.Process(

240 target=run_startup_in_subprocess,

241 args=(

vllm/benchmarks/sweep/cli.py

+2 -1

@@ -3,6 +3,7 @@

3 import argparse

5 from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG

6 +from vllm.utils.argparse_utils import FlexibleArgumentParser

8 from .plot import SweepPlotArgs

9 from .plot import main as plot_main

@@ -24,7 +25,7 @@ SUBCOMMANDS = (

25 )

27 -def add_cli_args(parser: argparse.ArgumentParser):

28 +def add_cli_args(parser: FlexibleArgumentParser):

29 subparsers = parser.add_subparsers(required=True, dest="sweep_type")

31 for cmd, entrypoint in SUBCOMMANDS:

vllm/benchmarks/sweep/param_sweep.py

+2 -3

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import json

4 import os

5 -from typing import Any

7 class ParameterSweep(list["ParameterSweepItem"]):

@@ -62,8 +61,8 @@ class ParameterSweepItem(dict[str, object]):

62 return cls(record)

65 - def __or__(self, other: dict[str, Any]):

66 - return type(self)(super().__or__(other))

64 + def __or__(self, other: dict[str, object], /) -> "ParameterSweepItem": # type: ignore[override]

65 + return ParameterSweepItem(super().__or__(other))

67 @property

68 def name(self) -> str:

vllm/benchmarks/sweep/plot.py

+6 -3

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, ClassVar

13 from typing_extensions import Self, override

15 +from vllm.utils.argparse_utils import FlexibleArgumentParser

16 from vllm.utils.collection_utils import full_groupby

17 from vllm.utils.import_utils import PlaceholderModule

@@ -51,6 +52,7 @@ class PlotFilterBase(ABC):

52 class PlotEqualTo(PlotFilterBase):

53 @override

54 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":

55 + target: float | str

56 try:

57 target = float(self.target)

58 except ValueError:

@@ -63,6 +65,7 @@ class PlotEqualTo(PlotFilterBase):

65 class PlotNotEqualTo(PlotFilterBase):

66 @override

67 def apply(self, df: "pd.DataFrame") -> "pd.DataFrame":

68 + target: float | str

69 try:

70 target = float(self.target)

71 except ValueError:

@@ -182,7 +185,7 @@ def _convert_inf_nan_strings(data: list[dict[str, object]]) -> list[dict[str, ob

185 """

186 converted_data = []

187 for record in data:

185 - converted_record = {}

188 + converted_record: dict[str, object] = {}

189 for key, value in record.items():

190 if isinstance(value, str):

191 if value in ["inf", "-inf", "nan"]:

@@ -531,7 +534,7 @@ class SweepPlotArgs:

534 )

535

536 @classmethod

534 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:

537 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser:

538 parser.add_argument(

539 "EXPERIMENT_DIR",

540 type=str,

@@ -682,7 +685,7 @@ def main(args: argparse.Namespace):

685

686

687 if __name__ == "__main__":

685 - parser = argparse.ArgumentParser(description=SweepPlotArgs.parser_help)

688 + parser = FlexibleArgumentParser(description=SweepPlotArgs.parser_help)

689 SweepPlotArgs.add_cli_args(parser)

690

691 main(parser.parse_args())

vllm/benchmarks/sweep/plot_pareto.py

+3 -2

@@ -8,6 +8,7 @@ from functools import partial

8 from pathlib import Path

9 from typing import TYPE_CHECKING, ClassVar

11 +from vllm.utils.argparse_utils import FlexibleArgumentParser

12 from vllm.utils.collection_utils import full_groupby

13 from vllm.utils.import_utils import PlaceholderModule

@@ -342,7 +343,7 @@ class SweepPlotParetoArgs:

343 )

344

345 @classmethod

345 - def add_cli_args(cls, parser: argparse.ArgumentParser):

346 + def add_cli_args(cls, parser: FlexibleArgumentParser):

347 parser.add_argument(

348 "EXPERIMENT_DIR",

349 type=str,

@@ -394,7 +395,7 @@ def main(args: argparse.Namespace):

395

396

397 if __name__ == "__main__":

397 - parser = argparse.ArgumentParser(description=SweepPlotParetoArgs.parser_help)

398 + parser = FlexibleArgumentParser(description=SweepPlotParetoArgs.parser_help)

399 SweepPlotParetoArgs.add_cli_args(parser)

400

401 main(parser.parse_args())

vllm/benchmarks/sweep/serve.py

+3 -2

@@ -10,6 +10,7 @@ from datetime import datetime

10 from pathlib import Path

11 from typing import ClassVar

13 +from vllm.utils.argparse_utils import FlexibleArgumentParser

14 from vllm.utils.import_utils import PlaceholderModule

16 from .param_sweep import ParameterSweep, ParameterSweepItem

@@ -368,7 +369,7 @@ class SweepServeArgs:

369 )

370

371 @classmethod

371 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:

372 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser:

373 parser.add_argument(

374 "--serve-cmd",

375 type=str,

@@ -531,7 +532,7 @@ def main(args: argparse.Namespace):

532

533

534 if __name__ == "__main__":

534 - parser = argparse.ArgumentParser(description=SweepServeArgs.parser_help)

535 + parser = FlexibleArgumentParser(description=SweepServeArgs.parser_help)

536 SweepServeArgs.add_cli_args(parser)

537

538 main(parser.parse_args())

vllm/benchmarks/sweep/serve_workload.py

+3 -2

@@ -10,6 +10,7 @@ import numpy as np

10 from typing_extensions import assert_never

12 from vllm.benchmarks.datasets import DEFAULT_NUM_PROMPTS

13 +from vllm.utils.argparse_utils import FlexibleArgumentParser

14 from vllm.utils.import_utils import PlaceholderModule

16 from .param_sweep import ParameterSweep, ParameterSweepItem

@@ -273,7 +274,7 @@ class SweepServeWorkloadArgs(SweepServeArgs):

274 )

275

276 @classmethod

276 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:

277 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser:

278 parser = super().add_cli_args(parser)

279

280 workload_group = parser.add_argument_group("workload options")

@@ -322,7 +323,7 @@ def main(args: argparse.Namespace):

323

324

325 if __name__ == "__main__":

325 - parser = argparse.ArgumentParser(description=SweepServeWorkloadArgs.parser_help)

326 + parser = FlexibleArgumentParser(description=SweepServeWorkloadArgs.parser_help)

327 SweepServeWorkloadArgs.add_cli_args(parser)

328

329 main(parser.parse_args())

vllm/benchmarks/sweep/startup.py

+2 -2

@@ -317,7 +317,7 @@ class SweepStartupArgs:

317 )

318

319 @classmethod

320 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:

320 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> FlexibleArgumentParser:

321 parser.add_argument(

322 "--startup-cmd",

323 type=str,

@@ -438,6 +438,6 @@ def main(args: argparse.Namespace):

438

439

440 if __name__ == "__main__":

441 - parser = argparse.ArgumentParser(description=SweepStartupArgs.parser_help)

441 + parser = FlexibleArgumentParser(description=SweepStartupArgs.parser_help)

442 SweepStartupArgs.add_cli_args(parser)

443 main(parser.parse_args())

vllm/benchmarks/throughput.py

+44 -17

@@ -18,6 +18,7 @@ from transformers import AutoModelForCausalLM, PreTrainedTokenizerBase

18 from vllm.benchmarks.datasets import (

19 AIMODataset,

20 ASRDataset,

21 + BenchmarkDataset,

22 BurstGPTDataset,

23 ConversationDataset,

24 InstructCoderDataset,

@@ -41,6 +42,7 @@ from vllm.outputs import RequestOutput

42 from vllm.platforms import current_platform

43 from vllm.sampling_params import BeamSearchParams

44 from vllm.tokenizers import TokenizerLike, get_tokenizer

45 +from vllm.utils.argparse_utils import FlexibleArgumentParser

46 from vllm.utils.async_utils import merge_async_iterators

@@ -102,13 +104,15 @@ def _run_vllm_requests(

104

105 prompts: list[TextPrompt | TokensPrompt] = []

106 sampling_params: list[SamplingParams] = []

105 - lora_requests: list[LoRARequest] | None = [] if enable_lora else None

107 + lora_requests: list[LoRARequest | None] | None = [] if enable_lora else None

108 for request in requests:

107 - prompt = (

108 - TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"])

109 - if "prompt_token_ids" in request.prompt

110 - else TextPrompt(prompt=request.prompt)

111 - )

109 + if isinstance(request.prompt, dict) and "prompt_token_ids" in request.prompt:

110 + prompt_token_ids = request.prompt["prompt_token_ids"]

111 + assert isinstance(prompt_token_ids, list)

112 + prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)

113 + else:

114 + assert isinstance(request.prompt, str)

115 + prompt = TextPrompt(prompt=request.prompt)

116 if request.multi_modal_data:

117 assert isinstance(request.multi_modal_data, dict)

118 prompt["multi_modal_data"] = request.multi_modal_data

@@ -159,7 +163,20 @@ def _run_vllm_requests(

163 end = time.perf_counter()

164 else:

165 assert lora_requests is None, "BeamSearch API does not support LoRA"

162 - beam_prompts = [request.prompt for request in requests]

166 + beam_prompts: list[TextPrompt | TokensPrompt] = []

167 + for request in requests:

168 + if isinstance(request.prompt, str):

169 + beam_prompts.append(TextPrompt(prompt=request.prompt))

170 + elif (

171 + isinstance(request.prompt, dict)

172 + and "prompt_token_ids" in request.prompt

173 + ):

174 + token_ids = request.prompt["prompt_token_ids"]

175 + assert isinstance(token_ids, list)

176 + beam_prompts.append(TokensPrompt(prompt_token_ids=token_ids))

177 + else:

178 + # Fallback: convert to string

179 + beam_prompts.append(TextPrompt(prompt=str(request.prompt)))

180 # output_len should be the same for all requests.

181 output_len = requests[0].expected_output_len

182 for request in requests:

@@ -268,7 +285,7 @@ def _run_vllm_chat_requests(

285 llm.wake_up(tags=["scheduling"])

286 outputs = llm.wait_for_completion(output_type=RequestOutput, use_tqdm=True)

287 else:

271 - outputs = llm.chat(prompts, sampling_params, use_tqdm=True)

288 + outputs = llm.chat(prompts, sampling_params, use_tqdm=True) # type: ignore[arg-type]

289

290 if do_profile:

291 llm.stop_profile()

@@ -339,11 +356,13 @@ async def _run_vllm_async_requests(

356 sampling_params: list[SamplingParams] = []

357 lora_requests: list[LoRARequest | None] = []

358 for request in requests:

342 - prompt = (

343 - TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"])

344 - if "prompt_token_ids" in request.prompt

345 - else TextPrompt(prompt=request.prompt)

346 - )

359 + if isinstance(request.prompt, dict) and "prompt_token_ids" in request.prompt:

360 + prompt_token_ids = request.prompt["prompt_token_ids"]

361 + assert isinstance(prompt_token_ids, list)

362 + prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)

363 + else:

364 + assert isinstance(request.prompt, str)

365 + prompt = TextPrompt(prompt=request.prompt)

366

367 if request.multi_modal_data:

368 assert isinstance(request.multi_modal_data, dict)

@@ -366,9 +385,11 @@ async def _run_vllm_async_requests(

385 start = time.perf_counter()

386 if do_profile:

387 await llm.start_profile()

369 - for i, (prompt, sp, lr) in enumerate(zip(prompts, sampling_params, lora_requests)):

388 + for i, (prompt_item, sp, lr) in enumerate(

389 + zip(prompts, sampling_params, lora_requests)

390 + ):

391 generator = llm.generate(

371 - prompt, sp, lora_request=lr, request_id=f"{request_id_prefix}{i}"

392 + prompt_item, sp, lora_request=lr, request_id=f"{request_id_prefix}{i}"

393 )

394 generators.append(generator)

395 all_gens = merge_async_iterators(*generators)

@@ -440,6 +461,7 @@ def _run_hf_requests(

461 prompt_len = requests[i].prompt_len

462 output_len = requests[i].expected_output_len

463 # Add the prompt to the batch.

464 + assert isinstance(prompt, str), "Prompt must be a string for HF backend"

465 batch.append(prompt)

466 max_prompt_len = max(max_prompt_len, prompt_len)

467 max_output_len = max(max_output_len, output_len)

@@ -500,6 +522,7 @@ def save_to_pytorch_benchmark_format(

522

523 def get_requests(args, tokenizer):

524 # Common parameters for all dataset types.

525 + dataset_cls: type[BenchmarkDataset]

526 common_kwargs = {

527 "dataset_path": args.dataset_path,

528 "random_seed": args.seed,

@@ -875,7 +898,7 @@ def validate_args(args):

898 )

899

900

878 -def add_cli_args(parser: argparse.ArgumentParser):

901 +def add_cli_args(parser: FlexibleArgumentParser):

902 parser.add_argument(

903 "--backend",

904 type=str,

@@ -1178,7 +1201,11 @@ def main(args: argparse.Namespace):

1201 total_prompt_tokens += (

1202 len(ro.prompt_token_ids) if ro.prompt_token_ids else 0

1203 )

1181 - total_output_tokens += sum(len(o.token_ids) for o in ro.outputs if o)

1204 + total_output_tokens += sum(

1205 + len(o.token_ids)

1206 + for o in ro.outputs

1207 + if o is not None and o.token_ids is not None

1208 + )

1209 total_num_tokens = total_prompt_tokens + total_output_tokens

1210 else:

1211 total_num_tokens = sum(r.prompt_len + r.expected_output_len for r in requests)

vllm/entrypoints/cli/benchmark/base.py

+2 -1

@@ -3,6 +3,7 @@

3 import argparse

5 from vllm.entrypoints.cli.types import CLISubcommand

6 +from vllm.utils.argparse_utils import FlexibleArgumentParser

9 class BenchmarkSubcommandBase(CLISubcommand):

@@ -11,7 +12,7 @@ class BenchmarkSubcommandBase(CLISubcommand):

12 help: str

14 @classmethod

14 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

15 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

16 """Add the CLI arguments to the parser."""

17 raise NotImplementedError

vllm/entrypoints/cli/benchmark/latency.py

+2 -1

@@ -4,6 +4,7 @@ import argparse

5 from vllm.benchmarks.latency import add_cli_args, main

6 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase

7 +from vllm.utils.argparse_utils import FlexibleArgumentParser

10 class BenchmarkLatencySubcommand(BenchmarkSubcommandBase):

@@ -13,7 +14,7 @@ class BenchmarkLatencySubcommand(BenchmarkSubcommandBase):

14 help = "Benchmark the latency of a single batch of requests."

16 @classmethod

16 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

17 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

18 add_cli_args(parser)

20 @staticmethod

vllm/entrypoints/cli/benchmark/mm_processor.py

+2 -1

@@ -4,6 +4,7 @@ import argparse

5 from vllm.benchmarks.mm_processor import add_cli_args, main

6 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase

7 +from vllm.utils.argparse_utils import FlexibleArgumentParser

10 class BenchmarkMMProcessorSubcommand(BenchmarkSubcommandBase):

@@ -13,7 +14,7 @@ class BenchmarkMMProcessorSubcommand(BenchmarkSubcommandBase):

14 help = "Benchmark multimodal processor latency across different configurations."

16 @classmethod

16 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

17 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

18 add_cli_args(parser)

20 @staticmethod

vllm/entrypoints/cli/benchmark/serve.py

+2 -1

@@ -4,6 +4,7 @@ import argparse

5 from vllm.benchmarks.serve import add_cli_args, main

6 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase

7 +from vllm.utils.argparse_utils import FlexibleArgumentParser

10 class BenchmarkServingSubcommand(BenchmarkSubcommandBase):

@@ -13,7 +14,7 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):

14 help = "Benchmark the online serving throughput."

16 @classmethod

16 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

17 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

18 add_cli_args(parser)

20 @staticmethod

vllm/entrypoints/cli/benchmark/startup.py

+2 -1

@@ -4,6 +4,7 @@ import argparse

5 from vllm.benchmarks.startup import add_cli_args, main

6 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase

7 +from vllm.utils.argparse_utils import FlexibleArgumentParser

10 class BenchmarkStartupSubcommand(BenchmarkSubcommandBase):

@@ -13,7 +14,7 @@ class BenchmarkStartupSubcommand(BenchmarkSubcommandBase):

14 help = "Benchmark the startup time of vLLM models."

16 @classmethod

16 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

17 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

18 add_cli_args(parser)

20 @staticmethod

vllm/entrypoints/cli/benchmark/sweep.py

+2 -1

@@ -4,6 +4,7 @@ import argparse

5 from vllm.benchmarks.sweep.cli import add_cli_args, main

6 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase

7 +from vllm.utils.argparse_utils import FlexibleArgumentParser

10 class BenchmarkSweepSubcommand(BenchmarkSubcommandBase):

@@ -13,7 +14,7 @@ class BenchmarkSweepSubcommand(BenchmarkSubcommandBase):

14 help = "Benchmark for a parameter sweep."

16 @classmethod

16 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

17 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

18 add_cli_args(parser)

20 @staticmethod

vllm/entrypoints/cli/benchmark/throughput.py

+2 -1

@@ -4,6 +4,7 @@ import argparse

5 from vllm.benchmarks.throughput import add_cli_args, main

6 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase

7 +from vllm.utils.argparse_utils import FlexibleArgumentParser

10 class BenchmarkThroughputSubcommand(BenchmarkSubcommandBase):

@@ -13,7 +14,7 @@ class BenchmarkThroughputSubcommand(BenchmarkSubcommandBase):

14 help = "Benchmark offline inference throughput."

16 @classmethod

16 - def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:

17 + def add_cli_args(cls, parser: FlexibleArgumentParser) -> None:

18 add_cli_args(parser)

20 @staticmethod