207 lines
7.8 KiB
Python
207 lines
7.8 KiB
Python
from typing import Any, Optional, Sequence, Union
|
|
|
|
import cloudpickle
|
|
from pydantic import ValidationError
|
|
|
|
from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
|
|
from vllm.engine.arg_utils import HfOverrides
|
|
from vllm.entrypoints.llm import LLM
|
|
from vllm.inputs import PromptType
|
|
from vllm.logger import init_logger
|
|
from vllm.plugins.io_processors import get_io_processor
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.usage.usage_lib import UsageContext
|
|
from vllm.utils import Counter
|
|
from vllm.v1.engine.llm_engine import LLMEngine
|
|
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
|
from vllm_omni.engine.output_processor import MultimodalOutputProcessor
|
|
from vllm_omni.engine.processor import OmniProcessor
|
|
from vllm_omni.entrypoints.omni_stage import OmniStage
|
|
from vllm_omni.entrypoints.utils import load_stage_configs_from_model
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OmniLLM:
|
|
def __init__(
|
|
self, model: str, stage_configs=None, log_stats: bool = False, **kwargs
|
|
):
|
|
if stage_configs is None:
|
|
self.initialize_stage_configs(model)
|
|
else:
|
|
self.stage_configs = stage_configs
|
|
|
|
self.stage_list = []
|
|
self.initialize_stages(model)
|
|
|
|
def initialize_stage_configs(self, model: str):
|
|
self.stage_configs = load_stage_configs_from_model(model)
|
|
|
|
def initialize_stages(self, model: str):
|
|
for stage_config in self.stage_configs:
|
|
stage = OmniStage(stage_config)
|
|
stage_llm = OmniStageLLM(model=model, **stage_config.engine_args)
|
|
stage.set_engine(stage_llm)
|
|
self.stage_list.append(stage)
|
|
|
|
def generate(
|
|
self,
|
|
prompts: Union[PromptType, Sequence[PromptType]],
|
|
sampling_params_list: Optional[
|
|
Union[SamplingParams, Sequence[SamplingParams]]
|
|
] = None,
|
|
) -> list[OmniRequestOutput]:
|
|
"""Generate text outputs for the given prompts."""
|
|
final_outputs: list[OmniRequestOutput] = []
|
|
if len(sampling_params_list) != len(self.stage_list):
|
|
raise ValueError(
|
|
f"Expected {len(self.stage_list)} sampling params, "
|
|
f"got {len(sampling_params_list)}"
|
|
)
|
|
for stage_id, stage in enumerate(self.stage_list):
|
|
if stage_id > 0:
|
|
engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
|
|
else:
|
|
engine_inputs = prompts
|
|
engine_outputs = self._run_generation(
|
|
stage, sampling_params_list[stage_id], engine_inputs
|
|
)
|
|
stage.set_engine_outputs(engine_outputs)
|
|
if hasattr(stage, "final_output") and stage.final_output:
|
|
final_outputs.append(
|
|
OmniRequestOutput(
|
|
stage_id=stage_id,
|
|
final_output_type=stage.final_output_type,
|
|
request_output=engine_outputs,
|
|
)
|
|
)
|
|
return final_outputs
|
|
|
|
def _run_generation(
|
|
self,
|
|
stage: OmniStage,
|
|
sampling_params: SamplingParams,
|
|
prompts: Union[PromptType, Sequence[PromptType]],
|
|
):
|
|
engine_outputs = []
|
|
for ro in stage.engine.generate(prompts, sampling_params):
|
|
engine_outputs.append(ro)
|
|
return engine_outputs
|
|
|
|
|
|
class OmniStageLLM(LLM):
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
compilation_config: Optional[
|
|
Union[int, dict[str, Any], CompilationConfig]
|
|
] = None,
|
|
hf_overrides: Optional[HfOverrides] = None,
|
|
structured_outputs_config: Optional[
|
|
Union[dict[str, Any], StructuredOutputsConfig]
|
|
] = None,
|
|
**kwargs,
|
|
):
|
|
"""LLM constructor."""
|
|
if "disable_log_stats" not in kwargs:
|
|
kwargs["disable_log_stats"] = True
|
|
|
|
if "worker_cls" in kwargs:
|
|
worker_cls = kwargs["worker_cls"]
|
|
# if the worker_cls is not qualified string name,
|
|
# we serialize it using cloudpickle to avoid pickling issues
|
|
if isinstance(worker_cls, type):
|
|
kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
|
|
|
|
if "kv_transfer_config" in kwargs and isinstance(
|
|
kwargs["kv_transfer_config"], dict
|
|
):
|
|
from vllm.config.kv_transfer import KVTransferConfig
|
|
|
|
raw_config_dict = kwargs["kv_transfer_config"]
|
|
try:
|
|
kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
|
|
except ValidationError as e:
|
|
logger.error(
|
|
"Failed to convert 'kv_transfer_config' dict to "
|
|
"KVTransferConfig object. Dict: %s. Error: %s",
|
|
raw_config_dict,
|
|
e,
|
|
)
|
|
# Consider re-raising a more specific vLLM error or ValueError
|
|
# to provide better context to the user.
|
|
raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
|
|
|
|
if hf_overrides is None:
|
|
hf_overrides = {}
|
|
|
|
if compilation_config is not None:
|
|
if isinstance(compilation_config, int):
|
|
compilation_config_instance = CompilationConfig(
|
|
level=compilation_config
|
|
)
|
|
elif isinstance(compilation_config, dict):
|
|
compilation_config_instance = CompilationConfig(
|
|
**{
|
|
k: v
|
|
for k, v in compilation_config.items()
|
|
if is_init_field(CompilationConfig, k)
|
|
}
|
|
)
|
|
else:
|
|
compilation_config_instance = compilation_config
|
|
else:
|
|
compilation_config_instance = CompilationConfig()
|
|
|
|
if structured_outputs_config is not None:
|
|
if isinstance(structured_outputs_config, dict):
|
|
structured_outputs_instance = StructuredOutputsConfig(
|
|
**{
|
|
k: v
|
|
for k, v in structured_outputs_config.items()
|
|
if is_init_field(StructuredOutputsConfig, k)
|
|
}
|
|
)
|
|
else:
|
|
structured_outputs_instance = structured_outputs_config
|
|
else:
|
|
structured_outputs_instance = StructuredOutputsConfig()
|
|
|
|
engine_args = OmniEngineArgs(
|
|
model=model,
|
|
hf_overrides=hf_overrides,
|
|
compilation_config=compilation_config_instance,
|
|
structured_outputs_config=structured_outputs_instance,
|
|
**kwargs,
|
|
)
|
|
|
|
# Create the Engine (autoselects V0 vs V1)
|
|
self.llm_engine = LLMEngine.from_engine_args(
|
|
engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
|
|
)
|
|
self.llm_engine.output_processor = MultimodalOutputProcessor(
|
|
tokenizer=self.llm_engine.tokenizer,
|
|
log_stats=self.llm_engine.log_stats,
|
|
engine_core_output_type=engine_args.engine_output_type,
|
|
)
|
|
self.llm_engine.processor = OmniProcessor(
|
|
vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
|
|
)
|
|
self.engine_class = type(self.llm_engine)
|
|
|
|
self.request_counter = Counter()
|
|
self.default_sampling_params: Union[dict[str, Any], None] = None
|
|
|
|
supported_tasks = self.llm_engine.get_supported_tasks() # type: ignore
|
|
|
|
logger.info("Supported_tasks: %s", supported_tasks)
|
|
|
|
self.supported_tasks = supported_tasks
|
|
|
|
# Load the Input/Output processor plugin if any
|
|
io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
|
|
self.io_processor = get_io_processor(
|
|
self.llm_engine.vllm_config, io_processor_plugin
|
|
)
|