| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | import time |
| 6 | from contextlib import ExitStack |
| 7 | from unittest.mock import AsyncMock, MagicMock, call |
| 8 | |
| 9 | import pytest |
| 10 | |
| 11 | import vllm.v1.engine.async_llm as async_llm_module |
| 12 | from vllm import SamplingParams |
| 13 | from vllm.assets.image import ImageAsset |
| 14 | from vllm.config import LoggingConfig, VllmConfig |
| 15 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 16 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 17 | ChatCompletionRequest, |
| 18 | ChatCompletionResponse, |
| 19 | ) |
| 20 | from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat |
| 21 | from vllm.entrypoints.openai.models.protocol import BaseModelPath |
| 22 | from vllm.entrypoints.openai.models.serving import OpenAIServingModels |
| 23 | from vllm.exceptions import VLLMValidationError |
| 24 | from vllm.inputs import PromptType |
| 25 | from vllm.outputs import RequestOutput |
| 26 | from vllm.platforms import current_platform |
| 27 | from vllm.sampling_params import RequestOutputKind |
| 28 | from vllm.utils.torch_utils import set_default_torch_num_threads |
| 29 | from vllm.v1.engine.async_llm import AsyncLLM |
| 30 | from vllm.v1.metrics.loggers import ( |
| 31 | AggregatedLoggingStatLogger, |
| 32 | LoggingStatLogger, |
| 33 | PerEngineStatLoggerAdapter, |
| 34 | PrometheusStatLogger, |
| 35 | ) |
| 36 | |
| 37 | if not current_platform.is_cuda(): |
| 38 | pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True) |
| 39 | |
| 40 | TEXT_ENGINE_ARGS = AsyncEngineArgs( |
| 41 | model="meta-llama/Llama-3.2-1B-Instruct", |
| 42 | enforce_eager=True, |
| 43 | ) |
| 44 | |
| 45 | VISION_ENGINE_ARGS = AsyncEngineArgs( |
| 46 | model="Qwen/Qwen2-VL-2B-Instruct", enforce_eager=True |
| 47 | ) |
| 48 | |
| 49 | TEXT_PROMPT = "Hello my name is Robert and" |
| 50 | |
| 51 | VISION_PROMPT_TEMPLATE = ( |
| 52 | "<|im_start|>system\nYou are a helpful assistant.<|im_end|>" |
| 53 | "\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" |
| 54 | "What is in the image?<|im_end|>\n" |
| 55 | "<|im_start|>assistant\n" |
| 56 | ) |
| 57 | VISION_PROMPT = { |
| 58 | "prompt": VISION_PROMPT_TEMPLATE, |
| 59 | "multi_modal_data": {"image": ImageAsset("stop_sign").pil_image}, |
| 60 | } |
| 61 | |
| 62 | |
| 63 | def _mock_async_llm_dependencies(monkeypatch: pytest.MonkeyPatch): |
| 64 | renderer = MagicMock() |
| 65 | engine_core = MagicMock() |
| 66 | engine_core.profile_async = AsyncMock() |
| 67 | monkeypatch.setattr( |
| 68 | async_llm_module, "maybe_register_config_serialize_by_value", MagicMock() |
| 69 | ) |
| 70 | monkeypatch.setattr( |
| 71 | async_llm_module, |
| 72 | "load_stat_logger_plugin_factories", |
| 73 | MagicMock(return_value=[]), |
| 74 | ) |
| 75 | monkeypatch.setattr( |
| 76 | async_llm_module, |
| 77 | "renderer_from_config", |
| 78 | MagicMock(return_value=renderer), |
| 79 | ) |
| 80 | monkeypatch.setattr(async_llm_module, "InputProcessor", MagicMock()) |
| 81 | monkeypatch.setattr(async_llm_module, "OutputProcessor", MagicMock()) |
| 82 | monkeypatch.setattr( |
| 83 | async_llm_module.EngineCoreClient, |
| 84 | "make_async_mp_client", |
| 85 | MagicMock(return_value=engine_core), |
| 86 | ) |
| 87 | return engine_core |
| 88 | |
| 89 | |
| 90 | def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): |
| 91 | vllm_config = MagicMock() |
| 92 | vllm_config.logging_config = LoggingConfig() |
| 93 | vllm_config.observability_config.otlp_traces_endpoint = None |
| 94 | vllm_config.scheduler_config.stream_interval = 1 |
| 95 | vllm_config.profiler_config.profiler = "cuda" |
| 96 | vllm_config.profiler_config.ignore_frontend = False |
| 97 | engine_core = _mock_async_llm_dependencies(monkeypatch) |
| 98 | |
| 99 | engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) |
| 100 | |
| 101 | async def profile(): |
| 102 | await engine.start_profile() |
| 103 | await engine.stop_profile() |
| 104 | |
| 105 | asyncio.run(profile()) |
| 106 | |
| 107 | assert engine.profiler is None |
| 108 | engine_core.profile_async.assert_has_awaits([call(True, None), call(False)]) |
| 109 | |
| 110 | |
| 111 | def test_cuda_only_torch_profiler_skips_frontend_cpu_trace( |
| 112 | monkeypatch: pytest.MonkeyPatch, |
| 113 | ): |
| 114 | vllm_config = MagicMock() |
| 115 | vllm_config.logging_config = LoggingConfig() |
| 116 | vllm_config.observability_config.otlp_traces_endpoint = None |
| 117 | vllm_config.scheduler_config.stream_interval = 1 |
| 118 | vllm_config.profiler_config.profiler = "torch" |
| 119 | vllm_config.profiler_config.ignore_frontend = False |
| 120 | vllm_config.profiler_config.torch_profiler_activities = ["CUDA"] |
| 121 | _mock_async_llm_dependencies(monkeypatch) |
| 122 | profiler = MagicMock() |
| 123 | monkeypatch.setattr(async_llm_module, "TorchProfilerWrapper", profiler) |
| 124 | |
| 125 | engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) |
| 126 | |
| 127 | assert engine.profiler is None |
| 128 | profiler.assert_not_called() |
| 129 | |
| 130 | |
| 131 | async def generate( |
| 132 | engine: AsyncLLM, |
| 133 | request_id: str, |
| 134 | prompt: PromptType, |
| 135 | output_kind: RequestOutputKind, |
| 136 | max_tokens: int, |
| 137 | n: int = 1, |
| 138 | prompt_logprobs: int | None = None, |
| 139 | cancel_after: int | None = None, |
| 140 | ) -> tuple[int, str]: |
| 141 | # Ensure generate doesn't complete too fast for cancellation test. |
| 142 | await asyncio.sleep(0.2) |
| 143 | |
| 144 | count = 0 |
| 145 | sampling_params = SamplingParams( |
| 146 | max_tokens=max_tokens, |
| 147 | ignore_eos=True, |
| 148 | output_kind=output_kind, |
| 149 | temperature=0.5, |
| 150 | seed=33, |
| 151 | n=n, |
| 152 | prompt_logprobs=prompt_logprobs, |
| 153 | ) |
| 154 | async for out in engine.generate( |
| 155 | request_id=request_id, prompt=prompt, sampling_params=sampling_params |
| 156 | ): |
| 157 | num_tokens = sum(len(output.token_ids) for output in out.outputs) |
| 158 | if output_kind == RequestOutputKind.DELTA: |
| 159 | count += num_tokens |
| 160 | else: |
| 161 | count = num_tokens |
| 162 | |
| 163 | if cancel_after is not None and count >= cancel_after: |
| 164 | return count, request_id |
| 165 | |
| 166 | await asyncio.sleep(0.0) |
| 167 | |
| 168 | return count, request_id |
| 169 | |
| 170 | |
| 171 | @pytest.mark.parametrize( |
| 172 | "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY] |
| 173 | ) |
| 174 | @pytest.mark.parametrize( |
| 175 | "engine_args,prompt", |
| 176 | [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)], |
| 177 | ) |
| 178 | @pytest.mark.asyncio |
| 179 | async def test_load( |
| 180 | output_kind: RequestOutputKind, |
| 181 | engine_args: AsyncEngineArgs, |
| 182 | prompt: PromptType, |
| 183 | ): |
| 184 | with ExitStack() as after: |
| 185 | with set_default_torch_num_threads(1): |
| 186 | engine = AsyncLLM.from_engine_args(engine_args) |
| 187 | after.callback(engine.shutdown) |
| 188 | |
| 189 | NUM_REQUESTS = 100 |
| 190 | NUM_EXPECTED_TOKENS = 10 |
| 191 | |
| 192 | request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] |
| 193 | |
| 194 | # Create concurrent requests. |
| 195 | tasks = [] |
| 196 | for request_id in request_ids: |
| 197 | tasks.append( |
| 198 | asyncio.create_task( |
| 199 | generate( |
| 200 | engine, request_id, prompt, output_kind, NUM_EXPECTED_TOKENS |
| 201 | ) |
| 202 | ) |
| 203 | ) |
| 204 | |
| 205 | # Confirm that we got all the EXPECTED tokens from the requests. |
| 206 | done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) |
| 207 | for task in pending: |
| 208 | task.cancel() |
| 209 | for task in done: |
| 210 | num_generated_tokens, request_id = await task |
| 211 | assert num_generated_tokens == NUM_EXPECTED_TOKENS, ( |
| 212 | f"{request_id} generated {num_generated_tokens} but " |
| 213 | f"expected {NUM_EXPECTED_TOKENS}" |
| 214 | ) |
| 215 | |
| 216 | assert not engine.output_processor.has_unfinished_requests() |
| 217 | |
| 218 | |
| 219 | @pytest.mark.parametrize( |
| 220 | "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY] |
| 221 | ) |
| 222 | @pytest.mark.parametrize( |
| 223 | "engine_args,prompt", |
| 224 | [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)], |
| 225 | ) |
| 226 | @pytest.mark.asyncio |
| 227 | async def test_abort( |
| 228 | output_kind: RequestOutputKind, |
| 229 | engine_args: AsyncEngineArgs, |
| 230 | prompt: PromptType, |
| 231 | ): |
| 232 | with ExitStack() as after: |
| 233 | with set_default_torch_num_threads(1): |
| 234 | engine = AsyncLLM.from_engine_args(engine_args) |
| 235 | after.callback(engine.shutdown) |
| 236 | |
| 237 | NUM_REQUESTS = 100 |
| 238 | NUM_EXPECTED_TOKENS = 100 |
| 239 | NUM_EXPECTED_TOKENS_LONG = 50000 |
| 240 | REQUEST_IDS_TO_ABORT = range(1, 100, 10) |
| 241 | PARALLEL_SAMPLE_REQ_IDS = range(1, 100, 15) |
| 242 | |
| 243 | request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] |
| 244 | |
| 245 | # Create concurrent requests. |
| 246 | tasks: list[asyncio.Task] = [] |
| 247 | for idx, request_id in enumerate(request_ids): |
| 248 | max_tokens = ( |
| 249 | NUM_EXPECTED_TOKENS_LONG |
| 250 | if (idx in REQUEST_IDS_TO_ABORT) |
| 251 | else NUM_EXPECTED_TOKENS |
| 252 | ) |
| 253 | n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1 |
| 254 | tasks.append( |
| 255 | asyncio.create_task( |
| 256 | generate(engine, request_id, prompt, output_kind, max_tokens, n) |
| 257 | ) |
| 258 | ) |
| 259 | |
| 260 | # API server cancels requests when they disconnect. |
| 261 | for idx in REQUEST_IDS_TO_ABORT: |
| 262 | tasks[idx].cancel() |
| 263 | await asyncio.sleep(0.1) |
| 264 | |
| 265 | # Confirm the other requests are okay. |
| 266 | for idx, task in enumerate(tasks): |
| 267 | # Confirm that it was actually canceled. |
| 268 | if idx in REQUEST_IDS_TO_ABORT: |
| 269 | with pytest.raises(asyncio.CancelledError): |
| 270 | await task |
| 271 | else: |
| 272 | # Otherwise, make sure the request was not impacted. |
| 273 | num_generated_tokens, request_id = await task |
| 274 | n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1 |
| 275 | expected_tokens = NUM_EXPECTED_TOKENS * n |
| 276 | assert num_generated_tokens == expected_tokens, ( |
| 277 | f"{request_id} generated {num_generated_tokens} but " |
| 278 | f"expected {expected_tokens}" |
| 279 | ) |
| 280 | |
| 281 | # Make sure all aborted requests were really aborted. |
| 282 | assert not engine.output_processor.has_unfinished_requests() |
| 283 | |
| 284 | # Confirm we can do another generation. |
| 285 | request_id = f"request-{REQUEST_IDS_TO_ABORT[0]}" |
| 286 | task = asyncio.create_task( |
| 287 | generate(engine, request_id, prompt, output_kind, NUM_EXPECTED_TOKENS) |
| 288 | ) |
| 289 | num_generated_tokens, request_id = await task |
| 290 | assert num_generated_tokens == NUM_EXPECTED_TOKENS |
| 291 | assert not engine.output_processor.has_unfinished_requests() |
| 292 | |
| 293 | |
| 294 | @pytest.mark.parametrize( |
| 295 | "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY] |
| 296 | ) |
| 297 | @pytest.mark.asyncio |
| 298 | async def test_multi_abort(output_kind: RequestOutputKind): |
| 299 | with ExitStack() as after: |
| 300 | with set_default_torch_num_threads(1): |
| 301 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 302 | after.callback(engine.shutdown) |
| 303 | |
| 304 | NUM_REQUESTS = 50 |
| 305 | NUM_EXPECTED_TOKENS = 100 |
| 306 | NUM_EXPECTED_TOKENS_LONG = 50000 |
| 307 | REQUEST_IDS_TO_ABORT = [5, 10, 15, 20, 25] |
| 308 | PARALLEL_SAMPLE_REQ_IDS = [5, 15, 30, 35] |
| 309 | |
| 310 | request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] |
| 311 | |
| 312 | # Create concurrent requests. |
| 313 | tasks: list[asyncio.Task] = [] |
| 314 | for idx, request_id in enumerate(request_ids): |
| 315 | max_tokens = ( |
| 316 | NUM_EXPECTED_TOKENS_LONG |
| 317 | if (idx in REQUEST_IDS_TO_ABORT) |
| 318 | else NUM_EXPECTED_TOKENS |
| 319 | ) |
| 320 | n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1 |
| 321 | tasks.append( |
| 322 | asyncio.create_task( |
| 323 | generate( |
| 324 | engine, request_id, TEXT_PROMPT, output_kind, max_tokens, n |
| 325 | ) |
| 326 | ) |
| 327 | ) |
| 328 | |
| 329 | # Let requests start generating, use a longer sleep to ensure all |
| 330 | # requests have exited prefill and produced at least one |
| 331 | # decode token before we abort. |
| 332 | await asyncio.sleep(1.0) |
| 333 | |
| 334 | # Use multi-abort to abort multiple requests at once |
| 335 | abort_request_ids = [request_ids[i] for i in REQUEST_IDS_TO_ABORT] |
| 336 | await engine.abort(abort_request_ids, internal=False) |
| 337 | |
| 338 | # Wait for all tasks to complete |
| 339 | results = await asyncio.gather(*tasks, return_exceptions=True) |
| 340 | |
| 341 | # Verify results |
| 342 | for idx, result in enumerate(results): |
| 343 | if idx in REQUEST_IDS_TO_ABORT: |
| 344 | # Aborted requests should return partial results |
| 345 | assert isinstance(result, tuple), ( |
| 346 | f"Request {idx} should have completed with partial results" |
| 347 | ) |
| 348 | num_generated_tokens, request_id = result |
| 349 | # Should have generated some tokens before abort |
| 350 | assert num_generated_tokens > 0, ( |
| 351 | f"Aborted request {request_id} should have generated some tokens" |
| 352 | ) |
| 353 | else: |
| 354 | # Non-aborted requests should complete normally |
| 355 | assert isinstance(result, tuple), ( |
| 356 | f"Request {idx} should have completed successfully" |
| 357 | ) |
| 358 | num_generated_tokens, request_id = result |
| 359 | n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1 |
| 360 | expected_tokens = NUM_EXPECTED_TOKENS * n |
| 361 | assert num_generated_tokens == expected_tokens, ( |
| 362 | f"{request_id} generated {num_generated_tokens} but " |
| 363 | f"expected {expected_tokens}" |
| 364 | ) |
| 365 | |
| 366 | # Make sure all aborted requests were cleaned up |
| 367 | assert not engine.output_processor.has_unfinished_requests() |
| 368 | |
| 369 | |
| 370 | @pytest.mark.parametrize("n", [1, 3]) |
| 371 | @pytest.mark.parametrize( |
| 372 | "engine_args,prompt", |
| 373 | [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)], |
| 374 | ) |
| 375 | @pytest.mark.asyncio |
| 376 | async def test_finished_flag( |
| 377 | n: int, |
| 378 | engine_args: AsyncEngineArgs, |
| 379 | prompt: PromptType, |
| 380 | ): |
| 381 | with ExitStack() as after: |
| 382 | with set_default_torch_num_threads(1): |
| 383 | engine = AsyncLLM.from_engine_args(engine_args) |
| 384 | after.callback(engine.shutdown) |
| 385 | |
| 386 | sampling_params = SamplingParams( |
| 387 | max_tokens=100, |
| 388 | output_kind=RequestOutputKind.DELTA, |
| 389 | temperature=1.0, |
| 390 | seed=33, |
| 391 | n=n, |
| 392 | ) |
| 393 | outputs = [ |
| 394 | out |
| 395 | async for out in engine.generate( |
| 396 | request_id="request-33", prompt=prompt, sampling_params=sampling_params |
| 397 | ) |
| 398 | ] |
| 399 | |
| 400 | # Assert only the last output has the finished flag set |
| 401 | assert all(not out.finished for out in outputs[:-1]) |
| 402 | assert outputs[-1].finished |
| 403 | |
| 404 | |
| 405 | @pytest.mark.parametrize( |
| 406 | "engine_args,prompt", |
| 407 | [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)], |
| 408 | ) |
| 409 | @pytest.mark.asyncio |
| 410 | async def test_mid_stream_cancellation( |
| 411 | engine_args: AsyncEngineArgs, prompt: PromptType |
| 412 | ): |
| 413 | """Test that requests can be cancelled mid-stream.""" |
| 414 | with ExitStack() as after: |
| 415 | with set_default_torch_num_threads(1): |
| 416 | engine = AsyncLLM.from_engine_args(engine_args) |
| 417 | after.callback(engine.shutdown) |
| 418 | |
| 419 | NUM_REQUESTS = 100 |
| 420 | NUM_TOKENS = 1000 |
| 421 | NUM_EXPECTED_TOKENS = 20 |
| 422 | |
| 423 | request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] |
| 424 | |
| 425 | # Create concurrent requests that will be cancelled mid-stream |
| 426 | tasks = [] |
| 427 | for request_id in request_ids: |
| 428 | tasks.append( |
| 429 | asyncio.create_task( |
| 430 | generate( |
| 431 | engine, |
| 432 | request_id, |
| 433 | prompt, |
| 434 | RequestOutputKind.DELTA, |
| 435 | NUM_TOKENS, |
| 436 | cancel_after=NUM_EXPECTED_TOKENS, |
| 437 | ) |
| 438 | ) |
| 439 | ) |
| 440 | |
| 441 | # Wait for all tasks to complete |
| 442 | results = await asyncio.gather(*tasks) |
| 443 | |
| 444 | # Verify all tasks were cancelled at the expected point. |
| 445 | # Uses >= because the cancel check is `count >= cancel_after`. |
| 446 | for num_generated_tokens, request_id in results: |
| 447 | assert num_generated_tokens >= NUM_EXPECTED_TOKENS, ( |
| 448 | f"{request_id} generated {num_generated_tokens} tokens but " |
| 449 | f"expected to cancel after {NUM_EXPECTED_TOKENS}" |
| 450 | ) |
| 451 | |
| 452 | # Make sure no requests are left hanging |
| 453 | assert not engine.output_processor.has_unfinished_requests() |
| 454 | |
| 455 | # Confirm we can reuse the request id after the cancellations. |
| 456 | request_id = request_ids[0] |
| 457 | task = asyncio.create_task( |
| 458 | generate( |
| 459 | engine, request_id, prompt, RequestOutputKind.DELTA, NUM_EXPECTED_TOKENS |
| 460 | ) |
| 461 | ) |
| 462 | num_generated_tokens, request_id = await task |
| 463 | assert num_generated_tokens == NUM_EXPECTED_TOKENS |
| 464 | assert not engine.output_processor.has_unfinished_requests() |
| 465 | |
| 466 | |
| 467 | class MockLoggingStatLogger(LoggingStatLogger): |
| 468 | def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): |
| 469 | super().__init__(vllm_config, engine_index) |
| 470 | self.log = MagicMock() |
| 471 | |
| 472 | |
| 473 | class MockAggregatedStatLogger(AggregatedLoggingStatLogger): |
| 474 | def __init__(self, vllm_config: VllmConfig, engine_indexes: list[int]): |
| 475 | super().__init__(vllm_config, engine_indexes) |
| 476 | self.log = MagicMock() |
| 477 | |
| 478 | |
| 479 | @pytest.mark.asyncio |
| 480 | async def test_customize_loggers(monkeypatch): |
| 481 | """Test that we can customize the loggers. |
| 482 | If a customized logger is provided at the init, it should |
| 483 | be added to the default loggers. |
| 484 | """ |
| 485 | with ExitStack() as after: |
| 486 | with set_default_torch_num_threads(1): |
| 487 | engine = AsyncLLM.from_engine_args( |
| 488 | TEXT_ENGINE_ARGS, |
| 489 | stat_loggers=[MockLoggingStatLogger], |
| 490 | ) |
| 491 | after.callback(engine.shutdown) |
| 492 | |
| 493 | await engine.do_log_stats() |
| 494 | |
| 495 | stat_loggers = engine.logger_manager.stat_loggers |
| 496 | assert ( |
| 497 | len(stat_loggers) == 3 |
| 498 | ) # MockLoggingStatLogger + LoggingStatLogger + Promethus Logger |
| 499 | print(f"{stat_loggers=}") |
| 500 | stat_loggers[0].per_engine_stat_loggers[0].log.assert_called_once() |
| 501 | assert isinstance(stat_loggers[1], PerEngineStatLoggerAdapter) |
| 502 | assert isinstance(stat_loggers[1].per_engine_stat_loggers[0], LoggingStatLogger) |
| 503 | assert isinstance(stat_loggers[2], PrometheusStatLogger) |
| 504 | |
| 505 | |
| 506 | @pytest.mark.asyncio |
| 507 | async def test_customize_aggregated_loggers(): |
| 508 | """Test that we can customize the aggregated loggers. |
| 509 | If a customized logger is provided at the init, it should |
| 510 | be added to the default loggers. |
| 511 | """ |
| 512 | with ExitStack() as after: |
| 513 | with set_default_torch_num_threads(1): |
| 514 | engine = AsyncLLM.from_engine_args( |
| 515 | TEXT_ENGINE_ARGS, |
| 516 | stat_loggers=[MockLoggingStatLogger, MockAggregatedStatLogger], |
| 517 | ) |
| 518 | after.callback(engine.shutdown) |
| 519 | |
| 520 | await engine.do_log_stats() |
| 521 | |
| 522 | stat_loggers = engine.logger_manager.stat_loggers |
| 523 | assert len(stat_loggers) == 4 |
| 524 | # MockLoggingStatLogger + MockAggregatedStatLogger |
| 525 | # + LoggingStatLogger + PrometheusStatLogger |
| 526 | stat_loggers[0].per_engine_stat_loggers[0].log.assert_called_once() |
| 527 | stat_loggers[1].log.assert_called_once() |
| 528 | assert isinstance(stat_loggers[2], PerEngineStatLoggerAdapter) |
| 529 | assert isinstance(stat_loggers[2].per_engine_stat_loggers[0], LoggingStatLogger) |
| 530 | assert isinstance(stat_loggers[3], PrometheusStatLogger) |
| 531 | |
| 532 | |
| 533 | @pytest.mark.asyncio(scope="module") |
| 534 | async def test_dp_rank_argument(): |
| 535 | with ExitStack() as after: |
| 536 | with set_default_torch_num_threads(1): |
| 537 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 538 | after.callback(engine.shutdown) |
| 539 | |
| 540 | sampling_params = SamplingParams( |
| 541 | max_tokens=100, |
| 542 | output_kind=RequestOutputKind.DELTA, |
| 543 | temperature=1.0, |
| 544 | seed=33, |
| 545 | ) |
| 546 | |
| 547 | # Test with valid DP rank. |
| 548 | async for _ in engine.generate( |
| 549 | request_id="request-34", |
| 550 | prompt=TEXT_PROMPT, |
| 551 | sampling_params=sampling_params, |
| 552 | data_parallel_rank=0, |
| 553 | ): |
| 554 | pass |
| 555 | |
| 556 | # Test with out-of-range DP rank. |
| 557 | with pytest.raises(VLLMValidationError): |
| 558 | async for _ in engine.generate( |
| 559 | request_id="request-35", |
| 560 | prompt=TEXT_PROMPT, |
| 561 | sampling_params=sampling_params, |
| 562 | data_parallel_rank=1, |
| 563 | ): |
| 564 | pass |
| 565 | |
| 566 | |
| 567 | @pytest.mark.asyncio(scope="module") |
| 568 | async def test_header_dp_rank_argument(): |
| 569 | with ExitStack() as after: |
| 570 | with set_default_torch_num_threads(1): |
| 571 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 572 | after.callback(engine.shutdown) |
| 573 | |
| 574 | MODEL_NAME = "test-model" |
| 575 | BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)] |
| 576 | |
| 577 | # Create models first |
| 578 | models = OpenAIServingModels( |
| 579 | engine_client=engine, |
| 580 | base_model_paths=BASE_MODEL_PATHS, |
| 581 | ) |
| 582 | |
| 583 | # Create render serving instance (required by OpenAIServingChat) |
| 584 | from vllm.renderers.online_renderer import OnlineRenderer |
| 585 | |
| 586 | online_renderer = OnlineRenderer( |
| 587 | model_config=engine.model_config, |
| 588 | renderer=engine.renderer, |
| 589 | request_logger=None, |
| 590 | chat_template=None, |
| 591 | chat_template_content_format="auto", |
| 592 | ) |
| 593 | |
| 594 | # Create serving chat instance |
| 595 | serving_chat = OpenAIServingChat( |
| 596 | engine_client=engine, |
| 597 | models=models, |
| 598 | response_role="assistant", |
| 599 | online_renderer=online_renderer, |
| 600 | chat_template=None, |
| 601 | chat_template_content_format="auto", |
| 602 | request_logger=None, |
| 603 | ) |
| 604 | # Create a chat completion request |
| 605 | req = ChatCompletionRequest( |
| 606 | model=MODEL_NAME, |
| 607 | messages=[{"role": "user", "content": TEXT_PROMPT}], |
| 608 | max_tokens=100, |
| 609 | temperature=1.0, |
| 610 | seed=33, |
| 611 | ) |
| 612 | # Test 1: Valid DP rank (0) |
| 613 | mock_raw_request = MagicMock() |
| 614 | mock_raw_request.headers = {"X-data-parallel-rank": "0"} |
| 615 | mock_raw_request.state = MagicMock() |
| 616 | |
| 617 | # Should succeed with valid rank |
| 618 | response = await serving_chat.create_chat_completion(req, mock_raw_request) |
| 619 | assert isinstance(response, ChatCompletionResponse), ( |
| 620 | "Expected a ChatCompletionResponse for valid DP rank" |
| 621 | ) |
| 622 | |
| 623 | # Test 2: Out-of-range DP rank (1) |
| 624 | mock_raw_request.headers = {"X-data-parallel-rank": "1"} |
| 625 | |
| 626 | # should raise VLLMValidationError for out-of-range rank |
| 627 | with pytest.raises(VLLMValidationError): |
| 628 | await serving_chat.create_chat_completion(req, mock_raw_request) |
| 629 | |
| 630 | |
| 631 | @pytest.mark.asyncio |
| 632 | async def test_check_health(): |
| 633 | """Test that check_health returns normally for healthy engine |
| 634 | and raises EngineDeadError when the engine is dead. |
| 635 | """ |
| 636 | from unittest.mock import patch |
| 637 | |
| 638 | from vllm.v1.engine.exceptions import EngineDeadError |
| 639 | |
| 640 | with ExitStack() as after: |
| 641 | with set_default_torch_num_threads(1): |
| 642 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 643 | after.callback(engine.shutdown) |
| 644 | |
| 645 | # Test 1: Healthy engine should not raise any exception |
| 646 | await engine.check_health() |
| 647 | |
| 648 | # Test 2: Mock the errored property to simulate a dead engine |
| 649 | with ( |
| 650 | patch.object( |
| 651 | type(engine), |
| 652 | "errored", |
| 653 | new_callable=lambda: property(lambda self: True), |
| 654 | ), |
| 655 | pytest.raises(EngineDeadError), |
| 656 | ): |
| 657 | await engine.check_health() |
| 658 | |
| 659 | # Test 3: Verify healthy engine still works after mock |
| 660 | await engine.check_health() |
| 661 | |
| 662 | |
| 663 | @pytest.mark.parametrize( |
| 664 | "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY] |
| 665 | ) |
| 666 | @pytest.mark.asyncio |
| 667 | async def test_abort_final_output(output_kind: RequestOutputKind): |
| 668 | """Test that abort() returns a final output with correct information.""" |
| 669 | with ExitStack() as after: |
| 670 | with set_default_torch_num_threads(1): |
| 671 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 672 | after.callback(engine.shutdown) |
| 673 | |
| 674 | request_id = "test-abort-final-output" |
| 675 | |
| 676 | # Start a long-running request |
| 677 | sampling_params = SamplingParams( |
| 678 | max_tokens=3000, # Long enough to allow abort |
| 679 | ignore_eos=True, |
| 680 | output_kind=output_kind, |
| 681 | temperature=0.5, |
| 682 | seed=42, |
| 683 | ) |
| 684 | |
| 685 | outputs: list[RequestOutput] = [] |
| 686 | generated = asyncio.create_task( |
| 687 | collect_outputs(engine, request_id, TEXT_PROMPT, sampling_params, outputs) |
| 688 | ) |
| 689 | |
| 690 | # Let it generate some tokens |
| 691 | await asyncio.sleep(0.5) |
| 692 | |
| 693 | # Abort the request |
| 694 | await engine.abort(request_id, internal=False) |
| 695 | |
| 696 | # Wait for generation to complete and return final output |
| 697 | final_output = await generated |
| 698 | |
| 699 | # Verify we got a final output |
| 700 | assert final_output is not None |
| 701 | assert final_output.finished |
| 702 | assert len(final_output.outputs) == 1 |
| 703 | |
| 704 | assert final_output.outputs[0].finish_reason == "abort" |
| 705 | assert final_output.outputs[0].stop_reason is None |
| 706 | |
| 707 | # Verify num_cached_tokens is set correctly |
| 708 | assert hasattr(final_output, "num_cached_tokens") |
| 709 | assert final_output.num_cached_tokens >= 0 |
| 710 | |
| 711 | # If we got intermediate outputs, verify they are consistent |
| 712 | if output_kind == RequestOutputKind.DELTA: |
| 713 | # For DELTA, sum all intermediate tokens should <= final tokens |
| 714 | token_count = sum(len(output.outputs[0].token_ids) for output in outputs) |
| 715 | assert token_count > 0 |
| 716 | # This would ordinarily be 0, but could end up > 0 if the |
| 717 | # final abort is coalesced with another chunk in the output queue. |
| 718 | assert len(final_output.outputs[0].token_ids) >= 0 |
| 719 | else: |
| 720 | # For FINAL_ONLY, we should only get the final output |
| 721 | assert len(outputs) == 0 |
| 722 | assert len(final_output.outputs[0].token_ids) > 0 |
| 723 | |
| 724 | assert not engine.output_processor.has_unfinished_requests() |
| 725 | |
| 726 | |
| 727 | async def collect_outputs( |
| 728 | engine: AsyncLLM, |
| 729 | request_id: str, |
| 730 | prompt: PromptType, |
| 731 | sampling_params: SamplingParams, |
| 732 | outputs_list: list[RequestOutput], |
| 733 | ) -> RequestOutput | None: |
| 734 | """Helper to collect outputs and return the final one.""" |
| 735 | final_output: RequestOutput | None = None |
| 736 | async for output in engine.generate( |
| 737 | request_id=request_id, prompt=prompt, sampling_params=sampling_params |
| 738 | ): |
| 739 | if not output.finished: |
| 740 | outputs_list.append(output) |
| 741 | final_output = output |
| 742 | return final_output |
| 743 | |
| 744 | |
| 745 | # ============================================================================= |
| 746 | # Pause/Resume Tests |
| 747 | # ============================================================================= |
| 748 | |
| 749 | |
| 750 | @pytest.mark.asyncio |
| 751 | async def test_pause_resume_basic(): |
| 752 | """Test basic pause/resume flag behavior and idempotency. |
| 753 | |
| 754 | Tests: |
| 755 | - pause_generation sets the paused flag |
| 756 | - resume_generation clears the paused flag |
| 757 | - calling pause when already paused is a no-op |
| 758 | - calling resume when not paused is safe |
| 759 | - all pause modes work with no requests in flight |
| 760 | - rapid pause/resume cycles don't break the engine |
| 761 | """ |
| 762 | with ExitStack() as after: |
| 763 | with set_default_torch_num_threads(1): |
| 764 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 765 | after.callback(engine.shutdown) |
| 766 | |
| 767 | # Initially not paused |
| 768 | assert not await engine.is_paused() |
| 769 | |
| 770 | # Resume when not paused should be safe |
| 771 | await engine.resume_generation() |
| 772 | assert not await engine.is_paused() |
| 773 | |
| 774 | # Pause sets flag |
| 775 | await engine.pause_generation(mode="abort") |
| 776 | assert await engine.is_paused() |
| 777 | |
| 778 | # Pause when already paused is a no-op |
| 779 | await engine.pause_generation(mode="abort") |
| 780 | assert await engine.is_paused() |
| 781 | |
| 782 | # Resume clears flag |
| 783 | await engine.resume_generation() |
| 784 | assert not await engine.is_paused() |
| 785 | |
| 786 | # Test all modes with no requests in flight |
| 787 | for mode in ("abort", "wait", "keep"): |
| 788 | await engine.pause_generation(mode=mode) |
| 789 | assert await engine.is_paused() |
| 790 | await engine.resume_generation() |
| 791 | assert not await engine.is_paused() |
| 792 | |
| 793 | # Concurrent pause/resume race conditions - should not deadlock or raise |
| 794 | await asyncio.gather( |
| 795 | engine.pause_generation(mode="abort"), |
| 796 | engine.resume_generation(), |
| 797 | engine.pause_generation(mode="abort"), |
| 798 | engine.resume_generation(), |
| 799 | ) |
| 800 | |
| 801 | # Ensure we end in a known state |
| 802 | await engine.resume_generation() |
| 803 | assert not await engine.is_paused() |
| 804 | |
| 805 | # Engine should still work after all cycles |
| 806 | sampling_params = SamplingParams(max_tokens=5) |
| 807 | async for out in engine.generate( |
| 808 | request_id="post-cycles", |
| 809 | prompt=TEXT_PROMPT, |
| 810 | sampling_params=sampling_params, |
| 811 | ): |
| 812 | pass |
| 813 | assert out.finished |
| 814 | |
| 815 | |
| 816 | @pytest.mark.asyncio |
| 817 | async def test_pause_abort(): |
| 818 | """Test that mode='abort' aborts in-flight requests immediately.""" |
| 819 | with ExitStack() as after: |
| 820 | with set_default_torch_num_threads(1): |
| 821 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 822 | after.callback(engine.shutdown) |
| 823 | |
| 824 | # Start a long-running request |
| 825 | sampling_params = SamplingParams(max_tokens=1000, ignore_eos=True) |
| 826 | outputs: list[RequestOutput] = [] |
| 827 | |
| 828 | async def gen(): |
| 829 | async for out in engine.generate( |
| 830 | request_id="test-abort-pause", |
| 831 | prompt=TEXT_PROMPT, |
| 832 | sampling_params=sampling_params, |
| 833 | ): |
| 834 | outputs.append(out) |
| 835 | return outputs[-1] if outputs else None |
| 836 | |
| 837 | # Start generation task |
| 838 | gen_task = asyncio.create_task(gen()) |
| 839 | |
| 840 | # Wait for some tokens to be generated |
| 841 | while len(outputs) < 3: |
| 842 | await asyncio.sleep(0.01) |
| 843 | |
| 844 | # Pause with abort mode |
| 845 | await engine.pause_generation(mode="abort") |
| 846 | |
| 847 | # Wait for task to complete (should be aborted) |
| 848 | final_output = await gen_task |
| 849 | |
| 850 | # Request should be finished (aborted) |
| 851 | assert final_output is not None |
| 852 | assert final_output.finished |
| 853 | assert final_output.outputs[0].finish_reason == "abort" |
| 854 | |
| 855 | # Also test that new requests are blocked while paused, then resume |
| 856 | assert await engine.is_paused() |
| 857 | |
| 858 | request_completed = False |
| 859 | |
| 860 | async def gen_blocked(): |
| 861 | nonlocal request_completed |
| 862 | async for out in engine.generate( |
| 863 | request_id="test-blocked", |
| 864 | prompt=TEXT_PROMPT, |
| 865 | sampling_params=SamplingParams(max_tokens=5), |
| 866 | ): |
| 867 | pass |
| 868 | request_completed = True |
| 869 | return out |
| 870 | |
| 871 | # Start a request (should block) |
| 872 | gen_task2 = asyncio.create_task(gen_blocked()) |
| 873 | |
| 874 | # Wait a bit - request should not have completed |
| 875 | await asyncio.sleep(0.3) |
| 876 | assert not request_completed, "Request should be blocked while paused" |
| 877 | |
| 878 | # Resume |
| 879 | await engine.resume_generation() |
| 880 | |
| 881 | # Now request should complete |
| 882 | final_output2 = await asyncio.wait_for(gen_task2, timeout=10.0) |
| 883 | assert request_completed |
| 884 | assert final_output2.finished |
| 885 | |
| 886 | |
| 887 | @pytest.mark.asyncio |
| 888 | async def test_pause_then_abort_queued_request(): |
| 889 | """Test that aborting a request that was submitted while paused (in |
| 890 | _paused_adds_queue) aborts it and notifies the client; the request does |
| 891 | not run after resume. |
| 892 | """ |
| 893 | with ExitStack() as after: |
| 894 | with set_default_torch_num_threads(1): |
| 895 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 896 | after.callback(engine.shutdown) |
| 897 | |
| 898 | request_id = "abort-queued-request" |
| 899 | sampling_params = SamplingParams(max_tokens=20, ignore_eos=True) |
| 900 | outputs: list[RequestOutput] = [] |
| 901 | |
| 902 | # Pause first so the next add goes to _paused_adds_queue |
| 903 | await engine.pause_generation(mode="keep") |
| 904 | assert await engine.is_paused() |
| 905 | |
| 906 | async def gen(): |
| 907 | async for out in engine.generate( |
| 908 | request_id=request_id, |
| 909 | prompt=TEXT_PROMPT, |
| 910 | sampling_params=sampling_params, |
| 911 | ): |
| 912 | outputs.append(out) |
| 913 | return outputs[-1] if outputs else None |
| 914 | |
| 915 | gen_task = asyncio.create_task(gen()) |
| 916 | |
| 917 | # Give the request time to reach the engine and sit in _paused_adds_queue |
| 918 | await asyncio.sleep(0.2) |
| 919 | |
| 920 | # Abort the queued request |
| 921 | await engine.abort(request_id, internal=False) |
| 922 | |
| 923 | # Resume so the engine can process and deliver the abort output |
| 924 | await engine.resume_generation() |
| 925 | |
| 926 | final_output = await asyncio.wait_for(gen_task, timeout=10.0) |
| 927 | assert final_output is not None |
| 928 | assert final_output.finished |
| 929 | assert final_output.outputs[0].finish_reason == "abort" |
| 930 | # Request was never run, so no tokens |
| 931 | assert len(final_output.outputs[0].token_ids) == 0 |
| 932 | |
| 933 | |
| 934 | @pytest.mark.asyncio |
| 935 | async def test_pause_wait(): |
| 936 | """Test that mode='wait' waits for in-flight requests to complete.""" |
| 937 | with ExitStack() as after: |
| 938 | with set_default_torch_num_threads(1): |
| 939 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 940 | after.callback(engine.shutdown) |
| 941 | |
| 942 | # Start a request - use fewer tokens since wait mode waits for completion |
| 943 | sampling_params = SamplingParams(max_tokens=10, ignore_eos=True) |
| 944 | got_first_token = asyncio.Event() |
| 945 | request_completed = False |
| 946 | |
| 947 | async def gen(): |
| 948 | nonlocal request_completed |
| 949 | async for out in engine.generate( |
| 950 | request_id="test-wait", |
| 951 | prompt=TEXT_PROMPT, |
| 952 | sampling_params=sampling_params, |
| 953 | ): |
| 954 | got_first_token.set() |
| 955 | request_completed = True |
| 956 | return out |
| 957 | |
| 958 | # Start generation |
| 959 | gen_task = asyncio.create_task(gen()) |
| 960 | |
| 961 | # Wait for generation to start (event-driven) |
| 962 | await asyncio.wait_for(got_first_token.wait(), timeout=30.0) |
| 963 | |
| 964 | # Pause with wait mode - should wait for request to finish |
| 965 | await engine.pause_generation(mode="wait") |
| 966 | |
| 967 | # By now the request should be done (wait mode waits for completion) |
| 968 | assert request_completed, "Request should have completed during wait" |
| 969 | |
| 970 | final_output = gen_task.result() |
| 971 | assert final_output.finished |
| 972 | # Should complete normally, not aborted |
| 973 | assert final_output.outputs[0].finish_reason != "eos" |
| 974 | |
| 975 | |
| 976 | @pytest.mark.asyncio |
| 977 | async def test_pause_keep_single_request(): |
| 978 | """Test that mode='keep' freezes a single request and resumes with timing gap.""" |
| 979 | with ExitStack() as after: |
| 980 | with set_default_torch_num_threads(1): |
| 981 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 982 | after.callback(engine.shutdown) |
| 983 | |
| 984 | sampling_params = SamplingParams(max_tokens=30, ignore_eos=True) |
| 985 | token_times: list[tuple[int, float]] = [] |
| 986 | pause_duration = 5.0 |
| 987 | pause_token_idx = 0 |
| 988 | |
| 989 | async def generator_task(): |
| 990 | """Generate tokens and record timestamps.""" |
| 991 | async for output in engine.generate( |
| 992 | request_id="test-keep-single", |
| 993 | prompt=TEXT_PROMPT, |
| 994 | sampling_params=sampling_params, |
| 995 | ): |
| 996 | token_count = len(output.outputs[0].token_ids) |
| 997 | token_times.append((token_count, time.monotonic())) |
| 998 | return output |
| 999 | |
| 1000 | async def controller_task(): |
| 1001 | """Pause and resume the engine.""" |
| 1002 | nonlocal pause_token_idx |
| 1003 | # Wait for some tokens (event-driven, handles slow token generation) |
| 1004 | while len(token_times) < 5: |
| 1005 | await asyncio.sleep(0.01) |
| 1006 | |
| 1007 | # Pause with keep mode |
| 1008 | await engine.pause_generation(mode="keep") |
| 1009 | pause_token_idx = len(token_times) |
| 1010 | |
| 1011 | # Sleep while paused |
| 1012 | await asyncio.sleep(pause_duration) |
| 1013 | |
| 1014 | # Resume |
| 1015 | await engine.resume_generation() |
| 1016 | |
| 1017 | # Run both tasks with timeout for slow generation |
| 1018 | gen_task = asyncio.create_task(generator_task()) |
| 1019 | ctrl_task = asyncio.create_task(controller_task()) |
| 1020 | |
| 1021 | final_output, _ = await asyncio.wait_for( |
| 1022 | asyncio.gather(gen_task, ctrl_task), timeout=60.0 |
| 1023 | ) |
| 1024 | |
| 1025 | # Request should complete with all tokens |
| 1026 | assert final_output.finished |
| 1027 | assert len(final_output.outputs[0].token_ids) == 30 |
| 1028 | |
| 1029 | # Check the gap at the recorded pause index matches the pause duration |
| 1030 | pause_gap = ( |
| 1031 | token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1] |
| 1032 | ) |
| 1033 | assert pause_gap >= pause_duration * 0.8, ( |
| 1034 | f"Expected gap of ~{pause_duration}s after pause, got {pause_gap:.3f}s" |
| 1035 | ) |
| 1036 | |
| 1037 | |
| 1038 | @pytest.mark.asyncio |
| 1039 | async def test_pause_keep_multi_request(): |
| 1040 | """Test that mode='keep' freezes multiple concurrent requests and all resume.""" |
| 1041 | with ExitStack() as after: |
| 1042 | with set_default_torch_num_threads(1): |
| 1043 | engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS) |
| 1044 | after.callback(engine.shutdown) |
| 1045 | |
| 1046 | num_requests = 3 |
| 1047 | sampling_params = SamplingParams(max_tokens=10, ignore_eos=True) |
| 1048 | completed_requests: list[str] = [] |
| 1049 | any_token_generated = asyncio.Event() |
| 1050 | |
| 1051 | async def gen_multi(request_id: str): |
| 1052 | async for out in engine.generate( |
| 1053 | request_id=request_id, |
| 1054 | prompt=TEXT_PROMPT, |
| 1055 | sampling_params=sampling_params, |
| 1056 | ): |
| 1057 | any_token_generated.set() |
| 1058 | completed_requests.append(request_id) |
| 1059 | return out |
| 1060 | |
| 1061 | # Start multiple requests |
| 1062 | tasks = [ |
| 1063 | asyncio.create_task(gen_multi(f"req-multi-{i}")) |
| 1064 | for i in range(num_requests) |
| 1065 | ] |
| 1066 | |
| 1067 | # Wait for at least one token across any request (event-driven) |
| 1068 | await asyncio.wait_for(any_token_generated.wait(), timeout=30.0) |
| 1069 | |
| 1070 | # Pause with keep mode |
| 1071 | await engine.pause_generation(mode="keep") |
| 1072 | |
| 1073 | # Wait while paused |
| 1074 | await asyncio.sleep(0.5) |
| 1075 | |
| 1076 | # Resume |
| 1077 | await engine.resume_generation() |
| 1078 | |
| 1079 | # All requests should complete |
| 1080 | results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=60.0) |
| 1081 | |
| 1082 | assert len(completed_requests) == num_requests |
| 1083 | for result in results: |
| 1084 | assert result.finished |
| 1085 | assert len(result.outputs[0].token_ids) == 10 |