class PoolingIOProcessor:
"""Processor for handling preprocessing & postprocessing ops for pooling requests.
This class manages both online (serving) and offline (batch) processing of pooling
requests, handling chat and completion formats.
"""
name: str
def __init__(
self,
vllm_config: VllmConfig,
renderer: BaseRenderer,
chat_template_config: ChatTemplateConfig,
):
self.vllm_config = vllm_config
self.model_config = vllm_config.model_config
self.renderer = renderer
self.chat_template = chat_template_config.chat_template
self.chat_template_content_format: Final = (
chat_template_config.chat_template_content_format
)
self.trust_request_chat_template = (
chat_template_config.trust_request_chat_template
)
#######################################
# online APIs
def create_pooling_params(self, request):
return request.to_pooling_params()
def pre_process_online(self, ctx: PoolingServeContext):
request = ctx.request
if isinstance(request, PoolingChatLikeRequest):
self._validate_chat_template(
request_chat_template=request.chat_template,
chat_template_kwargs=request.chat_template_kwargs,
trust_request_chat_template=self.trust_request_chat_template,
)
_, engine_inputs = self._preprocess_chat_online(
request,
request.messages,
default_template=self.chat_template,
default_template_content_format=self.chat_template_content_format,
default_template_kwargs=None,
)
elif isinstance(request, PoolingCompletionLikeRequest):
engine_inputs = self._preprocess_cmpl_online(
request,
prompt_input=request.input,
prompt_embeds=None,
)
else:
raise ValueError(f"Invalid {self.name} request type")
ctx.engine_inputs = engine_inputs
def post_process_online(
self,
ctx: PoolingServeContext,
):
pass
#######################################
# offline APIs
def get_request_factory_offline(
self, ctx: ALLOfflineInputsContext
) -> tuple[RequestFactory, int]:
assert isinstance(ctx, OfflineEncodeInputsContext)
prompts_seq = prompt_to_seq(ctx.prompts)
num_requests = len(prompts_seq)
pooling_task = ctx.pooling_task
parsed_prompts = [
(
prompt
if isinstance(prompt, bytes)
else parse_model_prompt(self.model_config, prompt)
)
for prompt in prompts_seq
]
tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
**(ctx.tokenization_kwargs or {})
)
pooling_params: PoolingParams | Sequence[PoolingParams]
if ctx.pooling_params is None:
pooling_params = PoolingParams()
else:
pooling_params = ctx.pooling_params
params_seq = self._params_to_seq(pooling_params, num_requests)
for param in params_seq:
if param.task is None:
param.task = pooling_task
elif pooling_task == "plugin":
# `plugin` task uses io_processor.parse_request to verify inputs.
# We actually allow plugin to overwrite pooling_task.
pass
elif param.task != pooling_task:
msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
raise ValueError(msg)
seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
def request_factory() -> RequestGenerator:
for i in range(num_requests):
yield EncodeCMPLRenderParams(
prompts=parsed_prompts[i],
tok_params=tok_params,
prompt_extras=None,
skip_mm_cache=False,
params=params_seq[i],
lora_requests=seq_lora_requests[i],
priorities=seq_priority[i],
)
return request_factory, num_requests
def post_process_offline(
self,
ctx: OfflineOutputsContext,
) -> list[PoolingRequestOutput]:
return ctx.outputs
#######################################
# helpers
def render(
self,
render_params: EncodeCMPLRenderParams
| EncodeChatRenderParams
| ScoringRenderParams,
) -> PoolingEngineInput:
if "conversations" in render_params:
render_params = cast(EncodeChatRenderParams, render_params)
(_,), engine_input = self.renderer.render_chat(
conversations=[render_params["conversations"]],
chat_params=render_params["chat_params"],
tok_params=render_params["tok_params"],
prompt_extras=render_params["prompt_extras"],
skip_mm_cache=render_params["skip_mm_cache"],
)
elif "prompts" in render_params:
render_params = cast(EncodeCMPLRenderParams, render_params)
engine_input = self.renderer.render_cmpl(
prompts=[render_params["prompts"]],
tok_params=render_params["tok_params"],
prompt_extras=render_params["prompt_extras"],
skip_mm_cache=render_params["skip_mm_cache"],
)
else:
raise ValueError(
f"Unsupported render_params type {render_params.__class__.__name__}"
)
return PoolingEngineInput(
prompts=engine_input[0],
params=render_params["params"],
lora_requests=render_params["lora_requests"],
priorities=render_params["priorities"],
)
def _preprocess_cmpl_online(
self,
request: RendererRequest,
prompt_input: str | list[str] | list[int] | list[list[int]] | None,
prompt_embeds: bytes | list[bytes] | None,
) -> list[EngineInput]:
renderer = self.renderer
model_config = self.model_config
prompts = list[SingletonPrompt | bytes]()
if prompt_embeds is not None: # embeds take higher priority
prompts.extend(prompt_to_seq(prompt_embeds))
if prompt_input is not None:
prompts.extend(prompt_to_seq(prompt_input))
parsed_prompts = [
(
prompt
if isinstance(prompt, bytes)
else parse_model_prompt(model_config, prompt)
)
for prompt in prompts
]
tok_params = request.build_tok_params(model_config)
return renderer.render_cmpl(
parsed_prompts,
tok_params,
prompt_extras={
k: v
for k in ("mm_processor_kwargs", "cache_salt")
if (v := getattr(request, k, None)) is not None
},
)
def _preprocess_chat_online(
self,
request: RendererChatRequest,
messages: list[ChatCompletionMessageParam],
default_template: str | None,
default_template_content_format: ChatTemplateContentFormatOption,
default_template_kwargs: dict[str, Any] | None,
tool_dicts: list[dict[str, Any]] | None = None,
tool_parser: type[ToolParser] | None = None,
) -> tuple[list[ConversationMessage], list[EngineInput]]:
renderer = self.renderer
default_template_kwargs = merge_kwargs(
default_template_kwargs,
dict(
tools=tool_dicts,
tokenize=is_mistral_tokenizer(renderer.tokenizer),
),
)
mm_config = self.model_config.multimodal_config
tok_params = request.build_tok_params(self.model_config)
chat_params = request.build_chat_params(
default_template, default_template_content_format
).with_defaults(
default_template_kwargs,
default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
)
(conversation,), (engine_input,) = renderer.render_chat(
[messages],
chat_params,
tok_params,
prompt_extras={
k: v
for k in ("mm_processor_kwargs", "cache_salt")
if (v := getattr(request, k, None)) is not None
},
)
return conversation, [engine_input]
def _validate_chat_template(
self,
request_chat_template: str | None,
chat_template_kwargs: dict[str, Any] | None,
trust_request_chat_template: bool,
):
if not trust_request_chat_template and (
request_chat_template is not None
or (
chat_template_kwargs
and chat_template_kwargs.get("chat_template") is not None
)
):
raise ValueError(
"Chat template is passed with request, but "
"--trust-request-chat-template is not set. "
"Refused request with untrusted chat template."
)
return None
def _params_to_seq(
self,
params: PoolingParams | Sequence[PoolingParams],
num_requests: int,
) -> Sequence[PoolingParams]:
if isinstance(params, Sequence):
if len(params) != num_requests:
raise ValueError(
f"The lengths of prompts ({num_requests}) "
f"and params ({len(params)}) must be the same."
)
return params
return [params] * num_requests
def _lora_request_to_seq(
self,
lora_request: LoRARequest | None | Sequence[LoRARequest | None],
num_requests: int,
) -> Sequence[LoRARequest | None]:
if isinstance(lora_request, Sequence):
if len(lora_request) != num_requests:
raise ValueError(
f"The lengths of prompts ({num_requests}) "
f"and lora_request ({len(lora_request)}) must be the same."
)
return lora_request
return [lora_request] * num_requests
def _priority_to_seq(
self,
priority: Sequence[int] | None,
num_requests: int,
) -> Sequence[int]:
if priority is not None:
if len(priority) != num_requests:
raise ValueError(
f"The lengths of prompts ({num_requests}) "
f"and priority ({len(priority)}) must be the same."
)
return priority
return [0] * num_requests