xinference 0.8.1__py3-none-any.whl → 0.8.3__py3-none-any.whl

xinference/client/handlers.py

@@ -1,3 +1,6 @@
+from .restful.restful_client import (  # noqa: F401
+    RESTfulAudioModelHandle as AudioModelHandle,
+)
 from .restful.restful_client import (  # noqa: F401
     RESTfulChatglmCppChatModelHandle as ChatglmCppChatModelHandle,
 )

xinference/client/restful/restful_client.py

@@ -400,19 +400,17 @@ class RESTfulChatModelHandle(RESTfulGenerateModelHandle):
         return response_data
 
 
-class RESTfulMultimodalModelHandle(RESTfulModelHandle):
+class RESTfulChatglmCppChatModelHandle(RESTfulModelHandle):
     def chat(
         self,
-        prompt: Any,
+        prompt: str,
         system_prompt: Optional[str] = None,
         chat_history: Optional[List["ChatCompletionMessage"]] = None,
         tools: Optional[List[Dict]] = None,
-        generate_config: Optional[
-            Union["LlamaCppGenerateConfig", "PytorchGenerateConfig"]
-        ] = None,
+        generate_config: Optional["ChatglmCppGenerateConfig"] = None,
     ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
         """
-        Given a list of messages comprising a conversation, the model will return a response via RESTful APIs.
+        Given a list of messages comprising a conversation, the ChatGLM model will return a response via RESTful APIs.
 
         Parameters
         ----------
@@ -424,10 +422,8 @@ class RESTfulMultimodalModelHandle(RESTfulModelHandle):
             A list of messages comprising the conversation so far.
         tools: Optional[List[Dict]]
             A tool list.
-        generate_config: Optional[Union["LlamaCppGenerateConfig", "PytorchGenerateConfig"]]
-            Additional configuration for the chat generation.
-            "LlamaCppGenerateConfig" -> configuration for ggml model
-            "PytorchGenerateConfig" -> configuration for pytorch model
+        generate_config: Optional["ChatglmCppGenerateConfig"]
+            Additional configuration for ChatGLM chat generation.
 
         Returns
         -------
@@ -451,7 +447,6 @@ class RESTfulMultimodalModelHandle(RESTfulModelHandle):
         if chat_history and chat_history[0]["role"] == "system":
             if system_prompt is not None:
                 chat_history[0]["content"] = system_prompt
-
         else:
             if system_prompt is not None:
                 chat_history.insert(0, {"role": "system", "content": system_prompt})
@@ -463,8 +458,7 @@ class RESTfulMultimodalModelHandle(RESTfulModelHandle):
             "messages": chat_history,
         }
         if tools is not None:
-            raise RuntimeError("Multimodal does not support function call.")
-
+            request_body["tools"] = tools
         if generate_config is not None:
             for key, value in generate_config.items():
                 request_body[key] = value
@@ -486,67 +480,51 @@ class RESTfulMultimodalModelHandle(RESTfulModelHandle):
         return response_data
 
 
-class RESTfulChatglmCppChatModelHandle(RESTfulModelHandle):
-    def chat(
+class RESTfulChatglmCppGenerateModelHandle(RESTfulChatglmCppChatModelHandle):
+    def generate(
         self,
         prompt: str,
-        chat_history: Optional[List["ChatCompletionMessage"]] = None,
-        tools: Optional[List[Dict]] = None,
         generate_config: Optional["ChatglmCppGenerateConfig"] = None,
-    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+    ) -> Union["Completion", Iterator["CompletionChunk"]]:
         """
-        Given a list of messages comprising a conversation, the ChatGLM model will return a response via RESTful APIs.
+        Given a prompt, the ChatGLM model will generate a response via RESTful APIs.
 
         Parameters
         ----------
         prompt: str
             The user's input.
-        chat_history: Optional[List["ChatCompletionMessage"]]
-            A list of messages comprising the conversation so far.
-        tools: Optional[List[Dict]]
-            A tool list.
         generate_config: Optional["ChatglmCppGenerateConfig"]
             Additional configuration for ChatGLM chat generation.
 
         Returns
         -------
-        Union["ChatCompletion", Iterator["ChatCompletionChunk"]]
+        Union["Completion", Iterator["CompletionChunk"]]
             Stream is a parameter in generate_config.
-            When stream is set to True, the function will return Iterator["ChatCompletionChunk"].
-            When stream is set to False, the function will return "ChatCompletion".
+            When stream is set to True, the function will return Iterator["CompletionChunk"].
+            When stream is set to False, the function will return "Completion".
 
         Raises
         ------
         RuntimeError
-            Report the failure to generate the chat from the server. Detailed information provided in error message.
+            Report the failure to generate the content from the server. Detailed information provided in error message.
 
         """
 
-        url = f"{self._base_url}/v1/chat/completions"
-
-        if chat_history is None:
-            chat_history = []
-
-        chat_history.append({"role": "user", "content": prompt})
+        url = f"{self._base_url}/v1/completions"
 
-        request_body: Dict[str, Any] = {
-            "model": self._model_uid,
-            "messages": chat_history,
-        }
-        if tools is not None:
-            request_body["tools"] = tools
+        request_body: Dict[str, Any] = {"model": self._model_uid, "prompt": prompt}
         if generate_config is not None:
             for key, value in generate_config.items():
                 request_body[key] = value
 
         stream = bool(generate_config and generate_config.get("stream"))
+
         response = requests.post(
             url, json=request_body, stream=stream, headers=self.auth_headers
         )
-
         if response.status_code != 200:
             raise RuntimeError(
-                f"Failed to generate chat completion, detail: {_get_error_string(response)}"
+                f"Failed to generate completion, detail: {response.json()['detail']}"
             )
 
         if stream:
@@ -556,56 +534,111 @@ class RESTfulChatglmCppChatModelHandle(RESTfulModelHandle):
         return response_data
 
 
-class RESTfulChatglmCppGenerateModelHandle(RESTfulChatglmCppChatModelHandle):
-    def generate(
+class RESTfulAudioModelHandle(RESTfulModelHandle):
+    def transcriptions(
         self,
-        prompt: str,
-        generate_config: Optional["ChatglmCppGenerateConfig"] = None,
-    ) -> Union["Completion", Iterator["CompletionChunk"]]:
+        audio: bytes,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: Optional[str] = "json",
+        temperature: Optional[float] = 0,
+    ):
         """
-        Given a prompt, the ChatGLM model will generate a response via RESTful APIs.
+        Transcribes audio into the input language.
 
         Parameters
         ----------
-        prompt: str
-            The user's input.
-        generate_config: Optional["ChatglmCppGenerateConfig"]
-            Additional configuration for ChatGLM chat generation.
+
+        audio: bytes
+            The audio file object (not file name) to transcribe, in one of these formats: flac, mp3, mp4, mpeg,
+            mpga, m4a, ogg, wav, or webm.
+        language: Optional[str]
+            The language of the input audio. Supplying the input language in ISO-639-1
+            (https://en.wikipedia.org/wiki/List_of_ISO_639_language_codes) format will improve accuracy and latency.
+        prompt: Optional[str]
+            An optional text to guide the model's style or continue a previous audio segment.
+            The prompt should match the audio language.
+        response_format: Optional[str], defaults to json
+            The format of the transcript output, in one of these options: json, text, srt, verbose_json, or vtt.
+        temperature: Optional[float], defaults to 0
+            The sampling temperature, between 0 and 1. Higher values like 0.8 will make the output more random,
+            while lower values like 0.2 will make it more focused and deterministic.
+            If set to 0, the model will use log probability to automatically increase the temperature
+            until certain thresholds are hit.
 
         Returns
         -------
-        Union["Completion", Iterator["CompletionChunk"]]
-            Stream is a parameter in generate_config.
-            When stream is set to True, the function will return Iterator["CompletionChunk"].
-            When stream is set to False, the function will return "Completion".
+            The transcribed text.
+        """
+        url = f"{self._base_url}/v1/audio/transcriptions"
+        params = {
+            "model": self._model_uid,
+            "language": language,
+            "prompt": prompt,
+            "response_format": response_format,
+            "temperature": temperature,
+        }
+        files: List[Any] = []
+        for key, value in params.items():
+            files.append((key, (None, value)))
+        files.append(("file", ("file", audio, "application/octet-stream")))
+        response = requests.post(url, files=files, headers=self.auth_headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to transcribe the audio, detail: {_get_error_string(response)}"
+            )
 
-        Raises
-        ------
-        RuntimeError
-            Report the failure to generate the content from the server. Detailed information provided in error message.
+        response_data = response.json()
+        return response_data
 
+    def translations(
+        self,
+        audio: bytes,
+        prompt: Optional[str] = None,
+        response_format: Optional[str] = "json",
+        temperature: Optional[float] = 0,
+    ):
         """
+        Translates audio into English.
 
-        url = f"{self._base_url}/v1/completions"
-
-        request_body: Dict[str, Any] = {"model": self._model_uid, "prompt": prompt}
-        if generate_config is not None:
-            for key, value in generate_config.items():
-                request_body[key] = value
+        Parameters
+        ----------
 
-        stream = bool(generate_config and generate_config.get("stream"))
+        audio: bytes
+            The audio file object (not file name) to transcribe, in one of these formats: flac, mp3, mp4, mpeg,
+            mpga, m4a, ogg, wav, or webm.
+        prompt: Optional[str]
+            An optional text to guide the model's style or continue a previous audio segment.
+            The prompt should match the audio language.
+        response_format: Optional[str], defaults to json
+            The format of the transcript output, in one of these options: json, text, srt, verbose_json, or vtt.
+        temperature: Optional[float], defaults to 0
+            The sampling temperature, between 0 and 1. Higher values like 0.8 will make the output more random,
+            while lower values like 0.2 will make it more focused and deterministic.
+            If set to 0, the model will use log probability to automatically increase the temperature
+            until certain thresholds are hit.
 
-        response = requests.post(
-            url, json=request_body, stream=stream, headers=self.auth_headers
-        )
+        Returns
+        -------
+            The translated text.
+        """
+        url = f"{self._base_url}/v1/audio/translations"
+        params = {
+            "model": self._model_uid,
+            "prompt": prompt,
+            "response_format": response_format,
+            "temperature": temperature,
+        }
+        files: List[Any] = []
+        for key, value in params.items():
+            files.append((key, (None, value)))
+        files.append(("file", ("file", audio, "application/octet-stream")))
+        response = requests.post(url, files=files, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
-                f"Failed to generate completion, detail: {response.json()['detail']}"
+                f"Failed to translate the audio, detail: {_get_error_string(response)}"
             )
 
-        if stream:
-            return streaming_response_iterator(response.iter_lines())
-
         response_data = response.json()
         return response_data
 
@@ -889,8 +922,8 @@ class Client:
             return RESTfulRerankModelHandle(
                 model_uid, self.base_url, auth_headers=self._headers
             )
-        elif desc["model_type"] == "multimodal":
-            return RESTfulMultimodalModelHandle(
+        elif desc["model_type"] == "audio":
+            return RESTfulAudioModelHandle(
                 model_uid, self.base_url, auth_headers=self._headers
             )
         else:

xinference/constants.py

@@ -18,8 +18,12 @@ from pathlib import Path
 XINFERENCE_ENV_ENDPOINT = "XINFERENCE_ENDPOINT"
 XINFERENCE_ENV_MODEL_SRC = "XINFERENCE_MODEL_SRC"
 XINFERENCE_ENV_HOME_PATH = "XINFERENCE_HOME"
-XINFERENCE_ENV_HEALTH_CHECK_ATTEMPTS = "XINFERENCE_HEALTH_CHECK_ATTEMPTS"
+XINFERENCE_ENV_HEALTH_CHECK_FAILURE_THRESHOLD = (
+    "XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD"
+)
 XINFERENCE_ENV_HEALTH_CHECK_INTERVAL = "XINFERENCE_HEALTH_CHECK_INTERVAL"
+XINFERENCE_ENV_HEALTH_CHECK_TIMEOUT = "XINFERENCE_HEALTH_CHECK_TIMEOUT"
+XINFERENCE_ENV_DISABLE_HEALTH_CHECK = "XINFERENCE_DISABLE_HEALTH_CHECK"
 XINFERENCE_ENV_DISABLE_VLLM = "XINFERENCE_DISABLE_VLLM"
 
 
@@ -47,10 +51,16 @@ XINFERENCE_DEFAULT_ENDPOINT_PORT = 9997
 XINFERENCE_DEFAULT_LOG_FILE_NAME = "xinference.log"
 XINFERENCE_LOG_MAX_BYTES = 100 * 1024 * 1024
 XINFERENCE_LOG_BACKUP_COUNT = 30
-XINFERENCE_HEALTH_CHECK_ATTEMPTS = int(
-    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_ATTEMPTS, 3)
+XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD = int(
+    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_FAILURE_THRESHOLD, 5)
 )
 XINFERENCE_HEALTH_CHECK_INTERVAL = int(
-    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_INTERVAL, 3)
+    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_INTERVAL, 5)
+)
+XINFERENCE_HEALTH_CHECK_TIMEOUT = int(
+    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_TIMEOUT, 10)
+)
+XINFERENCE_DISABLE_HEALTH_CHECK = bool(
+    int(os.environ.get(XINFERENCE_ENV_DISABLE_HEALTH_CHECK, 0))
 )
 XINFERENCE_DISABLE_VLLM = bool(int(os.environ.get(XINFERENCE_ENV_DISABLE_VLLM, 0)))

xinference/core/cache_tracker.py

@@ -0,0 +1,102 @@
+# Copyright 2022-2024 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from logging import getLogger
+from typing import Dict, List, Optional
+
+import xoscar as xo
+
+logger = getLogger(__name__)
+
+
+class CacheTrackerActor(xo.Actor):
+    def __init__(self):
+        super().__init__()
+        self._model_name_to_version_info: Dict[str, List[Dict]] = {}
+
+    @classmethod
+    def uid(cls) -> str:
+        return "cache_tracker"
+
+    @staticmethod
+    def _map_address_to_file_location(
+        model_version: Dict[str, List[Dict]], address: str
+    ):
+        for model_name, model_versions in model_version.items():
+            for info_dict in model_versions:
+                info_dict["model_file_location"] = (
+                    {address: info_dict["model_file_location"]}
+                    if info_dict["cache_status"]
+                    else None
+                )
+
+    @staticmethod
+    def _update_file_location(data: Dict, origin_version_info: Dict):
+        if origin_version_info["model_file_location"] is None:
+            origin_version_info["model_file_location"] = data
+        else:
+            assert isinstance(origin_version_info["model_file_location"], dict)
+            origin_version_info["model_file_location"].update(data)
+
+    def record_model_version(self, version_info: Dict[str, List[Dict]], address: str):
+        self._map_address_to_file_location(version_info, address)
+        for model_name, model_versions in version_info.items():
+            if model_name not in self._model_name_to_version_info:
+                self._model_name_to_version_info[model_name] = model_versions
+            else:
+                assert len(model_versions) == len(
+                    self._model_name_to_version_info[model_name]
+                ), "Model version info inconsistency between supervisor and worker"
+                for version, origin_version in zip(
+                    model_versions, self._model_name_to_version_info[model_name]
+                ):
+                    if (
+                        version["cache_status"]
+                        and version["model_file_location"] is not None
+                    ):
+                        origin_version["cache_status"] = True
+                        self._update_file_location(
+                            version["model_file_location"], origin_version
+                        )
+
+    def update_cache_status(
+        self,
+        address: str,
+        model_name: str,
+        model_version: Optional[str],
+        model_path: str,
+    ):
+        if model_name not in self._model_name_to_version_info:
+            logger.warning(f"Not record version info for {model_name} for now.")
+        else:
+            for version_info in self._model_name_to_version_info[model_name]:
+                if model_version is None:  # image model
+                    self._update_file_location({address: model_path}, version_info)
+                    version_info["cache_status"] = True
+                else:
+                    if version_info["model_version"] == model_version:
+                        self._update_file_location({address: model_path}, version_info)
+                        version_info["cache_status"] = True
+
+    def unregister_model_version(self, model_name: str):
+        self._model_name_to_version_info.pop(model_name, None)
+
+    def get_model_versions(self, model_name: str) -> List[Dict]:
+        if model_name not in self._model_name_to_version_info:
+            logger.warning(f"Not record version info for model_name: {model_name}")
+            return []
+        else:
+            return self._model_name_to_version_info[model_name]
+
+    def get_model_version_count(self, model_name: str) -> int:
+        return len(self.get_model_versions(model_name))

xinference/core/chat_interface.py

@@ -27,7 +27,6 @@ from ..client.restful.restful_client import (
     RESTfulChatglmCppChatModelHandle,
     RESTfulChatModelHandle,
     RESTfulGenerateModelHandle,
-    RESTfulMultimodalModelHandle,
 )
 from ..types import ChatCompletionMessage
 
@@ -66,7 +65,7 @@ class GradioInterface:
         )
 
     def build(self) -> "gr.Blocks":
-        if self.model_type == "multimodal":
+        if "vision" in self.model_ability:
             interface = self.build_chat_vl_interface()
         elif "chat" in self.model_ability:
             interface = self.build_chat_interface()
@@ -99,9 +98,16 @@ class GradioInterface:
             return flat_list
 
         def to_chat(lst: List[str]) -> List[ChatCompletionMessage]:
+            from ..model.llm import BUILTIN_LLM_PROMPT_STYLE
+
             res = []
+            prompt_style = BUILTIN_LLM_PROMPT_STYLE.get(self.model_name)
+            if prompt_style is None:
+                roles = ["assistant", "user"]
+            else:
+                roles = prompt_style.roles
             for i in range(len(lst)):
-                role = "assistant" if i % 2 == 1 else "user"
+                role = roles[0] if i % 2 == 1 else roles[1]
                 res.append(ChatCompletionMessage(role=role, content=lst[i]))
             return res
 
@@ -191,7 +197,7 @@ class GradioInterface:
             client = RESTfulClient(self.endpoint)
             client._set_token(self._access_token)
             model = client.get_model(self.model_uid)
-            assert isinstance(model, RESTfulMultimodalModelHandle)
+            assert isinstance(model, RESTfulChatModelHandle)
 
             prompt = history[-1]
             assert prompt["role"] == "user"

xinference/core/event.py

@@ -0,0 +1,56 @@
+# Copyright 2022-2024 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import queue
+from collections import defaultdict
+from enum import Enum
+from typing import Dict, List, TypedDict
+
+import xoscar as xo
+
+MAX_EVENT_COUNT_PER_MODEL = 100
+
+
+class EventType(Enum):
+    INFO = 1
+    WARNING = 2
+    ERROR = 3
+
+
+class Event(TypedDict):
+    event_type: EventType
+    event_ts: int
+    event_content: str
+
+
+class EventCollectorActor(xo.StatelessActor):
+    def __init__(self):
+        super().__init__()
+        self._model_uid_to_events: Dict[str, queue.Queue] = defaultdict(
+            lambda: queue.Queue(maxsize=MAX_EVENT_COUNT_PER_MODEL)
+        )
+
+    @classmethod
+    def uid(cls) -> str:
+        return "event_collector"
+
+    def get_model_events(self, model_uid: str) -> List[Dict]:
+        event_queue = self._model_uid_to_events.get(model_uid)
+        if event_queue is None:
+            return []
+        else:
+            return [dict(e, event_type=e["event_type"].name) for e in event_queue.queue]
+
+    def report_event(self, model_uid: str, event: Event):
+        self._model_uid_to_events[model_uid].put(event)

xinference/core/model.py

@@ -426,6 +426,50 @@ class ModelActor(xo.StatelessActor):
             )
         raise AttributeError(f"Model {self._model.model_spec} is not for reranking.")
 
+    @log_async(logger=logger, args_formatter=lambda _, kwargs: kwargs.pop("audio"))
+    @request_limit
+    async def transcriptions(
+        self,
+        audio: bytes,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        temperature: float = 0,
+    ):
+        if hasattr(self._model, "transcriptions"):
+            return await self._call_wrapper(
+                self._model.transcriptions,
+                audio,
+                language,
+                prompt,
+                response_format,
+                temperature,
+            )
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating transcriptions."
+        )
+
+    @log_async(logger=logger, args_formatter=lambda _, kwargs: kwargs.pop("audio"))
+    @request_limit
+    async def translations(
+        self,
+        audio: bytes,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        temperature: float = 0,
+    ):
+        if hasattr(self._model, "translations"):
+            return await self._call_wrapper(
+                self._model.translations,
+                audio,
+                prompt,
+                response_format,
+                temperature,
+            )
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating translations."
+        )
+
     @log_async(logger=logger)
     @request_limit
     async def text_to_image(

xinference/core/resource.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import Dict
+from typing import Dict, Union
 
 import psutil
 
+from .utils import get_nvidia_gpu_info
+
 
 @dataclass
 class ResourceStatus:
@@ -26,7 +28,14 @@ class ResourceStatus:
     memory_total: float
 
 
-def gather_node_info() -> Dict[str, ResourceStatus]:
+@dataclass
+class GPUStatus:
+    mem_total: float
+    mem_free: float
+    mem_used: float
+
+
+def gather_node_info() -> Dict[str, Union[ResourceStatus, GPUStatus]]:
     node_resource = dict()
     mem_info = psutil.virtual_memory()
     node_resource["cpu"] = ResourceStatus(
@@ -35,13 +44,11 @@ def gather_node_info() -> Dict[str, ResourceStatus]:
         memory_available=mem_info.available,
         memory_total=mem_info.total,
     )
-    # TODO: record GPU stats
-    # for idx, gpu_card_stat in enumerate(resource.cuda_card_stats()):
-    #     node_resource[f"gpu-{idx}"] = ResourceStatus(
-    #         available=gpu_card_stat.gpu_usage / 100.0,
-    #         total=1,
-    #         memory_available=gpu_card_stat.fb_mem_info.available,
-    #         memory_total=gpu_card_stat.fb_mem_info.total,
-    #     )
-
-    return node_resource
+    for gpu_idx, gpu_info in get_nvidia_gpu_info().items():
+        node_resource[gpu_idx] = GPUStatus(  # type: ignore
+            mem_total=gpu_info["total"],
+            mem_used=gpu_info["used"],
+            mem_free=gpu_info["free"],
+        )
+
+    return node_resource  # type: ignore

xinference/core/status_guard.py

@@ -33,6 +33,7 @@ class LaunchStatus(Enum):
 class InstanceInfo(BaseModel):
     model_name: str
     model_uid: str
+    model_version: Optional[str]
     model_ability: List[str]
     replica: int
     status: str
@@ -82,5 +83,8 @@ class StatusGuardActor(xo.StatelessActor):
             else self._drop_terminated_info(all_infos)
         )
 
+    def get_instance_count(self, model_name: str) -> int:
+        return len(self.get_instance_info(model_name=model_name))
+
     def update_instance_info(self, model_uid: str, info: Dict):
         self._model_uid_to_info[model_uid].update(**info)