Remove AsyncLLMEngine busy loop, shield background task (#1059)

This commit is contained in:
Antoni Baum
2023-09-17 00:29:08 -07:00
committed by GitHub
parent e3e79e9e8a
commit ff36139ffc
4 changed files with 154 additions and 18 deletions

View File

@@ -4,10 +4,25 @@ from vllm.engine.async_llm_engine import RequestTracker
from vllm.outputs import RequestOutput
class DummyEvent:
def __init__(self):
self._flag = False
def set(self):
self._flag = True
def clear(self):
self._flag = False
def test_request_tracker():
tracker = RequestTracker()
tracker.new_requests_event = DummyEvent()
stream_1 = tracker.add_request("1")
assert tracker.new_requests_event._flag
new, finished = tracker.get_new_and_finished_requests()
assert not tracker.new_requests_event._flag
assert len(new) == 1
assert new[0]["request_id"] == "1"
assert not finished
@@ -15,7 +30,9 @@ def test_request_tracker():
stream_2 = tracker.add_request("2")
stream_3 = tracker.add_request("3")
assert tracker.new_requests_event._flag
new, finished = tracker.get_new_and_finished_requests()
assert not tracker.new_requests_event._flag
assert len(new) == 2
assert new[0]["request_id"] == "2"
assert new[1]["request_id"] == "3"
@@ -26,6 +43,7 @@ def test_request_tracker():
# request_ids must be unique
with pytest.raises(KeyError):
tracker.add_request("1")
assert not tracker.new_requests_event._flag
tracker.abort_request("1")
new, finished = tracker.get_new_and_finished_requests()
@@ -36,6 +54,7 @@ def test_request_tracker():
stream_4 = tracker.add_request("4")
tracker.abort_request("4")
assert tracker.new_requests_event._flag
new, finished = tracker.get_new_and_finished_requests()
assert len(finished) == 1
assert "4" in finished
@@ -43,9 +62,11 @@ def test_request_tracker():
assert stream_4.finished
stream_5 = tracker.add_request("5")
assert tracker.new_requests_event._flag
tracker.process_request_output(
RequestOutput("2", "output", [], [], finished=True))
new, finished = tracker.get_new_and_finished_requests()
assert not tracker.new_requests_event._flag
assert len(finished) == 1
assert "2" in finished
assert len(new) == 1