diff --git a/assets/wechat.jpg b/assets/wechat.jpg index 40a9d988e..a50496546 100644 Binary files a/assets/wechat.jpg and b/assets/wechat.jpg differ diff --git a/pilot/componet.py b/pilot/componet.py index d88ad384f..2c3980cfc 100644 --- a/pilot/componet.py +++ b/pilot/componet.py @@ -3,12 +3,15 @@ from abc import ABC, abstractmethod from typing import Type, Dict, TypeVar, Optional, Union, TYPE_CHECKING from enum import Enum +import logging import asyncio # Checking for type hints during runtime if TYPE_CHECKING: from fastapi import FastAPI +logger = logging.getLogger(__name__) + class LifeCycle: """This class defines hooks for lifecycle events of a component.""" @@ -40,6 +43,7 @@ async def async_before_stop(self): class ComponetType(str, Enum): WORKER_MANAGER = "dbgpt_worker_manager" + MODEL_CONTROLLER = "dbgpt_model_controller" class BaseComponet(LifeCycle, ABC): @@ -92,6 +96,7 @@ def register_instance(self, instance: T): raise RuntimeError( f"Componse name {name} already exists: {self.componets[name]}" ) + logger.info(f"Register componet with name {name} and instance: {instance}") self.componets[name] = instance instance.init_app(self) diff --git a/pilot/model/cluster/controller/controller.py b/pilot/model/cluster/controller/controller.py index e1d55eab7..54360e477 100644 --- a/pilot/model/cluster/controller/controller.py +++ b/pilot/model/cluster/controller/controller.py @@ -4,6 +4,7 @@ from typing import List from fastapi import APIRouter, FastAPI +from pilot.componet import BaseComponet, ComponetType, SystemApp from pilot.model.base import ModelInstance from pilot.model.parameter import ModelControllerParameters from pilot.model.cluster.registry import EmbeddedModelRegistry, ModelRegistry @@ -14,7 +15,12 @@ ) -class BaseModelController(ABC): +class BaseModelController(BaseComponet, ABC): + name = ComponetType.MODEL_CONTROLLER + + def init_app(self, system_app: SystemApp): + pass + @abstractmethod async def register_instance(self, instance: ModelInstance) -> bool: """Register a given model instance""" @@ -25,7 +31,7 @@ async def deregister_instance(self, instance: ModelInstance) -> bool: @abstractmethod async def get_all_instances( - self, model_name: str, healthy_only: bool = False + self, model_name: str = None, healthy_only: bool = False ) -> List[ModelInstance]: """Fetch all instances of a given model. Optionally, fetch only the healthy instances.""" @@ -51,7 +57,7 @@ async def deregister_instance(self, instance: ModelInstance) -> bool: return await self.registry.deregister_instance(instance) async def get_all_instances( - self, model_name: str, healthy_only: bool = False + self, model_name: str = None, healthy_only: bool = False ) -> List[ModelInstance]: logging.info( f"Get all instances with {model_name}, healthy_only: {healthy_only}" @@ -94,7 +100,7 @@ async def get_all_model_instances(self) -> List[ModelInstance]: @sync_api_remote(path="/api/controller/models") def sync_get_all_instances( - self, model_name: str, healthy_only: bool = False + self, model_name: str = None, healthy_only: bool = False ) -> List[ModelInstance]: pass @@ -110,7 +116,7 @@ async def deregister_instance(self, instance: ModelInstance) -> bool: return await self.backend.deregister_instance(instance) async def get_all_instances( - self, model_name: str, healthy_only: bool = False + self, model_name: str = None, healthy_only: bool = False ) -> List[ModelInstance]: return await self.backend.get_all_instances(model_name, healthy_only) diff --git a/pilot/openapi/api_v1/api_v1.py b/pilot/openapi/api_v1/api_v1.py index 1507a9a1f..3af082026 100644 --- a/pilot/openapi/api_v1/api_v1.py +++ b/pilot/openapi/api_v1/api_v1.py @@ -18,6 +18,7 @@ from typing import List import tempfile +from pilot.componet import ComponetType from pilot.openapi.api_view_model import ( Result, ConversationVo, @@ -352,20 +353,17 @@ async def chat_completions(dialogue: ConversationVo = Body()): async def model_types(request: Request): print(f"/controller/model/types") try: - import httpx - - async with httpx.AsyncClient() as client: - base_url = request.base_url - response = await client.get( - f"{base_url}api/controller/models?healthy_only=true", - ) types = set() - if response.status_code == 200: - models = json.loads(response.text) - for model in models: - worker_type = model["model_name"].split("@")[1] - if worker_type == "llm": - types.add(model["model_name"].split("@")[0]) + from pilot.model.cluster.controller.controller import BaseModelController + + controller = CFG.SYSTEM_APP.get_componet( + ComponetType.MODEL_CONTROLLER, BaseModelController + ) + models = await controller.get_all_instances(healthy_only=True) + for model in models: + worker_name, worker_type = model.model_name.split("@") + if worker_type == "llm": + types.add(worker_name) return Result.succ(list(types)) except Exception as e: diff --git a/pilot/server/componet_configs.py b/pilot/server/componet_configs.py index 745937068..b68b052fa 100644 --- a/pilot/server/componet_configs.py +++ b/pilot/server/componet_configs.py @@ -9,10 +9,12 @@ def initialize_componets(system_app: SystemApp, embedding_model_name: str): from pilot.model.cluster import worker_manager + from pilot.model.cluster.controller.controller import controller system_app.register( RemoteEmbeddingFactory, worker_manager, model_name=embedding_model_name ) + system_app.register_instance(controller) class RemoteEmbeddingFactory(EmbeddingFactory):