xinference 1.0.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

xinference/core/supervisor.py

@@ -267,6 +267,14 @@ class SupervisorActor(xo.StatelessActor):
                 signal.SIGTERM, lambda: asyncio.create_task(signal_handler())
             )
 
+        from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
+        from ..model.llm.vllm.xavier.collective_manager import CollectiveManager
+
+        self._block_tracker_mapping: Dict[str, xo.ActorRefType[VLLMBlockTracker]] = {}
+        self._collective_manager_mapping: Dict[
+            str, xo.ActorRefType[CollectiveManager]
+        ] = {}
+
     @typing.no_type_check
     async def get_cluster_device_info(self, detailed: bool = False) -> List:
         import psutil
@@ -959,29 +967,83 @@ class SupervisorActor(xo.StatelessActor):
         if model_uid is None:
             model_uid = self._gen_model_uid(model_name)
 
+        # Xavier-related
+        enable_xavier: bool = (
+            bool(kwargs.pop("enable_xavier", False))
+            and model_engine is not None
+            and model_engine.lower() == "vllm"
+        )
+        store_address = None
+        store_port = None
+        world_size = None
+        if enable_xavier:
+            if replica <= 1:
+                logger.warning(f"Enabling xavier when `replica<=1` is meaningless.")
+                enable_xavier = False
+            else:
+                from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
+                from ..model.llm.vllm.xavier.collective_manager import CollectiveManager
+
+                self._block_tracker_mapping[model_uid] = await xo.create_actor(
+                    VLLMBlockTracker,
+                    address=self.address,
+                    uid=f"{VLLMBlockTracker.default_uid()}-{model_uid}",
+                )
+                world_size = replica + 1
+                logger.info(f"Going to start xavier with world size: {world_size}")
+                self._collective_manager_mapping[model_uid] = await xo.create_actor(
+                    CollectiveManager,
+                    address=self.address,
+                    uid=f"{CollectiveManager.default_uid()}-{model_uid}",
+                    model_uid=model_uid,
+                )
+                logger.info(f"Start collective manager for {model_uid} done.")
+
         model_size = str(model_size_in_billions) if model_size_in_billions else ""
         logger.debug(
             f"Enter launch_builtin_model, model_uid: {model_uid}, model_name: {model_name}, model_size: {model_size}, "
-            f"model_format: {model_format}, quantization: {quantization}, replica: {replica}, "
+            f"model_format: {model_format}, quantization: {quantization}, replica: {replica}, enable_xavier: {enable_xavier}, "
             f"kwargs: {kwargs}"
         )
 
-        async def _launch_one_model(_replica_model_uid):
+        async def _launch_one_model(worker_ref, _replica_model_uid, rank: int):
             if _replica_model_uid in self._replica_model_uid_to_worker:
                 raise ValueError(
                     f"Model is already in the model list, uid: {_replica_model_uid}"
                 )
+
+            nonlocal store_address
+            nonlocal store_port
+            xavier_config = (
+                {
+                    "block_tracker_uid": self._block_tracker_mapping[model_uid].uid,
+                    "block_tracker_address": self._block_tracker_mapping[
+                        model_uid
+                    ].address,
+                    "rank": rank,
+                    "world_size": world_size,
+                    "store_address": store_address,
+                    "store_port": store_port,
+                }
+                if enable_xavier
+                else None
+            )
+
+            if enable_xavier and rank == 0:
+                rank0_address, _port = await worker_ref.launch_rank0_model(
+                    _replica_model_uid, xavier_config
+                )
+                self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref
+                store_address = rank0_address.split(":")[0]
+                store_port = _port
+                return rank0_address
+
             replica_gpu_idx = assign_replica_gpu(_replica_model_uid, replica, gpu_idx)
             nonlocal model_type
 
-            worker_ref = (
-                target_ip_worker_ref
-                if target_ip_worker_ref is not None
-                else await self._choose_worker()
-            )
             # LLM as default for compatibility
             model_type = model_type or "LLM"
