[sleep functionality] code refactor about sleep/wake_up (#50431)

Signed-off-by: Andy Xie <andy.xning@gmail.com>

vllm/entrypoints/serve/dev/sleep/api_router.py

@@ -24,8 +24,6 @@ async def sleep(raw_request: Request):
     level = raw_request.query_params.get("level", "1")
     mode = raw_request.query_params.get("mode", "abort")
     await engine_client(raw_request).sleep(int(level), mode)
-    # FIXME: in v0 with frontend multiprocessing, the sleep command
-    # is sent but does not finish yet when we return a response.
     return Response(status_code=200)
 
 
@@ -37,8 +35,6 @@ async def wake_up(raw_request: Request):
         tags = None
     logger.info("wake up the engine with tags: %s", tags)
     await engine_client(raw_request).wake_up(tags)
-    # FIXME: in v0 with frontend multiprocessing, the wake-up command
-    # is sent but does not finish yet when we return a response.
     return Response(status_code=200)
 
 

vllm/v1/core/sched/scheduler.py

@@ -2476,6 +2476,7 @@ class Scheduler(SchedulerInterface):
         return self._pause_state
 
     def set_pause_state(self, pause_state: PauseState) -> None:
+        logger.info("setting pause state to %s", pause_state.name)
         self._pause_state = pause_state
 
     def _free_request_blocks(self, request: Request):

vllm/v1/engine/core.py

@@ -15,7 +15,7 @@ from functools import partial
 from inspect import isclass, signature
 from logging import DEBUG
 from multiprocessing.queues import Queue
-from typing import Any, TypeVar, cast
+from typing import Any, TypeVar, cast, get_args
 
 import msgspec
 import zmq
@@ -855,7 +855,7 @@ class EngineCore:
         - ``keep``: Set PAUSED_ALL; return a Future that completes when the
           output queue is empty.
         """
-        if mode not in ("keep", "abort", "wait"):
+        if mode not in get_args(PauseMode):
             raise ValueError(f"Invalid pause mode: {mode}")
         if mode == "wait":
             raise ValueError("'wait' mode can't be used in inproc-engine mode")
@@ -1933,7 +1933,7 @@ class EngineCoreProc(EngineCore):
         - ``keep``: Set PAUSED_ALL; return a Future that completes when the
           output queue is empty.
         """
-        if mode not in ("keep", "abort", "wait"):
+        if mode not in get_args(PauseMode):
             raise ValueError(f"Invalid pause mode: {mode}")
 
         def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None: