update

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/v1/engine/test_llm_engine.py

@@ -2,11 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import random
 from typing import TYPE_CHECKING
+from unittest.mock import MagicMock
 
 import pytest
 
 from vllm import LLM
 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
+from vllm.v1.engine.llm_engine import LLMEngine
 from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector
 
 if TYPE_CHECKING:
@@ -104,6 +106,55 @@ def _get_test_sampling_params(
     ], n_list
 
 
+def test_llm_engine_sleep_state_cache():
+    llm_engine = object.__new__(LLMEngine)
+    llm_engine.multiprocess_mode = True
+    llm_engine._sleeping_tags = set()
+    llm_engine.engine_core = MagicMock()
+    llm_engine.logger_manager = None
+
+    assert llm_engine.is_sleeping() is False
+    llm_engine.engine_core.is_sleeping.assert_not_called()
+
+    llm_engine.sleep(level=0, mode="keep")
+    llm_engine.engine_core.sleep.assert_called_once_with(0, "keep")
+    assert llm_engine.is_sleeping() is True
+    assert llm_engine._sleeping_tags == {"scheduling"}
+    llm_engine.engine_core.is_sleeping.assert_not_called()
+
+    llm_engine.wake_up(["scheduling"])
+    llm_engine.engine_core.wake_up.assert_called_once_with(["scheduling"])
+    assert llm_engine.is_sleeping() is False
+    assert llm_engine._sleeping_tags == set()
+    llm_engine.engine_core.is_sleeping.assert_not_called()
+
+    llm_engine.sleep(level=1, mode="keep")
+    assert llm_engine.is_sleeping() is True
+    assert llm_engine._sleeping_tags == {"scheduling", "weights", "kv_cache"}
+
+    llm_engine.wake_up(["scheduling"])
+    assert llm_engine.is_sleeping() is True
+    assert llm_engine._sleeping_tags == {"weights", "kv_cache"}
+
+    llm_engine.wake_up(["weights"])
+    assert llm_engine.is_sleeping() is True
+    assert llm_engine._sleeping_tags == {"kv_cache"}
+
+    llm_engine.wake_up(["kv_cache"])
+    assert llm_engine.is_sleeping() is False
+    assert llm_engine._sleeping_tags == set()
+    llm_engine.engine_core.is_sleeping.assert_not_called()
+
+    llm_engine = object.__new__(LLMEngine)
+    llm_engine.multiprocess_mode = False
+    llm_engine._sleeping_tags = set()
+    llm_engine.engine_core = MagicMock()
+    llm_engine.engine_core.is_sleeping.return_value = True
+
+    assert llm_engine.is_sleeping() is True
+    llm_engine.engine_core.is_sleeping.assert_called_once_with()
+
+
 def test_compatibility_with_skip_tokenizer_init(
     vllm_model_skip_tokenizer_init: VllmRunner,
     example_prompts: list[str],

vllm/v1/engine/llm_engine.py

@@ -86,6 +86,8 @@ class LLMEngine:
             self.dp_group = parallel_config.stateless_init_dp_group()
         else:
             self.dp_group = None
+        self.multiprocess_mode = multiprocess_mode
+        self._sleeping_tags: set[str] = set()
         self.should_execute_dummy_batch = False
 
         self.renderer = renderer = renderer_from_config(self.vllm_config)
@@ -352,17 +354,29 @@ class LLMEngine:
 
     def sleep(self, level: int = 1, mode: PauseMode = "abort"):
         self.engine_core.sleep(level, mode)
+        self._sleeping_tags.add("scheduling")
+        if level >= 1 and self._sleeping_tags == {"scheduling"}:
+            self._sleeping_tags.update({"weights", "kv_cache"})
 
         if self.logger_manager is not None:
             self.logger_manager.record_sleep_state(1, level)
 
     def wake_up(self, tags: list[str] | None = None):
         self.engine_core.wake_up(tags)
+        self._sleeping_tags.discard("scheduling")
+        if tags is None:
+            self._sleeping_tags.clear()
+        else:
+            self._sleeping_tags.difference_update(
+                tag for tag in tags if tag != "scheduling"
+            )
 
         if self.logger_manager is not None:
             self.logger_manager.record_sleep_state(0, 0)
 
     def is_sleeping(self) -> bool:
+        if self.multiprocess_mode:
+            return bool(self._sleeping_tags)
         return self.engine_core.is_sleeping()
 
     def get_metrics(self) -> list[Metric]: