xinference 0.15.3__py3-none-any.whl → 0.15.4__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-09-30T20:17:26+0800",
+ "date": "2024-10-12T18:28:41+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "00a9ee15279a60a6d75393c4720d8da5cbbf5796",
- "version": "0.15.3"
+ "full-revisionid": "c0be11504c70f6c392cbdb67c86cf12153353f70",
+ "version": "0.15.4"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -524,6 +524,16 @@ class RESTfulAPI:
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/requests/{request_id}/progress",
+            self.get_progress,
+            methods=["get"],
+            dependencies=(
+                [Security(self._auth_service, scopes=["models:read"])]
+                if self.is_authenticated()
+                else None
+            ),
+        )
         self._router.add_api_route(
             "/v1/images/generations",
             self.create_images,
@@ -1486,6 +1496,17 @@ class RESTfulAPI:
             await self._report_error_event(model_uid, str(e))
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def get_progress(self, request_id: str) -> JSONResponse:
+        try:
+            supervisor_ref = await self._get_supervisor_ref()
+            result = {"progress": await supervisor_ref.get_progress(request_id)}
+            return JSONResponse(content=result)
+        except KeyError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def create_images(self, request: Request) -> Response:
         body = TextToImageRequest.parse_obj(await request.json())
         model_uid = body.model
@@ -1853,10 +1874,16 @@ class RESTfulAPI:
             await self._report_error_event(model_uid, str(e))
             raise HTTPException(status_code=500, detail=str(e))
 
-        from ..model.llm.utils import GLM4_TOOL_CALL_FAMILY, QWEN_TOOL_CALL_FAMILY
+        from ..model.llm.utils import (
+            GLM4_TOOL_CALL_FAMILY,
+            LLAMA3_TOOL_CALL_FAMILY,
+            QWEN_TOOL_CALL_FAMILY,
+        )
 
         model_family = desc.get("model_family", "")
-        function_call_models = QWEN_TOOL_CALL_FAMILY + GLM4_TOOL_CALL_FAMILY
+        function_call_models = (
+            QWEN_TOOL_CALL_FAMILY + GLM4_TOOL_CALL_FAMILY + LLAMA3_TOOL_CALL_FAMILY
+        )
 
         if model_family not in function_call_models:
             if body.tools:

xinference/client/restful/restful_client.py

@@ -1385,6 +1385,16 @@ class Client:
         response_json = response.json()
         return response_json
 
+    def get_progress(self, request_id: str):
+        url = f"{self.base_url}/v1/requests/{request_id}/progress"
+        response = requests.get(url, headers=self._headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get progress, detail: {_get_error_string(response)}"
+            )
+        response_json = response.json()
+        return response_json
+
     def abort_cluster(self):
         url = f"{self.base_url}/v1/clusters"
         response = requests.delete(url, headers=self._headers)

xinference/constants.py

@@ -28,6 +28,7 @@ XINFERENCE_ENV_HEALTH_CHECK_TIMEOUT = "XINFERENCE_HEALTH_CHECK_TIMEOUT"
 XINFERENCE_ENV_DISABLE_HEALTH_CHECK = "XINFERENCE_DISABLE_HEALTH_CHECK"
 XINFERENCE_ENV_DISABLE_METRICS = "XINFERENCE_DISABLE_METRICS"
 XINFERENCE_ENV_TRANSFORMERS_ENABLE_BATCHING = "XINFERENCE_TRANSFORMERS_ENABLE_BATCHING"
+XINFERENCE_ENV_DOWNLOAD_MAX_ATTEMPTS = "XINFERENCE_DOWNLOAD_MAX_ATTEMPTS"
 
 
 def get_xinference_home() -> str:
@@ -82,3 +83,6 @@ XINFERENCE_DISABLE_METRICS = bool(
 XINFERENCE_TRANSFORMERS_ENABLE_BATCHING = bool(
     int(os.environ.get(XINFERENCE_ENV_TRANSFORMERS_ENABLE_BATCHING, 0))
 )
+XINFERENCE_DOWNLOAD_MAX_ATTEMPTS = int(
+    os.environ.get(XINFERENCE_ENV_DOWNLOAD_MAX_ATTEMPTS, 3)
+)

xinference/core/image_interface.py

@@ -16,6 +16,9 @@ import base64
 import io
 import logging
 import os
+import threading
+import time
+import uuid
 from typing import Dict, List, Optional, Union
 
 import gradio as gr
@@ -84,6 +87,7 @@ class ImageInterface:
             num_inference_steps: int,
             negative_prompt: Optional[str] = None,
             sampler_name: Optional[str] = None,
+            progress=gr.Progress(),
         ) -> PIL.Image.Image:
             from ..client import RESTfulClient
 
@@ -99,19 +103,43 @@ class ImageInterface:
             )
             sampler_name = None if sampler_name == "default" else sampler_name
 
