xinference 1.0.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

xinference/_compat.py

@@ -72,6 +72,7 @@ OpenAIChatCompletionToolParam = create_model_from_typeddict(ChatCompletionToolPa
 OpenAIChatCompletionNamedToolChoiceParam = create_model_from_typeddict(
     ChatCompletionNamedToolChoiceParam
 )
+from openai._types import Body
 
 
 class JSONSchema(BaseModel):
@@ -120,4 +121,5 @@ class CreateChatCompletionOpenAI(BaseModel):
     tools: Optional[Iterable[OpenAIChatCompletionToolParam]]  # type: ignore
     top_logprobs: Optional[int]
     top_p: Optional[float]
+    extra_body: Optional[Body]
     user: Optional[str]

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-11-29T16:57:04+0800",
+ "date": "2025-01-24T16:52:57+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "eb8ddd431f5c5fcb2216e25e0d43745f8455d9b9",
- "version": "1.0.1"
+ "full-revisionid": "a57b99b07b40d1082f69a8fc5b968d56bc3636bc",
+ "version": "1.2.1"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -94,9 +94,9 @@ class CreateCompletionRequest(CreateCompletion):
 
 class CreateEmbeddingRequest(BaseModel):
     model: str
-    input: Union[str, List[str], List[int], List[List[int]]] = Field(
-        description="The input to embed."
-    )
+    input: Union[
+        str, List[str], List[int], List[List[int]], Dict[str, str], List[Dict[str, str]]
+    ] = Field(description="The input to embed.")
     user: Optional[str] = None
 
     class Config:
@@ -1214,6 +1214,19 @@ class RESTfulAPI(CancelMixin):
     async def get_address(self) -> JSONResponse:
         return JSONResponse(content=self._supervisor_address)
 
+    async def _get_model_last_error(self, replica_model_uid: bytes, e: Exception):
+        if not isinstance(e, xo.ServerClosed):
+            return e
+        try:
+            model_status = await (await self._get_supervisor_ref()).get_model_status(
+                replica_model_uid.decode("utf-8")
+            )
+            if model_status is not None and model_status.last_error:
+                return Exception(model_status.last_error)
+        except Exception as ex:
+            return ex
+        return e
+
     async def create_completion(self, request: Request) -> Response:
         raw_body = await request.json()
         body = CreateCompletionRequest.parse_obj(raw_body)
@@ -1272,6 +1285,7 @@ class RESTfulAPI(CancelMixin):
                     )
                     return
                 except Exception as ex:
+                    ex = await self._get_model_last_error(model.uid, ex)
                     logger.exception("Completion stream got an error: %s", ex)
                     await self._report_error_event(model_uid, str(ex))
                     # https://github.com/openai/openai-python/blob/e0aafc6c1a45334ac889fe3e54957d309c3af93f/src/openai/_streaming.py#L107
@@ -1286,6 +1300,7 @@ class RESTfulAPI(CancelMixin):
                 data = await model.generate(body.prompt, kwargs, raw_params=raw_kwargs)
                 return Response(data, media_type="application/json")
             except Exception as e:
+                e = await self._get_model_last_error(model.uid, e)
                 logger.error(e, exc_info=True)
                 await self._report_error_event(model_uid, str(e))
                 self.handle_request_limit_error(e)
