xinference 1.2.2__py3-none-any.whl → 1.3.0__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-02-08T17:06:47+0800",
+ "date": "2025-02-21T22:33:35+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "ac97a13a831de6debda52e6fdb8c1bf9366be57c",
- "version": "1.2.2"
+ "full-revisionid": "bb778404b248ea91f4210e739e2b92b8a1be1455",
+ "version": "1.3.0"
 }
 '''  # END VERSION_JSON
 

xinference/client/restful/restful_client.py

@@ -917,11 +917,13 @@ class Client:
         model_format: Optional[str] = None,
         quantization: Optional[str] = None,
         replica: int = 1,
+        n_worker: int = 1,
         n_gpu: Optional[Union[int, str]] = "auto",
         peft_model_config: Optional[Dict] = None,
         request_limits: Optional[int] = None,
         worker_ip: Optional[str] = None,
         gpu_idx: Optional[Union[int, List[int]]] = None,
+        model_path: Optional[str] = None,
         **kwargs,
     ) -> str:
         """
@@ -945,8 +947,10 @@ class Client:
             The quantization of model.
         replica: Optional[int]
             The replica of model, default is 1.
+        n_worker: int
+            Number of workers to run.
         n_gpu: Optional[Union[int, str]],
-            The number of GPUs used by the model, default is "auto".
+            The number of GPUs used by the model, default is "auto". If n_worker>1, means number of GPUs per worker.
             ``n_gpu=None`` means cpu only, ``n_gpu=auto`` lets the system automatically determine the best number of GPUs to use.
         peft_model_config: Optional[Dict]
             - "lora_list": A List of PEFT (Parameter-Efficient Fine-Tuning) model and path.
@@ -959,6 +963,8 @@ class Client:
             Specify the worker ip where the model is located in a distributed scenario.
         gpu_idx: Optional[Union[int, List[int]]]
             Specify the GPU index where the model is located.
+        model_path: Optional[str]
+            Model path, if gguf format, should be the file path, otherwise, should be directory of the model.
         **kwargs:
             Any other parameters been specified.
 
@@ -985,10 +991,12 @@ class Client:
             "model_format": model_format,
             "quantization": quantization,
             "replica": replica,
+            "n_worker": n_worker,
             "n_gpu": n_gpu,
             "request_limits": request_limits,
             "worker_ip": worker_ip,
             "gpu_idx": gpu_idx,
+            "model_path": model_path,
         }
 
         for key, value in kwargs.items():

xinference/core/model.py

@@ -226,6 +226,9 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         model_description: Optional["ModelDescription"] = None,
         request_limits: Optional[int] = None,
         xavier_config: Optional[Dict] = None,
+        n_worker: Optional[int] = 1,
+        shard: Optional[int] = 0,
+        driver_info: Optional[dict] = None,  # for model across workers
     ):
         super().__init__()
         from ..model.llm.lmdeploy.core import LMDeployModel
@@ -263,6 +266,10 @@ class ModelActor(xo.StatelessActor, CancelMixin):
             "quantization": self._model_description.get("quantization", "none"),
         }
         self._loop: Optional[asyncio.AbstractEventLoop] = None
+        # model across workers
+        self._n_worker = n_worker
+        self._shard = shard
+        self._driver_info = driver_info
 
         self._scheduler_ref = None
         self._text_to_image_scheduler_ref = None
@@ -455,6 +462,8 @@ class ModelActor(xo.StatelessActor, CancelMixin):
             i += 1
             try:
                 self._model.load()
+                if hasattr(self._model, "driver_info"):
+                    self._driver_info = self._model.driver_info
                 break
             except Exception as e:
                 if (
@@ -477,6 +486,10 @@ class ModelActor(xo.StatelessActor, CancelMixin):
             )
         logger.info(f"{self} loaded")
 
+    async def wait_for_load(self):
+        if hasattr(self._model, "wait_for_load"):
+            self._model.wait_for_load()
+
     def model_uid(self):
         return (
             self._model.model_uid
@@ -488,6 +501,12 @@ class ModelActor(xo.StatelessActor, CancelMixin):
             )
         )
 
+    def get_driver_info(self):
+        # driver info is used for model across workers,
+        # the driver model actor(always the first worker)
+        # will hold driver information includes dist store etc.
+        return self._driver_info
+
     async def _handle_oom_error(self, ex):
         error_message = (
             f"Model actor is out of memory, model id: {self.model_uid()}, error: {ex}"

xinference/core/resource.py

@@ -17,7 +17,7 @@ from typing import Dict, Union
 
 import psutil
 
-from .utils import get_nvidia_gpu_info
+from ..device_utils import get_nvidia_gpu_info
 
 
 @dataclass
@@ -31,9 +31,12 @@ class ResourceStatus:
 
 @dataclass
 class GPUStatus:
+    name: str
     mem_total: float
     mem_free: float
     mem_used: float
+    mem_usage: float
+    gpu_util: float
 
 
 def gather_node_info() -> Dict[str, Union[ResourceStatus, GPUStatus]]:
@@ -48,9 +51,12 @@ def gather_node_info() -> Dict[str, Union[ResourceStatus, GPUStatus]]:
     )
     for gpu_idx, gpu_info in get_nvidia_gpu_info().items():
         node_resource[gpu_idx] = GPUStatus(  # type: ignore
+            name=gpu_info["name"],
             mem_total=gpu_info["total"],
             mem_used=gpu_info["used"],
             mem_free=gpu_info["free"],
+            mem_usage=gpu_info["used"] / gpu_info["total"],
+            gpu_util=gpu_info["util"],
         )
 
     return node_resource  # type: ignore

xinference/core/status_guard.py

@@ -39,6 +39,7 @@ class InstanceInfo(BaseModel):
     replica: int
     status: str
     instance_created_ts: int
+    n_worker: Optional[int] = 1
 
     def update(self, **kwargs):
         for field, value in kwargs.items():

xinference/core/supervisor.py

@@ -99,7 +99,11 @@ class SupervisorActor(xo.StatelessActor):
         self._worker_address_to_worker: Dict[str, xo.ActorRefType["WorkerActor"]] = {}  # type: ignore
         self._worker_status: Dict[str, WorkerStatus] = {}  # type: ignore
         self._replica_model_uid_to_worker: Dict[  # type: ignore
-            str, xo.ActorRefType["WorkerActor"]
+            str,
+            Union[
+                xo.ActorRefType["WorkerActor"],
+                Tuple[xo.ActorRefType["WorkerActor"], ...],
+            ],
         ] = {}
         self._model_uid_to_replica_info: Dict[str, ReplicaInfo] = {}  # type: ignore
         self._uptime = None
@@ -270,8 +274,8 @@ class SupervisorActor(xo.StatelessActor):
         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[
+        self._block_tracker_mapping: Dict[str, xo.ActorRefType[VLLMBlockTracker]] = {}  # type: ignore
+        self._collective_manager_mapping: Dict[  # type: ignore
             str, xo.ActorRefType[CollectiveManager]
         ] = {}
 
@@ -359,13 +363,16 @@ class SupervisorActor(xo.StatelessActor):
         worker_ref = await self._choose_worker()
         return await worker_ref.get_devices_count()
 
-    async def _choose_worker(self) -> xo.ActorRefType["WorkerActor"]:
+    async def _choose_worker(
+        self, available_workers: Optional[List[str]] = None
+    ) -> xo.ActorRefType["WorkerActor"]:
         # TODO: better allocation strategy.
         min_running_model_count = None
         target_worker = None
 
-        workers = list(self._worker_address_to_worker.values())
-        for worker in workers:
+        for worker_addr, worker in self._worker_address_to_worker.items():
+            if available_workers and worker_addr not in available_workers:
+                continue
             running_model_count = await worker.get_model_count()
             if (
                 min_running_model_count is None
@@ -911,6 +918,7 @@ class SupervisorActor(xo.StatelessActor):
         model_type: Optional[str],
         replica: int = 1,
         n_gpu: Optional[Union[int, str]] = "auto",
+        n_worker: Optional[int] = 1,
         request_limits: Optional[int] = None,
         wait_ready: bool = True,
         model_version: Optional[str] = None,
@@ -921,6 +929,35 @@ class SupervisorActor(xo.StatelessActor):
         model_path: Optional[str] = None,
         **kwargs,
     ) -> str:
+        if self.is_local_deployment() and n_worker > 1:  # type: ignore
+            # ignore n_worker > 1 if local deployment
+            logger.warning("Local deployment, ignore n_worker(%s)", n_worker)
+            n_worker = 1
+
+        if n_worker > 1:  # type: ignore
+            # distributed inference
+            return await self._launch_builtin_sharded_model(
+                model_uid,
+                model_name,
+                model_size_in_billions,
+                model_format,
+                quantization,
+                model_engine,
+                model_type,
+                replica=replica,
+                n_gpu=n_gpu,
+                n_worker=n_worker,
+                request_limits=request_limits,
+                wait_ready=wait_ready,
+                model_version=model_version,
+                peft_model_config=peft_model_config,
+                worker_ip=worker_ip,
+                gpu_idx=gpu_idx,
+                download_hub=download_hub,
+                model_path=model_path,
+                **kwargs,
+            )
+
         # search in worker first
         if not self.is_local_deployment():
             workers = list(self._worker_address_to_worker.values())
@@ -1157,6 +1194,150 @@ class SupervisorActor(xo.StatelessActor):
             task.add_done_callback(lambda _: callback_for_async_launch(model_uid))  # type: ignore
         return model_uid
 
+    async def _launch_builtin_sharded_model(
+        self,
+        model_uid: Optional[str],
+        model_name: str,
+        model_size_in_billions: Optional[Union[int, str]],
+        model_format: Optional[str],
+        quantization: Optional[str],
+        model_engine: Optional[str],
+        model_type: Optional[str],
+        replica: int = 1,
+        n_gpu: Optional[Union[int, str]] = "auto",
+        n_worker: Optional[int] = 1,
+        request_limits: Optional[int] = None,
+        wait_ready: bool = True,
+        model_version: Optional[str] = None,
+        peft_model_config: Optional[PeftModelConfig] = None,
+        worker_ip: Optional[str] = None,
+        gpu_idx: Optional[Union[int, List[int]]] = None,
+        download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+        model_path: Optional[str] = None,
+        **kwargs,
+    ):
+        available_workers = []
+        # search workers if registered
+        tasks = []
+        if not worker_ip:
+            all_workers = list(self._worker_address_to_worker)
+            for worker in all_workers:
+                tasks.append(
+                    self._worker_address_to_worker[worker].get_model_registration(
+                        model_type, model_name
+                    )
+                )
+            res = await asyncio.gather(*tasks)
+            for worker, res in zip(all_workers, res):
+                # check regi
+                if res:
+                    available_workers.append(worker)
+            if not available_workers:
+                # no registration, use all workers
+                available_workers = all_workers
+        else:
+            if isinstance(worker_ip, list):
+                available_workers.extend(worker_ip)
+            else:
+                available_workers.append(worker_ip)
+
+        async def _launch_model():
+            try:
+                for _idx, rep_model_uid in enumerate(
+                    iter_replica_model_uid(model_uid, replica)
+                ):
+                    replica_gpu_idx = assign_replica_gpu(
+                        rep_model_uid, replica, gpu_idx
+                    )
+                    # launch shard
+                    worker_refs = []
+                    driver_info = None
+                    for i_worker in range(n_worker):
+                        worker_ref = await self._choose_worker(available_workers)
+                        nonlocal model_type
+                        model_type = model_type or "LLM"
+                        if i_worker > 1:
+                            assert (
+                                driver_info is not None
+                            ), "driver info should be passed by first model shard"
+                        info = await worker_ref.launch_builtin_model(
+                            model_uid=rep_model_uid,
+                            model_name=model_name,
+                            model_size_in_billions=model_size_in_billions,
+                            model_format=model_format,
+                            quantization=quantization,
+                            model_engine=model_engine,
+                            model_type=model_type,
+                            n_gpu=n_gpu,
+                            request_limits=request_limits,
+                            peft_model_config=peft_model_config,
+                            gpu_idx=replica_gpu_idx,
+                            download_hub=download_hub,
+                            model_path=model_path,
+                            shard=i_worker,
+                            n_worker=n_worker,
+                            driver_info=driver_info,
+                            **kwargs,
+                        )
+                        if i_worker == 0:
+                            # info will be subpool address + driver info
+                            # for shard 0
+                            driver_info = info[1]
+                        worker_refs.append(worker_ref)
+                    self._replica_model_uid_to_worker[rep_model_uid] = worker_refs
+
+                    # for distributed inference,
+                    # launch will run asynchronously,
+                    # wait for load complete
+                    for worker_ref in worker_refs:
+                        await worker_ref.wait_for_load(rep_model_uid)
+            except:
+                # terminate_model will remove the replica info.
+                await self.terminate_model(model_uid, suppress_exception=True)
+                await self._status_guard_ref.update_instance_info(
+                    model_uid, {"status": LaunchStatus.ERROR.name}
+                )
+                raise
+
+        if model_uid is None:
+            model_uid = self._gen_model_uid(model_name)
+
+        if not is_valid_model_uid(model_uid):
+            raise ValueError(
+                "The model UID is invalid. Please specify the model UID by 0 < length <= 100."
+            )
+
+        if request_limits is not None and request_limits < 0:
+            raise ValueError(
+                "The `request_limits` parameter must be greater or equal than 0."
+            )
+
+        if model_uid in self._model_uid_to_replica_info:
+            raise ValueError(f"Model is already in the model list, uid: {model_uid}")
+
+        # Set replica info first for exception handler to terminate model.
+        self._model_uid_to_replica_info[model_uid] = ReplicaInfo(
+            replica=replica, scheduler=itertools.cycle(range(replica))
+        )
+        instance_info = InstanceInfo(
+            model_name=model_name,
+            model_uid=model_uid,
+            model_version=model_version,
+            model_ability=[],
+            replica=replica,
+            n_worker=n_worker,
+            status=LaunchStatus.CREATING.name,
+            instance_created_ts=int(time.time()),
+        )
+        await self._status_guard_ref.set_instance_info(model_uid, instance_info)
+        if wait_ready:
+            await _launch_model()
+        else:
+            task = asyncio.create_task(_launch_model())
+            ASYNC_LAUNCH_TASKS[model_uid] = task
+            task.add_done_callback(lambda _: callback_for_async_launch(model_uid))  # type: ignore
+        return model_uid
+
     async def get_instance_info(
         self, model_name: Optional[str], model_uid: Optional[str]
     ) -> List[Dict]:
@@ -1186,11 +1367,13 @@ class SupervisorActor(xo.StatelessActor):
                     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)
+                            worker_refs = self._replica_model_uid_to_worker[model_uid]
+                            if not isinstance(worker_refs, list):
+                                worker_refs = [worker_refs]
+                            for worker_ref in worker_refs:
+                                model_address = worker_ref.address
+                                if model_address == address:
+                                    dead_models.append(model_uid)
                         logger.error(
                             "Worker dead. address: %s, influenced models: %s",
                             address,
@@ -1222,13 +1405,18 @@ class SupervisorActor(xo.StatelessActor):
     @log_async(logger=logger)
     async def terminate_model(self, model_uid: str, suppress_exception=False):
         async def _terminate_one_model(_replica_model_uid):
-            worker_ref = self._replica_model_uid_to_worker.get(_replica_model_uid, None)
+            worker_refs = self._replica_model_uid_to_worker.get(
+                _replica_model_uid, None
+            )
+            if not isinstance(worker_refs, list):
+                worker_refs = [worker_refs]
 
-            if worker_ref is None:
-                raise ValueError(
-                    f"Model not found in the model list, uid: {_replica_model_uid}"
-                )
-            await worker_ref.terminate_model(model_uid=_replica_model_uid)
+            for worker_ref in worker_refs:
+                if worker_ref is None:
+                    raise ValueError(
+                        f"Model not found in the model list, uid: {_replica_model_uid}"
+                    )
+                await worker_ref.terminate_model(model_uid=_replica_model_uid)
             del self._replica_model_uid_to_worker[_replica_model_uid]
 
         replica_info = self._model_uid_to_replica_info.get(model_uid, None)
@@ -1290,6 +1478,9 @@ class SupervisorActor(xo.StatelessActor):
             raise ValueError(
                 f"Model not found in the model list, uid: {replica_model_uid}"
             )
+        if isinstance(worker_ref, list):
+            # get first worker to fetch information if model across workers
+            worker_ref = worker_ref[0]
         return await worker_ref.get_model(model_uid=replica_model_uid)
 
     @log_async(logger=logger)
@@ -1299,6 +1490,9 @@ class SupervisorActor(xo.StatelessActor):
             raise ValueError(
                 f"Model not found in the model list, uid: {replica_model_uid}"
             )
+        if isinstance(worker_ref, list):
+            # get status from first shard if model has multiple shards across workers
+            worker_ref = worker_ref[0]
         return await worker_ref.get_model_status(replica_model_uid)
 
     @log_async(logger=logger)
@@ -1314,6 +1508,9 @@ class SupervisorActor(xo.StatelessActor):
             raise ValueError(
                 f"Model not found in the model list, uid: {replica_model_uid}"
             )
+        if isinstance(worker_ref, list):
+            # get status from first shard if model has multiple shards across workers
+            worker_ref = worker_ref[0]
         info = await worker_ref.describe_model(model_uid=replica_model_uid)
         info["replica"] = replica_info.replica
         return info
@@ -1386,6 +1583,9 @@ class SupervisorActor(xo.StatelessActor):
             worker_ref = self._replica_model_uid_to_worker.get(rep_mid, None)
             if worker_ref is None:
                 continue
+            if isinstance(worker_ref, list):
+                # get status from first shard if model has multiple shards across workers
+                worker_ref = worker_ref[0]
             model_ref = await worker_ref.get_model(model_uid=rep_mid)
             result_info = await model_ref.abort_request(request_id, block_duration)
             res["msg"] = result_info
@@ -1415,8 +1615,17 @@ class SupervisorActor(xo.StatelessActor):
     async def remove_worker(self, worker_address: str):
         uids_to_remove = []
         for model_uid in self._replica_model_uid_to_worker:
-            if self._replica_model_uid_to_worker[model_uid].address == worker_address:
-                uids_to_remove.append(model_uid)
+            worker_refs = self._replica_model_uid_to_worker[model_uid]
+            if not isinstance(worker_refs, list):
+                worker_refs = [worker_refs]
+            for worker_ref in worker_refs:
+                model_address = worker_ref.address
+                if isinstance(model_address, str) and model_address == worker_address:
+                    uids_to_remove.append(model_uid)
+                elif (
+                    isinstance(model_address, list) and worker_address in model_address
+                ):
+                    uids_to_remove.append(model_uid)
 
         for replica_model_uid in uids_to_remove:
             model_uid, _ = parse_replica_model_uid(replica_model_uid)

xinference/core/utils.py

@@ -19,10 +19,9 @@ import string
 import uuid
 import weakref
 from enum import Enum
-from typing import Dict, Generator, List, Optional, Tuple, Union
+from typing import Generator, List, Optional, Tuple, Union
 
 import orjson
-from pynvml import nvmlDeviceGetCount, nvmlInit, nvmlShutdown
 
 from .._compat import BaseModel
 from ..constants import (
@@ -248,33 +247,6 @@ def parse_model_version(model_version: str, model_type: str) -> Tuple:
         raise ValueError(f"Not supported model_type: {model_type}")
 
 
-def _get_nvidia_gpu_mem_info(gpu_id: int) -> Dict[str, float]:
-    from pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
-
-    handler = nvmlDeviceGetHandleByIndex(gpu_id)
-    mem_info = nvmlDeviceGetMemoryInfo(handler)
-    return {"total": mem_info.total, "used": mem_info.used, "free": mem_info.free}
-
-
-def get_nvidia_gpu_info() -> Dict:
-    try:
-        nvmlInit()
-        device_count = nvmlDeviceGetCount()
-        res = {}
-        for i in range(device_count):
-            res[f"gpu-{i}"] = _get_nvidia_gpu_mem_info(i)
-        return res
-    except:
-        # TODO: add log here
-        # logger.debug(f"Cannot init nvml. Maybe due to lack of NVIDIA GPUs or incorrect installation of CUDA.")
-        return {}
-    finally:
-        try:
-            nvmlShutdown()
-        except:
-            pass
-
-
 def assign_replica_gpu(
     _replica_model_uid: str, replica: int, gpu_idx: Optional[Union[int, List[int]]]
 ) -> Optional[List[int]]:

xinference/core/worker.py

@@ -789,6 +789,9 @@ class WorkerActor(xo.StatelessActor):
         model_engine: Optional[str],
         model_type: str = "LLM",
         n_gpu: Optional[Union[int, str]] = "auto",
+        n_worker: Optional[int] = 1,
+        shard: Optional[int] = 0,
+        driver_info: Optional[dict] = None,
         peft_model_config: Optional[PeftModelConfig] = None,
         request_limits: Optional[int] = None,
         gpu_idx: Optional[Union[int, List[int]]] = None,
@@ -876,6 +879,18 @@ class WorkerActor(xo.StatelessActor):
                 xavier_config: Optional[Dict] = kwargs.pop("xavier_config", None)
                 if xavier_config is not None:
                     xavier_config["rank_address"] = subpool_address
+                model_kwargs = kwargs.copy()
+                if n_worker > 1:  # type: ignore
+                    # for model across workers,
+                    # add a few kwargs
+                    model_kwargs.update(
+                        dict(
+                            address=self.address,
+                            n_worker=n_worker,
+                            shard=shard,
+                            driver_info=driver_info,
+                        )
+                    )
                 model, model_description = await asyncio.to_thread(
                     create_model_instance,
                     subpool_address,
@@ -890,7 +905,7 @@ class WorkerActor(xo.StatelessActor):
                     peft_model_config,
                     download_hub,
                     model_path,
-                    **kwargs,
+                    **model_kwargs,
                 )
                 await self.update_cache_status(model_name, model_description)
                 model_ref = await xo.create_actor(
@@ -904,6 +919,9 @@ class WorkerActor(xo.StatelessActor):
                     model_description=model_description,
                     request_limits=request_limits,
                     xavier_config=xavier_config,
+                    n_worker=n_worker,
+                    shard=shard,
+                    driver_info=driver_info,
                 )
                 await model_ref.load()
             except:
@@ -933,7 +951,15 @@ class WorkerActor(xo.StatelessActor):
             origin_uid,
             {"model_ability": abilities, "status": LaunchStatus.READY.name},
         )
-        return subpool_address
+        if n_worker > 1 and shard == 0:  # type: ignore
+            return subpool_address, await model_ref.get_driver_info()
+        else:
+            return subpool_address
+
+    @log_async(logger=logger, level=logging.INFO)
+    async def wait_for_load(self, model_uid: str):
+        model_ref = self._model_uid_to_model[model_uid]
+        await model_ref.wait_for_load()
 
     @log_async(logger=logger, level=logging.INFO)
     async def terminate_model(self, model_uid: str, is_model_die=False):