[Feat][sleepmode] add omni sleepmode and ack protocol (#2022)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -27,7 +27,10 @@ from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
+    OmniACK,
     OmniDiffusionConfig,
+    OmniSleepTask,
+    OmniWakeTask,
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -77,6 +80,7 @@ class DiffusionWorker:
         self.model_runner: DiffusionModelRunner | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
+        self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
         # Create model runner
         self.model_runner = DiffusionModelRunner(
@@ -159,8 +163,10 @@ class DiffusionWorker:
     def _get_profiler(self) -> WorkerProfiler | None:
         return getattr(self, "profiler", None)
 
-    def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None:
+    def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None, **kwargs) -> None:
         """Load the diffusion model using DiffusionModelRunner."""
+        load_format = kwargs.get("load_format", load_format)
+        custom_pipeline_name = kwargs.get("custom_pipeline_name", custom_pipeline_name)
         with (
             set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config),
             set_current_vllm_config(self.vllm_config),
@@ -170,6 +176,8 @@ class DiffusionWorker:
                 load_format=load_format,
                 custom_pipeline_name=custom_pipeline_name,
             )
+            current_omni_platform.synchronize()
+            gc.collect()
         process_memory = get_process_gpu_memory(self.local_rank)
         if process_memory is not None:
             logger.info(
@@ -284,77 +292,176 @@ class DiffusionWorker:
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
-        process_memory_before_sleep = get_process_gpu_memory(self.local_rank)
-        free_bytes_before_sleep = None
-        if process_memory_before_sleep is None:
-            free_bytes_before_sleep = current_omni_platform.get_free_memory()
+        allocator = CuMemAllocator.get_instance()
+
+        usage_before = allocator.get_current_usage()
 
-        # Save the buffers before level 2 sleep
         if level == 2 and self.model_runner is not None:
+            if hasattr(self.model_runner, "graph_runners"):
+                self.model_runner.graph_runners.clear()
+                logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
             model = self.model_runner.pipeline
             self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
 
-        allocator = CuMemAllocator.get_instance()
-        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
-        process_memory_after_sleep = get_process_gpu_memory(self.local_rank)
-        if process_memory_before_sleep is not None and process_memory_after_sleep is not None:
-            freed_bytes = process_memory_before_sleep - process_memory_after_sleep
-            used_bytes = process_memory_after_sleep
-            accounting_scope = "process-scoped"
+        free_mem_before = current_omni_platform.get_free_memory()
+
+        # Level 1: Offload weights; Level 2: Total Discard
+        offload_tags = ("weights",) if level == 1 else tuple()
+        allocator.sleep(offload_tags=offload_tags)
+
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+        free_mem_after = current_omni_platform.get_free_memory()
+        try:
+            total_mem = current_omni_platform.get_device_total_memory()
+        except (NotImplementedError, AttributeError):
+            total_mem = torch.cuda.get_device_properties(self.device).total_memory
+
+        phys_freed_bytes = max(0, free_mem_after - free_mem_before)
+        phys_used_bytes = total_mem - free_mem_after
+
+        if usage_before > 0:
+            logger.info(
+                f"[Diffusion Worker {self.rank}] Sleep Level {level}: "
+                f"physically freed {phys_freed_bytes / GiB_bytes:.2f} GiB, "
+                f"{phys_used_bytes / GiB_bytes:.2f} GiB is still in use."
+            )
         else:
-            free_bytes_after_sleep = current_omni_platform.get_free_memory()
-            assert free_bytes_before_sleep is not None
-            device_id = self.device.index if self.device.index is not None else 0
-            total = current_omni_platform.get_device_total_memory(device_id)
-            freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
-            used_bytes = total - free_bytes_after_sleep
-            accounting_scope = "device-scoped fallback"
-        assert freed_bytes >= 0, "Memory usage increased after sleeping."
-        logger.info(
-            "Sleep mode (%s) freed %.2f GiB memory, %.2f GiB memory is still in use.",
-            accounting_scope,
-            freed_bytes / GiB_bytes,
-            used_bytes / GiB_bytes,
-        )
-        return True
+            logger.info(f"[Worker {self.rank}] Sleep Level {level} completed (GPU was already empty).")
+        logger.info(f"[Worker {self.rank}] Memory usage before sleep: {usage_before / GiB_bytes:.2f} GiB.")
+        return usage_before
 
     def wake_up(self, tags: list[str] | None = None) -> bool:
         """
-        Wake up the worker from sleep mode. See the sleep function
-        method for more details.
+        Wake up the worker from sleep mode.
+
+        Re-activates the memory allocator for the specified tags and restores
+        model buffers from CPU back to GPU if they were saved during Level 2 sleep.
 
         Args:
-            tags: An optional list of tags to reallocate the worker memory
-                for specific memory allocations. Values must be in
-                `("weights")`. If None, all memory is reallocated.
-                wake_up should be called with all tags (or None) before the
-                worker is used again.
+            tags: List of memory pool tags to re-activate (e.g., ["weights"]
+                  to match Level 1 sleep). If None, all pools are re-activated.
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
         allocator = CuMemAllocator.get_instance()
         allocator.wake_up(tags)
-
-        # Restore the buffers after level 2 sleep
+        current_omni_platform.synchronize()
         if len(self._sleep_saved_buffers) and self.model_runner is not None:
             model = self.model_runner.pipeline
             for name, buffer in model.named_buffers():
                 if name in self._sleep_saved_buffers:
                     buffer.data.copy_(self._sleep_saved_buffers[name].data)
             self._sleep_saved_buffers = {}
+            logger.info(f"[Worker {self.rank}] Buffers restored from CPU.")
+        logger.info(f"[Worker {self.rank}] Wake-up complete.")
         return True
 
+    def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK:
+        from vllm_omni.platforms import current_omni_platform
+
+        try:
+            if isinstance(task, dict):
+                task = OmniSleepTask(**task)
+            logger.info(f"[Worker {self.rank}] Handshake Received: Task {task.task_id}")
+
+            current_omni_platform.synchronize()
+            usage_before = current_omni_platform.get_current_memory_usage(self.device)
+            self.sleep(level=task.level)
+            current_omni_platform.synchronize()
+            usage_after = current_omni_platform.get_current_memory_usage(self.device)
+            real_freed = max(0, usage_before - usage_after)
+            logger.info(f"[Worker {self.rank}] Preparing ACK: freed_bytes={real_freed / GiB_bytes:.2f} GiB.")
+
+            # Ensure all ranks have completed sleep before measuring memory and sending ACK
+            if torch.distributed.is_initialized():
+                t_freed = torch.tensor([float(real_freed)], device=self.device)
+                torch.distributed.all_reduce(t_freed)
+                real_freed = int(t_freed.item())
+
+            if self.rank != 0:
+                return None
+
+            ack = OmniACK(
+                task_id=task.task_id,
+                status="SUCCESS",
+                stage_id=self.stage_id,
+                rank=self.rank,
+                freed_bytes=real_freed,
+                # return RL need metadata
+                metadata={
+                    "source": f"Platform_{current_omni_platform.get_device_name()}",
+                    "total_freed_gib": f"{real_freed / GiB_bytes:.2f}",
+                    "rank_residual_gib": f"{usage_after / GiB_bytes:.2f}",
+                },
+            )
+            logger.info(f"[Worker {self.rank}] ACK emitted. Freed {real_freed / GiB_bytes:.2f} GiB.")
+            return ack
+        except Exception as e:
+            logger.error(f"Sleep failed: {e}", exc_info=True)
+            if torch.distributed.is_initialized():
+                try:
+                    torch.distributed.barrier()
+                except Exception:
+                    pass
+            return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e))
+
+    def handle_wake_task(self, task: OmniWakeTask) -> OmniACK:
+        from vllm_omni.platforms import current_omni_platform
+
+        try:
+            if isinstance(task, dict):
+                task = OmniWakeTask(**task)
+            logger.info(f"[Worker {self.rank}] Responding to Wake-up Task: {task.task_id}")
+            self.wake_up(tags=task.tags)
+
+            logger.info(f"[Worker {self.rank}] wake_up logic finished, entering barrier...")
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+            current_omni_platform.synchronize()
+            usage_now = current_omni_platform.get_current_memory_usage(self.device)
+            current_used_gib = usage_now / (1024**3)
+
+            if self.rank != 0:
+                return None
+            logger.info(f"[Worker {self.rank}] PASSED barrier, about to return to loop.")
+
+            return OmniACK(
+                task_id=task.task_id,
+                status="SUCCESS",
+                stage_id=self.stage_id,
+                rank=self.rank,
+                metadata={
+                    "state": "WARM",
+                    "source": f"Platform_{current_omni_platform.get_device_name()}",
+                    "current_vram_gib": f"{current_used_gib:.2f}",
+                },
+            )
+        except Exception as e:
+            logger.error(f"Wake-up failed on Rank {self.rank}: {e}", exc_info=True)
+            if torch.distributed.is_initialized():
+                try:
+                    torch.distributed.barrier()
+                except Exception:
+                    pass
+            return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e))
+
     def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
         """Get memory pool context for sleep mode support."""
-        if self.od_config.enable_sleep_mode:
+        is_sleep_enabled = getattr(self.od_config, "enable_sleep_mode", False)
+        if is_sleep_enabled:
+            current_omni_platform.synchronize()
+            gc.collect()
             from vllm.device_allocator.cumem import CuMemAllocator
 
             allocator = CuMemAllocator.get_instance()
             if tag == "weights":
                 assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
+            logger.info(f"[Worker {self.rank}] Activating Diffusion CuMem pool for tag: {tag}")
             return allocator.use_memory_pool(tag=tag)
-        else:
-            return nullcontext()
+        return nullcontext()
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
@@ -395,11 +502,13 @@ class WorkerProc:
         od_config: OmniDiffusionConfig,
         gpu_id: int,
         broadcast_handle,
+        wake_event: mp.Event,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
     ):
         self.od_config = od_config
         self.gpu_id = gpu_id
+        self.wake_event = wake_event
 
         # Inter-process Communication
         self.context = zmq.Context(io_threads=2)
@@ -414,7 +523,13 @@ class WorkerProc:
         if gpu_id == 0:
             self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
             self.result_mq_handle = self.result_mq.export_handle()
+            WorkerProc._shared_result_handle = self.result_mq_handle
             logger.info(f"Worker {gpu_id} created result MessageQueue")
+        else:
+            handle = getattr(WorkerProc, "_shared_result_handle", None)
+            if handle:
+                self.result_mq = MessageQueue.create_from_handle(handle, gpu_id)
+                logger.info(f"Worker {gpu_id} attached to shared result MessageQueue")
 
         assert od_config.master_port is not None
 
@@ -438,13 +553,17 @@ class WorkerProc:
         )
         return wrapper
 
-    def return_result(self, output: object):
+    def return_result(self, output: Any):
         """Reply to client, only on rank 0."""
         if self.result_mq is not None:
+            if isinstance(output, OmniACK):
+                self.result_mq.enqueue(output)
+                return
             try:
                 pack_diffusion_output_shm(output)
             except Exception as e:
-                logger.warning("SHM pack failed, falling back to raw enqueue: %s", e)
+                if hasattr(output, "output"):
+                    logger.warning("SHM pack failed for model output: %s", e)
             self.result_mq.enqueue(output)
 
     def recv_message(self):
@@ -480,20 +599,31 @@ class WorkerProc:
         while self._running:
             msg = None
             try:
-                msg = self.recv_message()
-            except Exception as e:
-                logger.error(
-                    f"Error receiving message in worker loop: {e}",
-                    exc_info=True,
-                )
+                msg = self.mq.dequeue(timeout=1.0)
+            except Exception:
+                if self.wake_event and self.wake_event.is_set():
+                    self.wake_event.clear()
+                    logger.info(f"Worker {self.gpu_id} caught OOB POKE, forcing wake-up sequence.")
+                    msg = {"type": "wake_up", "task_id": "recovery-task", "tags": None}
+                else:
+                    continue
+            if msg is None:
                 continue
 
             if msg is None or len(msg) == 0:
                 logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id)
                 continue
 
+            if isinstance(msg, dict) and msg.get("type") == "sleep":
+                task = OmniSleepTask(level=msg.get("level", 2), task_id=msg.get("task_id", "local"))
+                ack = self.worker.handle_sleep_task(task)
+                self.return_result(ack)
+            elif isinstance(msg, dict) and msg.get("type") == "wake_up":
+                task = OmniWakeTask(tags=msg.get("tags"), task_id=msg.get("task_id", "local"))
+                ack = self.worker.handle_wake_task(task)
+                self.return_result(ack)
             # Route message based on type
-            if isinstance(msg, dict) and msg.get("type") == "rpc":
+            elif isinstance(msg, dict) and msg.get("type") == "rpc":
                 try:
                     result, should_reply = self.execute_rpc(msg)
                     if should_reply:
@@ -538,6 +668,7 @@ class WorkerProc:
         od_config: OmniDiffusionConfig,
         pipe_writer: mp.connection.Connection,
         broadcast_handle,
+        wake_event: mp.Event,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
     ) -> None:
@@ -549,6 +680,7 @@ class WorkerProc:
             od_config,
             gpu_id=rank,
             broadcast_handle=broadcast_handle,
+            wake_event=wake_event,
             worker_extension_cls=worker_extension_cls,
             custom_pipeline_args=custom_pipeline_args,
         )
@@ -574,6 +706,7 @@ class WorkerWrapperBase:
         gpu_id: int,
         od_config: OmniDiffusionConfig,
         base_worker_class: type = DiffusionWorker,
+        wake_event: mp.Event = None,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
     ):
@@ -734,6 +867,12 @@ class WorkerWrapperBase:
         """
         return self.worker.wake_up(tags)
 
+    def handle_sleep_task(self, task):
+        return self.worker.handle_sleep_task(task)
+
+    def handle_wake_task(self, task):
+        return self.worker.handle_wake_task(task)
+
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup resources."""
         return self.worker.shutdown()