vLLM/vllm/entrypoints/scale_out/derender/serving.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import time

import vllm.envs as envs
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionResponse,
    ChatCompletionStreamResponse,
)
from vllm.entrypoints.openai.completion.protocol import (
    CompletionResponse,
    CompletionStreamResponse,
)
from vllm.entrypoints.openai.models.serving import (
    OpenAIModelRegistry,
    OpenAIServingModels,
)
from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
from vllm.entrypoints.serve.engine.serving import BaseServing
from vllm.entrypoints.serve.utils.request_logger import RequestLogger
from vllm.inputs import EngineInput
from vllm.logger import init_logger
from vllm.renderers.online_derenderer import OnlineDerenderer

from ..token_in_token_out.mm_features import extract_mm_features
from ..token_in_token_out.protocol import (
    DerenderChatRequest,
    DerenderChatStreamRequest,
    DerenderCompletionRequest,
    DerenderCompletionStreamRequest,
    DerenderStreamState,
    GenerateResponse,
    MultiModalFeatures,
)

logger = init_logger(__name__)


class ServingDerender(BaseServing):
    def __init__(
        self,
        models: OpenAIServingModels | OpenAIModelRegistry,
        online_derenderer: "OnlineDerenderer",
        *,
        request_logger: RequestLogger | None = None,
    ) -> None:
        super().__init__(
            models=models,
            model_config=models.model_config,
            request_logger=request_logger,
        )

        self.online_derenderer = online_derenderer

    def _validate_derender_bounds(
        self,
        generate_responses: list[GenerateResponse],
    ) -> ErrorResponse | None:
        """Reject derender payloads that exceed resource bounds.

        Runs before any tokenizer.decode() or parser invocation to prevent
        CPU/memory exhaustion from oversized caller-supplied token structures.
        """
        max_n = envs.VLLM_MAX_N_SEQUENCES
        max_model_len = self.model_config.max_model_len
        # See ModelConfig.max_logprobs for semantics and default value.
        max_logprobs = self.model_config.max_logprobs

        if len(generate_responses) > max_n:
            return self.create_error_response(
                f"generate_responses count ({len(generate_responses)}) "
                f"exceeds server maximum ({max_n}). "
                f"Set VLLM_MAX_N_SEQUENCES to increase this limit."
            )

        for gen in generate_responses:
            if len(gen.choices) > max_n:
                return self.create_error_response(
                    f"choices count ({len(gen.choices)}) in response "
                    f"'{gen.request_id}' exceeds server maximum ({max_n})."
                )

            for choice in gen.choices:
                if choice.token_ids and len(choice.token_ids) > max_model_len:
                    return self.create_error_response(
                        f"token_ids length ({len(choice.token_ids)}) in "
                        f"choice {choice.index} exceeds "
                        f"max_model_len ({max_model_len})."
                    )
                if choice.logprobs and choice.logprobs.content:
                    if len(choice.logprobs.content) > max_model_len:
                        return self.create_error_response(
                            f"logprobs.content length "
                            f"({len(choice.logprobs.content)}) in "
                            f"choice {choice.index} exceeds "
                            f"max_model_len ({max_model_len})."
                        )
                    for entry in choice.logprobs.content:
                        if (
                            max_logprobs >= 0
                            and entry.top_logprobs
                            and len(entry.top_logprobs) > max_logprobs
                        ):
                            return self.create_error_response(
                                f"top_logprobs count "
                                f"({len(entry.top_logprobs)}) in "
                                f"choice {choice.index} exceeds "
                                f"max_logprobs ({max_logprobs})."
                            )

            if gen.prompt_logprobs and len(gen.prompt_logprobs) > max_model_len:
                return self.create_error_response(
                    f"prompt_logprobs length ({len(gen.prompt_logprobs)}) "
                    f"in response '{gen.request_id}' exceeds "
                    f"max_model_len ({max_model_len})."
                )

        return None

    async def derender_chat_response(
        self,
        request: DerenderChatRequest,
    ) -> ChatCompletionResponse | ErrorResponse:
        """Postprocess a GenerateResponse into a ChatCompletionResponse.

        Non-streaming only: expects the complete GenerateResponse with all
        token IDs present.  Uses ``parser.parse()`` for one-shot extraction.

        When ``request.chat_request`` is provided, the parser splits the
        output into (reasoning, content, tool_calls).  Otherwise falls
        back to plain detokenization.
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

        bounds_error = self._validate_derender_bounds([request.generate_response])
        if bounds_error is not None:
            return bounds_error

        if self.online_derenderer.parser is not None and request.chat_request is None:
            return self.create_error_response(
                "chat_request is required when a tool or reasoning parser is "
                "configured because plain detokenization would leak raw parser "
                "markup into content."
            )

        try:
            choices = await self.online_derenderer.derender_chat(
                request.generate_response, request.chat_request
            )
        except ValueError as exc:
            return self.create_error_response(str(exc))

        prompt_tokens = (
            request.prompt_tokens if request.prompt_tokens is not None else 0
        )
        gen = request.generate_response
        completion_tokens = sum(len(ch.token_ids) for ch in gen.choices if ch.token_ids)
        usage = UsageInfo(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )

        model_name = request.model or self.models.model_name()
        logger.debug(
            "derender_chat request_id=%s model=%s choices=%d completion_tokens=%d",
            gen.request_id,
            model_name,
            len(choices),
            completion_tokens,
        )
        return ChatCompletionResponse(
            id=gen.request_id,
            model=model_name,
            created=int(time.time()),
            choices=choices,
            usage=usage,
            prompt_logprobs=gen.prompt_logprobs,
            kv_transfer_params=gen.kv_transfer_params,
            metrics=gen.metrics,
        )

    async def derender_completion_response(
        self,
        request: DerenderCompletionRequest,
    ) -> CompletionResponse | ErrorResponse:
        """Postprocess a list of GenerateResponses into a CompletionResponse.

        Non-streaming only.  Mirrors the multi-prompt completions case: one
        GenerateResponse per prompt, parallel to the list[GenerateRequest]
        from /v1/completions/render.
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

        if not request.generate_responses:
            return self.create_error_response("generate_responses must not be empty")

        bounds_error = self._validate_derender_bounds(request.generate_responses)
        if bounds_error is not None:
            return bounds_error

        (
            choices,
            total_prompt_tokens,
            total_completion_tokens,
        ) = await self.online_derenderer.derender_completion(
            request.generate_responses,
            request.prompt_tokens,
            completion_request=request.completion_request,
        )

        first = request.generate_responses[0]
        kv_params = first.kv_transfer_params
        if any(
            r.kv_transfer_params != kv_params for r in request.generate_responses[1:]
        ):
            logger.warning(
                "derender_completion: kv_transfer_params differ across responses; "
                "setting to None on the aggregated response"
            )
            kv_params = None

        usage = UsageInfo(
            prompt_tokens=total_prompt_tokens,
            completion_tokens=total_completion_tokens,
            total_tokens=total_prompt_tokens + total_completion_tokens,
        )

        model_name = request.model or self.models.model_name()
        logger.debug(
            "derender_completion request_id=%s model=%s choices=%d"
            " completion_tokens=%d",
            first.request_id,
            model_name,
            len(choices),
            total_completion_tokens,
        )
        return CompletionResponse(
            id=first.request_id,
            model=model_name,
            created=int(time.time()),
            choices=choices,
            usage=usage,
            kv_transfer_params=kv_params,
            # Metrics describe one prompt. Parallel samples are already
            # suppressed by /generate; multi-prompt responses cannot be merged.
            metrics=first.metrics if len(request.generate_responses) == 1 else None,
        )

    async def derender_chat_stream_response(
        self,
        request: DerenderChatStreamRequest,
    ) -> tuple[ChatCompletionStreamResponse, DerenderStreamState] | ErrorResponse:
        """Streaming counterpart to ``derender_chat_response``.

        Processes one ``GenerateStreamResponse`` chunk and returns the
        derendered chunk together with the updated client carried state.
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

        if self.online_derenderer.parser is not None and request.chat_request is None:
            return self.create_error_response(
                "chat_request is required when a tool or reasoning parser is "
                "configured because plain detokenization would leak raw parser "
                "markup into content."
            )

        if (
            self.online_derenderer.parser is not None
            and request.prompt_token_ids is None
        ):
            return self.create_error_response(
                "prompt_token_ids is required when a tool or reasoning "
                "parser is configured. Without it parse_delta cannot tell "
                "whether the prompt left reasoning open (e.g. a chat "
                "template that pre-opens <think>) and would misclassify "
                "reasoning content as plain content."
            )

        # Each streamed chunk contains at most one choice per SSE event.
        # A single DerenderStreamState is threaded
        # through every choice in the chunk, so >1 would corrupt detok/parser
        # state across choices. See the matching check in
        # OnlineDerenderer.derender_chat_stream which this fails ahead of to
        # avoid touching the tokenizer at all on a malformed chunk.
        num_choices = len(request.generate_chunk.choices)
        if num_choices > 1:
            return self.create_error_response(
                f"derender_chat_stream expects at most one choice per chunk "
                f"(got {num_choices})."
            )

        # prompt_token_ids is folded in here too. It is caller supplied and
        # otherwise unbounded. A parser configured deployment rescans it
        # in full on every chunk (is_reasoning_end /
        # adjust_initial_state_from_prompt, see abstract_parser.parse_delta)
        # since parser state cannot be carried across calls.
        prompt_len = (
            len(request.prompt_token_ids) if request.prompt_token_ids is not None else 0
        )
        prior_len = (
            len(request.stream_state.output_token_ids)
            if request.stream_state is not None
            else 0
        )
        delta_len = sum(
            len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids
        )
        max_model_len = self.model_config.max_model_len
        if prompt_len + prior_len + delta_len > max_model_len:
            return self.create_error_response(
                f"prompt_token_ids length ({prompt_len}) plus "
                f"output_token_ids length ({prior_len}) plus delta "
                f"({delta_len}) exceeds max_model_len ({max_model_len})."
            )

        model_name = request.model or self.models.model_name()
        try:
            chunk, updated_state = await self.online_derenderer.derender_chat_stream(
                model=model_name,
                generate_chunk=request.generate_chunk,
                state=request.stream_state,
                chat_request=request.chat_request,
                prompt_tokens=request.prompt_tokens,
                prompt_token_ids=request.prompt_token_ids,
            )
        except Exception as exc:
            # The two ValueErrors derender_chat_stream can raise directly
            # (missing chat_request, >1 choice per chunk) are already
            # pre-checked above, so anything reaching here comes from the
            # parser replaying entirely client controlled state. Hermes
            # swallows its own streaming exceptions but finalize_generation
            # and engine based parsers don't, so an unexpected parser
            # failure on malformed (but well typed) state must surface as
            # a 400 here rather than an unhandled 500.
            logger.debug(
                "derender stream replay failed: %s: %s", type(exc).__name__, exc
            )
            return self.create_error_response("invalid stream_state: derender failed")

        logger.debug(
            "derender_chat_stream request_id=%s model=%s delta_tokens=%d",
            request.generate_chunk.request_id,
            model_name,
            sum(
                len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids
            ),
        )
        return chunk, updated_state

    async def derender_completion_stream_response(
        self,
        request: DerenderCompletionStreamRequest,
    ) -> tuple[CompletionStreamResponse, DerenderStreamState] | ErrorResponse:
        """Streaming counterpart to ``derender_completion_response``.

        Processes one ``GenerateStreamResponse`` chunk (one output sequence's
        delta) and returns the derendered chunk and updated state.
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

        model_name = request.model or self.models.model_name()
        try:
            (
                chunk,
                updated_state,
            ) = await self.online_derenderer.derender_completion_stream(
                model=model_name,
                generate_chunk=request.generate_chunk,
                state=request.stream_state,
                prompt_tokens=request.prompt_tokens,
                completion_request=request.completion_request,
            )
        except ValueError as exc:
            return self.create_error_response(str(exc))
        except (KeyError, IndexError) as exc:
            return self.create_error_response(
                f"invalid stream_state: detokenization failed ({exc})"
            )

        logger.debug(
            "derender_completion_stream request_id=%s model=%s delta_tokens=%d",
            request.generate_chunk.request_id,
            model_name,
            sum(
                len(c.token_ids) for c in request.generate_chunk.choices if c.token_ids
            ),
        )
        return chunk, updated_state

    @staticmethod
    def _extract_mm_features(
        engine_input: EngineInput,
    ) -> MultiModalFeatures | None:
        return extract_mm_features(engine_input)