-            await worker_ref.launch_builtin_model(
+            subpool_address = await worker_ref.launch_builtin_model(
                 model_uid=_replica_model_uid,
                 model_name=model_name,
                 model_size_in_billions=model_size_in_billions,
@@ -995,14 +1057,64 @@ class SupervisorActor(xo.StatelessActor):
                 gpu_idx=replica_gpu_idx,
                 download_hub=download_hub,
                 model_path=model_path,
+                xavier_config=xavier_config,
                 **kwargs,
             )
             self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref
+            return subpool_address
 
         async def _launch_model():
             try:
-                for rep_model_uid in iter_replica_model_uid(model_uid, replica):
-                    await _launch_one_model(rep_model_uid)
+                worker_refs = []
+                rank_addresses = []
+                for _idx, rep_model_uid in enumerate(
+                    iter_replica_model_uid(model_uid, replica)
+                ):
+                    worker_ref = (
+                        target_ip_worker_ref
+                        if target_ip_worker_ref is not None
+                        else await self._choose_worker()
+                    )
+                    if enable_xavier and _idx == 0:
+                        """
+                        Start the rank 0 model actor on the worker that holds the rank 1 replica,
+                        solely for constructing the collective communication world.
+                        """
+                        _uid = model_uid + "-rank0"
+                        rank0_address = await _launch_one_model(worker_ref, _uid, 0)
+                        worker_refs.append((worker_ref, _uid))
+                        rank_addresses.append(rank0_address)
+
+                    subpool_address = await _launch_one_model(
+                        worker_ref, rep_model_uid, _idx + 1
+                    )
+                    worker_refs.append((worker_ref, rep_model_uid))
+                    rank_addresses.append(subpool_address)
+
+                # For xavier, start all the vllm instances first,
+                # and then start the transfer component,
+                # because the transfer actor needs all the rank addresses used for collective communication
+                if enable_xavier:
+                    logger.debug(f"Init transfer component for xavier...")
+                    collective_manager_ref = self._collective_manager_mapping[model_uid]
+                    tasks = []
+                    for worker_ref, rep_model_uid in worker_refs:
+                        tasks.append(
+                            worker_ref.start_transfer_for_vllm(
+                                rep_model_uid, rank_addresses
+                            )
+                        )
+                    # Here you must use asyncio.gather, not a for loop,
+                    # or you will get stuck.
+                    await asyncio.gather(*tasks)
+
+                    # init collective_manager
+                    for idx, addr in enumerate(rank_addresses):
+                        await collective_manager_ref.register_rank(
+                            idx, addr, update=False
+                        )
+
+                    logger.debug(f"Init transfer component for xavier done.")
             except Exception:
                 # terminate_model will remove the replica info.
                 await self.terminate_model(model_uid, suppress_exception=True)
@@ -1131,6 +1243,38 @@ class SupervisorActor(xo.StatelessActor):
                     raise
         self._model_uid_to_replica_info.pop(model_uid, None)
 
+        # clear for xavier
+        rank0_uid = model_uid + "-rank0"
+        if rank0_uid in self._replica_model_uid_to_worker:
+            await _terminate_one_model(rank0_uid)
+
+        collective_manager_ref = self._collective_manager_mapping.pop(model_uid, None)
+        if collective_manager_ref is not None:
+            try:
+                await xo.destroy_actor(collective_manager_ref)
+            except Exception as e:
+                logger.debug(
+                    "Destroy collective_manager_ref failed, model uid: %s, error: %s",
+                    model_uid,
+                    e,
+                )
+            finally:
+                logger.debug(
+                    f"Destroy collective_manager_ref done. model uid: {model_uid}"
+                )
+        block_tracker_ref = self._block_tracker_mapping.pop(model_uid, None)
+        if block_tracker_ref is not None:
+            try:
+                await xo.destroy_actor(block_tracker_ref)
+            except Exception as e:
+                logger.debug(
+                    "Destroy block_tracker_ref failed, model uid: %s, error: %s",
+                    model_uid,
+                    e,
+                )
+            finally:
+                logger.debug(f"Destroy block_tracker_ref done. model uid: {model_uid}")
+
     @log_async(logger=logger)
     async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
         replica_info = self._model_uid_to_replica_info.get(model_uid, None)
@@ -1148,6 +1292,15 @@ class SupervisorActor(xo.StatelessActor):
             )
         return await worker_ref.get_model(model_uid=replica_model_uid)
 
