tcod/trinity/common/models/vllm_model.py

"""A wrapper around the vllm.AsyncEngine to handle async requests."""

import asyncio
import os
from collections import defaultdict
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch
from packaging.version import parse as parse_version
from PIL import Image
from transformers import AutoProcessor

from trinity.common.config import InferenceModelConfig
from trinity.common.experience import Experience
from trinity.common.models.mm_utils import (
    build_multi_modal_inputs,
    convert_messages_to_mm_format,
)
from trinity.common.models.model import BaseInferenceModel
from trinity.common.models.vllm_patch import get_vllm_version


# V0 engine is deprecated since vLLM v0.10.2, related code will be removed in the future.
class vLLMRolloutModel(BaseInferenceModel):
    """Wrapper around the vLLM engine to handle async requests.

    Args:
        config (Config): The config.
    """

    def __init__(
        self,
        config: InferenceModelConfig,
    ) -> None:
        super().__init__(config)

        import vllm
        from vllm.sampling_params import RequestOutputKind

        self.vllm_version = get_vllm_version()
        self.use_v1 = config.use_v1
        if config.tensor_parallel_size != 1:
            os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
            os.environ["VLLM_RAY_BUNDLE_INDICES"] = config.bundle_indices
        if self.vllm_version <= parse_version("0.11.0") and not vllm.envs.is_set("VLLM_USE_V1"):
            self.logger.info(f"Using vLLM v{int(config.use_v1)} engine")
            os.environ["VLLM_USE_V1"] = str(int(config.use_v1))
        if config.use_v1:
            os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
            os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(int(config.use_v1))
            os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
            os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
        if self.vllm_version >= parse_version("0.11.0"):
            os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
        if self.config.enable_runtime_lora_updating:
            os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "1"
        if not config.enforce_eager:
            # To avoid torch compile conflicts when multiple model are started simultaneously.
            # remove this when the following PR is released:
            # https://github.com/vllm-project/vllm/pull/27616
            os.environ["VLLM_CACHE_ROOT"] = os.path.expanduser(
                f"~/.cache/vllm/{config.bundle_indices}"
            )
        self.default_sampling_params = vllm.SamplingParams(
            n=1,
            temperature=config.temperature,
            max_tokens=config.max_response_tokens,
            min_tokens=config.min_response_tokens,
            truncate_prompt_tokens=(
                config.max_prompt_tokens if config.enable_prompt_truncation else None
            ),
            skip_special_tokens=True,
            include_stop_str_in_output=False,
            output_kind=RequestOutputKind.FINAL_ONLY,
            logprobs=config.logprobs,
            top_p=config.top_p,
            top_k=config.top_k,
            ignore_eos=config.ignore_eos,
        )
        self.ray_namespace = config.ray_namespace
        self.request_id = 0
        max_model_len = config.max_model_len
        self.enable_lora = config.enable_lora
        self.default_lora_path = config.lora_kwargs.pop("default_lora_path", None)
        if self.vllm_version >= parse_version("0.12.0"):
            rope_params = defaultdict(dict)
            if config.rope_scaling is not None:
                rope_params["rope_parameters"] = config.rope_scaling
            if config.rope_theta is not None:
                rope_params["rope_parameters"]["rope_theta"] = config.rope_theta
            if len(rope_params) > 0:
                rope_kwargs = {"hf_overrides": rope_params}
            else:
                rope_kwargs = {}
            self.logprobs_no_prefix_cache = True
        else:
            rope_kwargs = {
                key: getattr(config, key)
                for key in ["rope_scaling", "rope_theta"]
                if getattr(config, key) is not None
            }
            self.logprobs_no_prefix_cache = False
        engine_args = vllm.AsyncEngineArgs(
            model=config.model_path,
            enforce_eager=config.enforce_eager,
            worker_extension_cls="trinity.common.models.vllm_worker.WorkerExtension",
            tensor_parallel_size=config.tensor_parallel_size,
            seed=config.seed,
            distributed_executor_backend=("uni" if config.tensor_parallel_size == 1 else "ray"),
            max_model_len=max_model_len,
            enable_prefix_caching=config.enable_prefix_caching,
            enable_chunked_prefill=config.enable_chunked_prefill,
            dtype=config.dtype,
            trust_remote_code=True,
            gpu_memory_utilization=config.gpu_memory_utilization,
            override_generation_config={  # TODO: find a way to unittest this
                "temperature": config.temperature,
                "top_p": config.top_p,
                "top_k": config.top_k,
                "max_new_tokens": config.max_response_tokens,
                "repetition_penalty": config.repetition_penalty,
            },
            disable_log_stats=True,
            enable_lora=config.enable_lora,
            logprobs_mode="processed_logprobs",
            **rope_kwargs,
            **config.lora_kwargs,
        )
        if self.vllm_version > parse_version("0.10.0"):
            engine_args.enable_log_requests = config.enable_log_requests
        else:
            engine_args.disable_log_requests = not config.enable_log_requests
        if self.vllm_version >= parse_version("0.11.0"):
            engine_args.reasoning_parser = config.reasoning_parser
        if self.vllm_version >= parse_version("0.13.0"):
            engine_args.async_scheduling = False
        self.async_llm = vllm.AsyncLLMEngine.from_engine_args(engine_args)
        self.processor = None
        self.state_dict_meta = None
        self.model_version = 0  # TODO: resume the value from the checkpoint
        self.api_server_host = None
        self.api_server_port = None
        self.api_server = None
        self._prepared = False
        self.async_lock = asyncio.Lock()

    async def _initialize_tokenizer(self):
        if self.tokenizer is None:
            self.tokenizer = await self.async_llm.get_tokenizer()
        self.tokenizer.truncation_side = "left"

    def _initialize_processor(self):
        self.processor = AutoProcessor.from_pretrained(
            self.config.model_path, trust_remote_code=True
        )
        self.tokenizer = self.processor.tokenizer

    async def prepare(
        self,
    ) -> None:
        """Prepare the model for inference."""
        async with self.async_lock:
            if self._prepared:
                return
            await self._collective_rpc("apply_patches")
            await self.run_api_server()
            self._prepared = True

    async def chat(self, messages: List[Dict], lora_request=None, **kwargs) -> Sequence[Experience]:
        """Chat with the model with a list of messages in async.

        Args:
            messages (List[dict]): The input history messages.
            kwargs (dict): A dictionary of sampling parameters.

        Returns:
            A list of experiences.
        """
        if self.tokenizer is None:
            await self._initialize_tokenizer()

        prompt = self.apply_chat_template(self.tokenizer, messages)
        return await self.generate(prompt=prompt, lora_request=lora_request, **kwargs)

    async def generate(self, prompt: str, lora_request=None, **kwargs) -> Sequence[Experience]:
        """Generate a response from the provided prompt in async.

        Args:
            prompt (str): The input prompt.
            kwargs (dict): A dictionary of sampling parameters.

        Returns:
            A list of experiences.
        """
        if self.tokenizer is None:
            await self._initialize_tokenizer()

        token_ids, is_valid = self._handle_prompt_truncation(prompt, **kwargs)
        if not is_valid:
            return token_ids

        output = await self._generate_internal(
            prompt={"prompt_token_ids": token_ids}, lora_request=lora_request, **kwargs
        )
        experiences = [
            Experience(
                tokens=torch.cat(
                    (
                        torch.tensor(output.prompt_token_ids, dtype=torch.int32),
                        torch.tensor(output.outputs[i].token_ids, dtype=torch.int32),
                    )
                ),
                logprobs=torch.cat(
                    (
                        torch.tensor(
                            [
                                list(logprob_dict.values())[0].logprob
                                for logprob_dict in output.outputs[i].logprobs
                            ],
                            dtype=torch.float32,
                        ),
                    )
                ),
                prompt_length=len(output.prompt_token_ids),
                prompt_text=self.tokenizer.decode(output.prompt_token_ids),
                response_text=output.outputs[i].text,
            )
            for i in range(len(output.outputs))
        ]
        return experiences

    async def chat_mm(
        self, messages: List[Dict], images: List[Image.Image], videos: List[np.ndarray], **kwargs
    ) -> Sequence[Experience]:
        """Chat with the model with a list of messages in async.

        Args:
            messages (List[dict]): The input history messages.
            raw_mm_data (dict): The raw multi-modal data.
            kwargs (dict): A dictionary of sampling parameters.

        Returns:
            A list of experiences.
        """
        if self.processor is None:
            self._initialize_processor()
        messages = convert_messages_to_mm_format(messages)
        prompt = self.apply_chat_template(self.processor, messages)
        return await self.generate_mm(prompt=prompt, images=images, videos=videos, **kwargs)

    async def generate_mm(
        self,
        prompt: str = None,
        images: List[Image.Image] = None,
        videos: List[np.ndarray] = None,
        **kwargs,
    ) -> Sequence[Experience]:
        """Generate a response from the provided prompt in async.

        Args:
            prompt (str): The input prompt.
            images (List): The list of image inputs.
            videos (List): The list of video inputs.

        Returns:
            A list of experiences.
        """
        mm_inputs = build_multi_modal_inputs(
            prompt=prompt,
            images=images,
            videos=videos,
            processor=self.processor,
        )

        vllm_inputs = {
            "prompt": mm_inputs["prompt"],
            "multi_modal_data": mm_inputs["multi_modal_data"],
        }

        output = await self._generate_internal(prompt=vllm_inputs, **kwargs)
        experiences = [
            Experience(
                tokens=torch.cat(
                    (
                        torch.tensor(output.prompt_token_ids, dtype=torch.int32),
                        torch.tensor(output.outputs[i].token_ids, dtype=torch.int32),
                    )
                ),
                logprobs=torch.cat(
                    (
                        torch.tensor(
                            [
                                list(logprob_dict.values())[0].logprob
                                for logprob_dict in output.outputs[i].logprobs
                            ],
                            dtype=torch.float32,
                        ),
                    )
                ),
                prompt_length=len(output.prompt_token_ids),
                prompt_text=mm_inputs["prompt"],
                response_text=output.outputs[i].text,
                multi_modal_inputs=mm_inputs["multi_modal_inputs"],
            )
            for i in range(len(output.outputs))
        ]
        return experiences

    async def logprobs(  # type: ignore [override]
        self,
        token_ids: List[int],
        lora_request=None,
        temperature: Optional[float] = None,
    ) -> torch.Tensor:
        """Calculate the logprobs of the given tokens in async. Please slice the result carefully
        to align with the actual response length.

        Args:
            token_ids (List[int]): The input token ids (seq_length). Please make sure the length of
                it does not exceed `max_model_len - 1`.
            lora_request (LoRARequest, optional): The LoRA request. Defaults to None.
            temperature (float): The temperature for scaling logits.

        Returns:
            A tensor of logprobs (seq_length - 1).
        """
        temperature = temperature if temperature is not None else self.config.temperature
        if temperature is None:
            temperature = 1.0
        kwargs = {
            "n": 1,
            "max_tokens": 1,
            "prompt_logprobs": 0,  # vLLM return `prompt_logprobs + 1` logrpobs for each token
            "temperature": temperature,
        }
        # avoid using prefix cache when calculating logprobs, only for vLLM >= 0.12.0
        if self.logprobs_no_prefix_cache:
            kwargs["skip_reading_prefix_cache"] = True
        output = await self._generate_internal(
            prompt={"prompt_token_ids": token_ids},
            lora_request=lora_request,
            **kwargs,
        )
        return torch.tensor(
            [list(logprob_dict.values())[0].logprob for logprob_dict in output.prompt_logprobs[1:]],
            dtype=torch.float32,
        )

    async def add_lora_adapter(self, lora_request: Any) -> int:
        """Add a LoRA adapter to the vLLM engine.

        Args:
            lora_request (LoRARequest): The LoRA request.

        Returns:
            lora_id (int): The LoRA adapter ID.
        """
        lora_id = await self.async_llm.add_lora(lora_request)
        return lora_id

    async def remove_lora_adapter(self, lora_id: int) -> None:
        """Remove a LoRA adapter from the vLLM engine.

        Args:
            lora_id (int): The LoRA adapter ID.
        """
        await self.async_llm.remove_lora(lora_id)

    async def list_lora_adapters(self) -> Sequence[int]:
        """List all LoRA adapter IDs in the vLLM engine.

        Returns:
            lora_ids (List[int]): The list of LoRA adapter IDs.
        """
        lora_ids = await self.async_llm.list_loras()
        return list(lora_ids)

    async def sample(
        self,
        prompt: Any,
        num_samples: int,
        sampling_params: Any,
        include_prompt_logprobs: bool = False,
        topk_prompt_logprobs: int = 0,
        lora_request: Optional[Any] = None,
    ) -> Any:
        """Tinker compatible sampling interface.

        Args:
            prompt (ModelInput): The input prompt.
            num_samples (int): The number of samples to generate.
            sampling_params (SamplingParams): The sampling parameters.
            include_prompt_logprobs (bool): Whether to include prompt logprobs.
            topk_prompt_logprobs (int): The top-k prompt logprobs to include.
            lora_request (LoRARequest, optional): The LoRA request. Defaults to None.
        Returns:
            SampleResponse: The sample response.
        """
        from tinker.types import SampledSequence, SampleResponse

        params = {
            "max_tokens": sampling_params.max_tokens
            if sampling_params.max_tokens is not None
            else self.config.max_response_tokens,
            "seed": sampling_params.seed if sampling_params.seed is not None else self.config.seed,
            "top_k": sampling_params.top_k,
            "top_p": sampling_params.top_p,
            "temperature": sampling_params.temperature,
            "n": num_samples,
            "prompt_logprobs": (topk_prompt_logprobs if include_prompt_logprobs else None),
            # in vLLM, 0 means only return the chosen token's logprob
            "logprobs": 0,
        }
        if include_prompt_logprobs and self.logprobs_no_prefix_cache:
            params["skip_reading_prefix_cache"] = True
        if sampling_params.stop is not None:
            params["stop"] = sampling_params.stop
        req_output = await self._generate_internal(
            prompt={"prompt_token_ids": prompt.to_ints()},
            lora_request=lora_request,
            **params,
        )
        sequences = []
        # vLLM's prompt_logprobs output does not include a value for the first token.
        # Initialize with [None] to align with the prompt tokens.
        topk_prompt_logprobs_list: List[Optional[List[Tuple[int, float]]]] = [None]
        prompt_logprobs: List[Optional[float]] = [None]

        # collect prompt logprobs
        if include_prompt_logprobs:
            for logprob_dict in req_output.prompt_logprobs[1:]:
                prompt_logprobs.append(next(iter(logprob_dict.values())).logprob)
                if topk_prompt_logprobs > 0:
                    # collect top-k prompt logprobs
                    # logprob_dict: {token_id: Logprob(logprob, rank, ...), ...}
                    logprob_items = list(logprob_dict.items())
                    # sort by Logprob.rank
                    logprob_items_sorted = sorted(logprob_items, key=lambda x: x[1].rank)
                    # pick topk
                    topk = logprob_items_sorted[:topk_prompt_logprobs]
                    # record as (token_id, logprob)
                    topk_prompt_logprobs_list.append(
                        [(token_id, logprob.logprob) for token_id, logprob in topk]
                    )
        # collect response sequences
        for seq_output in req_output.outputs:
            seq = SampledSequence(
                stop_reason="length" if seq_output.finish_reason == "length" else "stop",
                tokens=seq_output.token_ids,
                logprobs=[
                    next(iter(logprob_dict.values())).logprob
                    for logprob_dict in seq_output.logprobs
                ],
            )
            sequences.append(seq)
        return SampleResponse(
            sequences=sequences,
            prompt_logprobs=prompt_logprobs if include_prompt_logprobs else None,
            topk_prompt_logprobs=(
                topk_prompt_logprobs_list
                if include_prompt_logprobs and topk_prompt_logprobs > 0
                else None
            ),
        )

    async def _generate_internal(self, prompt: Any, lora_request=None, **kwargs) -> Any:
        # Send the request to the LLM engine.
        self.request_id += 1
        stream = self.async_llm.generate(
            request_id=str(self.request_id),
            prompt=prompt,
            sampling_params=self._create_sampling_params(**kwargs),
            lora_request=lora_request,
        )

        # Consume the stream until the request is finished.
        async for request_output in stream:
            if request_output.finished:
                # Bypass the original full prompt.
                # request_output.prompt = request.prompt
                return request_output

        raise RuntimeError("[vLLM] The request is not finished. This should not happen.")

    async def shutdown(self):
        """Shutdown the vLLM v1 engine. This kills child processes forked
        by the vLLM engine. If not called, the child processes will be
        orphaned and will not be killed when the parent process exits,
        and they won't be able to be tracked by Ray anymore.
        """
        if self.api_server is not None:
            self.api_server.cancel()
            try:
                await self.api_server
            except asyncio.CancelledError:
                pass
            self.api_server = None
        if hasattr(self.async_llm, "shutdown"):
            self.logger.info("Shutting down vLLM engine")
            self.async_llm.shutdown()

    def _create_sampling_params(self, **kwargs):
        """Create sampling params."""
        if len(kwargs) == 0:
            return self.default_sampling_params
        params = self.default_sampling_params.clone()
        for k, v in kwargs.items():
            if hasattr(params, k):
                setattr(params, k, v)
        return params

    async def _collective_rpc(
        self,
        method: str,
        timeout: Optional[float] = None,
        args: tuple = (),
        kwargs: Optional[dict] = None,
    ):
        if self.use_v1:
            return await self.async_llm.collective_rpc(method, timeout, args, kwargs)
        else:
            return self.async_llm.engine.model_executor.collective_rpc(
                method, timeout, args, kwargs
            )

    async def sync_model(self, model_version: int, weight_source: str = "student") -> int:
        """Sync model weights to vLLM."""
        if self.enable_lora:
            # Revise the lora path; no need to sync weights manually.
            self.default_lora_path = self.default_lora_path.replace(
                f"global_step_{self.model_version}", f"global_step_{model_version}"
            )
            self.logger.info(
                f"Redirect `lora_path` from old_model_version={self.model_version} to {model_version=} successfully."
            )
            lora_int_ids = await self.async_llm.list_loras()
            for lora_id in lora_int_ids:
                await self.async_llm.remove_lora(lora_id)
            await self.async_llm.add_lora(self.get_lora_request(self.default_lora_path))
            self.model_version = model_version
            return model_version
        await self.async_llm.reset_prefix_cache()
        await self._collective_rpc("update_weight", kwargs={"source": weight_source})
        self.logger.info("Sync model weights to vLLM successfully.")
        self.model_version = model_version
        return model_version

    async def init_process_group(
        self,
        master_address: str,
        master_port: int,
        rank_offset: int,
        world_size: int,
        group_name: str,
        explorer_name: str,
        backend: str = "nccl",
        timeout: int = 1200,
        state_dict_meta: dict = None,
    ):
        return await self._collective_rpc(
            "init_process_group",
            args=(
                master_address,
                master_port,
                rank_offset,
                world_size,
                group_name,
                backend,
                timeout,
                state_dict_meta,
                explorer_name,
                self.ray_namespace,
            ),
        )

    async def run_api_server(self) -> bool:
        """Run the OpenAI API server in a Ray actor.

        Returns:
            success (bool): Whether the API server is started successfully.
        """
        if not self.config.enable_openai_api:
            self.logger.info("OpenAI API server is not enabled. Skipping...")
            return False  # Not enabled

        if self.api_server_host is not None and self.api_server_port is not None:
            self.logger.info("OpenAI API server is already running. Skipping...")
            return True  # already running

        api_server_host, api_server_port = self.get_available_address()
        from trinity.common.models.vllm_patch import get_api_server

        self.api_server = get_api_server(
            self.async_llm,
            host=api_server_host,
            port=api_server_port,
            config=self.config,
            logger=self.logger,
        )
        self.api_server_host = api_server_host
        self.api_server_port = api_server_port
        return True

    def get_api_server_url(self) -> Optional[str]:
        """Get the URL of the OpenAI API server.

        Returns:
            api_url (str): The URL of the OpenAI API server.
        """
        if not self._prepared:
            raise RuntimeError("Model is not prepared. Please call `prepare()` first.")
        if self.api_server_host is None or self.api_server_port is None:
            # openai api is not enabled
            return None
        return f"http://{self.api_server_host}:{self.api_server_port}"

    async def reset_prefix_cache(self) -> None:
        await self.async_llm.reset_prefix_cache()

    def get_model_version(self) -> int:
        return self.model_version

    def get_lora_request(self, lora_path: Optional[str] = None) -> Any:
        from vllm.lora.request import LoRARequest

        assert self.config.lora_modules is not None
        lora_request = LoRARequest(**self.config.lora_modules[0])
        if lora_path is not None:
            self.config.lora_modules[0]["lora_path"] = lora_path  # for consistency
            lora_request.lora_path = lora_path
        return lora_request

    async def get_message_token_len(self, messages) -> int:
        if self.tokenizer is None:
            await self._initialize_tokenizer()
        if self.chat_template is None:
            self.chat_template = self.tokenizer.get_chat_template()
        prompt = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            chat_template=self.chat_template,
            enable_thinking=self.enable_thinking,
        )
        prompt_token = self.tokenizer(  # type: ignore
            prompt, truncation=False, return_tensors="pt"
        )["input_ids"][0].tolist()
        return len(prompt_token)

    async def sleep(self, level: int = 1) -> None:
        await self.async_llm.sleep(level=level)

    async def wake_up(self) -> None:
        await self.async_llm.wake_up()