6cd4ae8acd · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -281,7 +281,7 @@ class LLM:

281 sampling_params: Optional[Union[SamplingParams,

282 Sequence[SamplingParams]]] = None,

283 *,

284 - use_tqdm: bool = True,

284 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

285 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

286 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

287 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -297,7 +297,7 @@ class LLM:

297 sampling_params: Optional[Union[SamplingParams,

298 list[SamplingParams]]] = None,

299 prompt_token_ids: Optional[list[int]] = None,

300 - use_tqdm: bool = True,

300 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

301 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

302 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

303 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -313,7 +313,7 @@ class LLM:

313 sampling_params: Optional[Union[SamplingParams,

314 list[SamplingParams]]] = None,

315 prompt_token_ids: Optional[list[list[int]]] = None,

316 - use_tqdm: bool = True,

316 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

317 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

318 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

319 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -330,7 +330,7 @@ class LLM:

330 list[SamplingParams]]] = None,

331 *,

332 prompt_token_ids: list[int],

333 - use_tqdm: bool = True,

333 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

334 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

335 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

336 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -347,7 +347,7 @@ class LLM:

347 list[SamplingParams]]] = None,

348 *,

349 prompt_token_ids: list[list[int]],

350 - use_tqdm: bool = True,

350 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

351 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

352 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

353 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -362,7 +362,7 @@ class LLM:

362 prompts: None,

363 sampling_params: None,

364 prompt_token_ids: Union[list[int], list[list[int]]],

365 - use_tqdm: bool = True,

365 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

366 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

367 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

368 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -382,7 +382,7 @@ class LLM:

382 sampling_params: Optional[Union[SamplingParams,

383 Sequence[SamplingParams]]] = None,

384 prompt_token_ids: Optional[Union[list[int], list[list[int]]]] = None,

385 - use_tqdm: bool = True,

385 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

386 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

387 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

388 guided_options_request: Optional[Union[LLMGuidedOptions,

@@ -404,7 +404,10 @@ class LLM:

404 When it is a single value, it is applied to every prompt.

405 When it is a list, the list must have the same length as the

406 prompts and it is paired one by one with the prompt.

407 - use_tqdm: Whether to use tqdm to display the progress bar.

407 + use_tqdm: If `True`, shows a tqdm progress bar.

408 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),

409 + it is used to create the progress bar.

410 + If `False`, no progress bar is created.

411 lora_request: LoRA request to use for generation, if any.

412 prompt_adapter_request: Prompt Adapter request to use for

413 generation, if any.

@@ -678,7 +681,7 @@ class LLM:

681 list[list[ChatCompletionMessageParam]]],

682 sampling_params: Optional[Union[SamplingParams,

683 list[SamplingParams]]] = None,

681 - use_tqdm: bool = True,

684 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

685 lora_request: Optional[LoRARequest] = None,

686 chat_template: Optional[str] = None,

687 chat_template_content_format: ChatTemplateContentFormatOption = "auto",

@@ -709,7 +712,10 @@ class LLM:

712 is a single value, it is applied to every prompt. When it

713 is a list, the list must have the same length as the

714 prompts and it is paired one by one with the prompt.

712 - use_tqdm: Whether to use tqdm to display the progress bar.

715 + use_tqdm: If `True`, shows a tqdm progress bar.

716 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),

717 + it is used to create the progress bar.

718 + If `False`, no progress bar is created.

719 lora_request: LoRA request to use for generation, if any.

720 chat_template: The template to use for structuring the chat.

721 If not provided, the model's default chat template will be used.

@@ -823,7 +829,7 @@ class LLM:

829 Sequence[PoolingParams]]] = None,

830 *,

831 truncate_prompt_tokens: Optional[int] = None,

826 - use_tqdm: bool = True,

832 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

833 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

834 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

835 ) -> list[PoolingRequestOutput]:

@@ -838,7 +844,7 @@ class LLM:

844 Sequence[PoolingParams]]] = None,

845 prompt_token_ids: Optional[list[int]] = None,

846 truncate_prompt_tokens: Optional[int] = None,

841 - use_tqdm: bool = True,

847 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

848 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

849 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

850 ) -> list[PoolingRequestOutput]:

@@ -853,7 +859,7 @@ class LLM:

859 Sequence[PoolingParams]]] = None,

860 prompt_token_ids: Optional[list[list[int]]] = None,

861 truncate_prompt_tokens: Optional[int] = None,

856 - use_tqdm: bool = True,

862 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

863 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

864 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

865 ) -> list[PoolingRequestOutput]:

@@ -869,7 +875,7 @@ class LLM:

875 *,

876 prompt_token_ids: list[int],

877 truncate_prompt_tokens: Optional[int] = None,

872 - use_tqdm: bool = True,

878 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

879 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

880 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

881 ) -> list[PoolingRequestOutput]:

@@ -885,7 +891,7 @@ class LLM:

891 *,

892 prompt_token_ids: list[list[int]],

893 truncate_prompt_tokens: Optional[int] = None,

888 - use_tqdm: bool = True,

894 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

895 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

896 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

897 ) -> list[PoolingRequestOutput]:

@@ -899,7 +905,7 @@ class LLM:

905 pooling_params: None,

906 prompt_token_ids: Union[list[int], list[list[int]]],

907 truncate_prompt_tokens: Optional[int] = None,

902 - use_tqdm: bool = True,

908 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

909 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

910 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

911 ) -> list[PoolingRequestOutput]:

@@ -918,7 +924,7 @@ class LLM:

924 Sequence[PoolingParams]]] = None,

925 prompt_token_ids: Optional[Union[list[int], list[list[int]]]] = None,

926 truncate_prompt_tokens: Optional[int] = None,

921 - use_tqdm: bool = True,

927 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

928 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

929 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

