| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from unittest.mock import AsyncMock, call |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.v1.engine.core_client import DPLBAsyncMPClient |
| 9 | |
| 10 | pytestmark = pytest.mark.cpu_test |
| 11 | |
| 12 | |
| 13 | @pytest.mark.asyncio |
| 14 | @pytest.mark.parametrize( |
| 15 | "results, expected", |
| 16 | [ |
| 17 | ([True, True], True), |
| 18 | ([True, False], False), |
| 19 | ([False, True], False), |
| 20 | ([False, False], False), |
| 21 | ], |
| 22 | ) |
| 23 | async def test_dp_wake_up_requires_all_engines_awake(results, expected): |
| 24 | client = object.__new__(DPLBAsyncMPClient) |
| 25 | client.core_engines = [b"engine-0", b"engine-1"] |
| 26 | client._call_utility_async = AsyncMock(side_effect=results) |
| 27 | |
| 28 | assert await client.wake_up_async(["weights"]) is expected |
| 29 | assert client._call_utility_async.await_args_list == [ |
| 30 | call("wake_up", ["weights"], engine=b"engine-0"), |
| 31 | call("wake_up", ["weights"], engine=b"engine-1"), |
| 32 | ] |