@@ -1317,14 +1332,11 @@ class RESTfulAPI(CancelMixin):
         try:
             embedding = await model.create_embedding(body.input, **kwargs)
             return Response(embedding, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def convert_ids_to_tokens(self, request: Request) -> Response:
@@ -1352,14 +1364,11 @@ class RESTfulAPI(CancelMixin):
         try:
             decoded_texts = await model.convert_ids_to_tokens(body.input, **kwargs)
             return Response(decoded_texts, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def rerank(self, request: Request) -> Response:
@@ -1393,14 +1402,11 @@ class RESTfulAPI(CancelMixin):
                 **parsed_kwargs,
             )
             return Response(scores, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_transcriptions(
@@ -1445,13 +1451,11 @@ class RESTfulAPI(CancelMixin):
                 **parsed_kwargs,
             )
             return Response(content=transcription, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_translations(
@@ -1496,13 +1500,11 @@ class RESTfulAPI(CancelMixin):
                 **parsed_kwargs,
             )
             return Response(content=translation, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_speech(
@@ -1558,14 +1560,11 @@ class RESTfulAPI(CancelMixin):
                 )
             else:
                 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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def get_progress(self, request_id: str) -> JSONResponse:
@@ -1611,14 +1610,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def sdapi_options(self, request: Request) -> Response:
@@ -1689,14 +1685,11 @@ class RESTfulAPI(CancelMixin):
                 **kwargs,
             )
             return Response(content=image_list, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def sdapi_img2img(self, request: Request) -> Response:
@@ -1723,14 +1716,11 @@ class RESTfulAPI(CancelMixin):
                 **kwargs,
             )
             return Response(content=image_list, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_variations(
@@ -1779,13 +1769,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_inpainting(
@@ -1841,13 +1829,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_ocr(
@@ -1887,13 +1873,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_flexible_infer(self, request: Request) -> Response:
@@ -1920,14 +1904,11 @@ class RESTfulAPI(CancelMixin):
         try:
             result = await model.infer(**kwargs)
             return Response(result, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_videos(self, request: Request) -> Response:
@@ -1952,14 +1933,11 @@ class RESTfulAPI(CancelMixin):
                 **kwargs,
             )
             return Response(content=video_list, media_type="application/json")
-        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:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_chat_completion(self, request: Request) -> Response:
@@ -2044,7 +2022,6 @@ class RESTfulAPI(CancelMixin):
                 )
         if body.tools and body.stream:
             is_vllm = await model.is_vllm_backend()
-
             if not (
                 (is_vllm and model_family in QWEN_TOOL_CALL_FAMILY)
                 or (not is_vllm and model_family in GLM4_TOOL_CALL_FAMILY)
@@ -2054,7 +2031,8 @@ class RESTfulAPI(CancelMixin):
                     detail="Streaming support for tool calls is available only when using "
                     "Qwen models with vLLM backend or GLM4-chat models without vLLM backend.",
                 )
-
+        if "skip_special_tokens" in raw_kwargs and await model.is_vllm_backend():
+            kwargs["skip_special_tokens"] = raw_kwargs["skip_special_tokens"]
         if body.stream:
 
             async def stream_results():
@@ -2084,6 +2062,7 @@ class RESTfulAPI(CancelMixin):
                     # TODO: Cannot yield here. Yield here would leads to error for the next streaming request.
                     return
                 except Exception as ex:
+                    ex = await self._get_model_last_error(model.uid, ex)
                     logger.exception("Chat completion stream got an error: %s", ex)
                     await self._report_error_event(model_uid, str(ex))
                     # https://github.com/openai/openai-python/blob/e0aafc6c1a45334ac889fe3e54957d309c3af93f/src/openai/_streaming.py#L107
@@ -2102,6 +2081,7 @@ class RESTfulAPI(CancelMixin):
                 )
                 return Response(content=data, media_type="application/json")
             except Exception as e:
+                e = await self._get_model_last_error(model.uid, e)
                 logger.error(e, exc_info=True)
                 await self._report_error_event(model_uid, str(e))
                 self.handle_request_limit_error(e)
@@ -2346,7 +2326,8 @@ class RESTfulAPI(CancelMixin):
     @staticmethod
     def extract_guided_params(raw_body: dict) -> dict:
         kwargs = {}
-        if raw_body.get("guided_json") is not None:
+        raw_extra_body: dict = raw_body.get("extra_body")  # type: ignore
+        if raw_body.get("guided_json"):
             kwargs["guided_json"] = raw_body.get("guided_json")
         if raw_body.get("guided_regex") is not None:
             kwargs["guided_regex"] = raw_body.get("guided_regex")
@@ -2362,6 +2343,31 @@ class RESTfulAPI(CancelMixin):
             kwargs["guided_whitespace_pattern"] = raw_body.get(
                 "guided_whitespace_pattern"
             )
+        # Parse OpenAI extra_body
+        if raw_extra_body is not None:
+            if raw_extra_body.get("guided_json"):
+                kwargs["guided_json"] = raw_extra_body.get("guided_json")
+            if raw_extra_body.get("guided_regex") is not None:
+                kwargs["guided_regex"] = raw_extra_body.get("guided_regex")
+            if raw_extra_body.get("guided_choice") is not None:
+                kwargs["guided_choice"] = raw_extra_body.get("guided_choice")
+            if raw_extra_body.get("guided_grammar") is not None:
+                kwargs["guided_grammar"] = raw_extra_body.get("guided_grammar")
+            if raw_extra_body.get("guided_json_object") is not None:
+                kwargs["guided_json_object"] = raw_extra_body.get("guided_json_object")
+            if raw_extra_body.get("guided_decoding_backend") is not None:
+                kwargs["guided_decoding_backend"] = raw_extra_body.get(
+                    "guided_decoding_backend"
+                )
+            if raw_extra_body.get("guided_whitespace_pattern") is not None:
+                kwargs["guided_whitespace_pattern"] = raw_extra_body.get(
+                    "guided_whitespace_pattern"
+                )
+            if raw_extra_body.get("platform") is not None:
+                kwargs["platform"] = raw_extra_body.get("platform")
+            if raw_extra_body.get("format") is not None:
+                kwargs["format"] = raw_extra_body.get("format")
+
         return kwargs
 
 

xinference/core/chat_interface.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import base64
+import html
 import logging
 import os
 from io import BytesIO
@@ -137,7 +138,11 @@ class GradioInterface:
                 if "content" not in delta:
                     continue
                 else:
-                    response_content += delta["content"]
+                    # some model like deepseek-r1-distill-qwen
+                    # will generate <think>...</think> ...
+                    # in gradio, no output will be rendered,
+                    # thus escape html tags in advance
+                    response_content += html.escape(delta["content"])
                     yield response_content
 
             yield response_content

xinference/core/model.py

@@ -35,6 +35,7 @@ from typing import (
     List,
     Optional,
     Union,
+    no_type_check,
 )
 
 import sse_starlette.sse
@@ -78,6 +79,9 @@ XINFERENCE_BATCHING_ALLOWED_VISION_MODELS = [
 ]
 
 XINFERENCE_TEXT_TO_IMAGE_BATCHING_ALLOWED_MODELS = ["FLUX.1-dev", "FLUX.1-schnell"]
+XINFERENCE_TEST_OUT_OF_MEMORY_ERROR = bool(
+    os.getenv("XINFERENCE_TEST_OUT_OF_MEMORY_ERROR", False)
+)
 
 
 def request_limit(fn):
@@ -118,20 +122,25 @@ def request_limit(fn):
 
 def oom_check(fn):
     @functools.wraps(fn)
-    def _wrapper(*args, **kwargs):
+    def _wrapper(self, *args, **kwargs):
         try:
-            return fn(*args, **kwargs)
-        except OutOfMemoryError:
-            logger.exception("Model actor is out of memory.")
-            os._exit(1)
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
+            return fn(self, *args, **kwargs)
+        except OutOfMemoryError as ex:
+            assert self._loop is not None
+            asyncio.run_coroutine_threadsafe(
+                self._handle_oom_error(ex), loop=self._loop
+            )
 
     @functools.wraps(fn)
-    async def _async_wrapper(*args, **kwargs):
+    async def _async_wrapper(self, *args, **kwargs):
         try:
-            return await fn(*args, **kwargs)
-        except OutOfMemoryError:
-            logger.exception("Model actor is out of memory.")
-            os._exit(1)
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
+            return await fn(self, *args, **kwargs)
+        except OutOfMemoryError as ex:
+            await self._handle_oom_error(ex)
 
     assert not inspect.isasyncgen(fn)
     assert not inspect.isgenerator(fn)
@@ -178,6 +187,16 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         if hasattr(self._model, "stop") and callable(self._model.stop):
             self._model.stop()
 
+        if isinstance(self._model, LLMVLLMModel):
+            if self._transfer_ref is not None:
+                try:
+                    await xo.destroy_actor(self._transfer_ref)
+                    del self._transfer_ref
+                except Exception as e:
+                    logger.debug(
+                        f"Destroy transfer actor failed, address: {self.address}, error: {e}"
+                    )
+
         if (
             isinstance(self._model, (LLMPytorchModel, LLMVLLMModel, SGLANGModel))
             and self._model.model_spec.model_format == "pytorch"
@@ -206,6 +225,7 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         replica_model_uid: str,
         model_description: Optional["ModelDescription"] = None,
         request_limits: Optional[int] = None,
+        xavier_config: Optional[Dict] = None,
     ):
         super().__init__()
         from ..model.llm.lmdeploy.core import LMDeployModel
@@ -247,6 +267,11 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         self._scheduler_ref = None
         self._text_to_image_scheduler_ref = None
 
+        if isinstance(self._model, VLLMModel):
+            self._xavier_config = xavier_config
+            self._model.set_xavier_config(xavier_config)
+            self._transfer_ref = None
+
     async def __post_create__(self):
         self._loop = asyncio.get_running_loop()
 
@@ -278,6 +303,29 @@ class ModelActor(xo.StatelessActor, CancelMixin):
     def decrease_serve_count(self):
         self._serve_count -= 1
 
+    @no_type_check
+    async def start_transfer_for_vllm(self, rank_addresses: List[str]):
+        from ..model.llm.vllm.core import VLLMModel
+        from ..model.llm.vllm.xavier.transfer import TransferActor
+
+        assert isinstance(self._model, VLLMModel)
+        rank = self._xavier_config.get("rank")  # type: ignore
+        self._transfer_ref = await xo.create_actor(
+            TransferActor,
+            address=self.address,
+            uid=f"{TransferActor.default_uid()}-{rank}",
+            rank=rank,
+            world_size=self._xavier_config.get("world_size"),  # type: ignore
+            rank_address=self._xavier_config.get("rank_address"),  # type: ignore
+            store_address=self._xavier_config.get("store_address"),  # type: ignore
+            store_port=self._xavier_config.get("store_port"),  # type: ignore
+            world_addresses=rank_addresses,
+        )
+        await self._model.init_xavier()
+        logger.debug(
+            f"Init transfer actor: {self._transfer_ref.address}, rank: {rank} done for vllm."  # type: ignore
+        )
+
     async def _record_completion_metrics(
         self, duration, completion_tokens, prompt_tokens
     ):
@@ -440,11 +488,24 @@ class ModelActor(xo.StatelessActor, CancelMixin):
             )
         )
 
+    async def _handle_oom_error(self, ex):
+        error_message = (
+            f"Model actor is out of memory, model id: {self.model_uid()}, error: {ex}"
+        )
+        logger.exception(error_message)
+        worker_ref = await self._get_worker_ref()
+        await worker_ref.update_model_status(
+            self._replica_model_uid, last_error=error_message
+        )
+        os._exit(1)
+
     def _to_generator(self, output_type: str, gen: types.GeneratorType):
         start_time = time.time()
         time_to_first_token = None
         final_usage = None
         try:
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
             for v in gen:
                 if time_to_first_token is None:
                     time_to_first_token = (time.time() - start_time) * 1000
@@ -456,11 +517,11 @@ class ModelActor(xo.StatelessActor, CancelMixin):
                         output_type == "binary"
                     ), f"Unknown output type '{output_type}'"
                 yield sse_starlette.sse.ensure_bytes(v, None)
-        except OutOfMemoryError:
-            logger.exception(
-                "Model actor is out of memory, model id: %s", self.model_uid()
+        except OutOfMemoryError as ex:
+            assert self._loop is not None
+            asyncio.run_coroutine_threadsafe(
+                self._handle_oom_error(ex), loop=self._loop
             )
-            os._exit(1)
         finally:
             if self._loop is not None and time_to_first_token is not None:
                 coro = self.record_metrics(
@@ -482,6 +543,8 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         time_to_first_token = None
         final_usage = None
         try:
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
             async for v in gen:
                 if time_to_first_token is None:
                     time_to_first_token = (time.time() - start_time) * 1000
@@ -494,11 +557,8 @@ class ModelActor(xo.StatelessActor, CancelMixin):
                         output_type == "binary"
                     ), f"Unknown output type '{output_type}'"
                 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)
+        except OutOfMemoryError as ex:
+            await self._handle_oom_error(ex)
         finally:
             coros = []
             if time_to_first_token is not None: