From 6771b32a5ed56e79fa65bcd41a3e6d0848c35117 Mon Sep 17 00:00:00 2001 From: Sourashis Roy Date: Tue, 24 Sep 2024 18:45:34 +0000 Subject: [PATCH 1/2] Fix test_schedule_swapped_simple --- tests/core/test_scheduler.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py index b3bc00280682c..8b4cc8bff439e 100644 --- a/tests/core/test_scheduler.py +++ b/tests/core/test_scheduler.py @@ -746,17 +746,20 @@ def test_schedule_decode_blocks_to_copy_update(use_v2_block_manager: bool): # applied. assert output.blocks_to_copy == [(2, 3)] - -def test_schedule_swapped_simple(): - scheduler = initialize_scheduler() +@pytest.mark.parametrize('use_v2_block_manager', [True, False]) +def test_schedule_swapped_simple(use_v2_block_manager:bool): + block_size=4 + scheduler = initialize_scheduler( + use_v2_block_manager=use_v2_block_manager, block_size=block_size) curr_loras = None blocks_to_swap_out: List[Tuple[int, int]] = [] - _, seq_group = create_dummy_prompt("1", prompt_length=60, best_of=2) + _, seq_group = create_dummy_prompt( + "1", prompt_length=4, best_of=2, block_size=block_size) scheduler._allocate_and_set_running(seq_group) - append_new_token_seq_group(60, seq_group, 1) + append_new_token_seq_group(4, seq_group, 1) scheduler._swap_out(seq_group, blocks_to_swap_out) scheduler._add_seq_group_to_swapped(seq_group) - + budget = create_token_budget() output = scheduler._schedule_swapped(budget, curr_loras) remaining_swapped = scheduler.swapped From 6d48951e98a51af897ce8e600e91ffc876cd1fdc Mon Sep 17 00:00:00 2001 From: Sourashis Roy Date: Tue, 24 Sep 2024 18:59:22 +0000 Subject: [PATCH 2/2] Fix format --- tests/core/test_scheduler.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py index 8b4cc8bff439e..88c6c3bb28e43 100644 --- a/tests/core/test_scheduler.py +++ b/tests/core/test_scheduler.py @@ -746,20 +746,23 @@ def test_schedule_decode_blocks_to_copy_update(use_v2_block_manager: bool): # applied. assert output.blocks_to_copy == [(2, 3)] + @pytest.mark.parametrize('use_v2_block_manager', [True, False]) -def test_schedule_swapped_simple(use_v2_block_manager:bool): - block_size=4 - scheduler = initialize_scheduler( - use_v2_block_manager=use_v2_block_manager, block_size=block_size) +def test_schedule_swapped_simple(use_v2_block_manager: bool): + block_size = 4 + scheduler = initialize_scheduler(use_v2_block_manager=use_v2_block_manager, + block_size=block_size) curr_loras = None blocks_to_swap_out: List[Tuple[int, int]] = [] - _, seq_group = create_dummy_prompt( - "1", prompt_length=4, best_of=2, block_size=block_size) + _, seq_group = create_dummy_prompt("1", + prompt_length=4, + best_of=2, + block_size=block_size) scheduler._allocate_and_set_running(seq_group) append_new_token_seq_group(4, seq_group, 1) scheduler._swap_out(seq_group, blocks_to_swap_out) scheduler._add_seq_group_to_swapped(seq_group) - + budget = create_token_budget() output = scheduler._schedule_swapped(budget, curr_loras) remaining_swapped = scheduler.swapped