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

xinference/core/model.py

@@ -13,21 +13,21 @@
 # limitations under the License.
 
 import asyncio
+import functools
 import inspect
 import json
 import os
-import uuid
+import time
+import types
+import weakref
 from typing import (
     TYPE_CHECKING,
-    Any,
     AsyncGenerator,
     Callable,
     Dict,
-    Generic,
     Iterator,
     List,
     Optional,
-    TypeVar,
     Union,
 )
 
@@ -35,8 +35,9 @@ import sse_starlette.sse
 import xoscar as xo
 
 if TYPE_CHECKING:
+    from .worker import WorkerActor
     from ..model.llm.core import LLM
-    from ..types import ChatCompletionChunk, CompletionChunk
+    from ..model.core import ModelDescription
     import PIL
 
 import logging
@@ -45,8 +46,6 @@ logger = logging.getLogger(__name__)
 
 from .utils import json_dumps, log_async
 
-T = TypeVar("T")
-
 try:
     from torch.cuda import OutOfMemoryError
 except ImportError:
@@ -88,38 +87,30 @@ def request_limit(fn):
     return wrapped_func
 
 
-class IteratorWrapper(Generic[T]):
-    def __init__(self, uid: str, model_actor_addr: str, model_actor_uid: str):
-        self._uid = uid
-        self._model_actor_addr = model_actor_addr
-        self._model_actor_uid = model_actor_uid
-        self._model_actor_ref: Optional[xo.ActorRefType["ModelActor"]] = None
-
-    async def destroy(self):
-        if self._model_actor_ref is None:
-            self._model_actor_ref = await xo.actor_ref(
-                address=self._model_actor_addr, uid=self._model_actor_uid
-            )
-        assert self._model_actor_ref is not None
-        return await self._model_actor_ref.destroy_generator(self._uid)
+def oom_check(fn):
+    @functools.wraps(fn)
+    def _wrapper(*args, **kwargs):
+        try:
+            return fn(*args, **kwargs)
+        except OutOfMemoryError:
+            logger.exception("Model actor is out of memory.")
+            os._exit(1)
 
-    def __aiter__(self):
-        return self
+    @functools.wraps(fn)
+    async def _async_wrapper(*args, **kwargs):
+        try:
+            return await fn(*args, **kwargs)
+        except OutOfMemoryError:
+            logger.exception("Model actor is out of memory.")
+            os._exit(1)
 
-    async def __anext__(self) -> T:
-        if self._model_actor_ref is None:
-            self._model_actor_ref = await xo.actor_ref(
-                address=self._model_actor_addr, uid=self._model_actor_uid
-            )
+    assert not inspect.isasyncgen(fn)
+    assert not inspect.isgenerator(fn)
 
-        try:
-            assert self._model_actor_ref is not None
-            return await self._model_actor_ref.next(self._uid)
-        except Exception as e:
-            if "StopIteration" in str(e):
-                raise StopAsyncIteration
-            else:
-                raise
+    if asyncio.iscoroutinefunction(fn):
+        return _async_wrapper
+    else:
+        return _wrapper
 
 
 class ModelActor(xo.StatelessActor):
@@ -152,22 +143,91 @@ class ModelActor(xo.StatelessActor):
             gc.collect()
             torch.cuda.empty_cache()
 
-    def __init__(self, model: "LLM", request_limits: Optional[int] = None):
+    def __init__(
+        self,
+        worker_address: str,
+        model: "LLM",
+        model_description: Optional["ModelDescription"] = None,
+        request_limits: Optional[int] = None,
+    ):
         super().__init__()
         from ..model.llm.pytorch.core import PytorchModel
         from ..model.llm.pytorch.spec_model import SpeculativeModel
         from ..model.llm.vllm.core import VLLMModel
 
+        self._worker_address = worker_address
         self._model = model
+        self._model_description = (
+            model_description.to_dict() if model_description else {}
+        )
         self._request_limits = request_limits
 
         self._generators: Dict[str, Union[Iterator, AsyncGenerator]] = {}
