176 lines
7.6 KiB
Python
176 lines
7.6 KiB
Python
|
|
|||
|
|
from collections.abc import Mapping
|
|||
|
|
from copy import copy
|
|||
|
|
from typing import Any, Callable, Optional, Union
|
|||
|
|
|
|||
|
|
from typing_extensions import TypeVar
|
|||
|
|
|
|||
|
|
import vllm.envs as envs
|
|||
|
|
from vllm.config import ParallelConfig, VllmConfig
|
|||
|
|
from vllm.distributed import stateless_destroy_torch_distributed_process_group
|
|||
|
|
from vllm.engine.arg_utils import EngineArgs
|
|||
|
|
from vllm.inputs import PromptType
|
|||
|
|
from vllm.logger import init_logger
|
|||
|
|
from vllm.lora.request import LoRARequest
|
|||
|
|
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
|
|||
|
|
from vllm.outputs import PoolingRequestOutput, RequestOutput
|
|||
|
|
from vllm.pooling_params import PoolingParams
|
|||
|
|
# from vllm.prompt_adapter.request import PromptAdapterRequest
|
|||
|
|
from vllm.sampling_params import SamplingParams
|
|||
|
|
from vllm.usage.usage_lib import UsageContext
|
|||
|
|
from vllm.utils import Device
|
|||
|
|
from vllm.v1.engine.core_client import EngineCoreClient
|
|||
|
|
from vllm.v1.engine.output_processor import OutputProcessor
|
|||
|
|
from vllm.v1.engine.parallel_sampling import ParentRequest
|
|||
|
|
from vllm.v1.engine.processor import Processor
|
|||
|
|
from vllm.v1.executor.abstract import Executor
|
|||
|
|
from vllm.v1.metrics.loggers import (PrometheusStatLogger, StatLoggerBase,
|
|||
|
|
StatLoggerFactory)
|
|||
|
|
from vllm.v1.metrics.reader import Metric, get_metrics_snapshot
|
|||
|
|
from vllm.v1.metrics.stats import IterationStats
|
|||
|
|
from vllm.v1.engine import EngineCoreRequest
|
|||
|
|
|
|||
|
|
logger = init_logger(__name__)
|
|||
|
|
class LLMEngine:
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_vllm_config(
|
|||
|
|
cls,
|
|||
|
|
vllm_config: VllmConfig,
|
|||
|
|
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
|
|||
|
|
stat_loggers: Optional[list[StatLoggerFactory]] = None,
|
|||
|
|
disable_log_stats: bool = False,
|
|||
|
|
) -> "LLMEngine":
|
|||
|
|
# vacc support spec_num = 1
|
|||
|
|
from .vllm_config_checker import check_spec_model
|
|||
|
|
check_spec_model(vllm_config)
|
|||
|
|
|
|||
|
|
from vllm.v1.engine.llm_engine import LLMEngine as DefaultLLM
|
|||
|
|
default_cls = DefaultLLM
|
|||
|
|
return default_cls(vllm_config=vllm_config,
|
|||
|
|
executor_class=Executor.get_class(vllm_config),
|
|||
|
|
log_stats=(not disable_log_stats),
|
|||
|
|
usage_context=usage_context,
|
|||
|
|
stat_loggers=stat_loggers,
|
|||
|
|
multiprocess_mode=envs.VLLM_ENABLE_V1_MULTIPROCESSING)
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_engine_args(
|
|||
|
|
cls,
|
|||
|
|
engine_args: EngineArgs,
|
|||
|
|
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
|
|||
|
|
stat_loggers: Optional[list[StatLoggerFactory]] = None,
|
|||
|
|
enable_multiprocessing: bool = False,
|
|||
|
|
) -> "LLMEngine":
|
|||
|
|
"""Creates an LLM engine from the engine arguments."""
|
|||
|
|
# Create the engine configs.
|
|||
|
|
vllm_config = engine_args.create_engine_config(usage_context)
|
|||
|
|
executor_class = Executor.get_class(vllm_config)
|
|||
|
|
|
|||
|
|
# vacc support spec_num = 1
|
|||
|
|
from .vllm_config_checker import check_spec_model
|
|||
|
|
check_spec_model(vllm_config)
|
|||
|
|
|
|||
|
|
if envs.VLLM_ENABLE_V1_MULTIPROCESSING:
|
|||
|
|
logger.debug("Enabling multiprocessing for LLMEngine.")
|
|||
|
|
enable_multiprocessing = True
|
|||
|
|
|
|||
|
|
# Create the LLMEngine.
|
|||
|
|
from vllm.v1.engine.llm_engine import LLMEngine as DefaultLLM
|
|||
|
|
default_cls = DefaultLLM
|
|||
|
|
return default_cls(vllm_config=vllm_config,
|
|||
|
|
executor_class=executor_class,
|
|||
|
|
log_stats=not engine_args.disable_log_stats,
|
|||
|
|
usage_context=usage_context,
|
|||
|
|
stat_loggers=stat_loggers,
|
|||
|
|
multiprocess_mode=enable_multiprocessing)
|
|||
|
|
|
|||
|
|
"""Legacy LLMEngine for backwards compatibility."""
|
|||
|
|
def add_requests(
|
|||
|
|
self,
|
|||
|
|
items: list[tuple[
|
|||
|
|
str, # request_id
|
|||
|
|
PromptType, # prompt
|
|||
|
|
Union[SamplingParams, PoolingParams], # params
|
|||
|
|
Optional[float], # arrival_time
|
|||
|
|
Optional[LoRARequest], # lora_request
|
|||
|
|
Optional[dict], # tokenization_kwargs
|
|||
|
|
Optional[dict], # trace_headers
|
|||
|
|
# Optional[PromptAdapterRequest], # prompt_adapter_request
|
|||
|
|
int, # priority
|
|||
|
|
]],
|
|||
|
|
) -> None:
|
|||
|
|
"""批量把请求送入 EngineCore，一次性触发 ADD_BULK。"""
|
|||
|
|
core_reqs: list[EngineCoreRequest] = []
|
|||
|
|
|
|||
|
|
for (request_id, prompt, params, arrival_time, lora_request,
|
|||
|
|
tokenization_kwargs, trace_headers,
|
|||
|
|
priority) in items:
|
|||
|
|
|
|||
|
|
# 复用现有逐条流程的解析入口
|
|||
|
|
prompt_str, request = self.processor.process_inputs(
|
|||
|
|
request_id=request_id,
|
|||
|
|
prompt=prompt,
|
|||
|
|
params=params,
|
|||
|
|
arrival_time=arrival_time,
|
|||
|
|
lora_request=lora_request,
|
|||
|
|
tokenization_kwargs=tokenization_kwargs,
|
|||
|
|
trace_headers=trace_headers,
|
|||
|
|
# prompt_adapter_request=prompt_adapter_request,
|
|||
|
|
priority=priority,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
n = params.n if isinstance(params, SamplingParams) else 1
|
|||
|
|
|
|||
|
|
if n == 1:
|
|||
|
|
# Make a new RequestState and queue.
|
|||
|
|
self.output_processor.add_request(request, prompt_str, None, 0)
|
|||
|
|
# Add the request to EngineCore.
|
|||
|
|
core_reqs.append(request)
|
|||
|
|
continue
|
|||
|
|
# self.engine_core.add_request(request)
|
|||
|
|
# return
|
|||
|
|
|
|||
|
|
# Fan out child requests (for n>1).
|
|||
|
|
parent_req = ParentRequest(request_id, params)
|
|||
|
|
for idx in range(n):
|
|||
|
|
request_id, params = parent_req.get_child_info(idx)
|
|||
|
|
child_request = request if idx == n - 1 else copy(request)
|
|||
|
|
child_request.request_id = request_id
|
|||
|
|
child_request.sampling_params = params
|
|||
|
|
|
|||
|
|
# Make a new RequestState and queue.
|
|||
|
|
self.output_processor.add_request(child_request, prompt_str,
|
|||
|
|
parent_req, idx)
|
|||
|
|
# Add the request to EngineCore.
|
|||
|
|
# self.engine_core.add_request(child_request)
|
|||
|
|
# print("add_requests: child_request id=", child_request.request_id)
|
|||
|
|
core_reqs.append(child_request)
|
|||
|
|
|
|||
|
|
# output_processor 需要为每个“实际进入引擎的 req_id”建索引。
|
|||
|
|
# 如果是 SamplingParams 且 n>1/best_of>1，要做 parent-children 拆分；
|
|||
|
|
# 否则直接登记单条。
|
|||
|
|
# if isinstance(params, SamplingParams) and (
|
|||
|
|
# (params.n is not None and params.n > 1) or
|
|||
|
|
# (getattr(params, "best_of", 1) and getattr(params, "best_of", 1) > 1)
|
|||
|
|
# ):
|
|||
|
|
# parent = self.parallel_sampler.create_parent(request_id, params)
|
|||
|
|
# # 注意：最后一个 child 可以直接复用 request，其余用 copy
|
|||
|
|
# children = self.parallel_sampler.materialize_children(parent, request)
|
|||
|
|
# for child_idx, child in enumerate(children):
|
|||
|
|
# self.output_processor.add_request(
|
|||
|
|
# request=child,
|
|||
|
|
# prompt_str=prompt_str,
|
|||
|
|
# parent=parent,
|
|||
|
|
# child_index=child_idx,
|
|||
|
|
# )
|
|||
|
|
# core_reqs.append(child)
|
|||
|
|
# else:
|
|||
|
|
# self.output_processor.add_request(request, prompt_str)
|
|||
|
|
# core_reqs.append(request)
|
|||
|
|
|
|||
|
|
# 关键：一次性下发给 Core。EngineCoreClient 会发送 ADD_BULK。
|
|||
|
|
print('self.engine_core', self.engine_core)
|
|||
|
|
self.engine_core.add_requests(core_reqs)
|
|||
|
|
|
|||
|
|
|