451 lines
16 KiB
Python
451 lines
16 KiB
Python
# 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)
|