+    @log_async(logger=logger)
+    async def get_model_status(self, replica_model_uid: str):
+        worker_ref = self._replica_model_uid_to_worker.get(replica_model_uid, None)
+        if worker_ref is None:
+            raise ValueError(
+                f"Model not found in the model list, uid: {replica_model_uid}"
+            )
+        return await worker_ref.get_model_status(replica_model_uid)
+
     @log_async(logger=logger)
     async def describe_model(self, model_uid: str) -> Dict[str, Any]:
         replica_info = self._model_uid_to_replica_info.get(model_uid, None)
@@ -1377,3 +1530,12 @@ class SupervisorActor(xo.StatelessActor):
 
     async def get_progress(self, request_id: str) -> float:
         return await self._progress_tracker.get_progress(request_id)
+
+    async def call_collective_manager(
+        self, model_uid: str, func_name: str, *args, **kwargs
+    ):
+        """
+        Used by worker.
+        """
+        collective_manager_ref = self._collective_manager_mapping[model_uid]
+        await getattr(collective_manager_ref, func_name)(*args, **kwargs)

xinference/core/utils.py

@@ -62,12 +62,16 @@ def log_async(
 
         @wraps(func)
         async def wrapped(*args, **kwargs):
-            try:
-                bound_args = sig.bind_partial(*args, **kwargs)
-                arguments = bound_args.arguments
-            except TypeError:
-                arguments = {}
-            request_id_str = arguments.get("request_id", "")
+            request_id_str = kwargs.get("request_id")
+            if not request_id_str:
+                # sometimes `request_id` not in kwargs
+                # we try to bind the arguments
+                try:
+                    bound_args = sig.bind_partial(*args, **kwargs)
+                    arguments = bound_args.arguments
+                except TypeError:
+                    arguments = {}
+                request_id_str = arguments.get("request_id", "")
             if not request_id_str:
                 request_id_str = uuid.uuid1()
                 if func_name == "text_to_image":
@@ -272,8 +276,8 @@ def get_nvidia_gpu_info() -> Dict:
 
 
 def assign_replica_gpu(
-    _replica_model_uid: str, replica: int, gpu_idx: Union[int, List[int]]
-) -> List[int]:
+    _replica_model_uid: str, replica: int, gpu_idx: Optional[Union[int, List[int]]]
+) -> Optional[List[int]]:
     model_uid, rep_id = parse_replica_model_uid(_replica_model_uid)
     rep_id, replica = int(rep_id), int(replica)
     if isinstance(gpu_idx, int):

xinference/core/worker.py

@@ -22,8 +22,9 @@ import signal
 import threading
 import time
 from collections import defaultdict
+from dataclasses import dataclass
 from logging import getLogger
-from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, no_type_check
 
 import xoscar as xo
 from async_timeout import timeout
@@ -58,6 +59,11 @@ else:
     MODEL_ACTOR_AUTO_RECOVER_LIMIT = None
 
 
+@dataclass
+class ModelStatus:
+    last_error: str = ""
+
+
 class WorkerActor(xo.StatelessActor):
     def __init__(
         self,
@@ -90,6 +96,7 @@ class WorkerActor(xo.StatelessActor):
         # attributes maintained after model launched:
         self._model_uid_to_model: Dict[str, xo.ActorRefType["ModelActor"]] = {}
         self._model_uid_to_model_spec: Dict[str, ModelDescription] = {}
+        self._model_uid_to_model_status: Dict[str, ModelStatus] = {}
         self._gpu_to_model_uid: Dict[int, str] = {}
         self._gpu_to_embedding_model_uids: Dict[int, Set[str]] = defaultdict(set)
         # Dict structure: gpu_index: {(replica_model_uid, model_type)}
@@ -177,12 +184,12 @@ class WorkerActor(xo.StatelessActor):
                             self._model_uid_to_recover_count[model_uid] = (
                                 recover_count - 1
                             )
-                            await self.launch_builtin_model(**launch_args)
+                            await self.recover_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)
+                        await self.recover_model(launch_args)
                 break
 
     @classmethod
@@ -866,6 +873,9 @@ class WorkerActor(xo.StatelessActor):
             )
 
             try:
+                xavier_config: Optional[Dict] = kwargs.pop("xavier_config", None)
+                if xavier_config is not None:
+                    xavier_config["rank_address"] = subpool_address
                 model, model_description = await asyncio.to_thread(
                     create_model_instance,
                     subpool_address,
@@ -893,6 +903,7 @@ class WorkerActor(xo.StatelessActor):
                     model=model,
                     model_description=model_description,
                     request_limits=request_limits,
+                    xavier_config=xavier_config,
                 )
                 await model_ref.load()
             except:
@@ -902,6 +913,7 @@ class WorkerActor(xo.StatelessActor):
                 raise
             self._model_uid_to_model[model_uid] = model_ref
             self._model_uid_to_model_spec[model_uid] = model_description
+            self._model_uid_to_model_status[model_uid] = ModelStatus()
             self._model_uid_to_addr[model_uid] = subpool_address
             self._model_uid_to_recover_count.setdefault(
                 model_uid, MODEL_ACTOR_AUTO_RECOVER_LIMIT
@@ -921,13 +933,18 @@ class WorkerActor(xo.StatelessActor):
             origin_uid,
             {"model_ability": abilities, "status": LaunchStatus.READY.name},
         )
+        return subpool_address
 
     @log_async(logger=logger, level=logging.INFO)
     async def terminate_model(self, model_uid: str, is_model_die=False):
         # Terminate model while its launching is not allow
         if model_uid in self._model_uid_launching_guard:
             raise ValueError(f"{model_uid} is launching")
-        origin_uid, _ = parse_replica_model_uid(model_uid)
+        # In special cases, if the suffix is `-rank0`, this is the Xavier's rank 0 model actor.
+        if model_uid.endswith("-rank0"):
+            origin_uid = model_uid.removesuffix("-rank0")
+        else:
+            origin_uid, _ = parse_replica_model_uid(model_uid)
         try:
             _ = await self.get_supervisor_ref()
             if self._event_collector_ref is not None:
@@ -976,6 +993,7 @@ class WorkerActor(xo.StatelessActor):
                 status = LaunchStatus.ERROR.name
             else:
                 status = LaunchStatus.TERMINATED.name
+                self._model_uid_to_model_status.pop(model_uid, None)
 
             if self._status_guard_ref is None:
                 _ = await self.get_supervisor_ref()
@@ -1010,6 +1028,9 @@ class WorkerActor(xo.StatelessActor):
 
     @log_sync(logger=logger)
     def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
+        model_status = self._model_uid_to_model_status.get(model_uid)
+        if model_status and model_status.last_error:
+            raise Exception(model_status.last_error)
         model_ref = self._model_uid_to_model.get(model_uid, None)
         if model_ref is None:
             raise ValueError(f"Model not found, uid: {model_uid}")
@@ -1138,6 +1159,83 @@ class WorkerActor(xo.StatelessActor):
         }
         return ret
 
+    def update_model_status(self, model_uid: str, **kwargs):
+        model_status = self._model_uid_to_model_status.get(model_uid)
+        if model_status is not None:
+            for k, v in kwargs.items():
+                setattr(model_status, k, v)
+
+    def get_model_status(self, model_uid: str):
+        return self._model_uid_to_model_status.get(model_uid)
+
     @staticmethod
     def record_metrics(name, op, kwargs):
         record_metrics(name, op, kwargs)
