| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Test for the fix in PR #29987: Eagerly abort cancelled final-step requests. |
| 5 | |
| 6 | This test verifies that when a request is aborted during its final execution |
| 7 | step (when it would naturally complete), it is properly marked as aborted |
| 8 | rather than being treated as normally completed. |
| 9 | |
| 10 | The test uses a dummy KV connector to verify that the connector receives |
| 11 | the correct finish status (FINISHED_ABORTED, not FINISHED_LENGTH_CAPPED). |
| 12 | """ |
| 13 | |
| 14 | import asyncio |
| 15 | import tempfile |
| 16 | import time |
| 17 | from pathlib import Path |
| 18 | from typing import Any |
| 19 | from unittest.mock import patch |
| 20 | |
| 21 | import pytest |
| 22 | |
| 23 | from vllm import SamplingParams |
| 24 | from vllm.config import KVTransferConfig, VllmConfig |
| 25 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 26 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 27 | KVConnectorBase_V1, |
| 28 | KVConnectorMetadata, |
| 29 | KVConnectorRole, |
| 30 | ) |
| 31 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 32 | from vllm.platforms import current_platform |
| 33 | from vllm.sampling_params import RequestOutputKind |
| 34 | from vllm.utils.torch_utils import set_default_torch_num_threads |
| 35 | from vllm.v1.core.sched.output import SchedulerOutput |
| 36 | from vllm.v1.engine.async_llm import AsyncLLM |
| 37 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 38 | from vllm.v1.request import Request |
| 39 | |
| 40 | if not current_platform.is_cuda(): |
| 41 | pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True) |
| 42 | |
| 43 | TEXT_PROMPT = "Hello" |
| 44 | |
| 45 | |
| 46 | class DummyKVConnectorMetadata(KVConnectorMetadata): |
| 47 | """Dummy metadata for the test connector.""" |
| 48 | |
| 49 | def __init__(self): |
| 50 | self.requests: list = [] |
| 51 | |
| 52 | |
| 53 | class DummyKVConnector(KVConnectorBase_V1): |
| 54 | """Dummy KV connector that captures request finish statuses to a file. |
| 55 | This is used to verify the fix - without the fix, a request aborted |
| 56 | during its final step would be captured as FINISHED_LENGTH_CAPPED |
| 57 | instead of FINISHED_ABORTED. |
| 58 | |
| 59 | The connector runs in a separate process, so we write statuses to a file |
| 60 | that can be read by the test process. |
| 61 | """ |
| 62 | |
| 63 | def __init__( |
| 64 | self, |
| 65 | vllm_config: VllmConfig, |
| 66 | role: KVConnectorRole, |
| 67 | kv_cache_config: KVCacheConfig, |
| 68 | ): |
| 69 | super().__init__(vllm_config, role, kv_cache_config) |
| 70 | # Get the status file path from extra config |
| 71 | extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config or {} |
| 72 | self.status_file = extra_config.get("status_file") |
| 73 | # Log that we were initialized |
| 74 | if self.status_file: |
| 75 | try: |
| 76 | with open(self.status_file, "a") as f: |
| 77 | f.write(f"INIT:{role.name}\n") |
| 78 | except Exception: |
| 79 | pass |
| 80 | |
| 81 | def get_num_new_matched_tokens( |
| 82 | self, |
| 83 | request: Request, |
| 84 | num_computed_tokens: int, |
| 85 | ) -> tuple[int | None, bool]: |
| 86 | return (0, False) |
| 87 | |
| 88 | def update_state_after_alloc( |
| 89 | self, |
| 90 | request: Request, |
| 91 | blocks: Any, |
| 92 | num_external_tokens: int, |
| 93 | ): |
| 94 | pass |
| 95 | |
| 96 | def build_connector_meta( |
| 97 | self, scheduler_output: SchedulerOutput |
| 98 | ) -> KVConnectorMetadata: |
| 99 | return DummyKVConnectorMetadata() |
| 100 | |
| 101 | def request_finished( |
| 102 | self, |
| 103 | request: Request, |
| 104 | block_ids: list[int], |
| 105 | ) -> tuple[bool, dict[str, Any] | None]: |
| 106 | """Capture the request status when finished by writing to a file.""" |
| 107 | if self.status_file: |
| 108 | try: |
| 109 | with open(self.status_file, "a") as f: |
| 110 | # Write the status name (e.g., "FINISHED_ABORTED") |
| 111 | f.write(f"{request.status.name}\n") |
| 112 | except Exception as e: |
| 113 | # Log but don't fail - this is just test instrumentation |
| 114 | print(f"[DummyKVConnector] Failed to write status: {e}") |
| 115 | return False, None |
| 116 | |
| 117 | def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None: |
| 118 | pass |
| 119 | |
| 120 | def wait_for_layer_load(self, layer_name: str) -> None: |
| 121 | pass |
| 122 | |
| 123 | def save_kv_layer( |
| 124 | self, |
| 125 | layer_name: str, |
| 126 | kv_layer: Any, |
| 127 | attn_metadata: Any, |
| 128 | **kwargs: Any, |
| 129 | ) -> None: |
| 130 | pass |
| 131 | |
| 132 | def wait_for_save(self): |
| 133 | pass |
| 134 | |
| 135 | |
| 136 | # Register the dummy connector |
| 137 | KVConnectorFactory.register_connector( |
| 138 | "DummyKVConnector", __name__, DummyKVConnector.__name__ |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | @pytest.mark.parametrize("async_scheduling", [False, True]) |
| 143 | @pytest.mark.asyncio |
| 144 | async def test_abort_during_final_step(async_scheduling: bool): |
| 145 | """Test that a request aborted during its final execution step is treated as |
| 146 | aborted rather than completed. |
| 147 | |
| 148 | This test: |
| 149 | 1. Monkeypatches execute_model to wait for a file to be deleted |
| 150 | 2. Configures a dummy KV connector to capture finish statuses |
| 151 | 3. Starts a request with max_tokens=1 (will complete on first decode step) |
| 152 | 4. Aborts the request, then deletes the file to unblock execute_model |
| 153 | 5. Verifies the KV connector received FINISHED_ABORTED not FINISHED_LENGTH_CAPPED |
| 154 | |
| 155 | See https://github.com/vllm-project/vllm/pull/29987. |
| 156 | |
| 157 | Without the fix, the KV connector would see FINISHED_LENGTH_CAPPED because |
| 158 | update_from_output() would mark the request as completed before processing |
| 159 | the abort. This causes KV cache blocks to not be freed properly in |
| 160 | disaggregated prefill scenarios. |
| 161 | |
| 162 | With the fix, _process_aborts_queue() runs before update_from_output(), so the |
| 163 | abort takes precedence and the KV connector sees FINISHED_ABORTED. |
| 164 | """ |
| 165 | # Create three temporary files: |
| 166 | # 1. ready_file: deleted by execute_model to signal it has started |
| 167 | # 2. block_file: execute_model waits for this to be deleted |
| 168 | # 3. status_file: KV connector writes finish statuses here |
| 169 | with tempfile.NamedTemporaryFile(delete=False) as f: |
| 170 | ready_file = Path(f.name) |
| 171 | with tempfile.NamedTemporaryFile(delete=False) as f2: |
| 172 | block_file = Path(f2.name) |
| 173 | with tempfile.NamedTemporaryFile(delete=False, mode="w") as f3: |
| 174 | status_file = Path(f3.name) |
| 175 | |
| 176 | try: |
| 177 | # Get the original execute_model method |
| 178 | from vllm.v1.worker.gpu_worker import Worker |
| 179 | |
| 180 | original_execute_model = Worker.execute_model |
| 181 | |
| 182 | def execute_model_with_wait(self, scheduler_output): |
| 183 | # V2's `gpu_worker.compile_or_warm_up_model` calls |
| 184 | # `warmup_kernels(...)` during engine init, which itself calls |
| 185 | # `Worker.execute_model` three times (prefill / decode / cleanup) |
| 186 | # to JIT compile triton kernels. None of those carry the test's |
| 187 | # request id, so we only stall when our actual request is being |
| 188 | # processed. |
| 189 | scheduled = scheduler_output.num_scheduled_tokens or {} |
| 190 | finished = scheduler_output.finished_req_ids or set() |
| 191 | |
| 192 | def is_target_request(req_ids): |
| 193 | return any( |
| 194 | rid == request_id or rid.startswith(f"{request_id}-") |
| 195 | for rid in req_ids |
| 196 | ) |
| 197 | |
| 198 | if is_target_request(scheduled) or is_target_request(finished): |
| 199 | # Signal that execute_model has been called by deleting ready_file |
| 200 | if ready_file.exists(): |
| 201 | ready_file.unlink() |
| 202 | |
| 203 | # Wait for the block file to be deleted (triggered from test after |
| 204 | # abort). This runs in the worker process (after fork), so we poll |
| 205 | # the filesystem. |
| 206 | while block_file.exists(): |
| 207 | time.sleep(0.01) |
| 208 | return original_execute_model(self, scheduler_output) |
| 209 | |
| 210 | # Patch execute_model to inject the wait |
| 211 | # This happens before the worker process is forked, so the patch applies there |
| 212 | with patch.object(Worker, "execute_model", execute_model_with_wait): |
| 213 | request_id = "test-abort-final-step" |
| 214 | |
| 215 | # Configure engine with dummy KV connector |
| 216 | # Pass the status file path so the connector can write to it |
| 217 | kv_transfer_config = KVTransferConfig( |
| 218 | kv_connector="DummyKVConnector", |
| 219 | kv_role="kv_both", |
| 220 | kv_connector_extra_config={"status_file": str(status_file)}, |
| 221 | ) |
| 222 | engine_args = AsyncEngineArgs( |
| 223 | model="meta-llama/Llama-3.2-1B-Instruct", |
| 224 | enforce_eager=True, |
| 225 | async_scheduling=async_scheduling, |
| 226 | kv_transfer_config=kv_transfer_config, |
| 227 | ) |
| 228 | |
| 229 | with set_default_torch_num_threads(1): |
| 230 | engine = AsyncLLM.from_engine_args(engine_args) |
| 231 | |
| 232 | try: |
| 233 | # Create a request that will complete after just 1 token |
| 234 | sampling_params = SamplingParams( |
| 235 | max_tokens=1, |
| 236 | ignore_eos=True, |
| 237 | output_kind=RequestOutputKind.DELTA, |
| 238 | ) |
| 239 | |
| 240 | # Start generation in a task |
| 241 | outputs = [] |
| 242 | |
| 243 | async def generate(): |
| 244 | async for output in engine.generate( |
| 245 | request_id=request_id, |
| 246 | prompt=TEXT_PROMPT, |
| 247 | sampling_params=sampling_params, |
| 248 | ): |
| 249 | outputs.append(output) |
| 250 | |
| 251 | gen_task = asyncio.create_task(generate()) |
| 252 | |
| 253 | # Wait for execute_model to signal it has started (with timeout) |
| 254 | timeout = 5.0 # 5 second timeout |
| 255 | start_time = time.time() |
| 256 | while ready_file.exists(): |
| 257 | if time.time() - start_time > timeout: |
| 258 | raise TimeoutError( |
| 259 | "Timeout waiting for execute_model to start. " |
| 260 | "The monkeypatch may not be working correctly, " |
| 261 | "for example if spawn was used instead of fork." |
| 262 | ) |
| 263 | await asyncio.sleep(0.01) |
| 264 | |
| 265 | # Abort the request while execute_model is blocked |
| 266 | await engine.abort(request_id) |
| 267 | |
| 268 | # Now unblock execute_model by deleting the file |
| 269 | # The abort should be processed before the model output |
| 270 | block_file.unlink() |
| 271 | |
| 272 | # Wait for generation to complete |
| 273 | await gen_task |
| 274 | |
| 275 | # Poll for the KV connector to record the finish status |
| 276 | timeout = 5.0 |
| 277 | start = time.time() |
| 278 | captured_statuses = [] |
| 279 | while time.time() - start < timeout: |
| 280 | with open(status_file) as f4: |
| 281 | status_lines = f4.read().strip().split("\n") |
| 282 | captured_statuses = [ |
| 283 | line |
| 284 | for line in status_lines |
| 285 | if line and line.startswith("FINISHED_") |
| 286 | ] |
| 287 | if captured_statuses: |
| 288 | break |
| 289 | await asyncio.sleep(0.05) |
| 290 | else: |
| 291 | raise TimeoutError( |
| 292 | "Timeout waiting for KV connector to record finish status." |
| 293 | ) |
| 294 | |
| 295 | # Verify we got output |
| 296 | assert len(outputs) > 0, "Should have received at least one output" |
| 297 | |
| 298 | # The final output should have finish_reason="abort" |
| 299 | final_output = outputs[-1] |
| 300 | assert final_output.finished, ( |
| 301 | "Final output should be marked as finished" |
| 302 | ) |
| 303 | assert final_output.outputs[0].finish_reason == "abort", ( |
| 304 | f"Expected finish_reason='abort' but got " |
| 305 | f"'{final_output.outputs[0].finish_reason}'. " |
| 306 | ) |
| 307 | |
| 308 | assert len(captured_statuses) >= 1, ( |
| 309 | f"Expected at least 1 captured finish status, got " |
| 310 | f"{len(captured_statuses)}. File content: {status_lines}" |
| 311 | ) |
| 312 | |
| 313 | assert "FINISHED_ABORTED" in captured_statuses, ( |
| 314 | f"KV connector should see FINISHED_ABORTED but got " |
| 315 | f"{captured_statuses}. " |
| 316 | ) |
| 317 | |
| 318 | # Verify cleanup |
| 319 | assert not engine.output_processor.has_unfinished_requests() |
| 320 | |
| 321 | finally: |
| 322 | # Shutdown the engine |
| 323 | engine.shutdown() |
| 324 | |
| 325 | finally: |
| 326 | # Clean up temporary files if they still exist |
| 327 | if ready_file.exists(): |
| 328 | ready_file.unlink() |
| 329 | if block_file.exists(): |
| 330 | block_file.unlink() |
| 331 | if status_file.exists(): |
| 332 | status_file.unlink() |