| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.v1.core.sched.scheduler import Scheduler |
| 10 | from vllm.v1.request import FinishReason, Request, RequestStatus |
| 11 | |
| 12 | from .utils import ( |
| 13 | create_model_runner_output, |
| 14 | create_request, |
| 15 | create_scheduler, |
| 16 | create_vllm_config, |
| 17 | ) |
| 18 | |
| 19 | pytestmark = pytest.mark.cpu_test |
| 20 | |
| 21 | |
| 22 | def _make_get_num_new_matched_tokens( |
| 23 | req_num_new_matched_tokens: dict[str, int], |
| 24 | async_load: bool, |
| 25 | ) -> Callable[[Request, int], tuple[int, bool]]: |
| 26 | def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]: |
| 27 | value = req_num_new_matched_tokens.get(request.request_id, 0) |
| 28 | return value, async_load |
| 29 | |
| 30 | return get_num_new_matched_tokens |
| 31 | |
| 32 | |
| 33 | @pytest.fixture |
| 34 | def fail_scheduler(): |
| 35 | """Scheduler with kv_load_failure_policy='fail'.""" |
| 36 | vllm_config = create_vllm_config() |
| 37 | vllm_config.kv_transfer_config.kv_load_failure_policy = "fail" |
| 38 | return create_scheduler(vllm_config) |
| 39 | |
| 40 | |
| 41 | def test_error_propagation_sync_load(fail_scheduler: Scheduler): |
| 42 | """Test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load).""" |
| 43 | num_prompt_blocks = 100 |
| 44 | num_external_computed_blocks = 99 |
| 45 | invalid_block_idx = 50 |
| 46 | |
| 47 | num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size |
| 48 | num_external_computed_tokens = ( |
| 49 | num_external_computed_blocks * fail_scheduler.block_size |
| 50 | ) |
| 51 | |
| 52 | request = create_request(num_tokens=num_prompt_tokens) |
| 53 | fail_scheduler.add_request(request=request) |
| 54 | |
| 55 | req_num_new_matched_tokens = { |
| 56 | request.request_id: num_external_computed_tokens, |
| 57 | } |
| 58 | |
| 59 | fail_scheduler.connector = Mock() |
| 60 | fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( |
| 61 | _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) |
| 62 | ) |
| 63 | fail_scheduler.connector.request_finished.return_value = (False, None) |
| 64 | fail_scheduler.connector.take_events.return_value = () |
| 65 | |
| 66 | scheduler_output = fail_scheduler.schedule() |
| 67 | |
| 68 | assert len(fail_scheduler.running) == 1 |
| 69 | assert len(scheduler_output.scheduled_new_reqs) == 1 |
| 70 | assert fail_scheduler.connector.get_num_new_matched_tokens.call_count == 1 |
| 71 | |
| 72 | req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0] |
| 73 | invalid_block_ids = {req_block_ids[invalid_block_idx]} |
| 74 | model_runner_output = create_model_runner_output( |
| 75 | [request], |
| 76 | invalid_block_ids=invalid_block_ids, |
| 77 | use_eos=True, |
| 78 | ) |
| 79 | |
| 80 | outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output) |
| 81 | |
| 82 | assert request.status == RequestStatus.FINISHED_ERROR |
| 83 | assert request.get_finished_reason() == FinishReason.ERROR |
| 84 | |
| 85 | assert len(outputs) == 1 |
| 86 | engine_outputs = next(iter(outputs.values())) |
| 87 | assert len(engine_outputs.outputs) == 1 |
| 88 | output = engine_outputs.outputs[0] |
| 89 | assert output.request_id == request.request_id |
| 90 | assert output.finish_reason == FinishReason.ERROR |
| 91 | |
| 92 | assert len(fail_scheduler.running) == 0 |
| 93 | |
| 94 | |
| 95 | def test_error_propagation_async_load(fail_scheduler: Scheduler): |
| 96 | """Test invalid_block_ids with fail policy -> FINISHED_ERROR (async load).""" |
| 97 | num_prompt_blocks = 100 |
| 98 | num_external_computed_blocks = 99 |
| 99 | invalid_block_idx = 50 |
| 100 | |
| 101 | num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size |
| 102 | num_external_computed_tokens = ( |
| 103 | num_external_computed_blocks * fail_scheduler.block_size |
| 104 | ) |
| 105 | |
| 106 | request = create_request(num_tokens=num_prompt_tokens) |
| 107 | fail_scheduler.add_request(request=request) |
| 108 | |
| 109 | req_num_new_matched_tokens = { |
| 110 | request.request_id: num_external_computed_tokens, |
| 111 | } |
| 112 | |
| 113 | fail_scheduler.connector = Mock() |
| 114 | fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( |
| 115 | _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) |
| 116 | ) |
| 117 | fail_scheduler.connector.request_finished.return_value = (False, None) |
| 118 | fail_scheduler.connector.take_events.return_value = () |
| 119 | |
| 120 | scheduler_output = fail_scheduler.schedule() |
| 121 | |
| 122 | assert len(fail_scheduler.kv_holding_waiting) == 1 |
| 123 | assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS |
| 124 | assert request.num_computed_tokens == num_external_computed_tokens |
| 125 | |
| 126 | (req_block_ids,) = fail_scheduler.kv_cache_manager.get_block_ids(request.request_id) |
| 127 | invalid_block_ids = {req_block_ids[invalid_block_idx]} |
| 128 | model_runner_output = create_model_runner_output( |
| 129 | reqs=[], |
| 130 | finished_recving=set(), |
| 131 | invalid_block_ids=invalid_block_ids, |
| 132 | use_eos=True, |
| 133 | ) |
| 134 | |
| 135 | outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output) |
| 136 | |
| 137 | assert request.status == RequestStatus.FINISHED_ERROR |
| 138 | assert request.get_finished_reason() == FinishReason.ERROR |
| 139 | |
| 140 | assert len(outputs) == 1 |
| 141 | engine_outputs = next(iter(outputs.values())) |
| 142 | assert len(engine_outputs.outputs) == 1 |
| 143 | output = engine_outputs.outputs[0] |
| 144 | assert output.request_id == request.request_id |
| 145 | assert output.finish_reason == FinishReason.ERROR |
| 146 | |
| 147 | assert len(fail_scheduler.waiting) == 0 |
| 148 | assert len(fail_scheduler.kv_holding_waiting) == 0 |
| 149 | assert not fail_scheduler.deferred_waiting |