Repository: SuperDuperDB/superduperdb
File: superduperdb/ext/llm/vllm.py (Python)

Code before:

    tensor_parallel_size: int = 1
    trust_remote_code: bool = True
    vllm_kwargs: dict = dc.field(default_factory=dict)

    def __post_init__(self):
        self.on_ray = self.on_ray or bool(self.ray_address)
        if "tensor_parallel_size" not in self.vllm_kwargs:
            self.vllm_kwargs["tensor_parallel_size"] = self.tensor_parallel_size

        if "trust_remote_code" not in self.vllm_kwargs:
            self.vllm_kwargs["trust_remote_code"] = self.trust_remote_code

        if "model" not in self.vllm_kwargs:
            self.vllm_kwargs["model"] = self.model_name

        super().__post_init__()

    def init(self):
        class _VLLMCore:
            """
            Wrapper for vllm model to support ray.
            Implementing the client in this way will no longer require vllm dependencies
            """

            def __init__(self, **kwargs) -> None:
                try:
                    from vllm import LLM
                except ImportError:
                    raise Exception(
                        "You must install vllm with command 'pip install vllm'"
                    )
                self.model = LLM(**kwargs)

            def generate(self, prompts: List[str], **kwargs) -> List[str]:
                from vllm import SamplingParams

                sampling_params = SamplingParams(**kwargs)
                results = self.model.generate(prompts, sampling_params, use_tqdm=False)
                results = [result.outputs[0].text for result in results]
                return results

        if self.on_ray:
            try:
                import ray
            except ImportError:
                raise Exception("You must install vllm with command 'pip install ray'")

            self.ray_config.setdefault("runtime_env", {"pip": ["vllm"]})

            if not ray.is_initialized():

Review comment (suggestion):
Change to AsyncLLMEngine because only AsyncLLMEngine supports real-time concurrent requests


related issues:
- https://github.com/vllm-project/vllm/issues/3248
- https://github.com/vllm-project/vllm/issues/1200

Code after:
        from vllm import SamplingParams
        from vllm.utils import random_uuid

        sampling_params = SamplingParams(**kwargs)
        request_id = random_uuid()
        results_generator = self.engine.generate(prompt, sampling_params, request_id)

        final_output = None
        async for request_output in results_generator:
            final_output = request_output

        assert final_output is not None
        prompt = final_output.prompt
        text_outputs = [output.text for output in final_output.outputs]
        n = sampling_params.n
        if n == 1:
            return text_outputs[0]
        return text_outputs

    async def abatch_predict(self, prompts: List[str], **kwargs):
        return await asyncio.gather(
            *[self.agenerate(prompt, **kwargs) for prompt in prompts]
        )

    def batch_predict(self, prompts: List[str], **kwargs):
        return asyncio.run(self.abatch_predict(prompts, **kwargs))


@public_api(stability='beta')
@dc.dataclass
class VllmModel(BaseLLMModel):
    """
    Load a large language model from VLLM.

    :param model_name: The name of the model to use.
    :param trust_remote_code: Whether to trust remote code.
    :param dtype: The data type to use.
    {parent_doc}
    """

    __doc__ = __doc__.format(parent_doc=BaseLLMModel.__doc__)

    tensor_parallel_size: int = 1
    trust_remote_code: bool = True
    vllm_kwargs: dict = dc.field(default_factory=dict)

    def __post_init__(self, artifacts):
        self.on_ray = self.on_ray or bool(self.ray_address)
        if "tensor_parallel_size" not in self.vllm_kwargs:
            self.vllm_kwargs["tensor_parallel_size"] = self.tensor_parallel_size