| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test for pause/resume with keep mode. |
| 4 | |
| 5 | This test uses concurrent tasks to verify the engine truly stops generating |
| 6 | during pause: |
| 7 | 1. Generator task: continuously generates and logs time between tokens |
| 8 | 2. Controller task: sends pause/resume commands |
| 9 | |
| 10 | If the engine properly pauses, we should see a gap in token timestamps |
| 11 | matching the pause duration. |
| 12 | """ |
| 13 | |
| 14 | import asyncio |
| 15 | import time |
| 16 | |
| 17 | from vllm import SamplingParams |
| 18 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 19 | from vllm.v1.engine.async_llm import AsyncLLM |
| 20 | |
| 21 | PAUSE_DURATION = 3.0 # seconds |
| 22 | |
| 23 | |
| 24 | async def main(): |
| 25 | # Create engine with a small model |
| 26 | engine_args = AsyncEngineArgs( |
| 27 | model="facebook/opt-125m", |
| 28 | enforce_eager=True, |
| 29 | ) |
| 30 | engine = AsyncLLM.from_engine_args(engine_args) |
| 31 | |
| 32 | prompt = "Write a story about a dragon. Once upon a time" |
| 33 | sampling_params = SamplingParams(max_tokens=30, ignore_eos=True) |
| 34 | |
| 35 | # Track token arrival times |
| 36 | token_times: list[tuple[int, float]] = [] # (token_count, timestamp) |
| 37 | pause_time: float = 0 |
| 38 | resume_time: float = 0 |
| 39 | pause_token_idx: int = 0 # Index in token_times when pause occurred |
| 40 | |
| 41 | async def generator_task(): |
| 42 | """Generate tokens and record timestamps.""" |
| 43 | async for output in engine.generate( |
| 44 | request_id="test-req", |
| 45 | prompt=prompt, |
| 46 | sampling_params=sampling_params, |
| 47 | ): |
| 48 | token_count = len(output.outputs[0].token_ids) |
| 49 | token_times.append((token_count, time.monotonic())) |
| 50 | print( |
| 51 | f"Token {token_count} arrived:" |
| 52 | f"T={token_times[-1][1] - token_times[0][1]:.3f}s" |
| 53 | ) |
| 54 | return output |
| 55 | |
| 56 | async def controller_task(): |
| 57 | """Pause and resume the engine after some tokens generated.""" |
| 58 | nonlocal pause_time, resume_time, pause_token_idx |
| 59 | |
| 60 | # Wait for some tokens to be generated |
| 61 | while len(token_times) < 5: |
| 62 | await asyncio.sleep(0.01) |
| 63 | |
| 64 | print(f"\nPausing engine (keep mode) at token {len(token_times)}") |
| 65 | pause_time = time.monotonic() |
| 66 | await engine.pause_generation(mode="keep") |
| 67 | pause_token_idx = len(token_times) |
| 68 | print(f"Paused! Sleeping for {PAUSE_DURATION}s...") |
| 69 | |
| 70 | # Sleep while paused - no tokens should be generated during this time |
| 71 | await asyncio.sleep(PAUSE_DURATION) |
| 72 | |
| 73 | print("Resuming engine...") |
| 74 | await engine.resume_generation() |
| 75 | resume_time = time.monotonic() |
| 76 | print("Resumed!\n") |
| 77 | |
| 78 | # Run both tasks concurrently |
| 79 | gen_task = asyncio.create_task(generator_task()) |
| 80 | ctrl_task = asyncio.create_task(controller_task()) |
| 81 | |
| 82 | final_output, _ = await asyncio.gather(gen_task, ctrl_task) |
| 83 | |
| 84 | # Verify the pause actually stopped generation. |
| 85 | # The gap after the pause token should be approximately the sleep duration. |
| 86 | pause_gap = token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1] |
| 87 | print( |
| 88 | f"\nGap after pause (token {pause_token_idx - 1} -> {pause_token_idx}): " |
| 89 | f"{pause_gap:.3f}s" |
| 90 | ) |
| 91 | if pause_gap >= PAUSE_DURATION * 0.9: |
| 92 | print(f"✓ Test passed! Engine paused for ~{pause_gap:.1f}s") |
| 93 | else: |
| 94 | print( |
| 95 | f"✗ Test failed! Expected ~{PAUSE_DURATION}s gap after pause, " |
| 96 | f"got {pause_gap:.3f}s" |
| 97 | ) |
| 98 | raise AssertionError("Engine did not properly pause") |
| 99 | |
| 100 | # Verify request completed |
| 101 | assert final_output.finished, "Request should have finished" |
| 102 | assert len(final_output.outputs[0].token_ids) == 30, "Should have all tokens" |
| 103 | |
| 104 | engine.shutdown() |
| 105 | |
| 106 | |
| 107 | if __name__ == "__main__": |
| 108 | asyncio.run(main()) |