# SPDX-License-Identifier: Apache-2.0 # 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, _pack_value_if_large, _unpack_if_shm_handle, pack_diffusion_output_shm, unpack_diffusion_output_shm, ) from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def _large_numel(dtype: torch.dtype) -> int: return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1 def _cleanup_shm_handle(value: object) -> None: if isinstance(value, dict) and value.get("__tensor_shm__"): with contextlib.suppress(FileNotFoundError): _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 output = DiffusionOutput(output={"image": image, "video": video, "metadata": {"keep": "inline"}}) pack_diffusion_output_shm(output) assert output.output["image"]["__tensor_shm__"] is True assert output.output["video"]["__tensor_shm__"] is True assert output.output["metadata"] == {"keep": "inline"} unpack_diffusion_output_shm(output) torch.testing.assert_close(output.output["image"], image) torch.testing.assert_close(output.output["video"], video) assert output.output["metadata"] == {"keep": "inline"} def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None: # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output. video = torch.arange(300_000, dtype=torch.float32) audio = torch.arange(300_000, dtype=torch.float32) * 3 output = DiffusionOutput(output=(video, audio)) pack_diffusion_output_shm(output) assert isinstance(output.output, tuple) assert output.output[0]["__tensor_shm__"] is True assert output.output[1]["__tensor_shm__"] is True unpack_diffusion_output_shm(output) assert isinstance(output.output, tuple) torch.testing.assert_close(output.output[0], video) torch.testing.assert_close(output.output[1], audio) def test_diffusion_output_list_tensors_round_trip_through_shm() -> None: frames = [torch.arange(300_000, dtype=torch.float32), torch.arange(300_000, dtype=torch.float32) + 1] output = DiffusionOutput(output=list(frames)) pack_diffusion_output_shm(output) assert isinstance(output.output, list) assert all(isinstance(item, dict) and item["__tensor_shm__"] is True for item in output.output) unpack_diffusion_output_shm(output) assert isinstance(output.output, list) torch.testing.assert_close(output.output[0], frames[0]) 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 = { "type": DIFFUSION_RPC_RESULT_ENVELOPE, "result": DiffusionOutput(output=tensor), "rank_statuses": [{"rank": 0, "ok": True}], } packed = pack_diffusion_output_shm(envelope) assert packed is envelope result = packed["result"] assert isinstance(result, DiffusionOutput) assert result.output["__tensor_shm__"] is True assert packed["rank_statuses"] == [{"rank": 0, "ok": True}] unpacked = unpack_diffusion_output_shm(packed) assert unpacked is envelope result = unpacked["result"] assert isinstance(result, DiffusionOutput) torch.testing.assert_close(result.output, tensor) 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 output = BatchRunnerOutput.from_list( [ RunnerOutput(request_id="req-0", finished=True, result=DiffusionOutput(output=first)), RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output={"image": second})), RunnerOutput(request_id="req-error", finished=True, result=DiffusionOutput(error="boom")), ] ) pack_diffusion_output_shm(output) assert output.runner_outputs[0].result.output["__tensor_shm__"] is True assert output.runner_outputs[1].result.output["image"]["__tensor_shm__"] is True assert output.runner_outputs[2].result.error == "boom" unpack_diffusion_output_shm(output) torch.testing.assert_close(output["req-0"].result.output, first) torch.testing.assert_close(output["req-1"].result.output["image"], second) assert output["req-error"].result.error == "boom" def test_per_request_views_of_a_batch_tensor_do_not_ship_whole_storage() -> None: """Regression: per-request slices used to pickle the whole batch storage. Each view is under the packing threshold, but pickle writes the entire shared storage per view, so a batch's message grew by the full batch tensor once per request and overflowed the ring buffer. """ import pickle num_requests = 16 # Each row stays under the threshold; the shared storage is well over it. row_numel = _SHM_TENSOR_THRESHOLD // (2 * torch.empty((), dtype=torch.float32).element_size()) batch = torch.arange(num_requests * row_numel, dtype=torch.float32).reshape(num_requests, row_numel) per_request = batch[0] assert per_request.nelement() * per_request.element_size() <= _SHM_TENSOR_THRESHOLD assert per_request.untyped_storage().nbytes() > _SHM_TENSOR_THRESHOLD output = BatchRunnerOutput.from_list( [ RunnerOutput( request_id=f"req-{i}", finished=True, result=DiffusionOutput(trajectory_latents=batch[i]), ) for i in range(num_requests) ] ) unpacked_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL)) pack_diffusion_output_shm(output) packed_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL)) assert output.runner_outputs[0].result.trajectory_latents["__tensor_shm__"] is True # Whole storage per request before; handles only after. storage_bytes = per_request.untyped_storage().nbytes() assert unpacked_size > num_requests * storage_bytes / 2 assert packed_size < _SHM_TENSOR_THRESHOLD unpack_diffusion_output_shm(output) for i in range(num_requests): torch.testing.assert_close(output[f"req-{i}"].result.trajectory_latents, batch[i]) def test_pack_value_keeps_tensor_at_threshold_inline() -> None: tensor = torch.arange( _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(), dtype=torch.float32, ) packed = _pack_value_if_large(tensor) assert packed is tensor def test_pack_value_packs_large_tensor_and_round_trips() -> None: tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32) packed = _pack_value_if_large(tensor) try: assert isinstance(packed, dict) assert packed["__tensor_shm__"] is True assert packed["shape"] == [tensor.numel()] assert packed["torch_dtype"] == "torch.float32" unpacked = _unpack_if_shm_handle(packed) assert isinstance(unpacked, torch.Tensor) torch.testing.assert_close(unpacked, tensor) finally: _cleanup_shm_handle(packed) def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_values() -> None: large = torch.arange(_large_numel(torch.float32), dtype=torch.float32) small = torch.arange(8, dtype=torch.float32) list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32) payload = { "media": { "large": large, "small": small, }, "list_value": [list_tensor], "metadata": {"prompt": "keep inline"}, } packed = _pack_value_if_large(payload) try: assert packed is not payload assert packed["media"] is not payload["media"] assert packed["media"]["large"]["__tensor_shm__"] is True assert packed["media"]["small"] is small # Lists are recursed too: the large tensor inside is packed and a new # list is returned, while the input payload is left untouched. assert packed["list_value"] is not payload["list_value"] assert packed["list_value"][0]["__tensor_shm__"] is True assert payload["list_value"][0] is list_tensor assert packed["metadata"] == {"prompt": "keep inline"} torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large) torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor) finally: if isinstance(packed, dict): _cleanup_shm_handle(packed.get("media", {}).get("large")) list_value = packed.get("list_value") if isinstance(list_value, list) and list_value: _cleanup_shm_handle(list_value[0]) def test_pack_value_preserves_dtype_shape_and_values_for_bfloat16() -> None: tensor = torch.arange(_large_numel(torch.bfloat16), dtype=torch.float32).to(torch.bfloat16).reshape(1, -1) packed = _pack_value_if_large(tensor) try: assert isinstance(packed, dict) assert packed["__tensor_shm__"] is True assert packed["shape"] == list(tensor.shape) assert packed["torch_dtype"] == "torch.bfloat16" assert packed["numpy_dtype"] == "float32" unpacked = _unpack_if_shm_handle(packed) assert isinstance(unpacked, torch.Tensor) assert unpacked.dtype == torch.bfloat16 torch.testing.assert_close(unpacked, tensor) finally: _cleanup_shm_handle(packed) def test_pack_value_packs_non_contiguous_large_tensor_values() -> None: tensor = torch.arange(_large_numel(torch.float32) * 2, dtype=torch.float32).reshape(-1, 2)[:, 0] assert not tensor.is_contiguous() packed = _pack_value_if_large(tensor) try: assert isinstance(packed, dict) assert packed["__tensor_shm__"] is True assert packed["shape"] == list(tensor.shape) unpacked = _unpack_if_shm_handle(packed) assert isinstance(unpacked, torch.Tensor) torch.testing.assert_close(unpacked, tensor) finally: _cleanup_shm_handle(packed)