-            response = model.text_to_image(
-                prompt=prompt,
-                n=n,
-                size=size,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                negative_prompt=negative_prompt,
-                sampler_name=sampler_name,
-                response_format="b64_json",
-            )
+            response = None
+            exc = None
+            request_id = str(uuid.uuid4())
+
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.text_to_image(
+                        request_id=request_id,
+                        prompt=prompt,
+                        n=n,
+                        size=size,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        negative_prompt=negative_prompt,
+                        sampler_name=sampler_name,
+                        response_format="b64_json",
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except (KeyError, RuntimeError):
+                    cur_progress = 0.0
+
+                progress(cur_progress, desc="Generating images")
+                time.sleep(1)
+
+            if exc:
+                raise exc
 
             images = []
-            for image_dict in response["data"]:
+            for image_dict in response["data"]:  # type: ignore
                 assert image_dict["b64_json"] is not None
                 image_data = base64.b64decode(image_dict["b64_json"])
                 image = PIL.Image.open(io.BytesIO(image_data))
@@ -184,6 +212,7 @@ class ImageInterface:
             num_inference_steps: int,
             padding_image_to_multiple: int,
             sampler_name: Optional[str] = None,
+            progress=gr.Progress(),
         ) -> PIL.Image.Image:
             from ..client import RESTfulClient
 
@@ -205,20 +234,44 @@ class ImageInterface:
             bio = io.BytesIO()
             image.save(bio, format="png")
 
