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

xinference/core/supervisor.py

@@ -21,10 +21,16 @@ from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Un
 
 import xoscar as xo
 
+from ..constants import (
+    XINFERENCE_DISABLE_HEALTH_CHECK,
+    XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD,
+    XINFERENCE_HEALTH_CHECK_INTERVAL,
+    XINFERENCE_HEALTH_CHECK_TIMEOUT,
+)
 from ..core import ModelActor
 from ..core.status_guard import InstanceInfo, LaunchStatus
 from .metrics import record_metrics
-from .resource import ResourceStatus
+from .resource import GPUStatus, ResourceStatus
 from .utils import (
     build_replica_model_uid,
     gen_random_string,
@@ -32,14 +38,15 @@ from .utils import (
     iter_replica_model_uid,
     log_async,
     log_sync,
+    parse_model_version,
     parse_replica_model_uid,
 )
 
 if TYPE_CHECKING:
+    from ..model.audio import AudioModelFamilyV1
     from ..model.embedding import EmbeddingModelSpec
     from ..model.image import ImageModelFamilyV1
     from ..model.llm import LLMFamilyV1
-    from ..model.multimodal import LVLMFamilyV1
     from ..model.rerank import RerankModelSpec
     from .worker import WorkerActor
 
@@ -47,7 +54,6 @@ if TYPE_CHECKING:
 logger = getLogger(__name__)
 
 
-DEFAULT_NODE_TIMEOUT = 60
 ASYNC_LAUNCH_TASKS = {}  # type: ignore
 
 
@@ -59,7 +65,8 @@ def callback_for_async_launch(model_uid: str):
 @dataclass
 class WorkerStatus:
     update_time: float
-    status: Dict[str, ResourceStatus]
+    failure_remaining_count: int
+    status: Dict[str, Union[ResourceStatus, GPUStatus]]
 
 
 @dataclass
@@ -86,9 +93,17 @@ class SupervisorActor(xo.StatelessActor):
 
     async def __post_create__(self):
         self._uptime = time.time()
-        # comment this line to avoid worker lost
-        # self._check_dead_nodes_task = asyncio.create_task(self._check_dead_nodes())
+        if not XINFERENCE_DISABLE_HEALTH_CHECK:
+            # Run _check_dead_nodes() in a dedicated thread.
+            from ..isolation import Isolation
+
+            self._isolation = Isolation(asyncio.new_event_loop(), threaded=True)
+            self._isolation.start()
+            asyncio.run_coroutine_threadsafe(
+                self._check_dead_nodes(), loop=self._isolation.loop
+            )
         logger.info(f"Xinference supervisor {self.address} started")
+        from .cache_tracker import CacheTrackerActor
         from .status_guard import StatusGuardActor
 
         self._status_guard_ref: xo.ActorRefType[
@@ -96,30 +111,98 @@ class SupervisorActor(xo.StatelessActor):
         ] = await xo.create_actor(
             StatusGuardActor, address=self.address, uid=StatusGuardActor.uid()
         )
+        self._cache_tracker_ref: xo.ActorRefType[
+            "CacheTrackerActor"
+        ] = await xo.create_actor(
+            CacheTrackerActor, address=self.address, uid=CacheTrackerActor.uid()
+        )
+
+        from .event import EventCollectorActor
+
+        self._event_collector_ref: xo.ActorRefType[
+            EventCollectorActor
+        ] = await xo.create_actor(
+            EventCollectorActor, address=self.address, uid=EventCollectorActor.uid()
+        )
 
         from ..model.embedding import (
             CustomEmbeddingModelSpec,
+            generate_embedding_description,
+            get_embedding_model_descriptions,
             register_embedding,
             unregister_embedding,
         )
-        from ..model.llm import register_llm, unregister_llm
-        from ..model.llm.llm_family import CustomLLMFamilyV1
-        from ..model.rerank.custom import (
+        from ..model.image import get_image_model_descriptions
+        from ..model.llm import (
+            CustomLLMFamilyV1,
+            generate_llm_description,
+            get_llm_model_descriptions,
+            register_llm,
+            unregister_llm,
+        )
+        from ..model.rerank import (
             CustomRerankModelSpec,
+            generate_rerank_description,
+            get_rerank_model_descriptions,
             register_rerank,
             unregister_rerank,
         )
 
         self._custom_register_type_to_cls: Dict[str, Tuple] = {
-            "LLM": (CustomLLMFamilyV1, register_llm, unregister_llm),
+            "LLM": (
+                CustomLLMFamilyV1,
+                register_llm,
+                unregister_llm,
+                generate_llm_description,
+            ),
             "embedding": (
                 CustomEmbeddingModelSpec,
                 register_embedding,
                 unregister_embedding,
+                generate_embedding_description,
+            ),
+            "rerank": (
+                CustomRerankModelSpec,
+                register_rerank,
+                unregister_rerank,
+                generate_rerank_description,
             ),
-            "rerank": (CustomRerankModelSpec, register_rerank, unregister_rerank),
         }
 
+        # record model version
+        model_version_infos: Dict[str, List[Dict]] = {}
+        model_version_infos.update(get_llm_model_descriptions())
+        model_version_infos.update(get_embedding_model_descriptions())
+        model_version_infos.update(get_rerank_model_descriptions())
+        model_version_infos.update(get_image_model_descriptions())
+        await self._cache_tracker_ref.record_model_version(
+            model_version_infos, self.address
+        )
+
+    async def get_cluster_device_info(self) -> List:
+        supervisor_device_info = {
+            "ip_address": self.address.split(":")[0],
+            "gpu_count": 0,
+            "gpu_vram_total": 0,
+        }
+        res = [{"node_type": "Supervisor", **supervisor_device_info}]
+        for worker_addr, worker_status in self._worker_status.items():
+            vram_total: float = sum(
+                [v.mem_total for k, v in worker_status.status.items() if k != "cpu"]  # type: ignore
+            )
+            total = (
+                vram_total if vram_total == 0 else f"{int(vram_total / 1024 / 1024)}MiB"
+            )
+            res.append(
+                {
+                    "node_type": "Worker",
+                    "ip_address": worker_addr.split(":")[0],
+                    "gpu_count": len(worker_status.status) - 1,
+                    "gpu_vram_total": total,
+                }
+            )
+        return res
+
     @staticmethod
     async def get_builtin_prompts() -> Dict[str, Any]:
         from ..model.llm.llm_family import BUILTIN_LLM_PROMPT_STYLE
@@ -180,99 +263,129 @@ class SupervisorActor(xo.StatelessActor):
             "workers": self._worker_status,
         }
 
-    def _to_llm_reg(
+    async def _to_llm_reg(
         self, llm_family: "LLMFamilyV1", is_builtin: bool
     ) -> Dict[str, Any]:
         from ..model.llm import get_cache_status
 
+        instance_cnt = await self.get_instance_count(llm_family.model_name)
+        version_cnt = await self.get_model_version_count(llm_family.model_name)
+
         if self.is_local_deployment():
             specs = []
             # TODO: does not work when the supervisor and worker are running on separate nodes.
             for spec in llm_family.model_specs:
                 cache_status = get_cache_status(llm_family, spec)
                 specs.append({**spec.dict(), "cache_status": cache_status})
-            return {**llm_family.dict(), "is_builtin": is_builtin, "model_specs": specs}
+            res = {**llm_family.dict(), "is_builtin": is_builtin, "model_specs": specs}
         else:
-            return {**llm_family.dict(), "is_builtin": is_builtin}
+            res = {**llm_family.dict(), "is_builtin": is_builtin}
+        res["model_version_count"] = version_cnt
+        res["model_instance_count"] = instance_cnt
+        return res
 
-    def _to_embedding_model_reg(
+    async def _to_embedding_model_reg(
         self, model_spec: "EmbeddingModelSpec", is_builtin: bool
     ) -> Dict[str, Any]:
         from ..model.embedding import get_cache_status
 
+        instance_cnt = await self.get_instance_count(model_spec.model_name)
+        version_cnt = await self.get_model_version_count(model_spec.model_name)
+
         if self.is_local_deployment():
             # TODO: does not work when the supervisor and worker are running on separate nodes.
             cache_status = get_cache_status(model_spec)
-            return {
+            res = {
                 **model_spec.dict(),
                 "cache_status": cache_status,
                 "is_builtin": is_builtin,
             }
         else:
-            return {
+            res = {
                 **model_spec.dict(),
                 "is_builtin": is_builtin,
             }
+        res["model_version_count"] = version_cnt
+        res["model_instance_count"] = instance_cnt
+        return res
 
-    def _to_rerank_model_reg(
+    async def _to_rerank_model_reg(
         self, model_spec: "RerankModelSpec", is_builtin: bool
     ) -> Dict[str, Any]:
         from ..model.rerank import get_cache_status
 
+        instance_cnt = await self.get_instance_count(model_spec.model_name)
+        version_cnt = await self.get_model_version_count(model_spec.model_name)
+
         if self.is_local_deployment():
             # TODO: does not work when the supervisor and worker are running on separate nodes.
             cache_status = get_cache_status(model_spec)
-            return {
+            res = {
                 **model_spec.dict(),
                 "cache_status": cache_status,
                 "is_builtin": is_builtin,
             }
         else:
-            return {
+            res = {
                 **model_spec.dict(),
                 "is_builtin": is_builtin,
             }
+        res["model_version_count"] = version_cnt
+        res["model_instance_count"] = instance_cnt
+        return res
 
-    def _to_image_model_reg(
+    async def _to_image_model_reg(
         self, model_family: "ImageModelFamilyV1", is_builtin: bool
     ) -> Dict[str, Any]:
         from ..model.image import get_cache_status
 
+        instance_cnt = await self.get_instance_count(model_family.model_name)
+        version_cnt = await self.get_model_version_count(model_family.model_name)
+
         if self.is_local_deployment():
             # TODO: does not work when the supervisor and worker are running on separate nodes.
             cache_status = get_cache_status(model_family)
-            return {
+            res = {
                 **model_family.dict(),
                 "cache_status": cache_status,
                 "is_builtin": is_builtin,
             }
         else:
-            return {
+            res = {
                 **model_family.dict(),
                 "is_builtin": is_builtin,
             }
+        res["model_version_count"] = version_cnt
+        res["model_instance_count"] = instance_cnt
+        return res
 
-    def _to_multimodal_reg(
-        self, model_family: "LVLMFamilyV1", is_builtin: bool
+    async def _to_audio_model_reg(
+        self, model_family: "AudioModelFamilyV1", is_builtin: bool
     ) -> Dict[str, Any]:
-        from ..model.llm import get_cache_status
+        from ..model.audio import get_cache_status
+
+        instance_cnt = await self.get_instance_count(model_family.model_name)
+        version_cnt = await self.get_model_version_count(model_family.model_name)
 
         if self.is_local_deployment():
-            specs = []
             # TODO: does not work when the supervisor and worker are running on separate nodes.
-            for spec in model_family.model_specs:
-                cache_status = get_cache_status(model_family, spec)
-                specs.append({**spec.dict(), "cache_status": cache_status})
-            return {
+            cache_status = get_cache_status(model_family)
+            res = {
                 **model_family.dict(),
+                "cache_status": cache_status,
                 "is_builtin": is_builtin,
-                "model_specs": specs,
             }
         else:
-            return {**model_family.dict(), "is_builtin": is_builtin}
+            res = {
+                **model_family.dict(),
+                "is_builtin": is_builtin,
+            }
+        res["model_version_count"] = version_cnt
+        res["model_instance_count"] = instance_cnt
+        return res
 
-    @log_sync(logger=logger)
-    def list_model_registrations(
+    @log_async(logger=logger)
+    async def list_model_registrations(
         self, model_type: str, detailed: bool = False
     ) -> List[Dict[str, Any]]:
         def sort_helper(item):
@@ -285,13 +398,13 @@ class SupervisorActor(xo.StatelessActor):
             ret = []
             for family in BUILTIN_LLM_FAMILIES:
                 if detailed:
-                    ret.append(self._to_llm_reg(family, True))
+                    ret.append(await self._to_llm_reg(family, True))
                 else:
                     ret.append({"model_name": family.model_name, "is_builtin": True})
 
             for family in get_user_defined_llm_families():
                 if detailed:
-                    ret.append(self._to_llm_reg(family, False))
+                    ret.append(await self._to_llm_reg(family, False))
                 else:
                     ret.append({"model_name": family.model_name, "is_builtin": False})
 
@@ -304,14 +417,16 @@ class SupervisorActor(xo.StatelessActor):
             ret = []
             for model_name, family in BUILTIN_EMBEDDING_MODELS.items():
                 if detailed:
-                    ret.append(self._to_embedding_model_reg(family, is_builtin=True))
+                    ret.append(
+                        await self._to_embedding_model_reg(family, is_builtin=True)
+                    )
                 else:
                     ret.append({"model_name": model_name, "is_builtin": True})
 
             for model_spec in get_user_defined_embeddings():
                 if detailed:
                     ret.append(
-                        self._to_embedding_model_reg(model_spec, is_builtin=False)
+                        await self._to_embedding_model_reg(model_spec, is_builtin=False)
                     )
                 else:
                     ret.append(
@@ -326,7 +441,19 @@ class SupervisorActor(xo.StatelessActor):
             ret = []
             for model_name, family in BUILTIN_IMAGE_MODELS.items():
                 if detailed:
-                    ret.append(self._to_image_model_reg(family, is_builtin=True))
+                    ret.append(await self._to_image_model_reg(family, is_builtin=True))
+                else:
+                    ret.append({"model_name": model_name, "is_builtin": True})
+
+            ret.sort(key=sort_helper)
+            return ret
+        elif model_type == "audio":
+            from ..model.audio import BUILTIN_AUDIO_MODELS
+
+            ret = []
+            for model_name, family in BUILTIN_AUDIO_MODELS.items():
+                if detailed:
+                    ret.append(await self._to_audio_model_reg(family, is_builtin=True))
                 else:
                     ret.append({"model_name": model_name, "is_builtin": True})
 
@@ -339,30 +466,20 @@ class SupervisorActor(xo.StatelessActor):
             ret = []
             for model_name, family in BUILTIN_RERANK_MODELS.items():
                 if detailed:
-                    ret.append(self._to_rerank_model_reg(family, is_builtin=True))
+                    ret.append(await self._to_rerank_model_reg(family, is_builtin=True))
                 else:
                     ret.append({"model_name": model_name, "is_builtin": True})
 
             for model_spec in get_user_defined_reranks():
                 if detailed:
-                    ret.append(self._to_rerank_model_reg(model_spec, is_builtin=False))
+                    ret.append(
+                        await self._to_rerank_model_reg(model_spec, is_builtin=False)
+                    )
                 else:
                     ret.append(
                         {"model_name": model_spec.model_name, "is_builtin": False}
                     )
 
-            ret.sort(key=sort_helper)
-            return ret
-        elif model_type == "multimodal":
-            from ..model.multimodal import BUILTIN_LVLM_FAMILIES
-
-            ret = []
-            for family in BUILTIN_LVLM_FAMILIES:
-                if detailed:
-                    ret.append(self._to_multimodal_reg(family, True))
-                else:
-                    ret.append({"model_name": family.model_name, "is_builtin": True})
-
             ret.sort(key=sort_helper)
             return ret
         else:
@@ -395,18 +512,18 @@ class SupervisorActor(xo.StatelessActor):
                 if f.model_name == model_name:
                     return f
             raise ValueError(f"Model {model_name} not found")
-        elif model_type == "rerank":
-            from ..model.rerank import BUILTIN_RERANK_MODELS
-            from ..model.rerank.custom import get_user_defined_reranks
+        elif model_type == "audio":
+            from ..model.audio import BUILTIN_AUDIO_MODELS
 
-            for f in list(BUILTIN_RERANK_MODELS.values()) + get_user_defined_reranks():
+            for f in BUILTIN_AUDIO_MODELS.values():
                 if f.model_name == model_name:
                     return f
             raise ValueError(f"Model {model_name} not found")
-        elif model_type == "multimodal":
-            from ..model.multimodal import BUILTIN_LVLM_FAMILIES
+        elif model_type == "rerank":
+            from ..model.rerank import BUILTIN_RERANK_MODELS
+            from ..model.rerank.custom import get_user_defined_reranks
 
-            for f in BUILTIN_LVLM_FAMILIES:
+            for f in list(BUILTIN_RERANK_MODELS.values()) + get_user_defined_reranks():
                 if f.model_name == model_name:
                     return f
             raise ValueError(f"Model {model_name} not found")
@@ -420,6 +537,7 @@ class SupervisorActor(xo.StatelessActor):
                 model_spec_cls,
                 register_fn,
                 unregister_fn,
+                generate_fn,
             ) = self._custom_register_type_to_cls[model_type]
 
             if not self.is_local_deployment():
@@ -430,6 +548,9 @@ class SupervisorActor(xo.StatelessActor):
             model_spec = model_spec_cls.parse_raw(model)
             try:
                 register_fn(model_spec, persist)
+                await self._cache_tracker_ref.record_model_version(
+                    generate_fn(model_spec), self.address
+                )
             except Exception as e:
                 unregister_fn(model_spec.model_name, raise_error=False)
                 raise e
@@ -439,8 +560,9 @@ class SupervisorActor(xo.StatelessActor):
     @log_async(logger=logger)
     async def unregister_model(self, model_type: str, model_name: str):
         if model_type in self._custom_register_type_to_cls:
-            _, _, unregister_fn = self._custom_register_type_to_cls[model_type]
+            _, _, unregister_fn, _ = self._custom_register_type_to_cls[model_type]
             unregister_fn(model_name)
+            await self._cache_tracker_ref.unregister_model_version(model_name)
 
             if not self.is_local_deployment():
                 workers = list(self._worker_address_to_worker.values())
@@ -457,6 +579,43 @@ class SupervisorActor(xo.StatelessActor):
         )
         return f"{model_name}-{gen_random_string(8)}"
 
+    async def get_model_versions(self, model_type: str, model_name: str) -> List[Dict]:
+        return await self._cache_tracker_ref.get_model_versions(model_name)
+
+    async def get_model_version_count(self, model_name: str) -> int:
+        return await self._cache_tracker_ref.get_model_version_count(model_name)
+
+    @log_async(logger=logger)
+    async def launch_model_by_version(
+        self,
+        model_uid: Optional[str],
+        model_type: str,
+        model_version: str,
+        replica: int = 1,
+        n_gpu: Optional[Union[int, str]] = "auto",
+        wait_ready: bool = True,
+    ):
+        parse_results = parse_model_version(model_version, model_type)
+
+        if model_type == "image" and len(parse_results) == 2:
+            kwargs = {"controlnet": parse_results[1]}
+        else:
+            kwargs = {}
+
+        return await self.launch_builtin_model(
+            model_uid=model_uid,
+            model_name=parse_results[0],
+            model_size_in_billions=parse_results[1] if model_type == "LLM" else None,
+            model_format=parse_results[2] if model_type == "LLM" else None,
+            quantization=parse_results[3] if model_type == "LLM" else None,
+            model_type=model_type,
+            replica=replica,
+            n_gpu=n_gpu,
+            wait_ready=wait_ready,
+            model_version=model_version,
+            **kwargs,
+        )
+
     async def launch_speculative_llm(
         self,
         model_uid: Optional[str],
@@ -529,6 +688,7 @@ class SupervisorActor(xo.StatelessActor):
         n_gpu: Optional[Union[int, str]] = "auto",
         request_limits: Optional[int] = None,
         wait_ready: bool = True,
+        model_version: Optional[str] = None,
         **kwargs,
     ) -> str:
         if model_uid is None:
@@ -601,6 +761,7 @@ class SupervisorActor(xo.StatelessActor):
         instance_info = InstanceInfo(
             model_name=model_name,
             model_uid=model_uid,
+            model_version=model_version,
             model_ability=[],
             replica=replica,
             status=LaunchStatus.CREATING.name,
@@ -623,29 +784,53 @@ class SupervisorActor(xo.StatelessActor):
         )
         return [info.dict() for info in sorted(infos, key=lambda info: info.model_uid)]
 
+    async def get_instance_count(self, model_name: str) -> int:
+        return await self._status_guard_ref.get_instance_count(model_name)
+
     async def _check_dead_nodes(self):
         while True:
-            dead_nodes = []
-            for address, status in self._worker_status.items():
-                if time.time() - status.update_time > DEFAULT_NODE_TIMEOUT:
-                    dead_models = []
-                    for model_uid in self._replica_model_uid_to_worker:
-                        if (
-                            self._replica_model_uid_to_worker[model_uid].address
-                            == address
-                        ):
-                            dead_models.append(model_uid)
-                    logger.error(
-                        "Worker timeout. address: %s, influenced models: %s",
-                        address,
-                        dead_models,
-                    )
-                    dead_nodes.append(address)
-
-            for address in dead_nodes:
-                self._worker_status.pop(address)
-                self._worker_address_to_worker.pop(address)
-            await asyncio.sleep(5)
+            try:
+                dead_nodes = []
+                for address, status in self._worker_status.items():
+                    if (
+                        time.time() - status.update_time
+                        > XINFERENCE_HEALTH_CHECK_TIMEOUT
+                    ):
+                        status.failure_remaining_count -= 1
+                    else:
+                        status.failure_remaining_count = (
+                            XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD
+                        )
+
+                    if status.failure_remaining_count <= 0:
+                        dead_models = []
+                        for model_uid in self._replica_model_uid_to_worker:
+                            if (
+                                self._replica_model_uid_to_worker[model_uid].address
+                                == address
+                            ):
+                                dead_models.append(model_uid)
+                        logger.error(
+                            "Worker dead. address: %s, influenced models: %s",
+                            address,
+                            dead_models,
+                        )
+                        dead_nodes.append(address)
+                    elif (
+                        status.failure_remaining_count
+                        != XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD
+                    ):
+                        logger.error(
+                            "Worker timeout. address: %s, check count remaining %s...",
+                            address,
+                            status.failure_remaining_count,
+                        )
+
+                for address in dead_nodes:
+                    self._worker_status.pop(address, None)
+                    self._worker_address_to_worker.pop(address, None)
+            finally:
+                await asyncio.sleep(XINFERENCE_HEALTH_CHECK_INTERVAL)
 
     @log_async(logger=logger)
     async def terminate_model(self, model_uid: str, suppress_exception=False):
@@ -744,13 +929,19 @@ class SupervisorActor(xo.StatelessActor):
             )
 
     async def report_worker_status(
-        self, worker_address: str, status: Dict[str, ResourceStatus]
+        self, worker_address: str, status: Dict[str, Union[ResourceStatus, GPUStatus]]
     ):
         if worker_address not in self._worker_status:
             logger.debug("Worker %s resources: %s", worker_address, status)
-        self._worker_status[worker_address] = WorkerStatus(
-            update_time=time.time(), status=status
-        )
+            self._worker_status[worker_address] = WorkerStatus(
+                update_time=time.time(),
+                failure_remaining_count=XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD,
+                status=status,
+            )
+        else:
+            worker_status = self._worker_status[worker_address]
+            worker_status.update_time = time.time()
+            worker_status.status = status
 
     @staticmethod
     def record_metrics(name, op, kwargs):