Skip to content

Commit

Permalink
[core] add wake_up doc and some sanity check (vllm-project#12361)
Browse files Browse the repository at this point in the history
Signed-off-by: youkaichao <[email protected]>
  • Loading branch information
youkaichao authored and rasmith committed Jan 30, 2025
1 parent 819fa05 commit 7db4b50
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 0 deletions.
3 changes: 3 additions & 0 deletions vllm/entrypoints/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -1157,6 +1157,9 @@ def sleep(self, level: int = 1):
self.llm_engine.sleep(level=level)

def wake_up(self):
"""
Wake up the engine from sleep mode. See the :meth:`sleep` method
for more details."""
self.llm_engine.wake_up()

# LEGACY
Expand Down
9 changes: 9 additions & 0 deletions vllm/executor/executor_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def __init__(
self.prompt_adapter_config = vllm_config.prompt_adapter_config
self.observability_config = vllm_config.observability_config
self._init_executor()
self.is_sleeping = False

@abstractmethod
def _init_executor(self) -> None:
Expand Down Expand Up @@ -194,10 +195,18 @@ def stop_profile(self) -> None:
self.collective_rpc("stop_profile")

def sleep(self, level: int = 1):
if self.is_sleeping:
logger.warning("Executor is already sleeping.")
return
self.collective_rpc("sleep", kwargs=dict(level=level))
self.is_sleeping = True

def wake_up(self):
if not self.is_sleeping:
logger.warning("Executor is not sleeping.")
return
self.collective_rpc("wake_up")
self.is_sleeping = False

def save_sharded_state(
self,
Expand Down

0 comments on commit 7db4b50

Please sign in to comment.