-            response = model.image_to_image(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                n=n,
-                image=bio.getvalue(),
-                size=size,
-                response_format="b64_json",
-                num_inference_steps=num_inference_steps,
-                padding_image_to_multiple=padding_image_to_multiple,
-                sampler_name=sampler_name,
-            )
+            response = None
+            exc = None
+            request_id = str(uuid.uuid4())
+
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.image_to_image(
+                        request_id=request_id,
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        n=n,
+                        image=bio.getvalue(),
+                        size=size,
+                        response_format="b64_json",
+                        num_inference_steps=num_inference_steps,
+                        padding_image_to_multiple=padding_image_to_multiple,
+                        sampler_name=sampler_name,
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except (KeyError, RuntimeError):
+                    cur_progress = 0.0
+
+                progress(cur_progress, desc="Generating images")
+                time.sleep(1)
+
+            if exc:
+                raise exc
 
             images = []
-            for image_dict in response["data"]:
+            for image_dict in response["data"]:  # type: ignore
                 assert image_dict["b64_json"] is not None
                 image_data = base64.b64decode(image_dict["b64_json"])
                 image = PIL.Image.open(io.BytesIO(image_data))

xinference/core/model.py

@@ -44,6 +44,7 @@ import xoscar as xo
 from ..constants import XINFERENCE_TRANSFORMERS_ENABLE_BATCHING
 
 if TYPE_CHECKING:
+    from .progress_tracker import ProgressTrackerActor
     from .worker import WorkerActor
     from ..model.llm.core import LLM
     from ..model.core import ModelDescription
@@ -177,6 +178,7 @@ class ModelActor(xo.StatelessActor):
 
     def __init__(
         self,
+        supervisor_address: str,
         worker_address: str,
         model: "LLM",
         model_description: Optional["ModelDescription"] = None,
@@ -188,6 +190,7 @@ class ModelActor(xo.StatelessActor):
         from ..model.llm.transformers.core import PytorchModel
         from ..model.llm.vllm.core import VLLMModel
 
+        self._supervisor_address = supervisor_address
         self._worker_address = worker_address
         self._model = model
         self._model_description = (
@@ -205,6 +208,7 @@ class ModelActor(xo.StatelessActor):
             else asyncio.locks.Lock()
         )
         self._worker_ref = None
+        self._progress_tracker_ref = None
         self._serve_count = 0
         self._metrics_labels = {
             "type": self._model_description.get("model_type", "unknown"),
@@ -275,6 +279,28 @@ class ModelActor(xo.StatelessActor):
             )
         return self._worker_ref
 
+    async def _get_progress_tracker_ref(
+        self,
+    ) -> xo.ActorRefType["ProgressTrackerActor"]:
+        from .progress_tracker import ProgressTrackerActor
+
+        if self._progress_tracker_ref is None:
+            self._progress_tracker_ref = await xo.actor_ref(
+                address=self._supervisor_address, uid=ProgressTrackerActor.default_uid()
+            )
+        return self._progress_tracker_ref
+
+    async def _get_progressor(self, request_id: str):
+        from .progress_tracker import Progressor
+
+        progressor = Progressor(
+            request_id,
+            await self._get_progress_tracker_ref(),
+            asyncio.get_running_loop(),
+        )
+        await progressor.start()
+        return progressor
+
     def is_vllm_backend(self) -> bool:
         from ..model.llm.vllm.core import VLLMModel
 
@@ -732,17 +758,20 @@ class ModelActor(xo.StatelessActor):
         *args,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         if hasattr(self._model, "text_to_image"):
-            return await self._call_wrapper_json(
-                self._model.text_to_image,
-                prompt,
-                n,
-                size,
-                response_format,
-                *args,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.text_to_image,
+                    prompt,
+                    n,
+                    size,
+                    response_format,
+                    *args,
+                    **kwargs,
+                )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating image."
         )
@@ -753,12 +782,15 @@ class ModelActor(xo.StatelessActor):
         self,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         if hasattr(self._model, "txt2img"):
-            return await self._call_wrapper_json(
-                self._model.txt2img,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.txt2img,
+                    **kwargs,
+                )
         raise AttributeError(f"Model {self._model.model_spec} is not for txt2img.")
 
     @log_async(
@@ -776,19 +808,22 @@ class ModelActor(xo.StatelessActor):
         *args,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         kwargs["negative_prompt"] = negative_prompt
         if hasattr(self._model, "image_to_image"):
-            return await self._call_wrapper_json(
-                self._model.image_to_image,
-                image,
-                prompt,
-                n,
-                size,
-                response_format,
-                *args,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.image_to_image,
+                    image,
+                    prompt,
+                    n,
+                    size,
+                    response_format,
+                    *args,
+                    **kwargs,
+                )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating image."
         )
@@ -799,12 +834,15 @@ class ModelActor(xo.StatelessActor):
         self,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         if hasattr(self._model, "img2img"):
-            return await self._call_wrapper_json(
-                self._model.img2img,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.img2img,
+                    **kwargs,
+                )
         raise AttributeError(f"Model {self._model.model_spec} is not for img2img.")
 
     @log_async(
@@ -823,20 +861,23 @@ class ModelActor(xo.StatelessActor):
         *args,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
+        kwargs["negative_prompt"] = negative_prompt
         if hasattr(self._model, "inpainting"):
-            return await self._call_wrapper_json(
-                self._model.inpainting,
-                image,
-                mask_image,
-                prompt,
-                negative_prompt,
-                n,
-                size,
-                response_format,
-                *args,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.inpainting,
+                    image,
+                    mask_image,
+                    prompt,
+                    n,
+                    size,
+                    response_format,
+                    *args,
+                    **kwargs,
+                )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating image."
         )

xinference/core/progress_tracker.py

@@ -0,0 +1,187 @@
+# 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 asyncio
+import dataclasses
+import logging
+import os
+import time
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import xoscar as xo
+
+TO_REMOVE_PROGRESS_INTERVAL = float(
+    os.getenv("XINFERENCE_REMOVE_PROGRESS_INTERVAL", 5 * 60)
+)  # 5min
+CHECK_PROGRESS_INTERVAL = float(
+    os.getenv("XINFERENCE_CHECK_PROGRESS_INTERVAL", 1 * 60)
+)  # 1min
+UPLOAD_PROGRESS_SPAN = float(
+    os.getenv("XINFERENCE_UPLOAD_PROGRESS_SPAN", 0.05)
+)  # not upload when change less than 0.1
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class _ProgressInfo:
+    progress: float
+    last_updated: float
+    info: Optional[str] = None
+
+
+class ProgressTrackerActor(xo.StatelessActor):
+    _request_id_to_progress: Dict[str, _ProgressInfo]
+
+    @classmethod
+    def default_uid(cls) -> str:
+        return "progress_tracker"
+
+    def __init__(
+        self,
+        to_remove_interval: float = TO_REMOVE_PROGRESS_INTERVAL,
+        check_interval: float = CHECK_PROGRESS_INTERVAL,
+    ):
+        super().__init__()
+
+        self._request_id_to_progress = {}
+        self._clear_finished_task = None
+        self._to_remove_interval = to_remove_interval
+        self._check_interval = check_interval
+
+    async def __post_create__(self):
+        self._clear_finished_task = asyncio.create_task(self._clear_finished())
+
+    async def __pre_destroy__(self):
+        if self._clear_finished_task:
+            self._clear_finished_task.cancel()
+
+    async def _clear_finished(self):
+        to_remove_request_ids = []
+        while True:
+            now = time.time()
+            for request_id, progress in self._request_id_to_progress.items():
+                if abs(progress.progress - 1.0) > 1e-5:
+                    continue
+
+                # finished
+                if now - progress.last_updated > self._to_remove_interval:
+                    to_remove_request_ids.append(request_id)
+
+            for rid in to_remove_request_ids:
+                del self._request_id_to_progress[rid]
+
+            if to_remove_request_ids:
+                logger.debug(
+                    "Remove requests %s due to it's finished for over %s seconds",
+                    to_remove_request_ids,
+                    self._to_remove_interval,
+                )
+
+            await asyncio.sleep(self._check_interval)
+
+    def start(self, request_id: str):
+        self._request_id_to_progress[request_id] = _ProgressInfo(
+            progress=0.0, last_updated=time.time()
+        )
+
+    def set_progress(self, request_id: str, progress: float):
+        assert progress <= 1.0
+        info = self._request_id_to_progress[request_id]
+        info.progress = progress
+        info.last_updated = time.time()
+        logger.debug(
+            "Setting progress, request id: %s, progress: %s", request_id, progress
+        )
+
+    def get_progress(self, request_id: str) -> float:
+        return self._request_id_to_progress[request_id].progress
+
+
+class Progressor:
+    _sub_progress_stack: List[Tuple[float, float]]
+
+    def __init__(
+        self,
+        request_id: str,
+        progress_tracker_ref: xo.ActorRefType["ProgressTrackerActor"],
+        loop: asyncio.AbstractEventLoop,
+        upload_span: float = UPLOAD_PROGRESS_SPAN,
+    ):
+        self.request_id = request_id
+        self.progress_tracker_ref = progress_tracker_ref
+        self.loop = loop
+        # uploading when progress changes over this span
+        # to prevent from frequently uploading
+        self._upload_span = upload_span
+
+        self._last_report_progress = 0.0
+        self._current_progress = 0.0
+        self._sub_progress_stack = [(0.0, 1.0)]
+        self._current_sub_progress_start = 0.0
+        self._current_sub_progress_end = 1.0
+
+    async def start(self):
+        if self.request_id:
+            await self.progress_tracker_ref.start(self.request_id)
+
+    def split_stages(self, n_stage: int, stage_weight: Optional[List[float]] = None):
+        if self.request_id:
+            if stage_weight is not None:
+                if len(stage_weight) != n_stage + 1:
+                    raise ValueError(
+                        f"stage_weight should have size {n_stage + 1}, got {len(stage_weight)}"
+                    )
+                progresses = stage_weight
+            else:
+                progresses = np.linspace(
+                    self._current_sub_progress_start,
+                    self._current_sub_progress_end,
+                    n_stage + 1,
+                )
+            spans = [(progresses[i], progresses[i + 1]) for i in range(n_stage)]
+            self._sub_progress_stack.extend(spans[::-1])
+
+    def __enter__(self):
+        if self.request_id:
+            (
+                self._current_sub_progress_start,
+                self._current_sub_progress_end,
+            ) = self._sub_progress_stack[-1]
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.request_id:
+            self._sub_progress_stack.pop()
+            # force to set progress to 1.0 for this sub progress
+            # nevertheless it is done or not
+            self.set_progress(1.0)
+        return False
+
+    def set_progress(self, progress: float):
+        if self.request_id:
+            self._current_progress = (
+                self._current_sub_progress_start
+                + (self._current_sub_progress_end - self._current_sub_progress_start)
+                * progress
+            )
+            if (
+                self._current_progress - self._last_report_progress >= self._upload_span
+                or 1.0 - progress < 1e-5
+            ):
+                set_progress = self.progress_tracker_ref.set_progress(
+                    self.request_id, self._current_progress
+                )
+                asyncio.run_coroutine_threadsafe(set_progress, self.loop)  # type: ignore
+                self._last_report_progress = self._current_progress

xinference/core/supervisor.py

@@ -130,6 +130,7 @@ class SupervisorActor(xo.StatelessActor):
             )
         logger.info(f"Xinference supervisor {self.address} started")
         from .cache_tracker import CacheTrackerActor
