refactor the solution of vllm integration (#60)

Author: depenglee1707

llmserve/backend/llm/engines/__init__.py

@@ -0,0 +1,29 @@
+from typing import TYPE_CHECKING, Type
+from llmserve.backend.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+if TYPE_CHECKING:
+    from ._base import LLMEngine
+
+from .generic import GenericEngine
+try:
+    from .vllm import VllmEngine
+except ImportError:
+    logger.info("Import vllm related stuff failed, please make sure 'vllm' is installed.")
+
+def get_engine_cls_by_name(name: str) -> Type["LLMEngine"]:
+    lowercase_globals = {k.lower(): v for k, v in globals().items()}
+    ret = lowercase_globals.get(
+        f"{name.lower()}engine", lowercase_globals.get(name.lower(), None)
+    )
+    assert ret
+    return ret
+
+
+__all__ = [
+    "get_engine_cls_by_name",
+    "GenericEngine",
+    "VllmEngine",
+]

llmserve/backend/llm/engines/_base.py

@@ -0,0 +1,66 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Any
+from ray.air import ScalingConfig
+from ray.util.placement_group import PlacementGroup
+from llmserve.backend.server.models import Prompt
+
+from llmserve.backend.logger import get_logger
+
+from typing import List, Optional
+from ray.air import ScalingConfig
+
+from llmserve.backend.logger import get_logger
+from llmserve.backend.server.models import Args, Prompt
+import asyncio
+
+logger = get_logger(__name__)
+
+class LLMEngine(ABC):
+    args: Args = None
+    """Initialize model and tokenizer and place them on the correct device.
+
+    Args:
+        device (torch.device): Device to place model and tokenizer on.
+        world_size (int): Number of GPUs to use.
+    """
+
+    def __init__(
+        self,
+        args: Args,
+
+    ):
+        self.args = args
+
+    @abstractmethod
+    async def launch_engine(
+            self, 
+            scaling_config: ScalingConfig,
+            placement_group: PlacementGroup,
+            scaling_options: dict,
+        ) -> Any:
+        """Load model.
+
+        Args:
+            model_id (str): Hugging Face model ID.
+        """
+        pass
+
+    @abstractmethod
+    async def predict(
+            self,
+            prompts: List[Prompt],
+            *,
+            timeout_s: float = 60,
+            start_timestamp: Optional[float] = None,
+            lock: asyncio.Lock,
+        ) -> List[str]:
+        """Load model.
+
+        Args:
+            model_id (str): Hugging Face model ID.
+        """
+        pass
+    
+    @abstractmethod
+    async def check_health(self):
+        pass