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

xinference/core/worker.py

@@ -15,7 +15,9 @@
 import asyncio
 import os
 import platform
+import queue
 import signal
+import threading
 from collections import defaultdict
 from logging import getLogger
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
@@ -25,8 +27,10 @@ from xoscar import MainActorPoolType
 
 from ..constants import XINFERENCE_CACHE_DIR
 from ..core import ModelActor
+from ..core.status_guard import LaunchStatus
 from ..model.core import ModelDescription, create_model_instance
 from ..utils import cuda_count
+from .metrics import launch_metrics_export_server, record_metrics
 from .resource import gather_node_info
 from .utils import log_async, log_sync, parse_replica_model_uid, purge_dir
 
@@ -34,6 +38,12 @@ logger = getLogger(__name__)
 
 
 DEFAULT_NODE_HEARTBEAT_INTERVAL = 5
+MODEL_ACTOR_AUTO_RECOVER_LIMIT: Optional[int]
+_MODEL_ACTOR_AUTO_RECOVER_LIMIT = os.getenv("XINFERENCE_MODEL_ACTOR_AUTO_RECOVER_LIMIT")
+if _MODEL_ACTOR_AUTO_RECOVER_LIMIT is not None:
+    MODEL_ACTOR_AUTO_RECOVER_LIMIT = int(_MODEL_ACTOR_AUTO_RECOVER_LIMIT)
+else:
+    MODEL_ACTOR_AUTO_RECOVER_LIMIT = None
 
 
 class WorkerActor(xo.StatelessActor):
@@ -42,6 +52,8 @@ class WorkerActor(xo.StatelessActor):
         supervisor_address: str,
         main_pool: MainActorPoolType,
         cuda_devices: List[int],
+        metrics_exporter_host: Optional[str] = None,
+        metrics_exporter_port: Optional[int] = None,
     ):
         super().__init__()
         # static attrs.
@@ -57,20 +69,71 @@ class WorkerActor(xo.StatelessActor):
         self._gpu_to_model_uid: Dict[int, str] = {}
         self._gpu_to_embedding_model_uids: Dict[int, Set[str]] = defaultdict(set)
         self._model_uid_to_addr: Dict[str, str] = {}
+        self._model_uid_to_recover_count: Dict[str, int] = {}
         self._model_uid_to_launch_args: Dict[str, Dict] = {}
 
+        # metrics export server.
+        if metrics_exporter_host is not None or metrics_exporter_port is not None:
+            logger.info(
+                f"Starting metrics export server at {metrics_exporter_host}:{metrics_exporter_port}"
+            )
+            q: queue.Queue = queue.Queue()
+            self._metrics_thread = threading.Thread(
+                name="Metrics Export Server",
+                target=launch_metrics_export_server,
+                args=(q, metrics_exporter_host, metrics_exporter_port),
+                daemon=True,
+            )
+            self._metrics_thread.start()
+            logger.info("Checking metrics export server...")
+            while self._metrics_thread.is_alive():
+                try:
+                    host, port = q.get(block=False)[:2]
+                    logger.info(f"Metrics server is started at: http://{host}:{port}")
+                    break
+                except queue.Empty:
+                    pass
+            else:
+                raise Exception("Metrics server thread exit.")
+
         self._lock = asyncio.Lock()
 
     async def recover_sub_pool(self, address):
-        logger.warning("Process %s is down, create model.", address)
+        logger.warning("Process %s is down.", address)
+        # Xoscar does not remove the address from sub_processes.
+        try:
+            await self._main_pool.remove_sub_pool(address)
+        except Exception:
+            pass
         for model_uid, addr in self._model_uid_to_addr.items():
             if addr == address:
                 launch_args = self._model_uid_to_launch_args.get(model_uid)
-                try:
-                    await self.terminate_model(model_uid)
-                except Exception:
-                    pass
-                await self.launch_builtin_model(**launch_args)
+                if launch_args is None:
+                    logger.warning(
+                        "Not recreate model because the it is down during launch."
+                    )
+                else:
+                    recover_count = self._model_uid_to_recover_count.get(model_uid)
+                    try:
+                        await self.terminate_model(model_uid)
+                    except Exception:
+                        pass
+                    if recover_count is not None:
+                        if recover_count > 0:
+                            logger.warning(
+                                "Recreating model actor %s, remain %s times ...",
+                                model_uid,
+                                recover_count - 1,
+                            )
+                            self._model_uid_to_recover_count[model_uid] = (
+                                recover_count - 1
+                            )
+                            await self.launch_builtin_model(**launch_args)
+                        else:
+                            logger.warning("Stop recreating model actor.")
+                    else:
+                        logger.warning("Recreating model actor %s ...", model_uid)
+                        await self.launch_builtin_model(**launch_args)
                 break
 
     @classmethod
