| |
| import asyncio |
| from types import SimpleNamespace |
|
|
| import pytest |
|
|
| from lmdeploy.messages import EngineOutput, ResponseType |
| from lmdeploy.pytorch.engine.engine import Engine |
| from lmdeploy.pytorch.engine.engine_instance import EngineInstance |
| from lmdeploy.pytorch.engine.request import RequestManager, RequestType, Response |
|
|
|
|
| class _FakeSequence: |
|
|
| def __init__(self, resp): |
| self.resp = resp |
|
|
|
|
| class _FakeSession: |
|
|
| def __init__(self, seq): |
| self.sequences = {0: seq} |
|
|
|
|
| class _FakeScheduler: |
|
|
| def __init__(self, session): |
| self.sessions = {1: session} |
| self.ended_sessions = [] |
|
|
| def end_session(self, session_id): |
| self.ended_sessions.append(session_id) |
| self.sessions.pop(session_id) |
|
|
|
|
| class _FakeEngineLoop: |
|
|
| def __init__(self, engine): |
| self.engine = engine |
| self.drained = False |
| self.resumed = False |
|
|
| async def drain_for_sleep(self): |
| assert self.engine.req_manager.is_request_blocked(RequestType.ADD_SESSION) |
| assert self.engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE) |
| self.drained = True |
|
|
| def resume_from_sleep(self): |
| self.resumed = True |
|
|
|
|
|
|
|
|
| class _FakeExecutor: |
|
|
| def __init__(self, engine): |
| self.engine = engine |
| self.sleep_calls = [] |
| self.wakeup_calls = [] |
|
|
| async def sleep(self, level=1): |
| assert self.engine.req_manager.is_request_blocked(RequestType.ADD_SESSION) |
| assert self.engine.scheduler.sessions == {} |
| self.sleep_calls.append(level) |
|
|
| def wakeup(self, tags=None): |
| self.wakeup_calls.append(tags) |
|
|
|
|
| @pytest.fixture |
| def event_loop(): |
| try: |
| old_loop = asyncio.get_event_loop() |
| except RuntimeError: |
| old_loop = None |
| new_loop = asyncio.new_event_loop() |
| try: |
| asyncio.set_event_loop(new_loop) |
| yield new_loop |
| finally: |
| pending = asyncio.all_tasks(new_loop) |
| for task in pending: |
| task.cancel() |
| if pending: |
| new_loop.run_until_complete( |
| asyncio.gather(*pending, return_exceptions=True)) |
| new_loop.run_until_complete(new_loop.shutdown_asyncgens()) |
| new_loop.stop() |
| new_loop.close() |
| asyncio.set_event_loop(old_loop) |
|
|
|
|
| def _build_sleeping_test_engine(event_loop): |
| engine = Engine.__new__(Engine) |
| engine.req_manager = RequestManager() |
| resp = Response(type=ResponseType.INTERNAL_ENGINE_ERROR, sender_id=0, event=asyncio.Event()) |
| seq = _FakeSequence(resp) |
| session = _FakeSession(seq) |
| engine.scheduler = _FakeScheduler(session) |
| engine._sleeping_tags = set() |
| engine._engine_loop = _FakeEngineLoop(engine) |
| engine.executor = _FakeExecutor(engine) |
| return engine, resp |
|
|
|
|
| def test_engine_sleep_blocks_inputs_cancels_sessions_then_sleeps(event_loop): |
| engine, resp = _build_sleeping_test_engine(event_loop) |
|
|
| event_loop.run_until_complete(engine.sleep(level=1)) |
|
|
| assert engine.req_manager.is_request_blocked(RequestType.ADD_SESSION) |
| assert engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE) |
| assert engine._engine_loop.drained |
| assert resp.type == ResponseType.CANCEL |
| assert resp.is_done |
| assert resp.event.is_set() |
| assert engine.scheduler.ended_sessions == [1] |
| assert engine.executor.sleep_calls == [1] |
|
|
|
|
| def test_engine_wakeup_reenables_inputs_only_after_all_tags(event_loop): |
| engine, _ = _build_sleeping_test_engine(event_loop) |
| engine.req_manager.block_request_types({RequestType.ADD_SESSION, RequestType.ADD_MESSAGE}) |
| engine._sleeping_tags = {'weights', 'kv_cache'} |
|
|
| engine.wakeup(['weights']) |
|
|
| assert engine.req_manager.is_request_blocked(RequestType.ADD_SESSION) |
| assert engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE) |
| assert not engine._engine_loop.resumed |
|
|
| engine.wakeup(['kv_cache']) |
|
|
| assert not engine.req_manager.is_request_blocked(RequestType.ADD_SESSION) |
| assert not engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE) |
| assert engine._engine_loop.resumed |
| assert engine.executor.wakeup_calls == [['weights'], ['kv_cache']] |
|
|
|
|
| def test_engine_instance_new_request_after_sleep_returns_cancel(event_loop): |
| engine = SimpleNamespace( |
| req_manager=RequestManager(), |
| max_session_len=8, |
| engine_config=SimpleNamespace(enable_transfer_obj_ref=False, distributed_executor_backend='uni'), |
| ) |
| engine.req_manager.block_request_types({RequestType.ADD_SESSION, RequestType.ADD_MESSAGE}) |
| inst = EngineInstance(engine) |
|
|
| async def __collect(): |
| outputs: list[EngineOutput] = [] |
| async for out in inst.async_stream_infer(1, [1, 2, 3]): |
| outputs.append(out) |
| return outputs |
|
|
| outputs = event_loop.run_until_complete(__collect()) |
|
|
| assert len(outputs) == 1 |
| assert outputs[0].status == ResponseType.CANCEL |
| assert engine.req_manager._loop_task is None |
|
|
