diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index 61089b295e..d1829bd068 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -932,7 +932,7 @@ async def _inner_call(i, generator, session): async for res in cleanup_generator: if await raw_request.is_disconnected(): # Abort the request if the client disconnects. - await VariableInterface.async_engine.stop_session(request.session_id) + await session.async_abort() return create_error_response(HTTPStatus.BAD_REQUEST, 'Client disconnected') final_res = res text += res.response @@ -960,7 +960,10 @@ async def _inner_call(i, generator, session): usage.completion_tokens += final_res.generate_token_len usage.total_tokens += total_tokens - await asyncio.gather(*[_inner_call(i, generators[i], sessions[i]) for i in range(len(generators))]) + inner_results = await asyncio.gather(*[_inner_call(i, generators[i], sessions[i]) for i in range(len(generators))]) + for inner_result in inner_results: + if inner_result is not None: + return inner_result response = CompletionResponse( id=request_id, diff --git a/tests/test_lmdeploy/serve/test_session_cleanup.py b/tests/test_lmdeploy/serve/test_session_cleanup.py index 8b8b1b0f3f..08ba65f740 100644 --- a/tests/test_lmdeploy/serve/test_session_cleanup.py +++ b/tests/test_lmdeploy/serve/test_session_cleanup.py @@ -257,3 +257,103 @@ async def _run_max_new_tokens_zero_cleans_up_session(): def test_max_new_tokens_zero_cleans_up_session(): asyncio.run(_run_max_new_tokens_zero_cleans_up_session()) + + +async def _run_completions_v1_disconnect_aborts_active_handle(): + from types import SimpleNamespace + + from lmdeploy.serve.core.async_engine import GenOut + from lmdeploy.serve.openai.api_server import VariableInterface, completions_v1 + from lmdeploy.serve.openai.protocol import CompletionRequest + + class _SpyHandle(_FakeHandle): + + def __init__(self): + self.cancelled_session_ids = [] + + async def async_cancel(self, session_id: int): + self.cancelled_session_ids.append(session_id) + + class _FakeAsyncEngine: + model_name = 'test-model' + epoch = 0 + + def __init__(self): + self.backend_config = SimpleNamespace() + self.session_mgr = SessionManager() + self.handle = _SpyHandle() + + async def generate(self, prompt, session, **kwargs): + # Mirror the real AsyncEngine: a handle is held for the duration of generation. + session._handle = self.handle + yield GenOut(response='hi', input_token_len=1, generate_token_len=1, finish_reason=None) + await asyncio.Event().wait() + + class _FakeRawRequest: + + async def json(self): + return {} + + async def is_disconnected(self): + return True + + engine = _FakeAsyncEngine() + origin_async_engine = VariableInterface.async_engine + VariableInterface.async_engine = engine + try: + request = CompletionRequest(model='test-model', prompt='hello', stream=False) + await completions_v1(request, _FakeRawRequest()) + finally: + VariableInterface.async_engine = origin_async_engine + + assert engine.handle.cancelled_session_ids == [0] + + +def test_completions_v1_disconnect_aborts_active_handle(): + asyncio.run(_run_completions_v1_disconnect_aborts_active_handle()) + + +async def _run_completions_v1_disconnect_returns_client_disconnected_response(): + from types import SimpleNamespace + + from fastapi.responses import JSONResponse + + from lmdeploy.serve.core.async_engine import GenOut + from lmdeploy.serve.openai.api_server import VariableInterface, completions_v1 + from lmdeploy.serve.openai.protocol import CompletionRequest + + class _FakeAsyncEngine: + model_name = 'test-model' + epoch = 0 + + def __init__(self): + self.backend_config = SimpleNamespace() + self.session_mgr = SessionManager() + + async def generate(self, prompt, session, **kwargs): + yield GenOut(response='hi', input_token_len=1, generate_token_len=1, finish_reason=None) + await asyncio.Event().wait() + + class _FakeRawRequest: + + async def json(self): + return {} + + async def is_disconnected(self): + return True + + origin_async_engine = VariableInterface.async_engine + VariableInterface.async_engine = _FakeAsyncEngine() + try: + request = CompletionRequest(model='test-model', prompt='hello', stream=False) + result = await completions_v1(request, _FakeRawRequest()) + finally: + VariableInterface.async_engine = origin_async_engine + + assert isinstance(result, JSONResponse) + assert result.status_code == 400 + assert b'Client disconnected' in result.body + + +def test_completions_v1_disconnect_returns_client_disconnected_response(): + asyncio.run(_run_completions_v1_disconnect_returns_client_disconnected_response())