@@ -78,8 +141,14 @@ class WorkerActor(xo.StatelessActor):
         return "worker"
 
     async def __post_create__(self):
+        from .status_guard import StatusGuardActor
         from .supervisor import SupervisorActor
 
+        self._status_guard_ref: xo.ActorRefType[
+            "StatusGuardActor"
+        ] = await xo.actor_ref(
+            address=self._supervisor_address, uid=StatusGuardActor.uid()
+        )
         self._supervisor_ref: xo.ActorRefType["SupervisorActor"] = await xo.actor_ref(
             address=self._supervisor_address, uid=SupervisorActor.uid()
         )
@@ -309,7 +378,12 @@ class WorkerActor(xo.StatelessActor):
 
         try:
             model_ref = await xo.create_actor(
-                ModelActor, address=subpool_address, uid=model_uid, model=model
+                ModelActor,
+                address=subpool_address,
+                uid=model_uid,
+                worker_address=self.address,
+                model=model,
+                model_description=model_description,
             )
             await model_ref.load()
         except:
@@ -324,6 +398,22 @@ class WorkerActor(xo.StatelessActor):
             self._gpu_to_model_uid[int(dev)] = model_uid
         self._model_uid_to_addr[model_uid] = subpool_address
 
+    async def _get_model_ability(self, model: Any, model_type: str) -> List[str]:
+        from ..model.llm.core import LLM
+
+        if model_type == "embedding":
+            return ["embed"]
+        elif model_type == "rerank":
+            return ["rerank"]
+        elif model_type == "image":
+            return ["text_to_image"]
+        elif model_type == "multimodal":
+            return ["multimodal"]
+        else:
+            assert model_type == "LLM"
+            assert isinstance(model, LLM)
+            return model.model_family.model_ability  # type: ignore
+
     @log_async(logger=logger)
     async def launch_builtin_model(
         self,
@@ -360,6 +450,7 @@ class WorkerActor(xo.StatelessActor):
         )
 
         try:
+            origin_uid, _, _ = parse_replica_model_uid(model_uid)
             model, model_description = await asyncio.to_thread(
                 create_model_instance,
                 subpool_address,
@@ -377,7 +468,9 @@ class WorkerActor(xo.StatelessActor):
                 ModelActor,
                 address=subpool_address,
                 uid=model_uid,
+                worker_address=self.address,
                 model=model,
+                model_description=model_description,
                 request_limits=request_limits,
             )
             await model_ref.load()
@@ -390,13 +483,27 @@ class WorkerActor(xo.StatelessActor):
         self._model_uid_to_model[model_uid] = model_ref
         self._model_uid_to_model_spec[model_uid] = model_description
         self._model_uid_to_addr[model_uid] = subpool_address
+        self._model_uid_to_recover_count.setdefault(
+            model_uid, MODEL_ACTOR_AUTO_RECOVER_LIMIT
+        )
         self._model_uid_to_launch_args[model_uid] = launch_args
 
+        # update status to READY
+        abilities = await self._get_model_ability(model, model_type)
+        await self._status_guard_ref.update_instance_info(
+            origin_uid,
+            {"model_ability": abilities, "status": LaunchStatus.READY.name},
+        )
+
     @log_async(logger=logger)
     async def terminate_model(self, model_uid: str):
+        origin_uid, _, _ = parse_replica_model_uid(model_uid)
+        await self._status_guard_ref.update_instance_info(
+            origin_uid, {"status": LaunchStatus.TERMINATING.name}
+        )
         model_ref = self._model_uid_to_model.get(model_uid, None)
         if model_ref is None:
-            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
+            logger.debug("Model not found, uid: %s", model_uid)
 
         try:
             await xo.destroy_actor(model_ref)
@@ -407,12 +514,20 @@ class WorkerActor(xo.StatelessActor):
         try:
             subpool_address = self._model_uid_to_addr[model_uid]
             await self._main_pool.remove_sub_pool(subpool_address)
+        except Exception as e:
+            logger.debug(
+                "Remove sub pool failed, model uid: %s, error: %s", model_uid, e
+            )
         finally:
-            del self._model_uid_to_model[model_uid]
-            del self._model_uid_to_model_spec[model_uid]
+            self._model_uid_to_model.pop(model_uid, None)
+            self._model_uid_to_model_spec.pop(model_uid, None)
             self.release_devices(model_uid)
-            del self._model_uid_to_addr[model_uid]
-            del self._model_uid_to_launch_args[model_uid]
+            self._model_uid_to_addr.pop(model_uid, None)
+            self._model_uid_to_recover_count.pop(model_uid, None)
+            self._model_uid_to_launch_args.pop(model_uid, None)
+            await self._status_guard_ref.update_instance_info(
+                origin_uid, {"status": LaunchStatus.TERMINATED.name}
+            )
 
     @log_async(logger=logger)
     async def list_models(self) -> Dict[str, Dict[str, Any]]:
@@ -427,7 +542,7 @@ class WorkerActor(xo.StatelessActor):
     def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
         model_ref = self._model_uid_to_model.get(model_uid, None)
         if model_ref is None:
-            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
+            raise ValueError(f"Model not found, uid: {model_uid}")
         return model_ref
 
     @log_sync(logger=logger)
@@ -460,3 +575,7 @@ class WorkerActor(xo.StatelessActor):
                 await asyncio.sleep(DEFAULT_NODE_HEARTBEAT_INTERVAL)
             except asyncio.CancelledError:  # pragma: no cover
                 break
+
+    @staticmethod
+    def record_metrics(name, op, kwargs):
+        record_metrics(name, op, kwargs)

xinference/deploy/cmdline.py

