File size: 5,802 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | # yapf: disable
import asyncio
import pytest
from lmdeploy.pytorch.engine.request import RequestManager, RequestType, ResponseType
# yapf: enable
class TestRequestHander:
@pytest.fixture
def event_loop(self):
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:
new_loop.stop()
new_loop.close()
asyncio.set_event_loop(old_loop)
@pytest.fixture
def manager(self):
yield RequestManager()
def test_bind(self, manager, event_loop):
def __stop_engine_callback(reqs, **kwargs):
for req in reqs:
resp = req.resp
resp.type = ResponseType.SUCCESS
resp.data = f'{req.data} success'
manager.response(resp)
async def __dummy_loop():
while True:
try:
await manager.step()
except Exception:
return
sender = manager.build_sender()
manager.set_main_loop_func(__dummy_loop)
# test not bind
resp = sender.send_async(RequestType.STOP_ENGINE, None)
resp = sender.recv(resp)
assert resp.type == ResponseType.HANDLER_NOT_EXIST
assert manager.is_loop_alive()
# test bind success
sender.send_async(RequestType.STOP_ENGINE, None)
manager.bind_func(RequestType.STOP_ENGINE, __stop_engine_callback)
resp = sender.send_async(RequestType.STOP_ENGINE, 'test')
resp = sender.recv(resp)
assert resp.data == 'test success'
# cleanup, cancel main task
task_to_cancel = manager._loop_task
manager.stop_loop()
event_loop.run_until_complete(asyncio.gather(task_to_cancel, return_exceptions=True))
def test_blocked_add_requests_are_cancelled_immediately(self, manager):
sender = manager.build_sender()
manager.block_request_types({RequestType.ADD_SESSION, RequestType.ADD_MESSAGE})
add_session_resp = sender.send_async(RequestType.ADD_SESSION, dict(session_id=1))
add_message_resp = sender.send_async(RequestType.ADD_MESSAGE, dict(session_id=1))
assert add_session_resp.type == ResponseType.CANCEL
assert add_session_resp.is_done
assert add_session_resp.event.is_set()
assert add_message_resp.type == ResponseType.CANCEL
assert add_message_resp.is_done
assert add_message_resp.event.is_set()
assert manager._loop_task is None
def test_cleanup_requests_are_allowed_while_add_requests_blocked(self, manager, event_loop):
def __success_callback(reqs, **kwargs):
for req in reqs:
req.resp.type = ResponseType.SUCCESS
manager.response(req.resp)
async def __dummy_loop():
while True:
try:
await manager.step()
except Exception:
return
sender = manager.build_sender()
manager.set_main_loop_func(__dummy_loop)
manager.bind_func(RequestType.STOP_SESSION, __success_callback)
manager.bind_func(RequestType.END_SESSION, __success_callback)
manager.block_request_types({RequestType.ADD_SESSION, RequestType.ADD_MESSAGE})
stop_resp = sender.send(RequestType.STOP_SESSION, dict(session_id=1))
end_resp = sender.send(RequestType.END_SESSION, dict(session_id=1))
assert stop_resp.type == ResponseType.SUCCESS
assert end_resp.type == ResponseType.SUCCESS
task_to_cancel = manager._loop_task
manager.stop_loop()
event_loop.run_until_complete(asyncio.gather(task_to_cancel, return_exceptions=True))
def test_queued_add_request_is_cancelled_when_blocked_before_processing(self, manager, event_loop):
async def __idle_loop():
await asyncio.Event().wait()
sender = manager.build_sender()
manager.set_main_loop_func(__idle_loop)
resp = sender.send_async(RequestType.ADD_SESSION, dict(session_id=1))
manager.block_request_types({RequestType.ADD_SESSION})
event_loop.run_until_complete(manager.step())
assert resp.type == ResponseType.CANCEL
assert resp.is_done
assert resp.event.is_set()
task_to_cancel = manager._loop_task
manager.stop_loop()
event_loop.run_until_complete(asyncio.gather(task_to_cancel, return_exceptions=True))
def test_unblock_request_types_restores_normal_processing(self, manager, event_loop):
def __success_callback(reqs, **kwargs):
for req in reqs:
req.resp.type = ResponseType.SUCCESS
manager.response(req.resp)
async def __dummy_loop():
while True:
try:
await manager.step()
except Exception:
return
sender = manager.build_sender()
manager.block_request_types({RequestType.ADD_SESSION})
blocked_resp = sender.send_async(RequestType.ADD_SESSION, dict(session_id=1))
assert blocked_resp.type == ResponseType.CANCEL
manager.unblock_request_types({RequestType.ADD_SESSION})
manager.set_main_loop_func(__dummy_loop)
manager.bind_func(RequestType.ADD_SESSION, __success_callback)
resp = sender.send(RequestType.ADD_SESSION, dict(session_id=1))
assert resp.type == ResponseType.SUCCESS
task_to_cancel = manager._loop_task
manager.stop_loop()
event_loop.run_until_complete(asyncio.gather(task_to_cancel, return_exceptions=True))
|