fix(model): use os._exit for OOM to trigger pool recovery

xinference/client/tests/test_async_client.py

@@ -616,6 +616,45 @@ async def test_model_error(set_test_oom_error, setup_cluster):
     await client.close()
 
 
+@pytest.mark.asyncio
+async def test_model_oom_triggers_restart(set_test_oom_error, setup_cluster):
+    # OOM must hard-exit the model subprocess via os._exit(1) so xoscar
+    # can detect the exit and recover it. With the old sys.exit(1) the
+    # SystemExit was swallowed by xoscar's _ErrorProcessor, the process
+    # never died, and this test would time out at the "subprocess did not
+    # exit" assert below.
+    endpoint, _ = setup_cluster
+    current_proc = psutil.Process()
+    before = set(current_proc.children(recursive=True))
+    client = AsyncRESTfulClient(endpoint)
+
+    model_uid = await client.launch_model(
+        model_name="qwen1.5-chat",
+        model_engine="llama.cpp",
+        model_size_in_billions="0_5",
+        quantization="q4_0",
+    )
+    assert len(await client.list_models()) == 1
+    model_proc = next(iter(set(current_proc.children(recursive=True)) - before))
+
+    model = await client.get_model(model_uid=model_uid)
+    assert isinstance(model, AsyncRESTfulChatModelHandle)
+
+    # Trigger OOM -> os._exit(1) on the model subprocess.
+    with pytest.raises(Exception):
+        await model.generate("Once upon a time, there was a very old computer")
+
+    # The old subprocess must actually exit (regression guard).
+    for _ in range(30):
+        if not model_proc.is_running() or model_proc.status() == psutil.STATUS_ZOMBIE:
+            break
+        await asyncio.sleep(1)
+    else:
+        assert False, "OOM subprocess did not exit (sys.exit swallowed by xoscar)"
+    await model.close()
+    await client.close()
+
+
 @pytest.mark.asyncio
 async def test_async_restful_chat_enable_thinking_injected():
     handle = AsyncRESTfulChatModelHandle.__new__(AsyncRESTfulChatModelHandle)

xinference/client/tests/test_client.py

@@ -545,6 +545,44 @@ def test_model_error(set_test_oom_error, setup_cluster):
         model.generate("Once upon a time, there was a very old computer")
 
 
+def test_model_oom_triggers_restart(set_test_oom_error, setup_cluster):
+    # OOM must hard-exit the model subprocess via os._exit(1) so xoscar
+    # can detect the exit and recover it. With the old sys.exit(1) the
+    # SystemExit was swallowed by xoscar's _ErrorProcessor, the process
+    # never died, and this test would time out at the "subprocess did not
+    # exit" assert below.
+    endpoint, _ = setup_cluster
+    current_proc = psutil.Process()
+    before = set(current_proc.children(recursive=True))
+    client = RESTfulClient(endpoint)
+
+    model_uid = client.launch_model(
+        model_name="qwen2.5-instruct",
+        model_engine="llama.cpp",
+        model_size_in_billions="0_5",
+        model_format="ggufv2",
+        quantization="q4_0",
+    )
+    assert len(client.list_models()) == 1
+    model_proc = next(iter(set(current_proc.children(recursive=True)) - before))
+
+    model = client.get_model(model_uid=model_uid)
+    assert isinstance(model, RESTfulChatModelHandle)
+
+    # Trigger OOM. The handler runs os._exit(1) on the model subprocess, so the
+    # client sees an error (OOM report or connection drop) -- either is fine.
+    with pytest.raises(Exception):
+        model.generate("Once upon a time, there was a very old computer")
+
+    # The old subprocess must actually exit (this is the regression guard).
+    for _ in range(30):
+        if not model_proc.is_running() or model_proc.status() == psutil.STATUS_ZOMBIE:
+            break
+        time.sleep(1)
+    else:
+        assert False, "OOM subprocess did not exit (sys.exit swallowed by xoscar)"
+
+
 def test_restful_chat_enable_thinking_injected():
     handle = RESTfulChatModelHandle("test-model", "http://localhost", {})
     dummy_session = _DummySession()

xinference/core/model.py

@@ -18,7 +18,6 @@
 import json
 import os
 import queue
-import sys
 import time
 import types
 import uuid
@@ -519,45 +518,64 @@ async def stop(self):
             await asyncio.to_thread(self._model.close)
 
     async def _handle_oom_error(self, ex):
+        # Reentrancy gate: on the sync batch path the OOM-triggering worker
+        # thread does not await this coroutine, so before os._exit(1) actually
+        # lands it can raise OOM again and concurrently schedule this handler.
+        # Only let the first invocation through to avoid re-entering stop()/del.
+        if getattr(self, "_oom_handling", False):
+            return
+        self._oom_handling = True
+
         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
-        )
-        # §4.1: Graceful cleanup instead of os._exit(1).
-        # 1) Stop model to release underlying resources (with timeout guard)
-        try:
-            await asyncio.wait_for(self.stop(), timeout=30)
-        except asyncio.TimeoutError:
-            logger.warning(
-                "Model stop() timed out after 30s for %s, proceeding with cleanup",
-                self.model_uid(),
+
+        async def _report_oom_status():
+            worker_ref = await self._get_worker_ref()
+            await worker_ref.update_model_status(
+                self._replica_model_uid, last_error=error_message
             )
+
+        try:
+            await xo.wait_for(_report_oom_status(), timeout=5)
+        except Exception:
+            logger.warning("Failed to report OOM status before exit", exc_info=True)
+
+        # Best-effort cleanup. At an OOM site this mostly cannot reclaim already
+        # reserved GPU fragments; failures are fine because os._exit(1) below is
+        # the real backstop. stop() timeout is 5s (not 30s): keep the exit delay
+        # bounded so xoscar's detection and recovery are not slowed down.
+        try:
+            await asyncio.wait_for(self.stop(), timeout=5)
         except Exception:
             logger.warning(
-                "Model stop() failed for %s, proceeding with cleanup",
-                self.model_uid(),
-                exc_info=True,
+                "Model stop() failed/timeout during OOM cleanup", exc_info=True
             )
-        # 2) Delete model reference to allow GC to reclaim GPU memory
         try:
             del self._model
         except AttributeError:
             pass
-        # 3) Release PyTorch GPU cache
         try:
             import torch
 
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
         except Exception:
             pass
-        # 4) Use sys.exit(1) instead of os._exit(1) so that Python cleanup hooks
-        #    run and xoscar can detect the sub-process exit, triggering recover_sub_pool.
-        sys.exit(1)
+
+        # Critical fix: must use os._exit(1), NOT sys.exit(1).
+        # sys.exit(1) only raises SystemExit, which xoscar's
+        # _ErrorProcessor.__exit__ (backends/pool.py) swallows with an
+        # unconditional `return True`. The process then stays alive, its
+        # returncode stays None, monitor_sub_pools keeps treating it as alive,
+        # and recover_sub_pool never fires. os._exit goes straight to the
+        # syscall so xoscar can detect the returncode change and rebuild the
+        # sub pool.
+        logger.error(
+            "Exiting model subprocess via os._exit(1) to trigger pool recovery"
+        )
+        os._exit(1)
 
     def _to_generator(self, output_type: str, gen: types.GeneratorType):
         start_time = time.time()