[Bugfix][v1] Prevent premature scheduler resumption during wake_up be…

vllm/v1/engine/core.py

@@ -176,7 +176,8 @@ def __init__(
         # Setup KV Caches and update CacheConfig after profiling.
         kv_cache_config = self._initialize_kv_caches(vllm_config)
         self.structured_output_manager = StructuredOutputManager(vllm_config)
-
+        # setup the guard for reload weight
+        self._needs_reload_weights: bool = False
         # Setup scheduler.
         Scheduler = vllm_config.scheduler_config.get_scheduler_cls()
 
@@ -1092,6 +1093,8 @@ def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None | Future:
         if level < 1:
             return pause_future
 
+        if level >= 2:
+            self._needs_reload_weights = True
         # Level 1+: Delegate to executor for GPU memory management
         model_executor = self.model_executor
         if pause_future is None:
@@ -1137,7 +1140,9 @@ def wake_up(self, tags: list[str] | None = None) -> bool:
 
         # Partial wakes intentionally keep the remaining allocations asleep.
         # Resume scheduling only once all executor memory is resident again.
-        fully_awake = not self.model_executor.is_sleeping
+        fully_awake = (
+            not self.model_executor.is_sleeping and not self._needs_reload_weights
+        )
         if fully_awake:
             self.resume_scheduler()
         return fully_awake
@@ -1217,6 +1222,12 @@ def collective_rpc(
         args: tuple = (),
         kwargs: dict[str, Any] | None = None,
     ) -> list[_R]:
+        if method == "reload_weights":
+            self._needs_reload_weights = False
+
+            if not self.model_executor.is_sleeping:
+                self.resume_scheduler()
+
         return self.model_executor.collective_rpc(method, timeout, args, kwargs)
 
     def set_weight_version(self, weight_version: str) -> None: