103 lines
4.1 KiB
Python
103 lines
4.1 KiB
Python
|
||
|
||
import itertools
|
||
import warnings
|
||
from collections.abc import Sequence
|
||
from contextlib import contextmanager
|
||
from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union,
|
||
cast, overload)
|
||
|
||
import cloudpickle
|
||
import torch.nn as nn
|
||
from pydantic import ValidationError
|
||
from tqdm.auto import tqdm
|
||
from typing_extensions import TypeVar, deprecated
|
||
|
||
from vllm.entrypoints.utils import _validate_truncation_size
|
||
from vllm.inputs import PromptType
|
||
from vllm.logger import init_logger
|
||
from vllm.lora.request import LoRARequest
|
||
from vllm.pooling_params import PoolingParams
|
||
from vllm.sampling_params import (RequestOutputKind, SamplingParams)
|
||
from vllm.platforms import current_platform
|
||
|
||
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
_R = TypeVar("_R", default=Any)
|
||
|
||
class LLM:
|
||
|
||
EPRECATE_LEGACY: ClassVar[bool] = True
|
||
def _validate_and_add_requests(
|
||
self,
|
||
prompts: Union[PromptType, Sequence[PromptType]],
|
||
params: Union[SamplingParams, Sequence[SamplingParams], PoolingParams,
|
||
Sequence[PoolingParams]],
|
||
*,
|
||
use_tqdm: Union[bool, Callable[..., tqdm]] = True,
|
||
lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],
|
||
priority: Optional[list[int]] = None,
|
||
) -> None:
|
||
|
||
if isinstance(prompts, (str, dict)):
|
||
# Convert a single prompt to a list.
|
||
prompts = [prompts]
|
||
|
||
num_requests = len(prompts)
|
||
if isinstance(params, Sequence) and len(params) != num_requests:
|
||
raise ValueError("The lengths of prompts and params "
|
||
"must be the same.")
|
||
if isinstance(lora_request,
|
||
Sequence) and len(lora_request) != num_requests:
|
||
raise ValueError("The lengths of prompts and lora_request "
|
||
"must be the same.")
|
||
|
||
for sp in params if isinstance(params, Sequence) else (params, ):
|
||
if isinstance(sp, SamplingParams):
|
||
# We only care about the final output
|
||
sp.output_kind = RequestOutputKind.FINAL_ONLY
|
||
|
||
# Add requests to the engine.
|
||
it = prompts
|
||
if use_tqdm:
|
||
tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
|
||
it = tqdm_func(it, desc="Adding requests")
|
||
|
||
if (hasattr(current_platform, 'supports_v1') and current_platform.supports_v1(current_platform)):
|
||
batch_items = []
|
||
model_config = self.llm_engine.model_config
|
||
for i, prompt in enumerate(it):
|
||
request_id = str(next(self.request_counter))
|
||
# print("requset_id===========", request_id)
|
||
param = params[i] if isinstance(params, Sequence) else params
|
||
tokenization_kwargs: dict[str, Any] = {}
|
||
_validate_truncation_size(model_config.max_model_len,
|
||
param.truncate_prompt_tokens,
|
||
tokenization_kwargs)
|
||
|
||
batch_items.append((
|
||
request_id,
|
||
prompt,
|
||
params[i] if isinstance(params, Sequence) else params,
|
||
None, # arrival_time，不用的话传 None
|
||
(lora_request[i] if isinstance(lora_request, Sequence)
|
||
else lora_request),
|
||
tokenization_kwargs,
|
||
None, # trace_headers（如无 APM/Tracing，None）
|
||
(priority[i] if priority else 0),
|
||
))
|
||
# 一次性下发给 EngineCore（走 ADD_BULK）
|
||
self.llm_engine.add_requests(batch_items)
|
||
else:
|
||
for i, prompt in enumerate(it):
|
||
self._add_request(
|
||
prompt,
|
||
params[i] if isinstance(params, Sequence) else params,
|
||
tokenization_kwargs=tokenization_kwargs,
|
||
lora_request=lora_request[i] if isinstance(
|
||
lora_request, Sequence) else lora_request,
|
||
priority=priority[i] if priority else 0,
|
||
)
|