汪志鹏andCursor
e37a89ffcb
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
115 lines
3.3 KiB
Python
115 lines
3.3 KiB
Python
import time
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def dummy_expand_func(prompt, sp0):
|
|
if prompt == "expand_me":
|
|
return [SimpleNamespace(prompt={"prompt": "neg"}, role="cfg_text", request_id_suffix="__cfg_text")]
|
|
return []
|
|
|
|
|
|
@pytest.fixture
|
|
def tracker():
|
|
sp0 = SimpleNamespace()
|
|
return CfgCompanionTracker(prompt_expand_func=dummy_expand_func, stage0_sampling_params=sp0, timeout_s=0.1)
|
|
|
|
|
|
def test_companion_tracker_initialization(tracker):
|
|
assert not tracker.is_active
|
|
assert tracker.num_companions == 0
|
|
|
|
|
|
def test_expand_prompts_registers_companions(tracker):
|
|
request_id_to_prompt = {"req1": "expand_me", "req2": "do_not_expand"}
|
|
|
|
pairs = tracker.expand_prompts(request_id_to_prompt)
|
|
|
|
assert len(pairs) == 1
|
|
companion_id, prompt = pairs[0]
|
|
assert companion_id == "req1__cfg_text"
|
|
assert prompt == {"prompt": "neg"}
|
|
|
|
assert tracker.is_active
|
|
assert tracker.num_companions == 1
|
|
assert tracker.is_companion("req1__cfg_text")
|
|
assert not tracker.is_companion("req2__cfg_text")
|
|
assert tracker.has_companions("req1")
|
|
assert not tracker.has_companions("req2")
|
|
|
|
comp_map = tracker.get_companion_request_ids("req1")
|
|
assert comp_map == {"cfg_text": "req1__cfg_text"}
|
|
|
|
|
|
def test_companion_lifecycle_success(tracker):
|
|
request_id_to_prompt = {"req1": "expand_me"}
|
|
tracker.expand_prompts(request_id_to_prompt)
|
|
|
|
# Defer parent
|
|
engine_outputs = {"out": 123}
|
|
tracker.defer_parent("req1", engine_outputs, stage_id=0)
|
|
|
|
# Initially not done
|
|
assert not tracker.all_companions_done("req1")
|
|
|
|
# Companion completes
|
|
parent_id = tracker.on_companion_completed("req1__cfg_text")
|
|
|
|
# Parent should be returned since all companions are done and it is pending
|
|
assert parent_id == "req1"
|
|
assert tracker.all_companions_done("req1")
|
|
|
|
# Pop pending parent
|
|
popped = tracker.pop_pending_parent("req1")
|
|
assert popped is not None
|
|
assert popped["engine_outputs"] == engine_outputs
|
|
assert popped["stage_id"] == 0
|
|
|
|
|
|
def test_companion_lifecycle_failure(tracker):
|
|
request_id_to_prompt = {"req1": "expand_me"}
|
|
tracker.expand_prompts(request_id_to_prompt)
|
|
|
|
tracker.defer_parent("req1", {"out": 123}, stage_id=0)
|
|
|
|
# Companion fails
|
|
parent_id, aborted = tracker.on_companion_error("req1__cfg_text")
|
|
|
|
assert parent_id == "req1"
|
|
assert aborted is True
|
|
assert tracker.is_parent_failed("req1")
|
|
|
|
# Parent should be removed from pending list
|
|
assert tracker.pop_pending_parent("req1") is None
|
|
|
|
# Consume failure
|
|
tracker.consume_parent_failure("req1")
|
|
assert not tracker.is_parent_failed("req1")
|
|
|
|
|
|
def test_companion_lifecycle_timeout(tracker):
|
|
request_id_to_prompt = {"req1": "expand_me"}
|
|
tracker.expand_prompts(request_id_to_prompt)
|
|
|
|
tracker.defer_parent("req1", {"out": 123}, stage_id=0)
|
|
|
|
# Initially no timeouts
|
|
timeouts = tracker.check_timeouts()
|
|
assert len(timeouts) == 0
|
|
|
|
# Wait for timeout
|
|
time.sleep(0.15)
|
|
|
|
# Check timeouts again
|
|
timeouts = tracker.check_timeouts()
|
|
assert len(timeouts) == 1
|
|
assert timeouts[0] == "req1"
|
|
|
|
# Should be removed from pending
|
|
assert tracker.pop_pending_parent("req1") is None
|