05b794f29b
Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: lishunyang12 <lishunyang03@gmail.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: lishunyang12 <lishunyang03@gmail.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: StormMelody <22371184@zju.edu.cn> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import queue
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.engine.messages import (
|
|
CollectiveRPCRequestMessage,
|
|
CollectiveRPCResultMessage,
|
|
ErrorMessage,
|
|
)
|
|
from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _request(rpc_id: str) -> CollectiveRPCRequestMessage:
|
|
return CollectiveRPCRequestMessage(
|
|
rpc_id=rpc_id,
|
|
method="health",
|
|
timeout=1,
|
|
args=(),
|
|
kwargs={},
|
|
stage_ids=None,
|
|
)
|
|
|
|
|
|
def test_correlated_rpc_client_unregisters_timeout_before_late_result() -> None:
|
|
request_queue: queue.Queue = queue.Queue()
|
|
result_queue: queue.Queue = queue.Queue()
|
|
client = CorrelatedRpcClient(request_queue, result_queue)
|
|
|
|
try:
|
|
with pytest.raises(TimeoutError, match="first timed out"):
|
|
client.execute(
|
|
("collective", "first"),
|
|
_request("first"),
|
|
timeout=0.01,
|
|
timeout_message="first timed out",
|
|
)
|
|
assert request_queue.get_nowait().rpc_id == "first"
|
|
|
|
result_queue.put(
|
|
CollectiveRPCResultMessage(
|
|
rpc_id="first",
|
|
method="health",
|
|
stage_ids=[0],
|
|
results=["late"],
|
|
)
|
|
)
|
|
result_queue.put(
|
|
CollectiveRPCResultMessage(
|
|
rpc_id="second",
|
|
method="health",
|
|
stage_ids=[0],
|
|
results=["current"],
|
|
)
|
|
)
|
|
current = client.execute(
|
|
("collective", "second"),
|
|
_request("second"),
|
|
timeout=1,
|
|
timeout_message="second timed out",
|
|
)
|
|
|
|
assert isinstance(current, CollectiveRPCResultMessage)
|
|
assert current.results == ["current"]
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_correlated_rpc_client_rejects_after_fatal_without_enqueuing() -> None:
|
|
request_queue: queue.Queue = queue.Queue()
|
|
result_queue: queue.Queue = queue.Queue()
|
|
client = CorrelatedRpcClient(request_queue, result_queue)
|
|
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
pending = executor.submit(
|
|
client.execute,
|
|
("collective", "pending"),
|
|
_request("pending"),
|
|
timeout=1,
|
|
timeout_message="unexpected timeout",
|
|
)
|
|
assert request_queue.get(timeout=1).rpc_id == "pending"
|
|
result_queue.put(ErrorMessage(error="orchestrator failed", fatal=True))
|
|
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
|
pending.result(timeout=1)
|
|
|
|
with pytest.raises(RuntimeError, match="orchestrator failed"):
|
|
client.execute(
|
|
("collective", "after-fatal"),
|
|
_request("after-fatal"),
|
|
timeout=1,
|
|
timeout_message="unexpected timeout",
|
|
)
|
|
assert request_queue.empty()
|
|
finally:
|
|
client.close()
|