1f60ddef5d
Signed-off-by: huanglong <huanglong@linux.alibaba.com> Signed-off-by: inkcherry <mingzhi.liu@amd.com> Co-authored-by: huanglong <huanglong@linux.alibaba.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: Huang Long <121648372+LLLL114@users.noreply.github.com>
85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
from sglang_router import mini_lb
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("result", "restored"),
|
|
[
|
|
(
|
|
{
|
|
"success": False,
|
|
"message": "instance is not idle",
|
|
"safe_to_restore": True,
|
|
},
|
|
True,
|
|
),
|
|
(
|
|
{
|
|
"success": False,
|
|
"message": "instance unhealthy, restart required",
|
|
"safe_to_restore": False,
|
|
},
|
|
False,
|
|
),
|
|
({"success": False, "message": "connection lost"}, False),
|
|
],
|
|
)
|
|
def test_failed_role_switch_restores_only_healthy_worker(monkeypatch, result, restored):
|
|
worker_url = "http://prefill:8000"
|
|
load_balancer = mini_lb.MiniLoadBalancer.__new__(mini_lb.MiniLoadBalancer)
|
|
load_balancer.timeout = 1
|
|
load_balancer.prefill_urls = [worker_url]
|
|
load_balancer.prefill_bootstrap_ports = [8998]
|
|
load_balancer.decode_urls = ["http://decode:8000"]
|
|
monkeypatch.setattr(mini_lb, "lb", load_balancer)
|
|
|
|
async def post_role_switch(*_args, **_kwargs):
|
|
return 400, result
|
|
|
|
monkeypatch.setattr(mini_lb, "_post_role_switch", post_role_switch)
|
|
response = asyncio.run(
|
|
mini_lb.pd_role_switch(
|
|
{
|
|
"worker_url": worker_url,
|
|
"new_role": "decode",
|
|
"drain_timeout_secs": 0,
|
|
}
|
|
)
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert (worker_url in load_balancer.prefill_urls) is restored
|
|
|
|
|
|
def test_role_switch_forwards_decode_graph_requirements(monkeypatch):
|
|
worker_url = "http://prefill:8000"
|
|
load_balancer = mini_lb.MiniLoadBalancer.__new__(mini_lb.MiniLoadBalancer)
|
|
load_balancer.timeout = 1
|
|
load_balancer.prefill_urls = [worker_url]
|
|
load_balancer.prefill_bootstrap_ports = [8998]
|
|
load_balancer.decode_urls = ["http://decode:8000"]
|
|
monkeypatch.setattr(mini_lb, "lb", load_balancer)
|
|
sent_body = {}
|
|
|
|
async def post_role_switch(_worker_url, body):
|
|
sent_body.update(body)
|
|
return 200, {"success": True, "message": "ok"}
|
|
|
|
monkeypatch.setattr(mini_lb, "_post_role_switch", post_role_switch)
|
|
response = asyncio.run(
|
|
mini_lb.pd_role_switch(
|
|
{
|
|
"worker_url": worker_url,
|
|
"new_role": "decode",
|
|
"decode_cuda_graph_bs": [1, 2, 4],
|
|
"decode_cuda_graph_memory_gb": 1.25,
|
|
}
|
|
)
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert sent_body["decode_cuda_graph_bs"] == [1, 2, 4]
|
|
assert sent_body["decode_cuda_graph_memory_gb"] == 1.25
|