diff --git a/xinference/client/tests/test_async_client.py b/xinference/client/tests/test_async_client.py index 36d9e61bea..bae4b6934f 100644 --- a/xinference/client/tests/test_async_client.py +++ b/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) diff --git a/xinference/client/tests/test_client.py b/xinference/client/tests/test_client.py index 6c9238d555..2385d3bd58 100644 --- a/xinference/client/tests/test_client.py +++ b/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() diff --git a/xinference/core/model.py b/xinference/core/model.py index a4937ea495..79cba41349 100644 --- a/xinference/core/model.py +++ b/xinference/core/model.py @@ -18,7 +18,6 @@ import json import os import queue -import sys import time import types import uuid @@ -519,35 +518,44 @@ 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 @@ -555,9 +563,19 @@ async def _handle_oom_error(self, ex): 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()