[Bugfix] Fix lost async diffusion outputs in request-level batching (#6023)

Signed-off-by: samithuang <285365963@qq.com>

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -65,6 +65,15 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
 logger = init_logger(__name__)
 
 _ASYNC_OUTPUT_THREAD_JOIN_TIMEOUT_S = 10.0
+# Maximum time (in seconds) to wait for pending background D2H / SHM packing
+# to drain before the worker executes memory-releasing lifecycle tasks
+# (e.g. during sleep/wake transitions). This barrier prevents device tensors
+# from being freed while the side CUDA stream is still actively reading them.
+_ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0
+
+# Worker entry points that release device memory. Background D2H/SHM packing
+# still reads model output tensors, so it must finish before these run.
+_MEMORY_RELEASING_METHODS = frozenset({"sleep", "handle_sleep_task"})
 
 
 @dataclass
@@ -850,6 +859,13 @@ class WorkerProc:
 
         self._async_output_queue: queue.Queue | None = None
         self._async_output_thread: threading.Thread | None = None
+        self._async_output_done = threading.Condition()
+        self._async_output_pending = 0
+        # MessageQueue.acquire_write() is single-writer: it reads current_idx,
+        # writes that block, then advances the index. The async output thread
+        # enqueues OUTPUT_READY while the main loop enqueues COMPUTE_DONE, so
+        # unsynchronized writers can target the same block and drop a message.
+        self._result_mq_lock = threading.Lock()
         if not self.od_config.step_execution:
             self._async_output_queue = queue.Queue()
             self._async_output_thread = threading.Thread(
@@ -882,25 +898,32 @@ class WorkerProc:
         )
         return wrapper
 
+    def _enqueue_result(self, msg: Any) -> None:
+        """Serialize writes to the single-writer result queue."""
+        with self._result_mq_lock:
+            self.result_mq.enqueue(msg)
+
     def _return_result(self, output: Any, rpc_id: str | None = None) -> None:
         """Reply to client, only on rank 0."""
         if self.result_mq is None:
             return
         if isinstance(output, OmniACK):
-            self.result_mq.enqueue(output)
+            self._enqueue_result(output)
             return
 
         # Async path: enqueue compute_done immediately, bg thread does D2H+SHM.
         if not self.od_config.step_execution and isinstance(output, (DiffusionOutput, BatchRunnerOutput)):
             async_output_id = WorkerProc._generate_async_output_id()
             gpu_event = current_omni_platform.record_device_event()
+            with self._async_output_done:
+                self._async_output_pending += 1
             self._async_output_queue.put((output, async_output_id, gpu_event))
             msg = AsyncDiffusionOutput(
                 kind=AsyncOutputKind.COMPUTE_DONE,
                 rpc_id=rpc_id,
                 async_output_id=async_output_id,
             )
-            self.result_mq.enqueue(msg)
+            self._enqueue_result(msg)
             return
 
         # Sync path (original, or async fallback).
@@ -909,7 +932,7 @@ class WorkerProc:
         except Exception as e:
             if hasattr(output, "output"):
                 logger.warning("SHM pack failed for model output: %s", e)
-        self.result_mq.enqueue(output)
+        self._enqueue_result(output)
 
     def _async_output_loop(self):
         """Background thread: D2H + SHM packing for async diffusion output.
@@ -932,7 +955,7 @@ class WorkerProc:
                 pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
                 d2h_stream.synchronize()
 
-                self.result_mq.enqueue(
+                self._enqueue_result(
                     AsyncDiffusionOutput(
                         kind=AsyncOutputKind.OUTPUT_READY,
                         async_output_id=async_output_id,
@@ -944,13 +967,38 @@ class WorkerProc:
                     "Async output packing failed for id '%s'; sending error",
                     async_output_id,
                 )
-                self.result_mq.enqueue(
+                self._enqueue_result(
                     AsyncDiffusionOutput(
                         kind=AsyncOutputKind.OUTPUT_READY,
                         async_output_id=async_output_id,
                         error="Background D2H/SHM packing failed",
                     )
                 )
+            finally:
+                with self._async_output_done:
+                    # Clamped: only items enqueued by _return_result are counted.
+                    self._async_output_pending = max(0, self._async_output_pending - 1)
+                    self._async_output_done.notify_all()
+
+    def drain_async_outputs(self, timeout: float = _ASYNC_OUTPUT_DRAIN_TIMEOUT_S) -> bool:
+        """Block until background D2H/SHM packing has no work left.
+
+        Returns False if outputs are still in flight when ``timeout`` expires.
+        """
+        with self._async_output_done:
+            if self._async_output_pending == 0:
+                return True
+            drained = self._async_output_done.wait_for(lambda: self._async_output_pending == 0, timeout=timeout)
+            pending = self._async_output_pending
+        if not drained:
+            logger.warning(
+                "Worker %d: %d async output(s) still in flight after %.1fs; "
+                "releasing device memory now may drop OUTPUT_READY messages",
+                self.gpu_id,
+                pending,
+                timeout,
+            )
+        return drained
 
     def shutdown(self) -> None:
         """Stop background work and release worker-owned IPC resources."""
@@ -1049,6 +1097,8 @@ class WorkerProc:
         }
 
         try:
+            if method in _MEMORY_RELEASING_METHODS:
+                self.drain_async_outputs()
             # Use execute_method from WorkerWrapperBase for consistent method resolution
             result = self.worker.execute_method(method, *args, **kwargs)
         except Exception as e:
@@ -1114,6 +1164,7 @@ class WorkerProc:
                 continue
 
             if isinstance(msg, dict) and msg.get("type") == "sleep":
+                self.drain_async_outputs()
                 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)
@@ -1140,7 +1191,7 @@ class WorkerProc:
                         if rpc_id is not None:
                             # Async RPC: must complete the executor's pending
                             # future so collective_rpc() doesn't hang.
-                            self.result_mq.enqueue(
+                            self._enqueue_result(
                                 AsyncDiffusionOutput(
                                     kind=AsyncOutputKind.RPC_RESULT,
                                     rpc_id=rpc_id,