xinference 0.8.4__py3-none-any.whl → 0.9.0__py3-none-any.whl

xinference/client/restful/restful_client.py

@@ -665,12 +665,16 @@ class Client:
     def _check_cluster_authenticated(self):
         url = f"{self.base_url}/v1/cluster/auth"
         response = requests.get(url)
-        if response.status_code != 200:
-            raise RuntimeError(
-                f"Failed to get cluster information, detail: {response.json()['detail']}"
-            )
-        response_data = response.json()
-        self._cluster_authed = bool(response_data["auth"])
+        # compatible with old version of xinference
+        if response.status_code == 404:
+            self._cluster_authed = False
+        else:
+            if response.status_code != 200:
+                raise RuntimeError(
+                    f"Failed to get cluster information, detail: {response.json()['detail']}"
+                )
+            response_data = response.json()
+            self._cluster_authed = bool(response_data["auth"])
 
     def login(self, username: str, password: str):
         if not self._cluster_authed:

xinference/core/model.py

@@ -44,6 +44,7 @@ import logging
 
 logger = logging.getLogger(__name__)
 
+from ..device_utils import empty_cache
 from .utils import json_dumps, log_async
 
 try:
@@ -130,7 +131,7 @@ class ModelActor(xo.StatelessActor):
             try:
                 import gc
 
-                import torch
+                import torch  # noqa: F401
             except ImportError:
                 error_message = "Failed to import module 'torch'"
                 installation_guide = [
@@ -141,7 +142,7 @@ class ModelActor(xo.StatelessActor):
 
             del self._model
             gc.collect()
-            torch.cuda.empty_cache()
+            empty_cache()
 
     def __init__(
         self,

xinference/core/resource.py

@@ -22,8 +22,9 @@ from .utils import get_nvidia_gpu_info
 
 @dataclass
 class ResourceStatus:
-    available: float
+    usage: float
     total: float
+    memory_used: float
     memory_available: float
     memory_total: float
 
@@ -39,8 +40,9 @@ def gather_node_info() -> Dict[str, Union[ResourceStatus, GPUStatus]]:
     node_resource = dict()
     mem_info = psutil.virtual_memory()
     node_resource["cpu"] = ResourceStatus(
-        available=psutil.cpu_percent() / 100.0,
+        usage=psutil.cpu_percent() / 100.0,
         total=psutil.cpu_count(),
+        memory_used=mem_info.used,
         memory_available=mem_info.available,
         memory_total=mem_info.total,
     )

xinference/core/status_guard.py

@@ -16,7 +16,8 @@ from logging import getLogger
 from typing import Dict, List, Optional
 
 import xoscar as xo
-from pydantic import BaseModel
+
+from .._compat import BaseModel
 
 logger = getLogger(__name__)
 

xinference/core/supervisor.py

@@ -15,6 +15,7 @@
 import asyncio
 import itertools
 import time
+import typing
 from dataclasses import dataclass
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union
@@ -179,12 +180,26 @@ class SupervisorActor(xo.StatelessActor):
             model_version_infos, self.address
         )
 
-    async def get_cluster_device_info(self) -> List:
+    @typing.no_type_check
+    async def get_cluster_device_info(self, detailed: bool = False) -> List:
+        import psutil
+
         supervisor_device_info = {
             "ip_address": self.address.split(":")[0],
             "gpu_count": 0,
             "gpu_vram_total": 0,
         }
+        if detailed:
+            supervisor_device_info["gpu_vram_total"] = 0
+            supervisor_device_info["gpu_vram_available"] = 0
+            supervisor_device_info["cpu_available"] = psutil.cpu_count() * (
+                1 - psutil.cpu_percent() / 100.0
+            )
+            supervisor_device_info["cpu_count"] = psutil.cpu_count()
+            mem_info = psutil.virtual_memory()
+            supervisor_device_info["mem_used"] = mem_info.used
+            supervisor_device_info["mem_available"] = mem_info.available
+            supervisor_device_info["mem_total"] = mem_info.total
         res = [{"node_type": "Supervisor", **supervisor_device_info}]
         for worker_addr, worker_status in self._worker_status.items():
             vram_total: float = sum(
@@ -193,14 +208,24 @@ class SupervisorActor(xo.StatelessActor):
             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,
-                }
-            )
+            info = {
+                "node_type": "Worker",
+                "ip_address": worker_addr.split(":")[0],
+                "gpu_count": len(worker_status.status) - 1,
+                "gpu_vram_total": total,
+            }
+            if detailed:
+                cpu_info = worker_status.status["cpu"]
+                info["cpu_available"] = cpu_info.total * (1 - cpu_info.usage)
+                info["cpu_count"] = cpu_info.total
+                info["mem_used"] = cpu_info.memory_used
+                info["mem_available"] = cpu_info.memory_available
+                info["mem_total"] = cpu_info.memory_total
+                info["gpu_vram_total"] = vram_total
+                info["gpu_vram_available"] = sum(
+                    [v.mem_free for k, v in worker_status.status.items() if k != "cpu"]
+                )
+            res.append(info)
         return res
 
     @staticmethod
