wehub-resource-sync
939 lines
42 KiB
Python
939 lines
42 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import asyncio
|
|
import inspect
|
|
import multiprocessing
|
|
import os
|
|
import time
|
|
import torch
|
|
from contextlib import contextmanager, nullcontext
|
|
from copy import copy, deepcopy
|
|
from packaging import version
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
from transformers import AutoConfig, GenerationConfig
|
|
from transformers.utils import is_torch_npu_available
|
|
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
|
|
|
|
from swift.metrics import Metric
|
|
from swift.model import get_processor
|
|
from swift.template import Template
|
|
from swift.utils import (disable_deepspeed_zero3, get_device, get_dist_setting, get_logger, is_dist,
|
|
safe_snapshot_download)
|
|
from .infer_engine import InferEngine
|
|
from .patch import patch_auto_tokenizer
|
|
from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
|
|
ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse,
|
|
EmbeddingResponseData, InferRequest, RequestConfig, random_uuid)
|
|
from .utils import AdapterRequest, InferStreamer, patch_npu_vllm, patch_vllm_memory_leak, patch_vllm_triton_device_guard
|
|
|
|
logger = get_logger()
|
|
try:
|
|
# After setting the environment variables, import vllm. This way of writing allows lint to pass.
|
|
os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
|
|
os.environ['VLLM_ENGINE_ITERATION_TIMEOUT_S'] = '86400'
|
|
import vllm
|
|
from vllm import AsyncEngineArgs, AsyncLLMEngine, EngineArgs, LLMEngine, SamplingParams
|
|
from vllm.pooling_params import PoolingParams
|
|
try:
|
|
# vLLM v0.12+ uses StructuredOutputsParams
|
|
from vllm.sampling_params import StructuredOutputsParams
|
|
except ImportError:
|
|
# Fallback for older vLLM versions
|
|
from vllm.sampling_params import GuidedDecodingParams as StructuredOutputsParams
|
|
except Exception:
|
|
raise
|
|
|
|
try:
|
|
from vllm.reasoning import ReasoningParserManager
|
|
except ImportError:
|
|
ReasoningParserManager = None
|
|
|
|
dtype_mapping = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}
|
|
|
|
|
|
def _patch_vllm_dp_coordinator_timeout():
|
|
# https://github.com/vllm-project/vllm/pull/37452 introduced a 30-second default timeout,
|
|
# which is prone to timing out in spawn scenarios. Patch it to 180 seconds here.
|
|
try:
|
|
from vllm.v1.engine import coordinator as coordinator_module
|
|
except ImportError:
|
|
return
|
|
|
|
coordinator_cls = coordinator_module.DPCoordinator
|
|
if not hasattr(coordinator_cls, '_wait_for_zmq_addrs'):
|
|
return
|
|
|
|
if getattr(coordinator_cls, '_swift_timeout_patched', False):
|
|
return
|
|
|
|
def _wait_for_zmq_addrs(self, zmq_addr_pipe):
|
|
t0 = time.monotonic()
|
|
try:
|
|
ready = multiprocessing.connection.wait([zmq_addr_pipe, self.proc.sentinel], timeout=180)
|
|
elapsed = time.monotonic() - t0
|
|
if not ready:
|
|
raise RuntimeError(f'DP Coordinator process failed to report ZMQ addresses '
|
|
f'within 180s (elapsed={elapsed:.1f}s).')
|
|
try:
|
|
return zmq_addr_pipe.recv()
|
|
except EOFError:
|
|
raise RuntimeError('DP Coordinator process failed during startup.') from None
|
|
finally:
|
|
zmq_addr_pipe.close()
|
|
|
|
coordinator_cls._wait_for_zmq_addrs = _wait_for_zmq_addrs
|
|
coordinator_cls._swift_timeout_patched = True
|
|
|
|
|
|
_patch_vllm_dp_coordinator_timeout()
|
|
|
|
|
|
@contextmanager
|
|
def _patch_rope_validation_ignore_keys():
|
|
"""Accept list-style RoPE validation ignore keys from older vLLM configs.
|
|
|
|
vLLM 0.18.x Qwen3.5 configs may pass ``ignore_keys_at_rope_validation``
|
|
as a list, while Transformers 5.x treats it as a set and performs a set
|
|
union during RoPE validation. vLLM release tags from 0.19.0 onward changed
|
|
the Qwen3.5 configs to set literals, but 0.18-based vLLM/vLLM-Ascend stacks
|
|
still need this compatibility layer. See vLLM PR:
|
|
https://github.com/vllm-project/vllm/pull/37338
|
|
"""
|
|
from transformers import PretrainedConfig
|
|
|
|
origin_convert = getattr(PretrainedConfig, 'convert_rope_params_to_dict', None)
|
|
if origin_convert is None:
|
|
yield
|
|
return
|
|
|
|
def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation=None, **kwargs):
|
|
if isinstance(ignore_keys_at_rope_validation, list):
|
|
ignore_keys_at_rope_validation = set(ignore_keys_at_rope_validation)
|
|
return origin_convert(self, ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs)
|
|
|
|
PretrainedConfig.convert_rope_params_to_dict = convert_rope_params_to_dict
|
|
try:
|
|
yield
|
|
finally:
|
|
PretrainedConfig.convert_rope_params_to_dict = origin_convert
|
|
|
|
|
|
class VllmEngine(InferEngine):
|
|
|
|
def __init__(
|
|
self,
|
|
model_id_or_path: str,
|
|
*,
|
|
template: Optional[Template] = None,
|
|
torch_dtype: Optional[torch.dtype] = None,
|
|
adapters: Optional[List[str]] = None,
|
|
use_async_engine: bool = False,
|
|
model_type: Optional[str] = None,
|
|
template_type: Optional[str] = None,
|
|
use_hf: Optional[bool] = None,
|
|
hub_token: Optional[str] = None,
|
|
revision: Optional[str] = None,
|
|
# engine_kwargs
|
|
gpu_memory_utilization: float = 0.9,
|
|
tensor_parallel_size: int = 1,
|
|
pipeline_parallel_size: int = 1,
|
|
enable_expert_parallel: bool = False,
|
|
max_model_len: Optional[int] = None,
|
|
max_num_seqs: int = 256,
|
|
disable_custom_all_reduce: bool = True,
|
|
enforce_eager: bool = False,
|
|
limit_mm_per_prompt: Optional[Dict[str, Any]] = None,
|
|
seed: Optional[int] = None,
|
|
task_type: Optional[str] = None, # embedding
|
|
disable_cascade_attn: bool = False,
|
|
load_format: str = 'auto',
|
|
mm_processor_cache_gb: Optional[float] = None,
|
|
logprobs_mode: Optional[str] = None,
|
|
speculative_config: Optional[Union[str, dict]] = None,
|
|
# lora
|
|
enable_lora: bool = False,
|
|
max_loras: int = 1,
|
|
max_lora_rank: int = 16,
|
|
enable_prefix_caching: Optional[bool] = None,
|
|
enable_sleep_mode: bool = False,
|
|
distributed_executor_backend: Optional[str] = None,
|
|
quantization: Optional[str] = None,
|
|
# reasoning parser
|
|
reasoning_parser: Optional[str] = None,
|
|
engine_kwargs: Optional[Dict[str, Any]] = None,
|
|
num_labels: Optional[int] = None,
|
|
reranker_use_activation: bool = True,
|
|
) -> None:
|
|
self.model_id_or_path = model_id_or_path
|
|
self.torch_dtype = torch_dtype
|
|
if isinstance(adapters, str):
|
|
adapters = [adapters]
|
|
self.default_adapter_request = None
|
|
if isinstance(adapters, list) and adapters:
|
|
assert len(adapters) == 1, 'Only one adapter is supported for now.'
|
|
enable_lora = True
|
|
self.default_adapter_request = AdapterRequest('default', adapters[0])
|
|
self.adapters = adapters or []
|
|
self.use_async_engine = use_async_engine
|
|
self.model_type = model_type
|
|
self.use_hf = use_hf
|
|
self.hub_token = hub_token
|
|
self.revision = revision
|
|
|
|
self.gpu_memory_utilization = gpu_memory_utilization
|
|
self.tensor_parallel_size = tensor_parallel_size
|
|
self.pipeline_parallel_size = pipeline_parallel_size
|
|
self.enable_expert_parallel = enable_expert_parallel
|
|
self.max_num_seqs = max_num_seqs
|
|
self.disable_custom_all_reduce = disable_custom_all_reduce
|
|
self.enforce_eager = enforce_eager
|
|
self.limit_mm_per_prompt = limit_mm_per_prompt
|
|
self.seed = seed
|
|
self.task_type = task_type
|
|
self.disable_cascade_attn = disable_cascade_attn
|
|
self.load_format = load_format
|
|
self.mm_processor_cache_gb = mm_processor_cache_gb
|
|
self.logprobs_mode = logprobs_mode
|
|
self.speculative_config = speculative_config
|
|
|
|
self.enable_lora = enable_lora
|
|
self.max_loras = max_loras
|
|
self.max_lora_rank = max_lora_rank
|
|
self.enable_prefix_caching = enable_prefix_caching
|
|
self.enable_sleep_mode = enable_sleep_mode
|
|
self.distributed_executor_backend = distributed_executor_backend
|
|
self.quantization = quantization
|
|
self.num_labels = num_labels
|
|
self.reranker_use_activation = reranker_use_activation
|
|
self._config_cls = None
|
|
|
|
patch_vllm_memory_leak()
|
|
patch_vllm_triton_device_guard()
|
|
self._adapters_pool = {}
|
|
if template is None:
|
|
processor = self._get_processor()
|
|
template = self._get_template(processor, template_type=template_type)
|
|
else:
|
|
safe_snapshot_download(
|
|
model_id_or_path,
|
|
revision=revision,
|
|
download_model=True,
|
|
use_hf=use_hf,
|
|
ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None),
|
|
hub_token=hub_token)
|
|
super().__init__(template)
|
|
if max_model_len is not None:
|
|
self.max_model_len = max_model_len
|
|
logger.info(f'Setting max_model_len: {max_model_len}')
|
|
self._prepare_engine_kwargs(max_model_len, engine_kwargs)
|
|
context = nullcontext()
|
|
if is_torch_npu_available() and (tensor_parallel_size == 1 or pipeline_parallel_size == 1):
|
|
colocate = (
|
|
getattr(self, '_swift_vllm_colocate_runtime', False)
|
|
or self.distributed_executor_backend == 'external_launcher')
|
|
context = patch_npu_vllm(get_device(), colocate=colocate)
|
|
with context:
|
|
self._prepare_engine()
|
|
self._load_generation_config()
|
|
self._fix_vllm_bug()
|
|
self.patch_remove_log()
|
|
self._request_count = 0
|
|
self._prepare_reasoning_parser(reasoning_parser)
|
|
|
|
def _get_processor(self):
|
|
return get_processor(
|
|
model_id_or_path=self.model_id_or_path,
|
|
torch_dtype=self.torch_dtype,
|
|
download_model=True,
|
|
model_type=self.model_type,
|
|
use_hf=self.use_hf,
|
|
hub_token=self.hub_token,
|
|
revision=self.revision,
|
|
num_labels=self.num_labels,
|
|
task_type=self.task_type)
|
|
|
|
def _prepare_engine(self) -> None:
|
|
with patch_auto_tokenizer(self.tokenizer), self._patch_auto_config(), \
|
|
_patch_rope_validation_ignore_keys(), disable_deepspeed_zero3():
|
|
llm_engine_cls = AsyncLLMEngine if self.use_async_engine else LLMEngine
|
|
engine = llm_engine_cls.from_engine_args(self.engine_args)
|
|
self.engine = engine
|
|
|
|
@contextmanager
|
|
def _patch_auto_config(self):
|
|
_old_from_pretrained = AutoConfig.from_pretrained
|
|
|
|
def _from_pretrained(*args, **kwargs):
|
|
config = deepcopy(self.config)
|
|
if self._version_ge('0.19'):
|
|
if self.model_type == 'deepseek_v4':
|
|
return _old_from_pretrained(*args, **kwargs)
|
|
if self._config_cls is None:
|
|
hf_config = _old_from_pretrained(*args, **kwargs)
|
|
self._config_cls = hf_config.__class__
|
|
if not isinstance(config, self._config_cls):
|
|
config.__class__ = self._config_cls
|
|
return config
|
|
|
|
AutoConfig.from_pretrained = _from_pretrained
|
|
try:
|
|
yield
|
|
finally:
|
|
AutoConfig.from_pretrained = _old_from_pretrained
|
|
|
|
def _prepare_engine_kwargs(self, max_model_len, engine_kwargs) -> None:
|
|
if engine_kwargs is None:
|
|
engine_kwargs = {}
|
|
if self.task_type == 'embedding':
|
|
self.task = 'embed'
|
|
elif self.task_type == 'seq_cls':
|
|
self.task = 'classify'
|
|
elif self.task_type in ('reranker', 'generative_reranker'):
|
|
self.task = 'score'
|
|
disable_log_stats = engine_kwargs.pop('disable_log_stats', True)
|
|
if self.use_async_engine:
|
|
engine_cls = AsyncEngineArgs
|
|
else:
|
|
engine_cls = EngineArgs
|
|
parameters = inspect.signature(engine_cls).parameters
|
|
if self.use_async_engine and 'disable_log_requests' in parameters:
|
|
engine_kwargs['disable_log_requests'] = True
|
|
if 'enable_lora' in parameters and self.enable_lora:
|
|
engine_kwargs['enable_lora'] = self.enable_lora
|
|
engine_kwargs['max_loras'] = self.max_loras
|
|
engine_kwargs['max_lora_rank'] = self.max_lora_rank
|
|
else:
|
|
assert not self.enable_lora, (
|
|
'The current version of vLLM does not support `enable_lora`. Please upgrade vLLM.')
|
|
|
|
if 'limit_mm_per_prompt' in parameters and self.limit_mm_per_prompt:
|
|
engine_kwargs['limit_mm_per_prompt'] = self.limit_mm_per_prompt
|
|
else:
|
|
assert not self.limit_mm_per_prompt, (
|
|
'The current version of vLLM does not support `limit_mm_per_prompt`. Please upgrade vLLM.')
|
|
for key in [
|
|
'enable_expert_parallel', 'enable_sleep_mode', 'disable_cascade_attn', 'load_format',
|
|
'mm_processor_cache_gb', 'speculative_config', 'logprobs_mode', 'quantization'
|
|
]:
|
|
if key in parameters:
|
|
value = getattr(self, key, None)
|
|
if value is not None:
|
|
engine_kwargs[key] = value
|
|
else:
|
|
logger.warning(f'The current version of vLLM does not support `{key}`. Ignored.')
|
|
for key in ['task', 'seed']:
|
|
val = getattr(self, key, None)
|
|
if val is not None:
|
|
engine_kwargs[key] = val
|
|
|
|
model_info = self.model_info
|
|
arch_mapping = {'deepseek_vl2': ['DeepseekVLV2ForCausalLM'], 'chatglm4v': ['GLM4VForCausalLM']}
|
|
if self.model_meta.model_type in arch_mapping:
|
|
architectures = arch_mapping[self.model_meta.model_type]
|
|
engine_kwargs['hf_overrides'] = {'architectures': architectures}
|
|
self.template.set_mode('vllm')
|
|
engine_kwargs.update(self.template.prepare_engine_kwargs())
|
|
if self.enable_prefix_caching is not None:
|
|
engine_kwargs['enable_prefix_caching'] = self.enable_prefix_caching
|
|
engine_args = engine_cls(
|
|
model=self.model_dir,
|
|
dtype=dtype_mapping[model_info.torch_dtype],
|
|
gpu_memory_utilization=self.gpu_memory_utilization,
|
|
tensor_parallel_size=self.tensor_parallel_size,
|
|
pipeline_parallel_size=self.pipeline_parallel_size,
|
|
max_model_len=max_model_len,
|
|
max_num_seqs=self.max_num_seqs,
|
|
disable_log_stats=disable_log_stats,
|
|
disable_custom_all_reduce=self.disable_custom_all_reduce,
|
|
enforce_eager=self.enforce_eager,
|
|
trust_remote_code=True,
|
|
distributed_executor_backend=self.distributed_executor_backend,
|
|
**engine_kwargs,
|
|
)
|
|
self.engine_args = engine_args
|
|
|
|
def _prepare_reasoning_parser(self, reasoning_parser: Optional[str]) -> None:
|
|
self.reasoning_parser = None
|
|
if not reasoning_parser:
|
|
return
|
|
|
|
# Validate reasoning_parser if provided
|
|
if ReasoningParserManager is None:
|
|
raise ImportError('the version of vLLM is too old, please upgrade vLLM')
|
|
|
|
valid_reasoning_parsers = list(ReasoningParserManager.reasoning_parsers.keys())
|
|
if reasoning_parser not in valid_reasoning_parsers:
|
|
raise ValueError(f'Invalid reasoning_parser: {reasoning_parser}. '
|
|
f'Available parsers: {valid_reasoning_parsers}')
|
|
logger.info(f'Using reasoning_parser: {reasoning_parser}')
|
|
|
|
reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(reasoning_parser)
|
|
self.reasoning_parser = reasoning_parser_cls(self.tokenizer)
|
|
|
|
def _fix_vllm_bug(self) -> None:
|
|
# fix vllm==0.4 bug (very slow)
|
|
tokenizer = self.tokenizer
|
|
if self._version_ge(
|
|
'0.4') and not self._version_ge('0.6') and not tokenizer.__class__.__name__.startswith('Cached'):
|
|
_tokenizer_len = len(tokenizer)
|
|
__old_len__ = tokenizer.__class__.__len__
|
|
|
|
def __len__(self) -> int:
|
|
if self is tokenizer:
|
|
return _tokenizer_len
|
|
else:
|
|
return __old_len__(self)
|
|
|
|
tokenizer.__class__.__len__ = __len__
|
|
|
|
def _load_generation_config(self) -> None:
|
|
generation_config_path = os.path.join(self.model_dir, 'generation_config.json')
|
|
if os.path.isfile(generation_config_path):
|
|
generation_config = GenerationConfig.from_pretrained(self.model_dir)
|
|
kwargs = generation_config.to_dict()
|
|
max_new_tokens = kwargs.get('max_new_tokens')
|
|
if max_new_tokens is not None:
|
|
kwargs['max_tokens'] = max_new_tokens
|
|
top_k = kwargs.get('top_k')
|
|
if top_k == 0:
|
|
kwargs['top_k'] = -1
|
|
parameters = inspect.signature(SamplingParams).parameters
|
|
for k, v in kwargs.copy().items():
|
|
if k not in parameters or v is None:
|
|
kwargs.pop(k)
|
|
self.generation_config = SamplingParams(**kwargs)
|
|
else:
|
|
self.generation_config = SamplingParams()
|
|
|
|
def _add_stop_words(self, generation_config: SamplingParams, request_config: RequestConfig) -> None:
|
|
template_meta = self.template.template_meta
|
|
stop_words = (request_config.stop or []) + (self.generation_config.stop or []) + template_meta.stop_words
|
|
generation_config.stop = self._get_stop_words(stop_words)
|
|
# stop parameter is not effective in v1 engine (test version: vllm 0.8.5.post)
|
|
generation_config.stop_token_ids = self._get_stop_token_ids(stop_words)
|
|
|
|
@staticmethod
|
|
def _version_ge(base_version: str):
|
|
vllm_version = vllm.__version__
|
|
if vllm_version is None or 'dev' in vllm_version:
|
|
return True
|
|
return version.parse(vllm_version) >= version.parse(base_version)
|
|
|
|
def _add_adapter(self, adapter_request: Optional[AdapterRequest] = None):
|
|
assert self.enable_lora, f'adapter_request: {adapter_request}, self.enable_lora: {self.enable_lora}'
|
|
from vllm.lora.request import LoRARequest
|
|
adapter_name = adapter_request.name
|
|
adapter_path = adapter_request.path
|
|
if adapter_name in self._adapters_pool:
|
|
lora_request = self._adapters_pool[adapter_name]
|
|
else:
|
|
lora_request = LoRARequest(
|
|
lora_name=adapter_name, lora_path=adapter_path, lora_int_id=len(self._adapters_pool) + 1)
|
|
self._adapters_pool[adapter_name] = lora_request
|
|
return lora_request
|
|
|
|
def _add_request(self,
|
|
inputs: Dict[str, Any],
|
|
generation_config: SamplingParams,
|
|
request_id: str,
|
|
adapter_request: Optional[AdapterRequest] = None):
|
|
kwargs = {}
|
|
adapter_request = adapter_request or self.default_adapter_request
|
|
if adapter_request:
|
|
kwargs['lora_request'] = self._add_adapter(adapter_request)
|
|
|
|
input_ids = inputs['input_ids']
|
|
if self._version_ge('0.4.3'):
|
|
llm_inputs = {'prompt_token_ids': input_ids}
|
|
mm_data = {}
|
|
for key in ['images', 'audios', 'videos']:
|
|
media_data = inputs.get(key) or []
|
|
if media_data:
|
|
if self._version_ge('0.6'):
|
|
|
|
mm_data[key.rstrip('s')] = media_data[0] if (
|
|
len(media_data) == 1 and
|
|
# compat qwen3_vl
|
|
not isinstance(media_data[0], tuple)) else media_data
|
|
else:
|
|
assert len(media_data) == 1, (
|
|
f'The current version of vllm only supports single {key}. Please upgrade to vllm >= 0.6.0')
|
|
mm_data[key.rstrip('s')] = media_data[0]
|
|
if mm_data:
|
|
llm_inputs['multi_modal_data'] = mm_data
|
|
mm_processor_kwargs = inputs.get('mm_processor_kwargs')
|
|
if mm_processor_kwargs:
|
|
llm_inputs['mm_processor_kwargs'] = mm_processor_kwargs
|
|
|
|
has_task_arg = 'task' in inspect.signature(PoolingParams).parameters
|
|
has_activation_arg = 'activation' in inspect.signature(PoolingParams).parameters
|
|
task_mapping = {
|
|
'embedding': 'embed',
|
|
'seq_cls': 'classify',
|
|
'reranker': 'score',
|
|
'generative_reranker': 'score',
|
|
}
|
|
if self.task_type in task_mapping:
|
|
pooling_kwargs = {}
|
|
if has_task_arg:
|
|
pooling_kwargs['task'] = task_mapping[self.task_type]
|
|
if self.task_type in ('reranker', 'generative_reranker') and \
|
|
has_activation_arg and self.reranker_use_activation:
|
|
pooling_kwargs['activation'] = True
|
|
pooling_params = PoolingParams(**pooling_kwargs)
|
|
return self.engine.encode(llm_inputs, pooling_params, request_id)
|
|
elif self.use_async_engine:
|
|
return self.engine.generate(llm_inputs, generation_config, request_id, **kwargs)
|
|
else:
|
|
return self.engine.add_request(request_id, llm_inputs, generation_config, **kwargs)
|
|
else:
|
|
if self.use_async_engine:
|
|
return self.engine.generate(None, generation_config, request_id, input_ids, **kwargs)
|
|
else:
|
|
return self.engine.add_request(request_id, None, generation_config, input_ids, **kwargs)
|
|
|
|
def _get_logprobs(self,
|
|
logprobs_list: Optional[List[Dict[int, float]]],
|
|
token_ids: List[int],
|
|
top_logprobs: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
if logprobs_list is None or len(token_ids) == 0:
|
|
return None
|
|
if len(token_ids) > 0:
|
|
logprobs_list = logprobs_list[-len(token_ids):]
|
|
for logprobs in logprobs_list:
|
|
for token_id, logprob in logprobs.items():
|
|
logprobs[token_id] = logprob.logprob
|
|
return super()._get_logprobs(logprobs_list, token_ids, top_logprobs)
|
|
|
|
def _prepare_generation_config(self, request_config: RequestConfig) -> SamplingParams:
|
|
kwargs = {'max_tokens': request_config.max_tokens}
|
|
for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']:
|
|
new_value = getattr(request_config, key)
|
|
if new_value is None:
|
|
kwargs[key] = getattr(self.generation_config, key)
|
|
else:
|
|
kwargs[key] = new_value
|
|
|
|
# Convert Swift's Chat Completions API style (logprobs: bool, top_logprobs: int)
|
|
# to vLLM's SamplingParams style (logprobs: int)
|
|
# vLLM semantics:
|
|
# - logprobs=None: no logprobs returned
|
|
# - logprobs=0: only sampled token's logprob
|
|
# - logprobs=N: top-N tokens + sampled token (up to N+1 total)
|
|
if request_config.logprobs:
|
|
# If logprobs=True, return log probabilities
|
|
if request_config.top_logprobs is not None and request_config.top_logprobs > 0:
|
|
# Return top_logprobs most likely tokens at each position
|
|
# (plus sampled token if not in top-N)
|
|
kwargs['logprobs'] = request_config.top_logprobs
|
|
else:
|
|
# Return only the sampled token's logprob
|
|
kwargs['logprobs'] = 0
|
|
|
|
if request_config.prompt_logprobs is not None:
|
|
kwargs['prompt_logprobs'] = request_config.prompt_logprobs
|
|
|
|
# TODO: beam search
|
|
for key in ['n', 'best_of', 'frequency_penalty', 'presence_penalty', 'seed']:
|
|
if hasattr(SamplingParams, key):
|
|
kwargs[key] = getattr(request_config, key)
|
|
|
|
# Handle structured outputs (guided decoding)
|
|
# vLLM v0.12+ uses 'structured_outputs' parameter, older versions use 'guided_decoding'
|
|
if request_config.structured_outputs_regex:
|
|
structured_outputs_param = StructuredOutputsParams(regex=request_config.structured_outputs_regex)
|
|
if hasattr(SamplingParams, 'structured_outputs'):
|
|
kwargs['structured_outputs'] = structured_outputs_param
|
|
else:
|
|
# Fallback for older vLLM versions
|
|
kwargs['guided_decoding'] = structured_outputs_param
|
|
|
|
res = SamplingParams(**kwargs)
|
|
|
|
if hasattr(res, 'output_kind') and res.n > 1:
|
|
# fix n > 1 in V1 Engine
|
|
from vllm.sampling_params import RequestOutputKind
|
|
res.output_kind = RequestOutputKind.FINAL_ONLY
|
|
return res
|
|
|
|
@property
|
|
def inner_model(self):
|
|
return self.engine.model_executor.driver_worker.worker.model_runner.model
|
|
|
|
@property
|
|
def inner_model_executor(self):
|
|
return self.engine.model_executor
|
|
|
|
async def _infer_stream_async(
|
|
self,
|
|
inputs: Dict[str, Any],
|
|
generation_config: SamplingParams,
|
|
adapter_request: Optional[AdapterRequest],
|
|
request_config: RequestConfig,
|
|
) -> AsyncIterator[ChatCompletionStreamResponse]:
|
|
request_id = random_uuid()
|
|
result_generator = self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
|
|
infer_streamers = [
|
|
InferStreamer(self.template, template_inputs=inputs['template_inputs']) for _ in range(generation_config.n)
|
|
]
|
|
token_idxs = [0 for _ in range(generation_config.n)]
|
|
async for result in result_generator:
|
|
res = self._create_chat_completion_stream_response(result, request_config, request_id, infer_streamers,
|
|
token_idxs)
|
|
if res is None:
|
|
continue
|
|
yield res
|
|
|
|
def _create_chat_completion_stream_response(self, result, request_config, request_id, infer_streamers,
|
|
token_idxs) -> Optional[ChatCompletionStreamResponse]:
|
|
is_diff = False
|
|
is_finished = False
|
|
for output in result.outputs:
|
|
output.token_ids = list(output.token_ids)
|
|
output.delta_text = infer_streamers[output.index].get_printable_text(output.token_ids, output.finished())
|
|
output.is_finished = output.finish_reason is not None
|
|
is_diff |= bool(output.delta_text)
|
|
is_finished |= output.is_finished
|
|
if not is_diff and not is_finished:
|
|
return
|
|
|
|
num_generated_tokens = sum(len(output.token_ids) for output in result.outputs)
|
|
usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens)
|
|
choices = []
|
|
previous_texts = [''] * len(result.outputs)
|
|
for output in result.outputs:
|
|
i = output.index
|
|
logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idxs[i]:],
|
|
request_config.top_logprobs)
|
|
|
|
# Handle reasoning content in streaming
|
|
delta_content = output.delta_text
|
|
delta_reasoning_content = None
|
|
|
|
if self.reasoning_parser and output.delta_text:
|
|
try:
|
|
# Get token IDs for the delta (new tokens in this step)
|
|
delta_token_ids = output.token_ids[token_idxs[i]:]
|
|
previous_token_ids = output.token_ids[:token_idxs[i]]
|
|
|
|
# Get current accumulated text for this output
|
|
previous_text = previous_texts[i]
|
|
current_text = previous_text + output.delta_text
|
|
previous_texts[i] = current_text
|
|
# Extract reasoning content from the delta
|
|
delta_message = self.reasoning_parser.extract_reasoning_content_streaming(
|
|
previous_text, current_text, output.delta_text, previous_token_ids, output.token_ids,
|
|
delta_token_ids)
|
|
|
|
if delta_message:
|
|
delta_reasoning_content = delta_message.reasoning_content
|
|
if delta_message.content:
|
|
delta_content = delta_message.content
|
|
else:
|
|
delta_content = None
|
|
|
|
except Exception as e:
|
|
logger.warning(f'Failed to extract reasoning content in streaming: {e}')
|
|
# Fallback to original delta_text
|
|
delta_content = output.delta_text
|
|
token_idxs[i] = len(output.token_ids)
|
|
|
|
toolcall = None
|
|
if output.is_finished:
|
|
toolcall = self._get_toolcall(
|
|
self.template.decode_generate_ids(output.token_ids, **infer_streamers[i].decode_kwargs))
|
|
|
|
choice = ChatCompletionResponseStreamChoice(
|
|
index=i,
|
|
delta=DeltaMessage(
|
|
role='assistant',
|
|
content=delta_content,
|
|
reasoning_content=delta_reasoning_content,
|
|
tool_calls=toolcall),
|
|
finish_reason=output.finish_reason,
|
|
logprobs=logprobs)
|
|
choices.append(choice)
|
|
return ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info, id=request_id)
|
|
|
|
@staticmethod
|
|
def _format_prompt_logprobs(prompt_logprobs):
|
|
if prompt_logprobs is None:
|
|
return None
|
|
result = []
|
|
for pos_lps in prompt_logprobs:
|
|
if pos_lps is None:
|
|
result.append(None)
|
|
else:
|
|
pos_dict = {}
|
|
for token_id, lp_obj in pos_lps.items():
|
|
pos_dict[str(token_id)] = {
|
|
'logprob': lp_obj.logprob,
|
|
'rank': getattr(lp_obj, 'rank', None),
|
|
'decoded_token': getattr(lp_obj, 'decoded_token', ''),
|
|
}
|
|
result.append(pos_dict)
|
|
return result
|
|
|
|
def _create_embedding_response(self, result, generation_config, request_id) -> EmbeddingResponse:
|
|
assert result is not None
|
|
embedding = result.outputs.data.cpu().numpy().tolist()
|
|
usage_info = self._get_usage_info(len(result.prompt_token_ids), 0)
|
|
return EmbeddingResponse(
|
|
model=self.model_name, data=[EmbeddingResponseData(embedding=embedding)], usage=usage_info, id=request_id)
|
|
|
|
def _create_chat_completion_response(
|
|
self,
|
|
result,
|
|
inputs,
|
|
request_config,
|
|
request_id,
|
|
) -> ChatCompletionResponse:
|
|
assert result is not None
|
|
num_generated_tokens = sum(len(output.token_ids) for output in result.outputs)
|
|
usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens)
|
|
choices = []
|
|
for output in result.outputs:
|
|
output.token_ids = list(output.token_ids)
|
|
response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])
|
|
|
|
# Extract reasoning content if reasoning_parser is enabled
|
|
reasoning_content = None
|
|
content = response
|
|
if self.reasoning_parser:
|
|
try:
|
|
reasoning_content, content = self.reasoning_parser.extract_reasoning_content(
|
|
response,
|
|
request=None # We don't have the original request here
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f'Failed to extract reasoning content: {e}')
|
|
# Fallback to original response
|
|
content = response
|
|
|
|
logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs)
|
|
toolcall = self._get_toolcall(content) # Use content instead of response for tool calls
|
|
token_ids = output.token_ids if request_config.return_details else None
|
|
choice = ChatCompletionResponseChoice(
|
|
index=output.index,
|
|
message=ChatMessage(
|
|
role='assistant', content=content, reasoning_content=reasoning_content, tool_calls=toolcall),
|
|
finish_reason=output.finish_reason,
|
|
logprobs=logprobs,
|
|
token_ids=token_ids)
|
|
choices.append(choice)
|
|
prompt_token_ids = None
|
|
images_size = None
|
|
if request_config.return_details:
|
|
prompt_token_ids = result.prompt_token_ids
|
|
images = inputs['template_inputs'].images
|
|
if all(isinstance(image, Image.Image) for image in images):
|
|
images_size = [image.size for image in images]
|
|
formatted_prompt_logprobs = None
|
|
if request_config.prompt_logprobs is not None:
|
|
formatted_prompt_logprobs = self._format_prompt_logprobs(result.prompt_logprobs)
|
|
return ChatCompletionResponse(
|
|
model=self.model_name,
|
|
choices=choices,
|
|
usage=usage_info,
|
|
id=request_id,
|
|
prompt_token_ids=prompt_token_ids,
|
|
prompt_logprobs=formatted_prompt_logprobs,
|
|
images_size=images_size)
|
|
|
|
def _create_seq_cls_response(
|
|
self,
|
|
result,
|
|
request_config,
|
|
request_id,
|
|
) -> ChatCompletionResponse:
|
|
assert result is not None
|
|
choices = []
|
|
preds = result.outputs.data
|
|
if preds.dim() == 1:
|
|
preds = preds.unsqueeze(0)
|
|
if self.task_type == 'seq_cls':
|
|
top_logprobs = request_config.top_logprobs or 20
|
|
preds, logprobs = self.template.decode_seq_cls(preds, top_logprobs)
|
|
else:
|
|
logprobs = [None] * len(preds)
|
|
num_prompt_token_ids = 0
|
|
num_generated_tokens = 0
|
|
for i, pred in enumerate(preds):
|
|
num_prompt_token_ids += len(result.prompt_token_ids)
|
|
num_generated_tokens += 1
|
|
if isinstance(pred, torch.Tensor):
|
|
pred = pred.tolist()
|
|
choices.append(
|
|
ChatCompletionResponseChoice(
|
|
index=0,
|
|
message=ChatMessage(role='assistant', content=pred, tool_calls=None),
|
|
finish_reason='stop',
|
|
logprobs=logprobs[i]))
|
|
usage_info = self._get_usage_info(num_prompt_token_ids, num_generated_tokens)
|
|
return ChatCompletionResponse(
|
|
model=self.model_name,
|
|
choices=choices,
|
|
usage=usage_info,
|
|
id=request_id,
|
|
prompt_token_ids=result.prompt_token_ids)
|
|
|
|
async def _infer_full_async(
|
|
self,
|
|
inputs: Dict[str, Any],
|
|
generation_config: SamplingParams,
|
|
adapter_request: Optional[AdapterRequest],
|
|
request_config: RequestConfig,
|
|
request_id: Optional[str] = None,
|
|
) -> Union[ChatCompletionResponse, EmbeddingResponse]:
|
|
if request_id is None:
|
|
request_id = random_uuid()
|
|
result_generator = self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
|
|
result = None
|
|
async for result in result_generator:
|
|
pass
|
|
if self.task_type == 'embedding':
|
|
return self._create_embedding_response(result, generation_config, request_id)
|
|
elif self.task_type in ('seq_cls', 'reranker', 'generative_reranker'):
|
|
return self._create_seq_cls_response(result, request_config, request_id)
|
|
else:
|
|
return self._create_chat_completion_response(result, inputs, request_config, request_id)
|
|
|
|
def _batch_infer_stream(self, *args, **kwargs):
|
|
if hasattr(self.engine, 'engine'):
|
|
self.engine.engine.model_executor.parallel_worker_tasks = None
|
|
return super()._batch_infer_stream(*args, **kwargs)
|
|
|
|
def infer(
|
|
self,
|
|
infer_requests: List[InferRequest],
|
|
request_config: Optional[RequestConfig] = None,
|
|
metrics: Optional[List[Metric]] = None,
|
|
*,
|
|
use_tqdm: Optional[bool] = None,
|
|
adapter_request: Optional[AdapterRequest] = None,
|
|
) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
|
|
if self.use_async_engine:
|
|
return super().infer(
|
|
infer_requests,
|
|
request_config,
|
|
metrics,
|
|
use_tqdm=use_tqdm,
|
|
adapter_request=adapter_request,
|
|
)
|
|
else:
|
|
request_config = deepcopy(request_config or RequestConfig())
|
|
if request_config.stream and len(infer_requests) > 1:
|
|
raise ValueError('If you want to use stream batch inference, you need to set use_async_engine to True.')
|
|
if use_tqdm is None:
|
|
use_tqdm = len(infer_requests) > 1
|
|
rank = get_dist_setting()[0]
|
|
if is_dist() and rank % self.engine_args.tensor_parallel_size != 0:
|
|
use_tqdm = False
|
|
self.template.set_mode('vllm')
|
|
batched_inputs, error_list = self._batch_encode(infer_requests, strict=getattr(self, 'strict', True))
|
|
request_id_list = []
|
|
for i, inputs in enumerate(batched_inputs):
|
|
request_id = str(self._request_count)
|
|
request_id_list.append(request_id)
|
|
self._request_count += 1
|
|
_request_config = deepcopy(request_config)
|
|
self.set_default_max_tokens(_request_config, inputs)
|
|
generation_config = self._prepare_generation_config(_request_config)
|
|
if generation_config.seed is not None:
|
|
generation_config.seed += i
|
|
self._add_stop_words(generation_config, _request_config)
|
|
self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
|
|
prog_bar = tqdm(total=len(batched_inputs), dynamic_ncols=True, disable=not use_tqdm)
|
|
outputs = {}
|
|
if request_config.stream:
|
|
|
|
def _gen_wrapper():
|
|
infer_streamers = [
|
|
InferStreamer(self.template, template_inputs=inputs['template_inputs'])
|
|
for _ in range(generation_config.n)
|
|
]
|
|
token_idxs = [0 for _ in range(generation_config.n)]
|
|
while self.engine.has_unfinished_requests():
|
|
result = self.engine.step()
|
|
if not result:
|
|
continue
|
|
result = result[0]
|
|
res = self._create_chat_completion_stream_response(result, request_config, request_id,
|
|
infer_streamers, token_idxs)
|
|
if res is None:
|
|
continue
|
|
yield res
|
|
if result.finished:
|
|
break
|
|
|
|
self._update_metrics(res, metrics)
|
|
|
|
return [_gen_wrapper()]
|
|
else:
|
|
while self.engine.has_unfinished_requests():
|
|
step_outputs = self.engine.step()
|
|
for output in step_outputs:
|
|
if output.finished:
|
|
outputs[output.request_id] = output
|
|
prog_bar.update()
|
|
prog_bar.close()
|
|
outputs = [outputs[request_id] for request_id in request_id_list]
|
|
res = [
|
|
self._create_chat_completion_response(result, inputs, request_config, request_id)
|
|
for request_id, inputs, result in zip(request_id_list, batched_inputs, outputs)
|
|
]
|
|
self._update_metrics(res, metrics)
|
|
return self._add_error_list(res, error_list)
|
|
|
|
async def infer_async(
|
|
self,
|
|
infer_request: InferRequest,
|
|
request_config: Optional[RequestConfig] = None,
|
|
*,
|
|
adapter_request: Optional[AdapterRequest] = None,
|
|
pre_infer_hook=None,
|
|
) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
|
|
if not self.use_async_engine:
|
|
raise ValueError('If you want to use `infer_async`, you need to pass `use_async_engine` as True.')
|
|
request_config = deepcopy(request_config or RequestConfig())
|
|
self.template.set_mode('vllm')
|
|
loop = asyncio.get_running_loop()
|
|
with torch.inference_mode():
|
|
inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)
|
|
self.set_default_max_tokens(request_config, inputs)
|
|
generation_config = self._prepare_generation_config(request_config)
|
|
self._add_stop_words(generation_config, request_config)
|
|
kwargs = {
|
|
'inputs': inputs,
|
|
'generation_config': generation_config,
|
|
'adapter_request': adapter_request,
|
|
'request_config': request_config,
|
|
}
|
|
if hasattr(infer_request, 'uuid') and infer_request.uuid:
|
|
# RolloutInferRequest
|
|
kwargs.update({'request_id': infer_request.uuid})
|
|
if pre_infer_hook:
|
|
kwargs = pre_infer_hook(kwargs)
|
|
if request_config.stream:
|
|
return self._infer_stream_async(**kwargs)
|
|
else:
|
|
return await self._infer_full_async(**kwargs)
|
|
|
|
@staticmethod
|
|
def patch_remove_log():
|
|
from vllm.engine import async_llm_engine
|
|
if not hasattr(async_llm_engine, '_log_task_completion'):
|
|
return
|
|
|
|
async_llm_engine._origin_log_task_completion = async_llm_engine._log_task_completion
|
|
|
|
def new_log_task_completion(task, error_callback) -> None:
|
|
try:
|
|
return_value = task.result()
|
|
raise AssertionError(f'The engine background task should never finish without an '
|
|
f'exception. {return_value}')
|
|
except asyncio.exceptions.CancelledError:
|
|
pass
|
|
|
|
async_llm_engine._log_task_completion = new_log_task_completion
|