930 ) -> list[PoolingRequestOutput]:

@@ -935,7 +941,10 @@ class LLM:

941 for more details about the format of each prompts.

942 pooling_params: The pooling parameters for pooling. If None, we

943 use the default pooling parameters.

938 - use_tqdm: Whether to use tqdm to display the progress bar.

944 + use_tqdm: If `True`, shows a tqdm progress bar.

945 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),

946 + it is used to create the progress bar.

947 + If `False`, no progress bar is created.

948 lora_request: LoRA request to use for generation, if any.

949 prompt_adapter_request: Prompt Adapter request to use for

950 generation, if any.

@@ -1005,7 +1014,7 @@ class LLM:

1014 /,

1015 *,

1016 truncate_prompt_tokens: Optional[int] = None,

1008 - use_tqdm: bool = True,

1017 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1018 pooling_params: Optional[Union[PoolingParams,

1019 Sequence[PoolingParams]]] = None,

1020 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

@@ -1024,7 +1033,10 @@ class LLM:

1033 for more details about the format of each prompts.

1034 pooling_params: The pooling parameters for pooling. If None, we

1035 use the default pooling parameters.

1027 - use_tqdm: Whether to use tqdm to display the progress bar.

1036 + use_tqdm: If `True`, shows a tqdm progress bar.

1037 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),

1038 + it is used to create the progress bar.

1039 + If `False`, no progress bar is created.

1040 lora_request: LoRA request to use for generation, if any.

1041 prompt_adapter_request: Prompt Adapter request to use for

1042 generation, if any.

@@ -1051,7 +1063,7 @@ class LLM:

1063 prompts: Union[PromptType, Sequence[PromptType]],

1064 /,

1065 *,

1054 - use_tqdm: bool = True,

1066 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1067 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1068 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1069 ) -> list[ClassificationRequestOutput]:

@@ -1066,7 +1078,10 @@ class LLM:

1078 prompts: The prompts to the LLM. You may pass a sequence of prompts

1079 for batch inference. See [PromptType][vllm.inputs.PromptType]

1080 for more details about the format of each prompts.

1069 - use_tqdm: Whether to use tqdm to display the progress bar.

1081 + use_tqdm: If `True`, shows a tqdm progress bar.

1082 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),

1083 + it is used to create the progress bar.

1084 + If `False`, no progress bar is created.

1085 lora_request: LoRA request to use for generation, if any.

1086 prompt_adapter_request: Prompt Adapter request to use for

1087 generation, if any.

@@ -1092,7 +1107,7 @@ class LLM:

1107 text_1: list[Union[str, TextPrompt, TokensPrompt]],

1108 text_2: list[Union[str, TextPrompt, TokensPrompt]],

1109 truncate_prompt_tokens: Optional[int] = None,

1095 - use_tqdm: bool = True,

1110 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1111 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1112 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1113 ) -> list[ScoringRequestOutput]:

@@ -1126,7 +1141,7 @@ class LLM:

1141 text_1: list[str],

1142 text_2: list[str],

1143 truncate_prompt_tokens: Optional[int] = None,

1129 - use_tqdm: bool = True,

1144 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1145 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1146 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1147 ) -> list[ScoringRequestOutput]:

@@ -1178,7 +1193,7 @@ class LLM:

1193 /,

1194 *,

1195 truncate_prompt_tokens: Optional[int] = None,

1181 - use_tqdm: bool = True,

1196 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1197 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1198 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1199 ) -> list[ScoringRequestOutput]:

@@ -1198,7 +1213,10 @@ class LLM:

1213 text_2: The texts to pair with the query to form the input

1214 to the LLM. See [PromptType][vllm.inputs.PromptType] for

1215 more details about the format of each prompts.

1201 - use_tqdm: Whether to use tqdm to display the progress bar.

1216 + use_tqdm: If `True`, shows a tqdm progress bar.

1217 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),

1218 + it is used to create the progress bar.

1219 + If `False`, no progress bar is created.

1220 lora_request: LoRA request to use for generation, if any.

1221 prompt_adapter_request: Prompt Adapter request to use for

1222 generation, if any.

@@ -1379,7 +1397,7 @@ class LLM:

1397 params: Union[SamplingParams, Sequence[SamplingParams], PoolingParams,

1398 Sequence[PoolingParams]],

1399 *,

1382 - use_tqdm: bool,

1400 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1401 lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],

1402 prompt_adapter_request: Optional[PromptAdapterRequest],

1403 tokenization_kwargs: Optional[dict[str, Any]] = None,

@@ -1417,7 +1435,8 @@ class LLM:

1435 # Add requests to the engine.

1436 it = prompts

1437 if use_tqdm:

1420 - it = tqdm(it, desc="Adding requests")

1438 + tqdm_func = use_tqdm if callable(use_tqdm) else tqdm

1439 + it = tqdm_func(it, desc="Adding requests")

1440

1441 for i, prompt in enumerate(it):

1442 self._add_request(

@@ -1474,12 +1493,15 @@ class LLM:

1493 return params

1494

1495 def _run_engine(

1477 - self, *, use_tqdm: bool

1496 + self,

1497 + *,

1498 + use_tqdm: Union[bool, Callable[..., tqdm]] = True

1499 ) -> list[Union[RequestOutput, PoolingRequestOutput]]:

1500 # Initialize tqdm.

1501 if use_tqdm:

1502 num_requests = self.llm_engine.get_num_unfinished_requests()

1482 - pbar = tqdm(

1503 + tqdm_func = use_tqdm if callable(use_tqdm) else tqdm

1504 + pbar = tqdm_func(

1505 total=num_requests,

1506 desc="Processed prompts",

1507 dynamic_ncols=True,

[Frontend] Add tqdm_leave_pbar to control progress bar visibility (#19357)