@@ -227,11 +252,11 @@ class SupervisorActor(xo.StatelessActor):
         }
 
     async def get_devices_count(self) -> int:
-        from ..utils import cuda_count
+        from ..device_utils import gpu_count
 
         if self.is_local_deployment():
-            return cuda_count()
-        # distributed deployment, choose a worker and return its cuda_count.
+            return gpu_count()
+        # distributed deployment, choose a worker and return its device_count.
         # Assume that each worker has the same count of cards.
         worker_ref = await self._choose_worker()
         return await worker_ref.get_devices_count()

xinference/core/utils.py

@@ -19,9 +19,10 @@ import string
 from typing import Dict, Generator, List, Tuple, Union
 
 import orjson
-from pydantic import BaseModel
 from pynvml import nvmlDeviceGetCount, nvmlInit, nvmlShutdown
 
+from .._compat import BaseModel
+
 logger = logging.getLogger(__name__)
 
 

xinference/core/worker.py

@@ -30,8 +30,8 @@ from xoscar import MainActorPoolType
 from ..constants import XINFERENCE_CACHE_DIR
 from ..core import ModelActor
 from ..core.status_guard import LaunchStatus
+from ..device_utils import gpu_count
 from ..model.core import ModelDescription, create_model_instance
-from ..utils import cuda_count
 from .event import Event, EventCollectorActor, EventType
 from .metrics import launch_metrics_export_server, record_metrics
 from .resource import gather_node_info
@@ -54,13 +54,13 @@ class WorkerActor(xo.StatelessActor):
         self,
         supervisor_address: str,
         main_pool: MainActorPoolType,
-        cuda_devices: List[int],
+        gpu_devices: List[int],
         metrics_exporter_host: Optional[str] = None,
         metrics_exporter_port: Optional[int] = None,
     ):
         super().__init__()
         # static attrs.
-        self._total_cuda_devices = cuda_devices
+        self._total_gpu_devices = gpu_devices
         self._supervisor_address = supervisor_address
         self._supervisor_ref = None
         self._main_pool = main_pool
@@ -244,9 +244,9 @@ class WorkerActor(xo.StatelessActor):
 
     @staticmethod
     def get_devices_count():
-        from ..utils import cuda_count
+        from ..device_utils import gpu_count
 
-        return cuda_count()
+        return gpu_count()
 
     @log_sync(logger=logger)
     def get_model_count(self) -> int:
@@ -263,7 +263,7 @@ class WorkerActor(xo.StatelessActor):
         we assume that embedding model only takes 1 GPU slot.
         """
         candidates = []
-        for _dev in self._total_cuda_devices:
+        for _dev in self._total_gpu_devices:
             if _dev not in self._gpu_to_model_uid:
                 candidates.append(_dev)
             else:
@@ -291,11 +291,11 @@ class WorkerActor(xo.StatelessActor):
         return device
 
     def allocate_devices(self, model_uid: str, n_gpu: int) -> List[int]:
-        if n_gpu > len(self._total_cuda_devices) - len(self._gpu_to_model_uid):
+        if n_gpu > len(self._total_gpu_devices) - len(self._gpu_to_model_uid):
             raise RuntimeError("No available slot found for the model")
 
         devices: List[int] = [
-            dev for dev in self._total_cuda_devices if dev not in self._gpu_to_model_uid
+            dev for dev in self._total_gpu_devices if dev not in self._gpu_to_model_uid
         ][:n_gpu]
         for dev in devices:
             self._gpu_to_model_uid[int(dev)] = model_uid
@@ -324,7 +324,7 @@ class WorkerActor(xo.StatelessActor):
     ) -> Tuple[str, List[str]]:
         env = {}
         devices = []
-        if isinstance(n_gpu, int) or (n_gpu == "auto" and cuda_count() > 0):
+        if isinstance(n_gpu, int) or (n_gpu == "auto" and gpu_count() > 0):
             # Currently, n_gpu=auto means using 1 GPU
             gpu_cnt = n_gpu if isinstance(n_gpu, int) else 1
             devices = (
@@ -396,10 +396,10 @@ class WorkerActor(xo.StatelessActor):
         n_gpu: Optional[Union[int, str]] = "auto",
     ):
         if n_gpu is not None:
-            if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > cuda_count()):
+            if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > gpu_count()):
                 raise ValueError(
                     f"The parameter `n_gpu` must be greater than 0 and "
-                    f"not greater than the number of GPUs: {cuda_count()} on the machine."
+                    f"not greater than the number of GPUs: {gpu_count()} on the machine."
                 )
             if isinstance(n_gpu, str) and n_gpu != "auto":
                 raise ValueError("Currently `n_gpu` only supports `auto`.")
@@ -504,10 +504,10 @@ class WorkerActor(xo.StatelessActor):
         launch_args.pop("kwargs")
         launch_args.update(kwargs)
         if n_gpu is not None:
-            if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > cuda_count()):
+            if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > gpu_count()):
                 raise ValueError(
                     f"The parameter `n_gpu` must be greater than 0 and "
-                    f"not greater than the number of GPUs: {cuda_count()} on the machine."
+                    f"not greater than the number of GPUs: {gpu_count()} on the machine."
                 )
             if isinstance(n_gpu, str) and n_gpu != "auto":
                 raise ValueError("Currently `n_gpu` only supports `auto`.")

xinference/deploy/worker.py

@@ -21,7 +21,7 @@ import xoscar as xo
 from xoscar import MainActorPoolType
 
 from ..core.worker import WorkerActor
-from ..utils import cuda_count
+from ..device_utils import gpu_count
 
 logger = logging.getLogger(__name__)
 
@@ -33,12 +33,12 @@ async def start_worker_components(
     metrics_exporter_host: Optional[str],
     metrics_exporter_port: Optional[int],
 ):
-    cuda_device_indices = []
-    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
-    if cuda_visible_devices:
-        cuda_device_indices.extend([int(i) for i in cuda_visible_devices.split(",")])
+    gpu_device_indices = []
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    if cuda_visible_devices is not None and cuda_visible_devices != "-1":
+        gpu_device_indices.extend([int(i) for i in cuda_visible_devices.split(",")])
     else:
-        cuda_device_indices = list(range(cuda_count()))
+        gpu_device_indices = list(range(gpu_count()))
 
     await xo.create_actor(
         WorkerActor,
@@ -46,7 +46,7 @@ async def start_worker_components(
         uid=WorkerActor.uid(),
         supervisor_address=supervisor_address,
         main_pool=main_pool,
-        cuda_devices=cuda_device_indices,
+        gpu_devices=gpu_device_indices,
         metrics_exporter_host=metrics_exporter_host,
         metrics_exporter_port=metrics_exporter_port,
     )

xinference/device_utils.py

@@ -0,0 +1,100 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import torch
+from typing_extensions import Literal, Union
+
+DeviceType = Literal["cuda", "mps", "xpu", "cpu"]
+
+
+def is_xpu_available() -> bool:
+    return hasattr(torch, "xpu") and torch.xpu.is_available()
+
+
+def get_available_device() -> DeviceType:
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    elif is_xpu_available():
+        return "xpu"
+    return "cpu"
+
+
+def is_device_available(device: str) -> bool:
+    if device == "cuda":
+        return torch.cuda.is_available()
+    elif device == "mps":
+        return torch.backends.mps.is_available()
+    elif device == "xpu":
+        return is_xpu_available()
+    elif device == "cpu":
+        return True
+
+    return False
+
+
+def move_model_to_available_device(model):
+    device = get_available_device()
+
+    if device == "cpu":
+        return model
+
+    return model.to(device)
+
+
+def get_device_preferred_dtype(device: str) -> Union[torch.dtype, None]:
+    if device == "cpu":
+        return torch.float32
+    elif device == "cuda" or device == "mps":
+        return torch.float16
+    elif device == "xpu":
+        return torch.bfloat16
+
+    return None
+
+
+def is_hf_accelerate_supported(device: str) -> bool:
+    return device == "cuda" or device == "xpu"
+
+
+def empty_cache():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+    if is_xpu_available():
+        torch.xpu.empty_cache()
+
+
+def gpu_count():
+    if torch.cuda.is_available():
+        cuda_visible_devices_env = os.getenv("CUDA_VISIBLE_DEVICES", None)
+
+        if cuda_visible_devices_env is None:
+            return torch.cuda.device_count()
+
+        cuda_visible_devices = (
+            cuda_visible_devices_env.split(",") if cuda_visible_devices_env else []
+        )
+
+        return min(torch.cuda.device_count(), len(cuda_visible_devices))
+    elif torch.backends.mps.is_available():
+        return 1
+    elif is_xpu_available():
+        return torch.xpu.device_count()
+    else:
+        return 0

xinference/fields.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pydantic import Field
+from ._compat import Field
 
 none_field = Field(None)
 

xinference/model/audio/core.py

@@ -16,8 +16,7 @@ import os
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple
 
-from pydantic import BaseModel
-
+from ..._compat import BaseModel
 from ...constants import XINFERENCE_CACHE_DIR
 from ..core import ModelDescription
 from ..utils import valid_model_revision

xinference/model/audio/whisper.py

@@ -14,6 +14,12 @@
 import logging
 from typing import TYPE_CHECKING, Dict, Optional
 
+from xinference.device_utils import (
+    get_available_device,
+    get_device_preferred_dtype,
+    is_device_available,
+)
+
 if TYPE_CHECKING:
     from .core import AudioModelFamilyV1
 
@@ -37,11 +43,15 @@ class WhisperModel:
         self._kwargs = kwargs
 
     def load(self):
-        import torch
         from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        if self._device is None:
+            self._device = get_available_device()
+        else:
+            if not is_device_available(self._device):
+                raise ValueError(f"Device {self._device} is not available!")
+
+        torch_dtype = get_device_preferred_dtype(self._device)
 
         model = AutoModelForSpeechSeq2Seq.from_pretrained(
             self._model_path,
@@ -49,7 +59,7 @@ class WhisperModel:
             low_cpu_mem_usage=True,
             use_safetensors=True,
         )
-        model.to(device)
+        model.to(self._device)
 
         processor = AutoProcessor.from_pretrained(self._model_path)
 
@@ -63,7 +73,7 @@ class WhisperModel:
             batch_size=16,
             return_timestamps=False,
             torch_dtype=torch_dtype,
-            device=device,
+            device=self._device,
         )
 
     def _call_model(
@@ -99,9 +109,11 @@ class WhisperModel:
             )
         return self._call_model(
             audio=audio,
-            generate_kwargs={"language": language, "task": "transcribe"}
-            if language is not None
-            else {"task": "transcribe"},
+            generate_kwargs=(
+                {"language": language, "task": "transcribe"}
+                if language is not None
+                else {"task": "transcribe"}
+            ),
             response_format=response_format,
         )
 

xinference/model/core.py

@@ -15,6 +15,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, List, Optional, Tuple
 
+from .._compat import BaseModel
+
 
 class ModelDescription(ABC):
     def __init__(
@@ -94,3 +96,10 @@ def create_model_instance(
         )
     else:
         raise ValueError(f"Unsupported model type: {model_type}.")
+
+
+class CacheableModelSpec(BaseModel):
+    model_name: str
+    model_id: str
+    model_revision: Optional[str]
+    model_hub: str = "huggingface"