+
+    async def start_transfer_for_vllm(
+        self, rep_model_uid: str, rank_addresses: List[str]
+    ):
+        model_ref = self._model_uid_to_model[rep_model_uid]
+        await model_ref.start_transfer_for_vllm(rank_addresses)
+
+    @log_async(logger=logger, level=logging.INFO)
+    async def launch_rank0_model(
+        self, rep_model_uid: str, xavier_config: Dict[str, Any]
+    ) -> Tuple[str, int]:
+        from ..model.llm.vllm.xavier.collective_manager import Rank0ModelActor
+
+        if os.name != "nt" and platform.system() != "Darwin":
+            # Linux
+            start_method = "forkserver"
+        else:
+            # Windows and macOS
+            start_method = "spawn"
+        subpool_address = await self._main_pool.append_sub_pool(
+            start_method=start_method
+        )
+
+        store_address = subpool_address.split(":")[0]
+        # Note that `store_port` needs to be generated on the worker,
+        # as the TCP store is on rank 0, not on the supervisor.
+        store_port = xo.utils.get_next_port()
+        self._model_uid_launching_guard[rep_model_uid] = True
+        try:
+            try:
+                xavier_config["rank_address"] = subpool_address
+                xavier_config["store_address"] = store_address
+                xavier_config["store_port"] = store_port
+                model_ref = await xo.create_actor(
+                    Rank0ModelActor,
+                    address=subpool_address,
+                    uid=rep_model_uid,
+                    xavier_config=xavier_config,
+                )
+            except:
+                await self._main_pool.remove_sub_pool(subpool_address)
+                raise
+            self._model_uid_to_model[rep_model_uid] = model_ref
+            self._model_uid_to_addr[rep_model_uid] = subpool_address
+        finally:
+            del self._model_uid_launching_guard[rep_model_uid]
+        return subpool_address, store_port
+
+    @no_type_check
+    async def recover_model(self, launch_args: Dict[str, Any]):
+        rep_model_uid = launch_args.get("model_uid")
+        origin_uid, _ = parse_replica_model_uid(rep_model_uid)
+        xavier_config: Optional[Dict[str, Any]] = launch_args.get("xavier_config", None)
+        is_xavier: bool = xavier_config is not None
+        supervisor_ref = await self.get_supervisor_ref(add_worker=False)
+        if is_xavier:
+            rank = xavier_config.get("rank")
+            await supervisor_ref.call_collective_manager(
+                origin_uid, "unregister_rank", rank
+            )
+        subpool_address = await self.launch_builtin_model(**launch_args)
+        if is_xavier:
+            model_ref = self._model_uid_to_model[rep_model_uid]
+            await model_ref.start_transfer_for_vllm([])
+            rank = xavier_config.get("rank")
+            await supervisor_ref.call_collective_manager(
+                origin_uid, "register_rank", rank, subpool_address, update=True
+            )

xinference/deploy/cmdline.py