+        self._current_generator = lambda: None
         self._lock = (
             None
             if isinstance(self._model, (PytorchModel, SpeculativeModel, VLLMModel))
             else asyncio.locks.Lock()
         )
+        self._worker_ref = None
         self._serve_count = 0
+        self._metrics_labels = {
+            "type": self._model_description.get("model_type", "unknown"),
+            "model": self.model_uid(),
+            "node": self._worker_address,
+            "format": self._model_description.get("model_format", "unknown"),
+            "quantization": self._model_description.get("quantization", "none"),
+        }
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+
+    async def __post_create__(self):
+        self._loop = asyncio.get_running_loop()
+
+    async def _record_completion_metrics(
+        self, duration, completion_tokens, prompt_tokens
+    ):
+        coros = []
+        if completion_tokens > 0:
+            coros.append(
+                self.record_metrics(
+                    "output_tokens_total_counter",
+                    "add",
+                    {
+                        "labels": self._metrics_labels,
+                        "value": completion_tokens,
+                    },
+                )
+            )
+        if prompt_tokens > 0:
+            coros.append(
+                self.record_metrics(
+                    "input_tokens_total_counter",
+                    "add",
+                    {"labels": self._metrics_labels, "value": prompt_tokens},
+                )
+            )
+        if completion_tokens > 0:
+            generate_throughput = completion_tokens / duration
+            coros.append(
+                self.record_metrics(
+                    "generate_throughput",
+                    "set",
+                    {
+                        "labels": self._metrics_labels,
+                        "value": generate_throughput,
+                    },
+                )
+            )
+        await asyncio.gather(*coros)
+
+    async def _get_worker_ref(self) -> xo.ActorRefType["WorkerActor"]:
+        from .worker import WorkerActor
+
+        if self._worker_ref is None:
+            self._worker_ref = await xo.actor_ref(
+                address=self._worker_address, uid=WorkerActor.uid()
+            )
+        return self._worker_ref
 
     def is_vllm_backend(self) -> bool:
         from ..model.llm.vllm.core import VLLMModel
@@ -188,106 +248,158 @@ class ModelActor(xo.StatelessActor):
             )
         )
 
-    def _wrap_generator(self, ret: Any):
-        if inspect.isgenerator(ret) or inspect.isasyncgen(ret):
-            if self._lock is not None and self._generators:
-                raise Exception("Parallel generation is not supported by ggml.")
-            generator_uid = str(uuid.uuid1())
-            self._generators[generator_uid] = ret
-
-            return IteratorWrapper(
-                uid=generator_uid,
-                model_actor_addr=self.address,
-                model_actor_uid=self.uid,
-            )
-        else:
-            return json_dumps(ret)
-
-    async def _call_wrapper(self, _wrapper: Callable):
+    def _to_json_generator(self, gen: types.GeneratorType):
+        start_time = time.time()
+        time_to_first_token = None
+        final_usage = None
         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)
+            for v in gen:
+                if time_to_first_token is None:
+                    time_to_first_token = (time.time() - start_time) * 1000
+                final_usage = v.pop("usage", None)
+                v = dict(data=json.dumps(v))
+                yield sse_starlette.sse.ensure_bytes(v, None)
         except OutOfMemoryError:
             logger.exception(
                 "Model actor is out of memory, model id: %s", self.model_uid()
             )
             os._exit(1)
+        finally:
+            if self._loop is not None and time_to_first_token is not None:
+                coro = self.record_metrics(
+                    "time_to_first_token",
+                    "set",
+                    {"labels": self._metrics_labels, "value": time_to_first_token},
+                )
+                asyncio.run_coroutine_threadsafe(coro, loop=self._loop)
+            if self._loop is not None and final_usage is not None:
+                coro = self._record_completion_metrics(
+                    time.time() - start_time,
+                    completion_tokens=final_usage["completion_tokens"],
+                    prompt_tokens=final_usage["prompt_tokens"],
+                )
+                asyncio.run_coroutine_threadsafe(coro, loop=self._loop)
 
-    async def _call_async_wrapper(self, _wrapper: Callable):
+    async def _to_json_async_gen(self, gen: types.AsyncGeneratorType):
+        start_time = time.time()
+        time_to_first_token = None
+        final_usage = None
         try:
-            return await asyncio.create_task(_wrapper())
+            async for v in gen:
+                if time_to_first_token is None:
+                    time_to_first_token = (time.time() - start_time) * 1000
+                final_usage = v.pop("usage", None)
+                v = await asyncio.to_thread(json.dumps, v)
+                v = dict(data=v)  # noqa: F821
+                yield await asyncio.to_thread(sse_starlette.sse.ensure_bytes, v, None)
         except OutOfMemoryError:
             logger.exception(
                 "Model actor is out of memory, model id: %s", self.model_uid()
             )
             os._exit(1)
+        finally:
+            coros = []
+            if time_to_first_token is not None:
+                coros.append(
+                    self.record_metrics(
+                        "time_to_first_token",
+                        "set",
+                        {"labels": self._metrics_labels, "value": time_to_first_token},
+                    )
+                )
+            if final_usage is not None:
+                coros.append(
+                    self._record_completion_metrics(
+                        time.time() - start_time,
+                        completion_tokens=final_usage["completion_tokens"],
+                        prompt_tokens=final_usage["prompt_tokens"],
+                    )
+                )
+            await asyncio.gather(*coros)
+
+    @oom_check
+    async def _call_wrapper(self, fn: Callable, *args, **kwargs):
+        if self._lock is None:
+            if inspect.iscoroutinefunction(fn):
+                ret = await fn(*args, **kwargs)
+            else:
+                ret = await asyncio.to_thread(fn, *args, **kwargs)
+        else:
+            async with self._lock:
+                if inspect.iscoroutinefunction(fn):
+                    ret = await fn(*args, **kwargs)
+                else:
+                    ret = await asyncio.to_thread(fn, *args, **kwargs)
+
+        if self._lock is not None and self._current_generator():
+            raise Exception("Parallel generation is not supported by ggml.")
+
+        if inspect.isgenerator(ret):
+            gen = self._to_json_generator(ret)
+            self._current_generator = weakref.ref(gen)
+            return gen
+        if inspect.isasyncgen(ret):
+            gen = self._to_json_async_gen(ret)
+            self._current_generator = weakref.ref(gen)
+            return gen
+        return await asyncio.to_thread(json_dumps, ret)
 
     @log_async(logger=logger)
     @request_limit
+    @xo.generator
     async def generate(self, prompt: str, *args, **kwargs):
-        if not hasattr(self._model, "generate") and not hasattr(
-            self._model, "async_generate"
-        ):
-            raise AttributeError(f"Model {self._model.model_spec} is not for generate.")
-
-        def _wrapper():
-            return self._wrap_generator(
-                getattr(self._model, "generate")(prompt, *args, **kwargs)
+        if hasattr(self._model, "generate"):
+            return await self._call_wrapper(
+                self._model.generate, prompt, *args, **kwargs
             )
-
-        async def _async_wrapper():
-            # for vLLM.
-            return self._wrap_generator(
-                await getattr(self._model, "async_generate")(prompt, *args, **kwargs)
+        if hasattr(self._model, "async_generate"):
+            return await self._call_wrapper(
+                self._model.async_generate, prompt, *args, **kwargs
             )
-
-        if hasattr(self._model, "generate"):
-            return await self._call_wrapper(_wrapper)
-        else:
-            return await self._call_async_wrapper(_async_wrapper)
+        raise AttributeError(f"Model {self._model.model_spec} is not for generate.")
 
     @log_async(logger=logger)
     @request_limit
+    @xo.generator
     async def chat(self, prompt: str, *args, **kwargs):
-        if not hasattr(self._model, "chat") and not hasattr(self._model, "async_chat"):
+        start_time = time.time()
+        response = None
+        try:
+            if hasattr(self._model, "chat"):
+                response = await self._call_wrapper(
+                    self._model.chat, prompt, *args, **kwargs
+                )
+                return response
+            if hasattr(self._model, "async_chat"):
+                response = await self._call_wrapper(
+                    self._model.async_chat, prompt, *args, **kwargs
+                )
+                return response
             raise AttributeError(f"Model {self._model.model_spec} is not for chat.")
-
-        def _wrapper():
-            return self._wrap_generator(
-                getattr(self._model, "chat")(prompt, *args, **kwargs)
-            )
-
-        async def _async_wrapper():
-            # for vLLM.
-            return self._wrap_generator(
-                await getattr(self._model, "async_chat")(prompt, *args, **kwargs)
-            )
-
-        if hasattr(self._model, "async_chat"):
-            return await self._call_async_wrapper(_async_wrapper)
-        else:
-            return await self._call_wrapper(_wrapper)
+        finally:
+            # For the non stream result.
+            if response is not None and isinstance(response, dict):
+                usage = response["usage"]
+                # Some backends may not have a valid usage, we just skip them.
+                completion_tokens = usage["completion_tokens"]
+                prompt_tokens = usage["prompt_tokens"]
+                await self._record_completion_metrics(
+                    time.time() - start_time,
+                    completion_tokens,
+                    prompt_tokens,
+                )
 
     @log_async(logger=logger)
     @request_limit
     async def create_embedding(self, input: Union[str, List[str]], *args, **kwargs):
-        if not hasattr(self._model, "create_embedding"):
-            raise AttributeError(
-                f"Model {self._model.model_spec} is not for creating embedding."
+        if hasattr(self._model, "create_embedding"):
+            return await self._call_wrapper(
+                self._model.create_embedding, input, *args, **kwargs
             )
 
-        def _wrapper():
-            data = getattr(self._model, "create_embedding")(input, *args, **kwargs)
-            return json_dumps(data)
-
-        return await self._call_wrapper(_wrapper)
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating embedding."
+        )
 
     @log_async(logger=logger)
     @request_limit
@@ -301,13 +413,9 @@ class ModelActor(xo.StatelessActor):
         *args,
         **kwargs,
     ):
-        if not hasattr(self._model, "rerank"):
-            raise AttributeError(
-                f"Model {self._model.model_spec} is not for reranking."
-            )
-
-        def _wrapper():
-            data = getattr(self._model, "rerank")(
+        if hasattr(self._model, "rerank"):
+            return await self._call_wrapper(
+                self._model.rerank,
                 documents,
                 query,
                 top_n,
@@ -316,9 +424,7 @@ class ModelActor(xo.StatelessActor):
                 *args,
                 **kwargs,
             )
-            return json_dumps(data)
-
-        return await self._call_wrapper(_wrapper)
+        raise AttributeError(f"Model {self._model.model_spec} is not for reranking.")
 
     @log_async(logger=logger)
     @request_limit
@@ -331,18 +437,19 @@ class ModelActor(xo.StatelessActor):
         *args,
         **kwargs,
     ):
-        if not hasattr(self._model, "text_to_image"):
-            raise AttributeError(
-                f"Model {self._model.model_spec} is not for creating image."
-            )
-
-        def _wrapper():
-            r = getattr(self._model, "text_to_image")(
-                prompt, n, size, response_format, *args, **kwargs
+        if hasattr(self._model, "text_to_image"):
+            return await self._call_wrapper(
+                self._model.text_to_image,
+                prompt,
+                n,
+                size,
+                response_format,
+                *args,
+                **kwargs,
             )
-            return json_dumps(r)
-
-        return await self._call_wrapper(_wrapper)
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating image."
+        )
 
     async def image_to_image(
         self,
@@ -355,13 +462,9 @@ class ModelActor(xo.StatelessActor):
         *args,
         **kwargs,
     ):
-        if not hasattr(self._model, "image_to_image"):
-            raise AttributeError(
-                f"Model {self._model.model_spec} is not for creating image."
-            )
-
-        def _wrapper():
-            r = getattr(self._model, "image_to_image")(
+        if hasattr(self._model, "image_to_image"):
+            return await self._call_wrapper(
+                self._model.image_to_image,
                 image,
                 prompt,
                 negative_prompt,
@@ -371,50 +474,10 @@ class ModelActor(xo.StatelessActor):
                 *args,
                 **kwargs,
             )
-            return json_dumps(r)
-
-        return await self._call_wrapper(_wrapper)
-
-    async def next(
-        self, generator_uid: str
-    ) -> Union["ChatCompletionChunk", "CompletionChunk"]:
-        assert generator_uid in self._generators
-        stop = object()
-        gen = self._generators[generator_uid]
-
-        def _wrapper():
-            try:
-                v = dict(data=json.dumps(next(gen)))
-                return sse_starlette.sse.ensure_bytes(v, None)
-            except StopIteration:
-                return stop
-
-        async def _async_wrapper():
-            try:
-                # anext is only available for Python >= 3.10
-                v = await gen.__anext__()
-                v = await asyncio.to_thread(json.dumps, v)
-                v = dict(data=v)  # noqa: F821
-                return await asyncio.to_thread(sse_starlette.sse.ensure_bytes, v, None)
-            except StopAsyncIteration:
-                return stop
-
-        if inspect.isgenerator(gen):
-            r = await self._call_wrapper(_wrapper)
-        elif inspect.isasyncgen(gen):
-            # for vLLM.
-            r = await self._call_async_wrapper(_async_wrapper)
-        else:
-            raise TypeError(
-                f"Unexpected type {type(gen)}, expecting generator or async generator"
-            )
-
-        if r is stop:
-            self._generators.pop(generator_uid, None)
-            raise Exception("StopIteration")
-        else:
-            return r
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating image."
+        )
 
-    @log_async(logger=logger)
-    async def destroy_generator(self, generator_uid: str):
-        self._generators.pop(generator_uid, None)
+    async def record_metrics(self, name, op, kwargs):
+        worker_ref = await self._get_worker_ref()
+        await worker_ref.record_metrics(name, op, kwargs)

xinference/core/status_guard.py

@@ -0,0 +1,86 @@
+# Copyright 2022-2024 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.
+from enum import Enum
+from logging import getLogger
+from typing import Dict, List, Optional
+
+import xoscar as xo
+from pydantic import BaseModel
+
+logger = getLogger(__name__)
+
+
+class LaunchStatus(Enum):
+    CREATING = 1
+    UPDATING = 2
+    TERMINATING = 3
+    TERMINATED = 4
+    READY = 5
+    ERROR = 6
+
+
+class InstanceInfo(BaseModel):
+    model_name: str
+    model_uid: str
+    model_ability: List[str]
+    replica: int
+    status: str
+    instance_created_ts: int
+
+    def update(self, **kwargs):
+        for field, value in kwargs.items():
+            setattr(self, field, value)
+
+
+class StatusGuardActor(xo.StatelessActor):
+    def __init__(self):
+        super().__init__()
+        self._model_uid_to_info: Dict[str, InstanceInfo] = {}
+
+    @classmethod
+    def uid(cls) -> str:
+        return "status_guard"
+
+    @staticmethod
+    def _drop_terminated_info(instance_infos: List[InstanceInfo]) -> List[InstanceInfo]:
+        return [
+            info
+            for info in instance_infos
+            if info.status != LaunchStatus.TERMINATED.name
+        ]
+
+    def set_instance_info(self, model_uid: str, info: InstanceInfo):
+        self._model_uid_to_info[model_uid] = info
+
+    def get_instance_info(
+        self, model_name: Optional[str] = None, model_uid: Optional[str] = None
+    ) -> List[InstanceInfo]:
+        if model_uid is not None:
+            return (
+                self._drop_terminated_info([self._model_uid_to_info[model_uid]])
+                if model_uid in self._model_uid_to_info
+                else []
+            )
+        all_infos: List[InstanceInfo] = list(self._model_uid_to_info.values())
+        filtered_infos: List[InstanceInfo] = list(
+            filter(lambda info: info.model_name == model_name, all_infos)
+        )
+        return (
+            self._drop_terminated_info(filtered_infos)
+            if model_name is not None
+            else self._drop_terminated_info(all_infos)
+        )
+
+    def update_instance_info(self, model_uid: str, info: Dict):
+        self._model_uid_to_info[model_uid].update(**info)