@@ -87,7 +87,12 @@ def get_stored_token(
 
 
 def start_local_cluster(
-    log_level: str, host: str, port: int, auth_config_file: Optional[str] = None
+    log_level: str,
+    host: str,
+    port: int,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
+    auth_config_file: Optional[str] = None,
 ):
     from .local import main
 
@@ -102,6 +107,8 @@ def start_local_cluster(
     main(
         host=host,
         port=port,
+        metrics_exporter_host=metrics_exporter_host,
+        metrics_exporter_port=metrics_exporter_port,
         logging_conf=dict_config,
         auth_config_file=auth_config_file,
     )
@@ -182,14 +189,41 @@ def cli(
     type=int,
     help="Specify the port number for the Xinference server.",
 )
+@click.option(
+    "--metrics-exporter-host",
+    "-MH",
+    default=None,
+    type=str,
+    help="Specify the host address for the Xinference metrics exporter server, default is the same as --host.",
+)
+@click.option(
+    "--metrics-exporter-port",
+    "-mp",
+    type=int,
+    help="Specify the port number for the Xinference metrics exporter server.",
+)
 @click.option(
     "--auth-config",
     type=str,
     help="Specify the auth config json file.",
 )
-def local(log_level: str, host: str, port: int, auth_config: Optional[str]):
+def local(
+    log_level: str,
+    host: str,
+    port: int,
+    metrics_exporter_host: Optional[str],
+    metrics_exporter_port: Optional[int],
+    auth_config: Optional[str],
+):
+    if metrics_exporter_host is None:
+        metrics_exporter_host = host
     start_local_cluster(
-        log_level=log_level, host=host, port=port, auth_config_file=auth_config
+        log_level=log_level,
+        host=host,
+        port=port,
+        metrics_exporter_host=metrics_exporter_host,
+        metrics_exporter_port=metrics_exporter_port,
+        auth_config_file=auth_config,
     )
 
 
@@ -276,8 +310,25 @@ def supervisor(
     type=int,
     help="Specify the port number for the Xinference worker.",
 )
+@click.option(
+    "--metrics-exporter-host",
+    "-MH",
+    default=XINFERENCE_DEFAULT_DISTRIBUTED_HOST,
+    type=str,
+    help="Specify the host address for the metrics exporter server.",
+)
+@click.option(
+    "--metrics-exporter-port",
+    type=int,
+    help="Specify the port number for the Xinference metrics exporter worker.",
+)
 def worker(
-    log_level: str, endpoint: Optional[str], host: str, worker_port: Optional[int]
+    log_level: str,
+    endpoint: Optional[str],
+    host: str,
+    worker_port: Optional[int],
+    metrics_exporter_host: Optional[str],
+    metrics_exporter_port: Optional[int],
 ):
     from ..deploy.worker import main
 
@@ -298,6 +349,8 @@ def worker(
     main(
         address=address,
         supervisor_address=supervisor_internal_addr,
+        metrics_exporter_host=metrics_exporter_host,
+        metrics_exporter_port=metrics_exporter_port,
         logging_conf=dict_config,
     )
 

xinference/deploy/local.py

@@ -35,6 +35,8 @@ logger = logging.getLogger(__name__)
 
 async def _start_local_cluster(
     address: str,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
     logging_conf: Optional[Dict] = None,
 ):
     from .utils import create_worker_actor_pool
@@ -50,7 +52,11 @@ async def _start_local_cluster(
             SupervisorActor, address=address, uid=SupervisorActor.uid()
         )
         await start_worker_components(
-            address=address, supervisor_address=address, main_pool=pool
+            address=address,
+            supervisor_address=address,
+            main_pool=pool,
+            metrics_exporter_host=metrics_exporter_host,
+            metrics_exporter_port=metrics_exporter_port,
         )
         await pool.join()
     except asyncio.CancelledError:
@@ -58,7 +64,12 @@ async def _start_local_cluster(
             await pool.stop()
 
 
-def run(address: str, logging_conf: Optional[Dict] = None):
+def run(
+    address: str,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
+    logging_conf: Optional[Dict] = None,
+):
     def sigterm_handler(signum, frame):
         sys.exit(0)
 
@@ -66,15 +77,26 @@ def run(address: str, logging_conf: Optional[Dict] = None):
 
     loop = asyncio.get_event_loop()
     task = loop.create_task(
-        _start_local_cluster(address=address, logging_conf=logging_conf)
+        _start_local_cluster(
+            address=address,
+            metrics_exporter_host=metrics_exporter_host,
+            metrics_exporter_port=metrics_exporter_port,
+            logging_conf=logging_conf,
+        )
     )
     loop.run_until_complete(task)
 
 
 def run_in_subprocess(
-    address: str, logging_conf: Optional[Dict] = None
+    address: str,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
+    logging_conf: Optional[Dict] = None,
 ) -> multiprocessing.Process:
-    p = multiprocessing.Process(target=run, args=(address, logging_conf))
+    p = multiprocessing.Process(
+        target=run,
+        args=(address, metrics_exporter_host, metrics_exporter_port, logging_conf),
+    )
     p.start()
     return p
 
@@ -82,11 +104,15 @@ def run_in_subprocess(
 def main(
     host: str,
     port: int,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
     logging_conf: Optional[Dict] = None,
     auth_config_file: Optional[str] = None,
 ):
     supervisor_address = f"{host}:{get_next_port()}"
-    local_cluster = run_in_subprocess(supervisor_address, logging_conf)
+    local_cluster = run_in_subprocess(
+        supervisor_address, metrics_exporter_host, metrics_exporter_port, logging_conf
+    )
 
     if not health_check(
         address=supervisor_address,

xinference/deploy/worker.py

@@ -27,7 +27,11 @@ logger = logging.getLogger(__name__)
 
 
 async def start_worker_components(
-    address: str, supervisor_address: str, main_pool: MainActorPoolType
+    address: str,
+    supervisor_address: str,
+    main_pool: MainActorPoolType,
+    metrics_exporter_host: Optional[str],
+    metrics_exporter_port: Optional[int],
 ):
     cuda_device_indices = []
     cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
@@ -43,24 +47,48 @@ async def start_worker_components(
         supervisor_address=supervisor_address,
         main_pool=main_pool,
         cuda_devices=cuda_device_indices,
+        metrics_exporter_host=metrics_exporter_host,
+        metrics_exporter_port=metrics_exporter_port,
     )
 
 
 async def _start_worker(
-    address: str, supervisor_address: str, logging_conf: Any = None
+    address: str,
+    supervisor_address: str,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
+    logging_conf: Any = None,
 ):
     from .utils import create_worker_actor_pool
 
     pool = await create_worker_actor_pool(address=address, logging_conf=logging_conf)
     await start_worker_components(
-        address=address, supervisor_address=supervisor_address, main_pool=pool
+        address=address,
+        supervisor_address=supervisor_address,
+        main_pool=pool,
+        metrics_exporter_host=metrics_exporter_host,
+        metrics_exporter_port=metrics_exporter_port,
     )
     await pool.join()
 
 
-def main(address: str, supervisor_address: str, logging_conf: Optional[dict] = None):
+def main(
+    address: str,
+    supervisor_address: str,
+    metrics_exporter_host: Optional[str] = None,
+    metrics_exporter_port: Optional[int] = None,
+    logging_conf: Optional[dict] = None,
+):
     loop = asyncio.get_event_loop()
-    task = loop.create_task(_start_worker(address, supervisor_address, logging_conf))
+    task = loop.create_task(
+        _start_worker(
+            address,
+            supervisor_address,
+            metrics_exporter_host,
+            metrics_exporter_port,
+            logging_conf,
+        )
+    )
 
     try:
         loop.run_until_complete(task)

xinference/fields.py

@@ -30,7 +30,10 @@ logprobs_field = Field(
 )
 
 max_tokens_field = Field(
-    default=128, ge=1, le=32768, description="The maximum number of tokens to generate."
+    default=1024,
+    ge=1,
+    le=32768,
+    description="The maximum number of tokens to generate.",
 )
 
 temperature_field = Field(

xinference/model/llm/__init__.py

@@ -21,6 +21,7 @@ from .llm_family import (
     BUILTIN_LLM_FAMILIES,
     BUILTIN_LLM_MODEL_CHAT_FAMILIES,
     BUILTIN_LLM_MODEL_GENERATE_FAMILIES,
+    BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES,
     BUILTIN_LLM_PROMPT_STYLE,
     BUILTIN_MODELSCOPE_LLM_FAMILIES,
     LLM_CLASSES,
@@ -47,6 +48,7 @@ def _install():
     from .pytorch.chatglm import ChatglmPytorchChatModel
     from .pytorch.core import PytorchChatModel, PytorchModel
     from .pytorch.falcon import FalconPytorchChatModel, FalconPytorchModel
+    from .pytorch.internlm2 import Internlm2PytorchChatModel
     from .pytorch.llama_2 import LlamaPytorchChatModel, LlamaPytorchModel
     from .pytorch.vicuna import VicunaPytorchChatModel
     from .vllm.core import VLLMChatModel, VLLMModel
@@ -79,6 +81,7 @@ def _install():
             LlamaPytorchChatModel,
             PytorchChatModel,
             FalconPytorchModel,
+            Internlm2PytorchChatModel,
             PytorchModel,
         ]
     )
@@ -102,6 +105,8 @@ def _install():
             BUILTIN_LLM_MODEL_CHAT_FAMILIES.add(model_spec.model_name)
         else:
             BUILTIN_LLM_MODEL_GENERATE_FAMILIES.add(model_spec.model_name)
+        if "tool_call" in model_spec.model_ability:
+            BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES.add(model_spec.model_name)
 
     modelscope_json_path = os.path.join(
         os.path.dirname(os.path.abspath(__file__)), "llm_family_modelscope.json"
@@ -123,6 +128,8 @@ def _install():
             BUILTIN_LLM_MODEL_CHAT_FAMILIES.add(model_spec.model_name)
         else:
             BUILTIN_LLM_MODEL_GENERATE_FAMILIES.add(model_spec.model_name)
+        if "tool_call" in model_spec.model_ability:
+            BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES.add(model_spec.model_name)
 
     from ...constants import XINFERENCE_MODEL_DIR
 

xinference/model/llm/ggml/llamacpp.py

@@ -306,7 +306,8 @@ class LlamaCppChatModel(LlamaCppModel, ChatModelMixin):
 
         generate_config = self._sanitize_generate_config(generate_config)
         # TODO(codingl2k1): qwen hacky to set stop for function call.
-        if tools and self.model_family.model_name == "qwen-chat":
+        model_family = self.model_family.model_family or self.model_family.model_name
+        if tools and "qwen-chat" == model_family:
             stop = generate_config.get("stop")
             if isinstance(stop, str):
                 generate_config["stop"] = [stop, "Observation:"]
@@ -326,6 +327,6 @@ class LlamaCppChatModel(LlamaCppModel, ChatModelMixin):
             assert not isinstance(c, Iterator)
             if tools:
                 return self._tool_calls_completion(
-                    self.model_family.model_name, self.model_uid, c, tools
+                    self.model_family, self.model_uid, c, tools
                 )
             return self._to_chat_completion(c)