xinference 0.11.3__py3-none-any.whl → 0.12.0__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-05-31T17:12:13+0800",
+ "date": "2024-06-07T15:04:33+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "69c09cd068a530cd2fdcac07e4e81f03d48f04f9",
- "version": "0.11.3"
+ "full-revisionid": "55c5636f2b6022842d1827eae373c8e5f162a1a3",
+ "version": "0.12.0"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -122,6 +122,14 @@ class TextToImageRequest(BaseModel):
     user: Optional[str] = None
 
 
+class SpeechRequest(BaseModel):
+    model: str
+    input: str
+    voice: Optional[str]
+    response_format: Optional[str] = "mp3"
+    speed: Optional[float] = 1.0
+
+
 class RegisterModelRequest(BaseModel):
     model: str
     persist: bool
@@ -337,6 +345,16 @@ class RESTfulAPI:
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/models/{model_uid}/requests/{request_id}/abort",
+            self.abort_request,
+            methods=["POST"],
+            dependencies=(
+                [Security(self._auth_service, scopes=["models:read"])]
+                if self.is_authenticated()
+                else None
+            ),
+        )
         self._router.add_api_route(
             "/v1/models/instance",
             self.launch_model_by_version,
@@ -418,6 +436,16 @@ class RESTfulAPI:
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/audio/speech",
+            self.create_speech,
+            methods=["POST"],
+            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,
@@ -1179,6 +1207,38 @@ class RESTfulAPI:
             await self._report_error_event(model_uid, str(e))
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def create_speech(self, request: Request) -> Response:
+        body = SpeechRequest.parse_obj(await request.json())
+        model_uid = body.model
+        try:
+            model = await (await self._get_supervisor_ref()).get_model(model_uid)
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            await self._report_error_event(model_uid, str(ve))
+            raise HTTPException(status_code=400, detail=str(ve))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            await self._report_error_event(model_uid, str(e))
+            raise HTTPException(status_code=500, detail=str(e))
+
+        try:
+            out = await model.speech(
+                input=body.input,
+                voice=body.voice,
+                response_format=body.response_format,
+                speed=body.speed,
+            )
+            return Response(media_type="application/octet-stream", content=out)
+        except RuntimeError as re:
+            logger.error(re, exc_info=True)
+            await self._report_error_event(model_uid, str(re))
+            self.handle_request_limit_error(re)
+            raise HTTPException(status_code=400, detail=str(re))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            await self._report_error_event(model_uid, str(e))
+            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
@@ -1518,6 +1578,15 @@ class RESTfulAPI:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def abort_request(self, model_uid: str, request_id: str) -> JSONResponse:
+        try:
+            supervisor_ref = await self._get_supervisor_ref()
+            res = await supervisor_ref.abort_request(model_uid, request_id)
+            return JSONResponse(content=res)
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def list_vllm_supported_model_families(self) -> JSONResponse:
         try:
             from ..model.llm.vllm.core import (

xinference/client/restful/restful_client.py

@@ -684,6 +684,49 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
         response_data = response.json()
         return response_data
 
+    def speech(
+        self,
+        input: str,
+        voice: str = "",
+        response_format: str = "mp3",
+        speed: float = 1.0,
+    ):
+        """
+        Generates audio from the input text.
+
+        Parameters
+        ----------
+
+        input: str
+            The text to generate audio for. The maximum length is 4096 characters.
+        voice: str
+            The voice to use when generating the audio.
+        response_format: str
+            The format to audio in.
+        speed: str
+            The speed of the generated audio.
+
+        Returns
+        -------
+        bytes
+            The generated audio binary.
+        """
+        url = f"{self._base_url}/v1/audio/speech"
+        params = {
+            "model": self._model_uid,
+            "input": input,
+            "voice": voice,
+            "response_format": response_format,
+            "speed": speed,
+        }
+        response = requests.post(url, json=params, headers=self.auth_headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to speech the text, detail: {_get_error_string(response)}"
+            )
+
+        return response.content
+
 
 class Client:
     def __init__(self, base_url, api_key: Optional[str] = None):
@@ -1181,3 +1224,30 @@ class Client:
 
         response_data = response.json()
         return response_data
+
+    def abort_request(self, model_uid: str, request_id: str):
+        """
+        Abort a request.
+        Abort a submitted request. If the request is finished or not found, this method will be a no-op.
+        Currently, this interface is only supported when batching is enabled for models on transformers backend.
+
+        Parameters
+        ----------
+        model_uid: str
+            Model uid.
+        request_id: str
+            Request id.
+        Returns
+        -------
+        Dict
+            Return empty dict.
+        """
+        url = f"{self.base_url}/v1/models/{model_uid}/requests/{request_id}/abort"
+        response = requests.post(url, headers=self._headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to abort request, detail: {_get_error_string(response)}"
+            )
+
+        response_data = response.json()
+        return response_data

xinference/constants.py

@@ -27,6 +27,7 @@ XINFERENCE_ENV_DISABLE_HEALTH_CHECK = "XINFERENCE_DISABLE_HEALTH_CHECK"
 XINFERENCE_ENV_DISABLE_VLLM = "XINFERENCE_DISABLE_VLLM"
 XINFERENCE_ENV_ENABLE_SGLANG = "XINFERENCE_ENABLE_SGLANG"
 XINFERENCE_ENV_DISABLE_METRICS = "XINFERENCE_DISABLE_METRICS"
+XINFERENCE_ENV_TRANSFORMERS_ENABLE_BATCHING = "XINFERENCE_TRANSFORMERS_ENABLE_BATCHING"
 
 
 def get_xinference_home() -> str:
@@ -70,3 +71,6 @@ XINFERENCE_ENABLE_SGLANG = bool(int(os.environ.get(XINFERENCE_ENV_ENABLE_SGLANG,
 XINFERENCE_DISABLE_METRICS = bool(
     int(os.environ.get(XINFERENCE_ENV_DISABLE_METRICS, 0))
 )
+XINFERENCE_TRANSFORMERS_ENABLE_BATCHING = bool(
+    int(os.environ.get(XINFERENCE_ENV_TRANSFORMERS_ENABLE_BATCHING, 0))
+)

xinference/core/model.py

@@ -20,9 +20,14 @@ import os
 import time
 import types
 import weakref
+from asyncio.queues import Queue
+from asyncio.tasks import wait_for
+from concurrent.futures import Future as ConcurrentFuture
 from typing import (
     TYPE_CHECKING,
+    Any,
     AsyncGenerator,
+    AsyncIterator,
     Callable,
     Dict,
     Generator,
@@ -35,6 +40,8 @@ from typing import (
 import sse_starlette.sse
 import xoscar as xo
 
+from ..constants import XINFERENCE_TRANSFORMERS_ENABLE_BATCHING
+
 if TYPE_CHECKING:
     from .worker import WorkerActor
     from ..model.llm.core import LLM
@@ -125,6 +132,16 @@ class ModelActor(xo.StatelessActor):
         from ..model.llm.pytorch.core import PytorchModel as LLMPytorchModel
         from ..model.llm.vllm.core import VLLMModel as LLMVLLMModel
 
+        if self.allow_batching():
+            try:
+                assert self._scheduler_ref is not None
+                await xo.destroy_actor(self._scheduler_ref)
+                del self._scheduler_ref
+            except Exception as e:
+                logger.debug(
+                    f"Destroy scheduler actor failed, address: {self.address}, error: {e}"
+                )
+
         if (
             isinstance(self._model, (LLMPytorchModel, LLMVLLMModel))
             and self._model.model_spec.model_format == "pytorch"
@@ -181,9 +198,20 @@ class ModelActor(xo.StatelessActor):
         }
         self._loop: Optional[asyncio.AbstractEventLoop] = None
 
+        self._scheduler_ref = None
+
     async def __post_create__(self):
         self._loop = asyncio.get_running_loop()
 
+        if self.allow_batching():
+            from .scheduler import SchedulerActor
+
+            self._scheduler_ref = await xo.create_actor(
+                SchedulerActor,
+                address=self.address,
+                uid=SchedulerActor.gen_uid(self.model_uid(), self._model.rep_id),
+            )
+
     async def _record_completion_metrics(
         self, duration, completion_tokens, prompt_tokens
     ):
@@ -235,8 +263,22 @@ class ModelActor(xo.StatelessActor):
 
         return isinstance(self._model, VLLMModel)
 
-    def load(self):
+    def allow_batching(self) -> bool:
+        from ..model.llm.pytorch.core import PytorchChatModel
+
+        return (
+            XINFERENCE_TRANSFORMERS_ENABLE_BATCHING
+            and isinstance(self._model, PytorchChatModel)
+            and self._model.__class__.__name__ == PytorchChatModel.__name__
+        )
+
+    async def load(self):
         self._model.load()
+        if self.allow_batching():
+            await self._scheduler_ref.set_model(self._model)
+            logger.debug(
+                f"Batching enabled for model: {self.model_uid()}, max_num_seqs: {self._model.get_max_num_seqs()}"
+            )
 
     def model_uid(self):
         return (
@@ -343,6 +385,8 @@ class ModelActor(xo.StatelessActor):
             gen = self._to_json_async_gen(ret)
             self._current_generator = weakref.ref(gen)
             return gen
+        if isinstance(ret, bytes):
+            return ret
         return await asyncio.to_thread(json_dumps, ret)
 
     @log_async(logger=logger)
@@ -359,6 +403,36 @@ class ModelActor(xo.StatelessActor):
             )
         raise AttributeError(f"Model {self._model.model_spec} is not for generate.")
 
+    async def _queue_consumer(
+        self, queue: Queue, timeout: Optional[float] = None
+    ) -> AsyncIterator[Any]:
+        from .scheduler import (
+            XINFERENCE_STREAMING_ABORT_FLAG,
+            XINFERENCE_STREAMING_DONE_FLAG,
+            XINFERENCE_STREAMING_ERROR_FLAG,
+        )
+
+        while True:
+            # TODO: timeout setting
+            res = await wait_for(queue.get(), timeout)
+            if res == XINFERENCE_STREAMING_DONE_FLAG:
+                break
+            elif res == XINFERENCE_STREAMING_ABORT_FLAG:
+                raise RuntimeError(
+                    f"This request has been cancelled by another `abort_request` request."
+                )
+            elif isinstance(res, str) and res.startswith(
+                XINFERENCE_STREAMING_ERROR_FLAG
+            ):
+                raise RuntimeError(res[len(XINFERENCE_STREAMING_ERROR_FLAG) :])
+            else:
+                yield res
+
+    @staticmethod
+    def get_stream_from_args(*args) -> bool:
+        assert args[2] is None or isinstance(args[2], dict)
+        return False if args[2] is None else args[2].get("stream", False)
+
     @log_async(logger=logger)
     @request_limit
     @xo.generator
@@ -366,17 +440,46 @@ class ModelActor(xo.StatelessActor):
         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.")
+            if self.allow_batching():
+                stream = self.get_stream_from_args(*args)
+                assert self._scheduler_ref is not None
+                if stream:
+                    assert self._scheduler_ref is not None
+                    queue: Queue[Any] = Queue()
+                    ret = self._queue_consumer(queue)
+                    await self._scheduler_ref.add_request(
+                        prompt, queue, *args, **kwargs
+                    )
+                    gen = self._to_json_async_gen(ret)
+                    self._current_generator = weakref.ref(gen)
+                    return gen
+                else:
+                    from .scheduler import XINFERENCE_NON_STREAMING_ABORT_FLAG
+
+                    assert self._loop is not None
+                    future = ConcurrentFuture()
+                    await self._scheduler_ref.add_request(
+                        prompt, future, *args, **kwargs
+                    )
+                    fut = asyncio.wrap_future(future, loop=self._loop)
+                    result = await fut
+                    if result == XINFERENCE_NON_STREAMING_ABORT_FLAG:
+                        raise RuntimeError(
+                            f"This request has been cancelled by another `abort_request` request."
+                        )
+                    return await asyncio.to_thread(json_dumps, result)
+            else:
+                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.")
         finally:
             # For the non stream result.
             record = None
@@ -395,6 +498,15 @@ class ModelActor(xo.StatelessActor):
                     prompt_tokens,
                 )
 
+    async def abort_request(self, request_id: str) -> str:
+        from .scheduler import AbortRequestMessage
+
+        if self.allow_batching():
+            if self._scheduler_ref is None:
+                return AbortRequestMessage.NOT_FOUND.name
+            return await self._scheduler_ref.abort_request(request_id)
+        return AbortRequestMessage.NO_OP.name
+
     @log_async(logger=logger)
     @request_limit
     async def create_embedding(self, input: Union[str, List[str]], *args, **kwargs):
@@ -482,6 +594,23 @@ class ModelActor(xo.StatelessActor):
             f"Model {self._model.model_spec} is not for creating translations."
         )
 
+    @log_async(logger=logger)
+    @request_limit
+    async def speech(
+        self, input: str, voice: str, response_format: str = "mp3", speed: float = 1.0
+    ):
+        if hasattr(self._model, "speech"):
+            return await self._call_wrapper(
+                self._model.speech,
+                input,
+                voice,
+                response_format,
+                speed,
+            )
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating speech."
+        )
+
     @log_async(logger=logger)
     @request_limit
     async def text_to_image(