BUG: request_limits does not work with streaming interfaces (#2571)

xorbitsai · Nov 25, 2024 · ddbe211 · ddbe211
1 parent f2b22bb
commit ddbe211
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 13 deletions.
diff --git a/xinference/api/restful_api.py b/xinference/api/restful_api.py
@@ -1274,6 +1274,8 @@ async def stream_results():
                     # https://github.com/openai/openai-python/blob/e0aafc6c1a45334ac889fe3e54957d309c3af93f/src/openai/_streaming.py#L107
                     yield dict(data=json.dumps({"error": str(ex)}))
                     return
+                finally:
+                    await model.decrease_serve_count()
 
             return EventSourceResponse(stream_results())
         else:
@@ -1540,8 +1542,16 @@ async def create_speech(
                 **parsed_kwargs,
             )
             if body.stream:
+
+                async def stream_results():
+                    try:
+                        async for item in out:
+                            yield item
+                    finally:
+                        await model.decrease_serve_count()
+
                 return EventSourceResponse(
-                    media_type="application/octet-stream", content=out
+                    media_type="application/octet-stream", content=stream_results()
                 )
             else:
                 return Response(media_type="application/octet-stream", content=out)
@@ -2072,6 +2082,8 @@ async def stream_results():
                     # https://github.com/openai/openai-python/blob/e0aafc6c1a45334ac889fe3e54957d309c3af93f/src/openai/_streaming.py#L107
                     yield dict(data=json.dumps({"error": str(ex)}))
                     return
+                finally:
+                    await model.decrease_serve_count()
 
             return EventSourceResponse(stream_results())
         else:

diff --git a/xinference/core/model.py b/xinference/core/model.py
@@ -91,21 +91,26 @@ async def wrapped_func(self, *args, **kwargs):
         logger.debug(
             f"Request {fn.__name__}, current serve request count: {self._serve_count}, request limit: {self._request_limits} for the model {self.model_uid()}"
         )
-        if self._request_limits is not None:
-            if 1 + self._serve_count <= self._request_limits:
-                self._serve_count += 1
-            else:
-                raise RuntimeError(
-                    f"Rate limit reached for the model. Request limit {self._request_limits} for the model: {self.model_uid()}"
-                )
+        if 1 + self._serve_count <= self._request_limits:
+            self._serve_count += 1
+        else:
+            raise RuntimeError(
+                f"Rate limit reached for the model. Request limit {self._request_limits} for the model: {self.model_uid()}"
+            )
+        ret = None
         try:
             ret = await fn(self, *args, **kwargs)
         finally:
-            if self._request_limits is not None:
+            if ret is not None and (
+                inspect.isasyncgen(ret) or inspect.isgenerator(ret)
+            ):
+                # stream case, let client call model_ref to decrease self._serve_count
+                pass
+            else:
                 self._serve_count -= 1
-            logger.debug(
-                f"After request {fn.__name__}, current serve request count: {self._serve_count} for the model {self.model_uid()}"
-            )
+                logger.debug(
+                    f"After request {fn.__name__}, current serve request count: {self._serve_count} for the model {self.model_uid()}"
+                )
         return ret
 
     return wrapped_func
@@ -215,7 +220,9 @@ def __init__(
         self._model_description = (
             model_description.to_dict() if model_description else {}
         )
-        self._request_limits = request_limits
+        self._request_limits = (
+            float("inf") if request_limits is None else request_limits
+        )
         self._pending_requests: asyncio.Queue = asyncio.Queue()
         self._handle_pending_requests_task = None
         self._lock = (
@@ -268,6 +275,9 @@ async def __post_create__(self):
     def __repr__(self) -> str:
         return f"ModelActor({self._replica_model_uid})"
 
+    def decrease_serve_count(self):
+        self._serve_count -= 1
+
     async def _record_completion_metrics(
         self, duration, completion_tokens, prompt_tokens
     ):