project_6_89d52222/upstream_ref/xllm/xllm/pybind/vlm.py

import os
import signal
import sys
from . import util
from typing import List, Optional, Union, Dict, Any
from xllm_export import (VLMMaster, Options, RequestOutput,
                         RequestParams, MMData)
from .errors import ValidationError
from .params import (
    SamplingParams,
    to_request_params_list,
)


class VLM:
    def __init__(
        self,
        model: str,
        task: str = "generate",
        devices: str = 'auto',
        draft_model: Optional[str] = None,
        draft_devices: Optional[str] = None,
        block_size: int = 128,
        max_cache_size: int = 0,
        max_memory_utilization: float = 0.9,
        disable_prefix_cache: bool = False,
        max_tokens_per_batch: int = 50000,
        max_seqs_per_batch: int = 256,
        max_tokens_per_chunk_for_prefill: int = 512,
        num_speculative_tokens: int = 0,
        num_request_handling_threads: int = 4,
        communication_backend: str = 'hccl',
        rank_tablefile: str = '',
        expert_parallel_degree: int = 0,
        disable_chunked_prefill: bool = False,
        enable_prefill_sp: bool = False,
        instance_role: str = 'DEFAULT',
        device_ip: str = '',
        transfer_listen_port: int = 26000,
        nnodes: int = 1,
        node_rank: int = 0,
        dp_size: int = 1,
        ep_size: int = 1,
        instance_name: str = '',
        enable_disagg_pd: bool = False,
        enable_schedule_overlap: bool = False,
        kv_cache_transfer_mode: str = 'PUSH',
        enable_shm: bool = False,
        is_local: bool = True,
        input_shm_size: int = 1024,
        output_shm_size: int = 128,
        **kwargs: Any,
    ) -> None:
        signal.signal(signal.SIGTERM, lambda s, f: sys.exit(0))
        signal.signal(signal.SIGINT, lambda s, f: sys.exit(0))

        if not os.path.exists(model):
            raise ValueError(f"model {model} not exists")
        self.model = model

        options = Options()
        options.model_path = model
        options.task_type = task
        options.devices = devices
        options.draft_model_path = draft_model
        options.draft_devices = draft_devices
        options.backend ="vlm"
        options.block_size = block_size
        options.max_cache_size = max_cache_size
        options.max_memory_utilization = max_memory_utilization
        if disable_prefix_cache:
            options.enable_prefix_cache = False
        else:
            options.enable_prefix_cache = True
        options.max_tokens_per_batch = max_tokens_per_batch
        options.max_seqs_per_batch = max_seqs_per_batch
        options.max_tokens_per_chunk_for_prefill = max_tokens_per_chunk_for_prefill
        options.num_speculative_tokens = num_speculative_tokens
        options.num_request_handling_threads = num_request_handling_threads
        options.communication_backend = communication_backend
        options.rank_tablefile = rank_tablefile
        options.expert_parallel_degree = expert_parallel_degree
        if disable_chunked_prefill:
            options.enable_chunked_prefill = False
        else:
            options.enable_chunked_prefill = True
        options.enable_prefill_sp = enable_prefill_sp
        free_port = util.get_free_port()
        options.master_node_addr = "127.0.0.1:" + str(free_port)
        options.device_ip = device_ip
        options.transfer_listen_port = transfer_listen_port
        options.nnodes = nnodes
        options.node_rank = node_rank
        options.dp_size = dp_size
        options.ep_size = ep_size
        options.instance_name = instance_name
        options.enable_disagg_pd = enable_disagg_pd
        options.enable_schedule_overlap = False
        options.kv_cache_transfer_mode = kv_cache_transfer_mode
        options.enable_offline_inference = True
        options.spawn_worker_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
        options.enable_shm = enable_shm
        options.is_local = is_local
        options.input_shm_size = input_shm_size
        options.output_shm_size = output_shm_size
        self.master = VLMMaster(options)

    def finish(self) -> None:
        try:
            #os.kill(os.getpid(), signal.SIGTERM)
            #os.kill(os.getpid(), signal.SIGKILL)
            util.terminate_process(os.getpid())
        except Exception as e:
            pass

    def generate(
        self,
        prompts: Union[
            str,
            List[str],
            Dict[str, Any],
            List[Dict[str, Any]],
        ],
        sampling_params: Optional[Union[
            SamplingParams,
            List[SamplingParams],
        ]] = None,
        wait_for_schedule: bool = True,
        **kwargs: Any,
    ) -> List[RequestOutput]:
        from . import mm_utils
        prompts, mm_datas, image_urls = mm_utils.normalize_vllm_style_inputs(prompts)

        request_params = kwargs.pop("request_params", None)
        if kwargs:
            unknown = ", ".join(kwargs.keys())
            raise TypeError(f"Unexpected keyword arguments: {unknown}")
        if request_params is None:
            request_params = sampling_params
        elif sampling_params is not None:
            raise ValueError("sampling_params and request_params cannot both be set")

        request_params_list = to_request_params_list(
            request_params, default_cls=SamplingParams
        )
        if len(request_params_list) not in (1, len(prompts)):
            raise ValueError(
                "The number of request_params must be 1 or equal to the number of prompts."
            )

        outputs = [None] * len(prompts)
        def callback(index: int, output: RequestOutput) -> bool:
            outputs[index] = output
            return True
        # schedule the batch requests
        if image_urls is not None:
            self.master.handle_batch_request_with_image_urls(
                prompts, image_urls, request_params_list, callback
            )
        else:
            self.master.handle_batch_request(
                prompts, mm_datas, request_params_list, callback
            )

        # wait for batch request to be scheduled
        if wait_for_schedule:
            pass

        # run until all scheduled requsts complete
        self.master.generate()

        # throw an exception if there is any error
        for index, output in enumerate(outputs):
            if output is None:
                raise RuntimeError("Request failed, no output received")
            if output.status is not None and not output.status.ok:
                raise ValidationError(output.status.code, output.status.message)
            # carry over the prompt to the output
            output.prompt = prompts[index]
        return outputs