xinference 0.6.5__py3-none-any.whl → 0.7.0__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-12-01T18:35:32+0800",
+ "date": "2023-12-08T13:45:18+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "909a428f4762a5e397123a42b8d2abff02eada83",
- "version": "0.6.5"
+ "full-revisionid": "b5a5f0a270f85e451591eba34fe615a0fc8ce4bf",
+ "version": "0.7.0"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -158,6 +158,9 @@ class RESTfulAPI:
         self._router.add_api_route(
             "/v1/models/prompts", self._get_builtin_prompts, methods=["GET"]
         )
+        self._router.add_api_route(
+            "/v1/cluster/devices", self._get_devices_count, methods=["GET"]
+        )
         self._router.add_api_route(
             "/v1/models/{model_uid}", self.describe_model, methods=["GET"]
         )
@@ -310,6 +313,17 @@ class RESTfulAPI:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def _get_devices_count(self) -> JSONResponse:
+        """
+        For internal usage
+        """
+        try:
+            data = await (await self._get_supervisor_ref()).get_devices_count()
+            return JSONResponse(content=data)
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def get_status(self) -> JSONResponse:
         try:
             data = await (await self._get_supervisor_ref()).get_status()

xinference/core/model.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import inspect
+import os
 import uuid
 from typing import (
     TYPE_CHECKING,
@@ -44,6 +45,15 @@ from .utils import json_dumps, log_async
 
 T = TypeVar("T")
 
+try:
+    from torch.cuda import OutOfMemoryError
+except ImportError:
+
+    class _OutOfMemoryError(Exception):
+        pass
+
+    OutOfMemoryError = _OutOfMemoryError
+
 
 def request_limit(fn):
     """
@@ -192,18 +202,30 @@ class ModelActor(xo.StatelessActor):
             return ret
 
     async def _call_wrapper(self, _wrapper: Callable):
-        assert not (
-            inspect.iscoroutinefunction(_wrapper)
-            or inspect.isasyncgenfunction(_wrapper)
-        )
-        if self._lock is None:
-            return await asyncio.to_thread(_wrapper)
-        else:
-            async with self._lock:
+        try:
+            assert not (
+                inspect.iscoroutinefunction(_wrapper)
+                or inspect.isasyncgenfunction(_wrapper)
+            )
+            if self._lock is None:
                 return await asyncio.to_thread(_wrapper)
+            else:
+                async with self._lock:
+                    return await asyncio.to_thread(_wrapper)
+        except OutOfMemoryError:
+            logger.exception(
+                "Model actor is out of memory, model id: %s", self.model_uid()
+            )
+            os._exit(1)
 
     async def _call_async_wrapper(self, _wrapper: Callable):
-        return await asyncio.create_task(_wrapper())
+        try:
+            return await asyncio.create_task(_wrapper())
+        except OutOfMemoryError:
+            logger.exception(
+                "Model actor is out of memory, model id: %s", self.model_uid()
+            )
+            os._exit(1)
 
     @log_async(logger=logger)
     @request_limit
@@ -365,7 +387,8 @@ class ModelActor(xo.StatelessActor):
 
         async def _async_wrapper():
             try:
-                return await anext(gen)  # noqa: F821
+                # anext is only available for Python >= 3.10
+                return await gen.__anext__()  # noqa: F821
             except StopAsyncIteration:
                 return stop
 

xinference/core/supervisor.py

@@ -106,6 +106,16 @@ class SupervisorActor(xo.StatelessActor):
             data[k] = v.dict()
         return data
 
+    async def get_devices_count(self) -> int:
+        from ..utils import cuda_count
+
+        if self.is_local_deployment():
+            return cuda_count()
+        # distributed deployment, choose a worker and return its cuda_count.
+        # Assume that each worker has the same count of cards.
+        worker_ref = await self._choose_worker()
+        return await worker_ref.get_devices_count()
+
     async def _choose_worker(self) -> xo.ActorRefType["WorkerActor"]:
         # TODO: better allocation strategy.
         min_running_model_count = None

xinference/core/worker.py

@@ -48,6 +48,7 @@ class WorkerActor(xo.StatelessActor):
         self._supervisor_address = supervisor_address
         self._supervisor_ref = None
         self._main_pool = main_pool
+        self._main_pool.recover_sub_pool = self.recover_sub_pool
 
         # internal states.
         self._model_uid_to_model: Dict[str, xo.ActorRefType["ModelActor"]] = {}
@@ -55,9 +56,22 @@ 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_launch_args: Dict[str, Dict] = {}
 
         self._lock = asyncio.Lock()
 
+    async def recover_sub_pool(self, address):
+        logger.warning("Process %s is down, create model.", address)
+        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)
+                break
+
     @classmethod
     def uid(cls) -> str:
         return "worker"
@@ -94,6 +108,12 @@ class WorkerActor(xo.StatelessActor):
     async def __pre_destroy__(self):
         self._upload_task.cancel()
 
+    @staticmethod
+    def get_devices_count():
+        from ..utils import cuda_count
+
+        return cuda_count()
+
     @log_sync(logger=logger)
     def get_model_count(self) -> int:
         return len(self._model_uid_to_model)
@@ -175,7 +195,7 @@ class WorkerActor(xo.StatelessActor):
             gpu_cnt = n_gpu if isinstance(n_gpu, int) else 1
             devices = (
                 [await self.allocate_devices_for_embedding(model_uid)]
-                if model_type == "embedding"
+                if model_type in ["embedding", "rerank"]
                 else self.allocate_devices(model_uid=model_uid, n_gpu=gpu_cnt)
             )
             env["CUDA_VISIBLE_DEVICES"] = ",".join([str(dev) for dev in devices])
@@ -278,7 +298,6 @@ class WorkerActor(xo.StatelessActor):
         for dev in devices:
             self._gpu_to_model_uid[int(dev)] = model_uid
         self._model_uid_to_addr[model_uid] = subpool_address
-        return model_ref
 
     @log_async(logger=logger)
     async def launch_builtin_model(
@@ -292,7 +311,9 @@ class WorkerActor(xo.StatelessActor):
         n_gpu: Optional[Union[int, str]] = "auto",
         request_limits: Optional[int] = None,
         **kwargs,
-    ) -> xo.ActorRefType["ModelActor"]:
+    ):
+        launch_args = locals()
+        launch_args.pop("self")
         if n_gpu is not None:
             if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > cuda_count()):
                 raise ValueError(
@@ -343,7 +364,7 @@ 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
-        return model_ref
+        self._model_uid_to_launch_args[model_uid] = launch_args
 
     @log_async(logger=logger)
     async def terminate_model(self, model_uid: str):
@@ -351,15 +372,21 @@ class WorkerActor(xo.StatelessActor):
         if model_ref is None:
             raise ValueError(f"Model not found in the model list, uid: {model_uid}")
 
-        await xo.destroy_actor(model_ref)
-        del self._model_uid_to_model[model_uid]
-        del self._model_uid_to_model_spec[model_uid]
-
-        self.release_devices(model_uid)
-
-        subpool_address = self._model_uid_to_addr[model_uid]
-        await self._main_pool.remove_sub_pool(subpool_address)
-        del self._model_uid_to_addr[model_uid]
+        try:
+            await xo.destroy_actor(model_ref)
+        except Exception as e:
+            logger.debug(
+                "Destroy model actor failed, model uid: %s, error: %s", model_uid, e
+            )
+        try:
+            subpool_address = self._model_uid_to_addr[model_uid]
+            await self._main_pool.remove_sub_pool(subpool_address)
+        finally:
+            del self._model_uid_to_model[model_uid]
+            del self._model_uid_to_model_spec[model_uid]
+            self.release_devices(model_uid)
+            del self._model_uid_to_addr[model_uid]
+            del self._model_uid_to_launch_args[model_uid]
 
     @log_async(logger=logger)
     async def list_models(self) -> Dict[str, Dict[str, Any]]:

xinference/deploy/utils.py

@@ -110,6 +110,7 @@ async def create_worker_actor_pool(
     return await xo.create_actor_pool(
         address=address,
         n_process=0,
+        auto_recover="process",
         subprocess_start_method=subprocess_start_method,
         logging_conf={"dict": logging_conf},
     )

xinference/model/embedding/core.py

@@ -220,6 +220,9 @@ class EmbeddingModel:
             ]
 
             raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")
+        from ..utils import patch_trust_remote_code
+
+        patch_trust_remote_code()
         self._model = SentenceTransformer(self._model_path, device=self._device)
 
     def create_embedding(self, sentences: Union[str, List[str]], **kwargs):

xinference/model/llm/vllm/core.py

@@ -319,7 +319,9 @@ class VLLMChatModel(VLLMModel, ChatModelMixin):
         if not generate_config:
             generate_config = {}
         if self.model_family.prompt_style:
-            if (not generate_config["stop"]) and self.model_family.prompt_style.stop:
+            if (
+                not generate_config.get("stop")
+            ) and self.model_family.prompt_style.stop:
                 generate_config["stop"] = self.model_family.prompt_style.stop.copy()
             if self.model_family.prompt_style.stop_token_ids:
                 generate_config.setdefault(

xinference/model/utils.py

@@ -211,3 +211,28 @@ def copy_from_src_to_dst(
             )
             if attempt + 1 == max_attempt:
                 raise
+
+
+def patch_trust_remote_code():
+    """sentence-transformers calls transformers without the trust_remote_code=True, some embedding
+    models will fail to load, e.g. jina-embeddings-v2-base-en
+
+    :return:
+    """
+    try:
+        from transformers.dynamic_module_utils import resolve_trust_remote_code
+    except ImportError:
+        logger.error("Patch transformers trust_remote_code failed.")
+    else:
+
+        def _patched_resolve_trust_remote_code(*args, **kwargs):
+            logger.info("Patched resolve_trust_remote_code: %s %s", args, kwargs)
+            return True
+
+        if (
+            resolve_trust_remote_code.__code__
+            != _patched_resolve_trust_remote_code.__code__
+        ):
+            resolve_trust_remote_code.__code__ = (
+                _patched_resolve_trust_remote_code.__code__
+            )

xinference/types.py

@@ -289,7 +289,7 @@ def get_pydantic_model_from_method(
         model.__fields__.pop(key)
     if exclude_fields is not None:
         for key in exclude_fields:
-            model.__fields__.pop(key)
+            model.__fields__.pop(key, None)
     if include_fields is not None:
         dummy_model = create_model("DummyModel", **include_fields)
         model.__fields__.update(dummy_model.__fields__)
@@ -307,10 +307,10 @@ def fix_forward_ref(model):
         if isinstance(field.annotation, ForwardRef):
             exclude_fields.append(key)
             include_fields[key] = (Optional[Any], None)
-    if exclude_fields is not None:
+    if exclude_fields:
         for key in exclude_fields:
-            model.__fields__.pop(key)
-    if include_fields is not None:
+            model.__fields__.pop(key, None)
+    if include_fields:
         dummy_model = create_model("DummyModel", **include_fields)
         model.__fields__.update(dummy_model.__fields__)
     return model

xinference/web/ui/build/asset-manifest.json

@@ -1,11 +1,11 @@
 {
   "files": {
-    "main.js": "./static/js/main.9b46963e.js",
+    "main.js": "./static/js/main.8126d441.js",
     "static/media/icon.webp": "./static/media/icon.4603d52c63041e5dfbfd.webp",
     "index.html": "./index.html",
-    "main.9b46963e.js.map": "./static/js/main.9b46963e.js.map"
+    "main.8126d441.js.map": "./static/js/main.8126d441.js.map"
   },
   "entrypoints": [
-    "static/js/main.9b46963e.js"
+    "static/js/main.8126d441.js"
   ]
 }

xinference/web/ui/build/index.html

@@ -1 +1 @@
-<!doctype html><html lang="en"><head><meta charset="utf-8"/><link rel="icon" href="./favicon.svg"/><meta name="viewport" content="width=device-width,initial-scale=1"/><meta name="theme-color" content="#000000"/><meta name="description" content="Web site created using create-react-app"/><link rel="apple-touch-icon" href="./logo192.png"/><link rel="manifest" href="./manifest.json"/><title>Xinference</title><script defer="defer" src="./static/js/main.9b46963e.js"></script></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root"></div></body></html>
+<!doctype html><html lang="en"><head><meta charset="utf-8"/><link rel="icon" href="./favicon.svg"/><meta name="viewport" content="width=device-width,initial-scale=1"/><meta name="theme-color" content="#000000"/><meta name="description" content="Web site created using create-react-app"/><link rel="apple-touch-icon" href="./logo192.png"/><link rel="manifest" href="./manifest.json"/><title>Xinference</title><script defer="defer" src="./static/js/main.8126d441.js"></script></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root"></div></body></html>