[Diffusion] Fix DLO DP concurrent request execution (#5864)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/test_diffusion_ipc.py

@@ -2,11 +2,19 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import contextlib
+import multiprocessing as mp
+import queue
+import threading
+import time
+from multiprocessing import shared_memory
 
+import numpy as np
 import pytest
 import torch
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import (
     _SHM_TENSOR_THRESHOLD,
     DIFFUSION_RPC_RESULT_ENVELOPE,
@@ -30,6 +38,129 @@ def _cleanup_shm_handle(value: object) -> None:
             _unpack_if_shm_handle(value)
 
 
+def _result_queue_worker(connection, rank: int) -> None:
+    """Send one nested NumPy result through a worker-owned MessageQueue."""
+    result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
+    try:
+        connection.send(result_mq.export_handle())
+        assert connection.recv() == "reader-ready"
+
+        frames = np.full(300_000, rank, dtype=np.float32)
+        audio = np.arange(300_000, dtype=np.float32) + rank
+        output = DiffusionOutput(
+            output={
+                "rank": rank,
+                "nested": {"frames": frames},
+                "audio": [audio],
+            }
+        )
+        pack_diffusion_output_shm(output)
+        connection.send(
+            [
+                output.output["nested"]["frames"]["name"],
+                output.output["audio"][0]["name"],
+            ]
+        )
+        result_mq.enqueue(output)
+        assert connection.recv() == "shutdown"
+    finally:
+        result_mq.shutdown()
+        connection.close()
+
+
+def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> None:
+    """Exercise the production per-worker queue/pump/SHM lifecycle."""
+    ctx = mp.get_context("spawn")
+    parent_connections = []
+    processes = []
+    result_mqs = []
+    shm_names = []
+    executor = None
+
+    try:
+        for rank in range(2):
+            parent_connection, child_connection = ctx.Pipe()
+            process = ctx.Process(
+                target=_result_queue_worker,
+                args=(child_connection, rank),
+                name=f"IPCResultWorker-{rank}",
+            )
+            process.start()
+            child_connection.close()
+            parent_connections.append(parent_connection)
+            processes.append(process)
+
+        for connection in parent_connections:
+            handle = connection.recv()
+            result_mqs.append(MessageQueue.create_from_handle(handle, 0))
+            connection.send("reader-ready")
+            shm_names.extend(connection.recv())
+
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._result_mqs = result_mqs
+        executor._result_mq = result_mqs[0]
+        executor._pump_running = False
+        executor._pump_stop = threading.Event()
+        executor._sync_result_buffer = queue.Queue()
+        executor._is_failed = False
+        executor._futures_lock = threading.RLock()
+        executor._rpc_futures = {}
+        executor._output_futures = {}
+        executor._completed_outputs = {}
+        executor._batch_split_map = {}
+        executor._start_result_pump()
+
+        deadline = time.monotonic() + 10.0
+        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
+            time.sleep(0.01)
+        assert executor._sync_result_buffer.qsize() == 2
+
+        outputs = []
+        for _ in range(2):
+            output = executor._sync_result_buffer.get_nowait()
+            unpack_diffusion_output_shm(output)
+            outputs.append(output)
+
+        assert sorted(output.output["rank"] for output in outputs) == [0, 1]
+        for output in outputs:
+            rank = output.output["rank"]
+            assert output.output["nested"]["frames"][0] == rank
+            assert output.output["audio"][0][0] == rank
+
+        executor._pump_stop.set()
+        for thread in executor._result_pump_threads:
+            thread.join(timeout=3.0)
+        assert all(not thread.is_alive() for thread in executor._result_pump_threads)
+    finally:
+        if executor is not None:
+            executor._pump_stop.set()
+            for thread in getattr(executor, "_result_pump_threads", []):
+                thread.join(timeout=3.0)
+
+        for connection in parent_connections:
+            with contextlib.suppress(BrokenPipeError, EOFError, OSError):
+                connection.send("shutdown")
+        for process in processes:
+            process.join(timeout=10.0)
+            if process.is_alive():
+                process.terminate()
+                process.join(timeout=5.0)
+        for result_mq in result_mqs:
+            with contextlib.suppress(Exception):
+                result_mq.shutdown()
+        for connection in parent_connections:
+            connection.close()
+
+    assert all(process.exitcode == 0 for process in processes)
+    for name in shm_names:
+        try:
+            leaked = shared_memory.SharedMemory(name=name)
+        except FileNotFoundError:
+            continue
+        leaked.close()
+        pytest.fail(f"shared memory segment {name} still exists after unpack")
+
+
 def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
     image = torch.arange(300_000, dtype=torch.float32)
     video = torch.arange(300_000, dtype=torch.float32) * 2
@@ -83,6 +214,29 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
     torch.testing.assert_close(output.output[1], frames[1])
 
 
+def test_diffusion_output_numpy_array_round_trips_through_shm() -> None:
+    frames = np.arange(300_000, dtype=np.float32)
+    output = DiffusionOutput(output=frames)
+
+    pack_diffusion_output_shm(output)
+
+    assert output.output["__ndarray_shm__"] is True
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, np.ndarray)
+    np.testing.assert_array_equal(output.output, frames)
+
+
+def test_numpy_object_array_stays_inline() -> None:
+    values = np.empty(200_000, dtype=object)
+    values[:] = "safe-to-pickle"
+
+    packed = _pack_value_if_large(values)
+
+    assert packed is values
+
+
 def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
     tensor = torch.arange(300_000, dtype=torch.float32)
     envelope = {
@@ -108,6 +262,25 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
     assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
 
 
+def test_rpc_result_envelope_dp_tagged_output_round_trips_through_shm() -> None:
+    frames = np.arange(300_000, dtype=np.float32)
+    envelope = {
+        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+        "result": {"dp_rank": 1, "output": DiffusionOutput(output=frames)},
+        "rank_statuses": [{"rank": 1, "ok": True}],
+    }
+
+    packed = pack_diffusion_output_shm(envelope)
+    tagged = packed["result"]
+    assert tagged["dp_rank"] == 1
+    assert tagged["output"].output["__ndarray_shm__"] is True
+
+    unpacked = unpack_diffusion_output_shm(packed)
+    tagged = unpacked["result"]
+    assert tagged["dp_rank"] == 1
+    np.testing.assert_array_equal(tagged["output"].output, frames)
+
+
 def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
     first = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
     second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1