+        from .progress_tracker import ProgressTrackerActor
         from .status_guard import StatusGuardActor
 
         self._status_guard_ref: xo.ActorRefType[  # type: ignore
@@ -142,6 +143,13 @@ class SupervisorActor(xo.StatelessActor):
         ] = await xo.create_actor(
             CacheTrackerActor, address=self.address, uid=CacheTrackerActor.default_uid()
         )
+        self._progress_tracker: xo.ActorRefType[  # type: ignore
+            "ProgressTrackerActor"
+        ] = await xo.create_actor(
+            ProgressTrackerActor,
+            address=self.address,
+            uid=ProgressTrackerActor.default_uid(),
+        )
 
         from .event import EventCollectorActor
 
@@ -1360,3 +1368,6 @@ class SupervisorActor(xo.StatelessActor):
     @staticmethod
     def record_metrics(name, op, kwargs):
         record_metrics(name, op, kwargs)
+
+    async def get_progress(self, request_id: str) -> float:
+        return await self._progress_tracker.get_progress(request_id)

xinference/core/worker.py

@@ -885,6 +885,7 @@ class WorkerActor(xo.StatelessActor):
                     ModelActor,
                     address=subpool_address,
                     uid=model_uid,
+                    supervisor_address=self._supervisor_address,
                     worker_address=self.address,
                     model=model,
                     model_description=model_description,

xinference/model/audio/chattts.py

@@ -53,7 +53,8 @@ class ChatTTSModel:
         torch._dynamo.config.suppress_errors = True
         torch.set_float32_matmul_precision("high")
         self._model = ChatTTS.Chat()
-        self._model.load(source="custom", custom_path=self._model_path, compile=True)
+        logger.info("Load ChatTTS model with kwargs: %s", self._kwargs)
+        self._model.load(source="custom", custom_path=self._model_path, **self._kwargs)
 
     def speech(
         self,

xinference/model/audio/core.py

@@ -25,8 +25,6 @@ from .fish_speech import FishSpeechModel
 from .funasr import FunASRModel
 from .whisper import WhisperModel
 
-MAX_ATTEMPTS = 3
-
 logger = logging.getLogger(__name__)
 
 # Used for check whether the model is cached.