Skip to content

Commit ac96fd4

Browse files
sroy745Alvant
authored andcommitted
Fix test_schedule_swapped_simple in test_scheduler.py (vllm-project#8780)
Signed-off-by: Alvant <alvasian@yandex.ru>
1 parent e710c5d commit ac96fd4

File tree

1 file changed

+10
-4
lines changed

1 file changed

+10
-4
lines changed

tests/core/test_scheduler.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -747,13 +747,19 @@ def test_schedule_decode_blocks_to_copy_update(use_v2_block_manager: bool):
747747
assert output.blocks_to_copy == [(2, 3)]
748748

749749

750-
def test_schedule_swapped_simple():
751-
scheduler = initialize_scheduler()
750+
@pytest.mark.parametrize('use_v2_block_manager', [True, False])
751+
def test_schedule_swapped_simple(use_v2_block_manager: bool):
752+
block_size = 4
753+
scheduler = initialize_scheduler(use_v2_block_manager=use_v2_block_manager,
754+
block_size=block_size)
752755
curr_loras = None
753756
blocks_to_swap_out: List[Tuple[int, int]] = []
754-
_, seq_group = create_dummy_prompt("1", prompt_length=60, best_of=2)
757+
_, seq_group = create_dummy_prompt("1",
758+
prompt_length=4,
759+
best_of=2,
760+
block_size=block_size)
755761
scheduler._allocate_and_set_running(seq_group)
756-
append_new_token_seq_group(60, seq_group, 1)
762+
append_new_token_seq_group(4, seq_group, 1)
757763
scheduler._swap_out(seq_group, blocks_to_swap_out)
758764
scheduler._add_seq_group_to_swapped(seq_group)
759765

0 commit comments

Comments
 (0)