fix(model): Fix the bug that the webserver cannot return model instances

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)
 

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)
 

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:

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):