A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 21:00 UTC
[RL] Add Pause and Resume Generation for Asynchronous RL Training (#28037)
Signed-off-by: SamitHuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Signed-off-by: samithuang <285365963@qq.com> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com>
Samit committed
Nov 20, 2025 at 19:01 UTC
371b1d4c61335ed4c1d7fb2acee75274cc6d4551
4 files changed
+182
vllm/engine/protocol.py
+27
@@ -149,6 +149,33 @@ class EngineClient(ABC):
149
"""Load a new LoRA adapter into the engine for future requests."""
150
...
151
152
+ @abstractmethod
153
+ async def pause_generation(
154
+ self,
155
+ *,
156
+ wait_for_inflight_requests: bool = False,
157
+ clear_cache: bool = True,
158
+ ) -> None:
159
+ """Pause new generation/encoding requests.
160
+
161
+ Args:
162
+ wait_for_inflight_requests: When ``True`` waits for in-flight requests
163
+ to finish before pausing. When ``False`` (default), aborts in-flight
164
+ requests immediately.
165
+ clear_cache: Whether to clear KV and prefix caches after draining.
166
+ """
167
+ ...
168
+
169
+ @abstractmethod
170
+ async def resume_generation(self) -> None:
171
+ """Resume accepting generation/encoding requests."""
172
+ ...
173
+
174
+ @abstractmethod
175
+ async def is_paused(self) -> bool:
176
+ """Return whether the engine is currently paused."""
177
+ ...
178
+
179
async def scale_elastic_ep(
180
self, new_data_parallel_size: int, drain_timeout: int = 300
181
) -> None:
vllm/entrypoints/openai/api_server.py
+78
@@ -394,6 +394,84 @@ async def get_server_load_metrics(request: Request):
394
return JSONResponse(content={"server_load": request.app.state.server_load_metrics})
395
396
397
+@router.post("/pause")
398
+async def pause_generation(
399
+ raw_request: Request,
400
+ wait_for_inflight_requests: bool = Query(False),
401
+ clear_cache: bool = Query(True),
402
+) -> JSONResponse:
403
+ """Pause generation requests to allow weight updates.
404
+
405
+ Args:
406
+ wait_for_inflight_requests: When ``True`` waits for in-flight
407
+ requests to finish before pausing. When ``False`` (default),
408
+ aborts any in-flight requests immediately.
409
+ clear_cache: Whether to clear KV/prefix caches after draining.
410
+ """
411
+
412
+ engine = engine_client(raw_request)
413
+
414
+ try:
415
+ await engine.pause_generation(
416
+ wait_for_inflight_requests=wait_for_inflight_requests,
417
+ clear_cache=clear_cache,
418
+ )
419
+ return JSONResponse(
420
+ content={"status": "paused"},
421
+ status_code=HTTPStatus.OK.value,
422
+ )
423
+
424
+ except ValueError as err:
425
+ return JSONResponse(
426
+ content={"error": str(err)},
427
+ status_code=HTTPStatus.BAD_REQUEST.value,
428
+ )
429
+ except Exception as err: # pragma: no cover - defensive
430
+ logger.exception("Failed to pause generation")
431
+ return JSONResponse(
432
+ content={"error": f"Failed to pause generation: {err}"},
433
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
434
+ )
435
+
436
+
437
+@router.post("/resume")
438
+async def resume_generation(raw_request: Request) -> JSONResponse:
439
+ """Resume generation after a pause."""
440
+
441
+ engine = engine_client(raw_request)
442
+
443
+ try:
444
+ await engine.resume_generation()
445
+ return JSONResponse(
446
+ content={"status": "resumed"},
447
+ status_code=HTTPStatus.OK.value,
448
+ )
449
+ except Exception as err: # pragma: no cover - defensive
450
+ logger.exception("Failed to resume generation")
451
+ return JSONResponse(
452
+ content={"error": f"Failed to resume generation: {err}"},
453
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
454
+ )
455
+
456
+
457
+@router.get("/is_paused")
458
+async def is_paused(raw_request: Request) -> JSONResponse:
459
+ """Return the current pause status."""
460
+
461
+ engine = engine_client(raw_request)
462
+
463
+ try:
464
+ paused = await engine.is_paused()
465
+ except Exception as err: # pragma: no cover - defensive
466
+ logger.exception("Failed to fetch pause status")
467
+ return JSONResponse(
468
+ content={"error": f"Failed to fetch pause status: {err}"},
469
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
470
+ )
471
+
472
+ return JSONResponse(content={"is_paused": paused})
473
+
474
+
475
@router.post(
476
"/tokenize",
477
dependencies=[Depends(validate_json_request)],
vllm/v1/engine/async_llm.py
+64
@@ -152,6 +152,10 @@ class AsyncLLM(EngineClient):
152
)
153
self.logger_manager.log_engine_initialized()
154
155
+ # Pause / resume state for async RL workflows.
156
+ self._pause_cond = asyncio.Condition()
157
+ self._paused = False
158
+
159
self.output_handler: asyncio.Task | None = None
160
try:
161
# Start output handler eagerly if we are in the asyncio eventloop.
@@ -404,6 +408,10 @@ class AsyncLLM(EngineClient):
408
# to handle startup failure gracefully in the OpenAI server.
409
self._run_output_handler()
410
411
+ # Wait until generation is resumed if the engine is paused.
412
+ async with self._pause_cond:
413
+ await self._pause_cond.wait_for(lambda: not self._paused)
414
+
415
if tokenization_kwargs is None:
416
tokenization_kwargs = {}
417
truncate_prompt_tokens = sampling_params.truncate_prompt_tokens
@@ -551,6 +559,58 @@ class AsyncLLM(EngineClient):
559
if self.log_requests:
560
logger.info("Aborted request(s) %s.", ",".join(request_ids))
561
562
+ async def pause_generation(
563
+ self,
564
+ *,
565
+ wait_for_inflight_requests: bool = False,
566
+ clear_cache: bool = True,
567
+ ) -> None:
568
+ """
569
+ Pause generation to allow model weight updates.
570
+
571
+ New generation/encoding requests are blocked until resume.
572
+
573
+ Args:
574
+ wait_for_inflight_requests: When ``True`` waits for in-flight
575
+ requests to finish before pausing. When ``False`` (default),
576
+ immediately aborts any in-flight requests.
577
+ clear_cache: Whether to clear KV cache and prefix cache after
578
+ draining. Set to ``False`` to preserve cache for faster resume.
579
+ Default is ``True`` (clear caches).
580
+ """
581
+
582
+ async with self._pause_cond:
583
+ if self._paused:
584
+ return
585
+ self._paused = True
586
+
587
+ if not wait_for_inflight_requests:
588
+ request_ids = list(self.output_processor.request_states.keys())
589
+ if request_ids:
590
+ await self.abort(request_ids)
591
+
592
+ # Wait for running requests to drain before clearing cache.
593
+ if self.output_processor.has_unfinished_requests():
594
+ await self.output_processor.wait_for_requests_to_drain()
595
+
596
+ # Clear cache
597
+ if clear_cache:
598
+ await self.reset_prefix_cache()
599
+ await self.reset_mm_cache()
600
+
601
+ async def resume_generation(self) -> None:
602
+ """Resume generation after :meth:`pause_generation`."""
603
+
604
+ async with self._pause_cond:
605
+ self._paused = False
606
+ self._pause_cond.notify_all() # Wake up all waiting requests
607
+
608
+ async def is_paused(self) -> bool:
609
+ """Return whether the engine is currently paused."""
610
+
611
+ async with self._pause_cond:
612
+ return self._paused
613
+
614
async def encode(
615
self,
616
prompt: PromptType,
@@ -582,6 +642,10 @@ class AsyncLLM(EngineClient):
642
# to handle startup failure gracefully in the OpenAI server.
643
self._run_output_handler()
644
645
+ # Respect pause state before accepting new requests.
646
+ async with self._pause_cond:
647
+ await self._pause_cond.wait_for(lambda: not self._paused)
648
+
649
if tokenization_kwargs is None:
650
tokenization_kwargs = {}
651
_validate_truncation_size(
vllm/v1/engine/output_processor.py
+13
@@ -350,6 +350,8 @@ class OutputProcessor:
350
self.parent_requests: dict[str, ParentRequest] = {}
351
self.lora_states = LoRARequestStates(log_stats)
352
self.tracer: Tracer | None = None
353
+ self._requests_drained = asyncio.Event()
354
+ self._requests_drained.set()
355
356
def get_num_unfinished_requests(self):
357
return len(self.request_states)
@@ -357,6 +359,11 @@ class OutputProcessor:
359
def has_unfinished_requests(self) -> bool:
360
return len(self.request_states) > 0
361
362
+ async def wait_for_requests_to_drain(self) -> None:
363
+ if not self.request_states:
364
+ return
365
+ await self._requests_drained.wait()
366
+
367
def propagate_error(self, e: Exception):
368
"""Propagate error to all generate() tasks."""
369
@@ -396,6 +403,8 @@ class OutputProcessor:
403
child_reqs = self.abort_requests(child_reqs)
404
request_ids_to_abort.extend(child_reqs)
405
self.parent_requests.pop(request_id, None)
406
+ if not self.request_states:
407
+ self._requests_drained.set()
408
return request_ids_to_abort
409
410
def add_request(
@@ -420,6 +429,8 @@ class OutputProcessor:
429
log_stats=self.log_stats,
430
stream_interval=self.stream_interval,
431
)
432
+ if self._requests_drained.is_set():
433
+ self._requests_drained.clear()
434
self.request_states[request_id] = req_state
435
if parent_req:
436
self.parent_requests[parent_req.request_id] = parent_req
@@ -511,6 +522,8 @@ class OutputProcessor:
522
parent_req = req_state.parent_req
523
if parent_req and not parent_req.child_requests:
524
self.parent_requests.pop(parent_req.request_id, None)
525
+ if not self.request_states:
526
+ self._requests_drained.set()
527
if not engine_core_output.finished:
528
# If req not finished in EngineCore, but Detokenizer
529
# detected stop string, abort needed in EngineCore.