53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
"""Unit tests for the device step of the diffusion pause barrier."""
|
||
|
|
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from unittest.mock import MagicMock
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||
|
|
|
||
|
|
|
||
|
|
def _make_worker(runner: object) -> DiffusionWorker:
|
||
|
|
worker = object.__new__(DiffusionWorker)
|
||
|
|
worker.model_runner = runner
|
||
|
|
worker.rank = 0
|
||
|
|
return worker
|
||
|
|
|
||
|
|
|
||
|
|
def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker):
|
||
|
|
platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||
|
|
order: list = []
|
||
|
|
manager = MagicMock()
|
||
|
|
manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True
|
||
|
|
platform.synchronize.side_effect = lambda: order.append("sync")
|
||
|
|
worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
|
||
|
|
|
||
|
|
worker.synchronize_device(timeout=2.5)
|
||
|
|
|
||
|
|
assert order == [("prefetch", 2.5), "sync"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker):
|
||
|
|
platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||
|
|
manager = MagicMock()
|
||
|
|
manager.wait_prefetch.return_value = False
|
||
|
|
worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
|
||
|
|
|
||
|
|
with pytest.raises(TimeoutError, match="prefetch"):
|
||
|
|
worker.synchronize_device(timeout=0.01)
|
||
|
|
|
||
|
|
platform.synchronize.assert_not_called()
|
||
|
|
|
||
|
|
|
||
|
|
def test_synchronize_device_without_transfer_manager_only_syncs(mocker):
|
||
|
|
platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||
|
|
worker = _make_worker(SimpleNamespace(kv_transfer_manager=None))
|
||
|
|
|
||
|
|
worker.synchronize_device(timeout=None)
|
||
|
|
|
||
|
|
platform.synchronize.assert_called_once_with()
|