@@ -846,7 +846,9 @@ def model_launch(
     kwargs = {}
     for i in range(0, len(ctx.args), 2):
         if not ctx.args[i].startswith("--"):
-            raise ValueError("You must specify extra kwargs with `--` prefix.")
+            raise ValueError(
+                f"You must specify extra kwargs with `--` prefix. There is an error in parameter passing that is {ctx.args[i]}."
+            )
         kwargs[ctx.args[i][2:]] = handle_click_args_type(ctx.args[i + 1])
     print(f"Launch model name: {model_name} with kwargs: {kwargs}", file=sys.stderr)
 

xinference/deploy/test/test_cmdline.py

@@ -23,6 +23,7 @@ from ..cmdline import (
     list_model_registrations,
     model_chat,
     model_generate,
+    model_launch,
     model_list,
     model_terminate,
     register_model,
@@ -311,3 +312,58 @@ def test_remove_cache(setup):
 
     assert result.exit_code == 0
     assert "Cache directory qwen1.5-chat has been deleted."
+
+
+def test_launch_error_in_passing_parameters():
+    runner = CliRunner()
+
+    # Known parameter but not provided with value.
+    result = runner.invoke(
+        model_launch,
+        [
+            "--model-engine",
+            "transformers",
+            "--model-name",
+            "qwen2.5-instruct",
+            "--model-uid",
+            "-s",
+            "0.5",
+            "-f",
+            "gptq",
+            "-q",
+            "INT4",
+            "111",
+            "-l",
+        ],
+    )
+    assert result.exit_code == 1
+    assert (
+        "You must specify extra kwargs with `--` prefix. There is an error in parameter passing that is 0.5."
+        in str(result)
+    )
+
+    # Unknown parameter
+    result = runner.invoke(
+        model_launch,
+        [
+            "--model-engine",
+            "transformers",
+            "--model-name",
+            "qwen2.5-instruct",
+            "--model-uid",
+            "123",
+            "-s",
+            "0.5",
+            "-f",
+            "gptq",
+            "-q",
+            "INT4",
+            "-l",
+            "111",
+        ],
+    )
+    assert result.exit_code == 1
+    assert (
+        "You must specify extra kwargs with `--` prefix. There is an error in parameter passing that is -l."
+        in str(result)
+    )

xinference/isolation.py

@@ -37,6 +37,30 @@ class Isolation:
         asyncio.set_event_loop(self._loop)
         self._stopped = asyncio.Event()
         self._loop.run_until_complete(self._stopped.wait())
+        self._cancel_all_tasks(self._loop)
+
+    @staticmethod
+    def _cancel_all_tasks(loop):
+        to_cancel = asyncio.all_tasks(loop)
+        if not to_cancel:
+            return
+
+        for task in to_cancel:
+            task.cancel()
+
+        loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
+
+        for task in to_cancel:
+            if task.cancelled():
+                continue
+            if task.exception() is not None:
+                loop.call_exception_handler(
+                    {
+                        "message": "unhandled exception during asyncio.run() shutdown",
+                        "exception": task.exception(),
+                        "task": task,
+                    }
+                )
 
     def start(self):
         if self._threaded:

xinference/model/audio/core.py

@@ -21,8 +21,11 @@ from ..core import CacheableModelSpec, ModelDescription
 from ..utils import valid_model_revision
 from .chattts import ChatTTSModel
 from .cosyvoice import CosyVoiceModel
+from .f5tts import F5TTSModel
+from .f5tts_mlx import F5TTSMLXModel
 from .fish_speech import FishSpeechModel
 from .funasr import FunASRModel
+from .melotts import MeloTTSModel
 from .whisper import WhisperModel
 from .whisper_mlx import WhisperMLXModel
 
@@ -46,6 +49,7 @@ class AudioModelFamilyV1(CacheableModelSpec):
     model_id: str
     model_revision: Optional[str]
     multilingual: bool
+    language: Optional[str]
     model_ability: Optional[str]
     default_model_config: Optional[Dict[str, Any]]
     default_transcription_config: Optional[Dict[str, Any]]
@@ -169,6 +173,9 @@ def create_audio_model_instance(
         ChatTTSModel,
         CosyVoiceModel,
         FishSpeechModel,
+        F5TTSModel,
+        F5TTSMLXModel,
+        MeloTTSModel,
     ],
     AudioModelDescription,
 ]:
@@ -182,6 +189,9 @@ def create_audio_model_instance(
         ChatTTSModel,
         CosyVoiceModel,
         FishSpeechModel,
+        F5TTSModel,
+        F5TTSMLXModel,
+        MeloTTSModel,
     ]
     if model_spec.model_family == "whisper":
         if not model_spec.engine:
@@ -196,6 +206,12 @@ def create_audio_model_instance(
         model = CosyVoiceModel(model_uid, model_path, model_spec, **kwargs)
     elif model_spec.model_family == "FishAudio":
         model = FishSpeechModel(model_uid, model_path, model_spec, **kwargs)
+    elif model_spec.model_family == "F5-TTS":
+        model = F5TTSModel(model_uid, model_path, model_spec, **kwargs)
+    elif model_spec.model_family == "F5-TTS-MLX":
+        model = F5TTSMLXModel(model_uid, model_path, model_spec, **kwargs)
+    elif model_spec.model_family == "MeloTTS":
+        model = MeloTTSModel(model_uid, model_path, model_spec, **kwargs)
     else:
         raise Exception(f"Unsupported audio model family: {model_spec.model_family}")
     model_description = AudioModelDescription(