7ae15c74c8
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com> Signed-off-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
606 lines
19 KiB
Python
606 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus
|
|
from vllm_omni.engine.membership_controller import MembershipController
|
|
from vllm_omni.engine.stage_pool import StagePool
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class FakeHub:
|
|
def __init__(self, snapshots=None):
|
|
self.snapshots = list(snapshots or [])
|
|
self.closed = False
|
|
|
|
def get_replica_list(self):
|
|
if self.snapshots:
|
|
return self.snapshots.pop(0)
|
|
return SimpleNamespace(replicas=[])
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
|
|
async def _wait_until(predicate, timeout=1):
|
|
async def _poll():
|
|
while not predicate():
|
|
await asyncio.sleep(0)
|
|
|
|
await asyncio.wait_for(_poll(), timeout=timeout)
|
|
|
|
|
|
class FakePool:
|
|
def __init__(self, stage_id: int):
|
|
self.stage_id = stage_id
|
|
self.clients = []
|
|
self.added = []
|
|
self.removed = []
|
|
self.invalidated = []
|
|
self.hub = None
|
|
self.lb = None
|
|
self.replica_ids = {}
|
|
|
|
def attach_hub(self, hub):
|
|
self.hub = hub
|
|
|
|
def attach_load_balancer(self, lb):
|
|
self.lb = lb
|
|
|
|
def add_client(self, input_addr, client, *, replica_id=None):
|
|
self.added.append((input_addr, client, replica_id))
|
|
self.clients.append(client)
|
|
resolved_replica_id = len(self.clients) - 1 if replica_id is None else replica_id
|
|
self.replica_ids[input_addr] = resolved_replica_id
|
|
return resolved_replica_id
|
|
|
|
def invalidate_addr(self, input_addr):
|
|
self.invalidated.append(input_addr)
|
|
return ["req-1", "req-2"]
|
|
|
|
def remove_client(self, input_addr):
|
|
self.removed.append(input_addr)
|
|
return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown"))
|
|
|
|
def get_client_by_addr(self, input_addr):
|
|
return SimpleNamespace() if input_addr in self.replica_ids else None
|
|
|
|
def get_replica_id_by_addr(self, input_addr):
|
|
return self.replica_ids.get(input_addr)
|
|
|
|
|
|
def _snapshot(*replicas):
|
|
return SimpleNamespace(replicas=list(replicas))
|
|
|
|
|
|
def _replica(stage_id: int, input_addr: str, status=ReplicaStatus.UP):
|
|
return SimpleNamespace(stage_id=stage_id, input_addr=input_addr, status=status)
|
|
|
|
|
|
def _controller(monkeypatch, pool, hub, remote_replica_factory=None):
|
|
import vllm_omni.engine.membership_controller as membership_mod
|
|
|
|
monkeypatch.setattr(membership_mod, "OmniCoordClientForHub", lambda _addr: hub)
|
|
if remote_replica_factory is None:
|
|
|
|
def remote_replica_factory(stage_id, replica_id):
|
|
return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"})
|
|
|
|
return MembershipController(
|
|
stage_pools=[pool],
|
|
coordinator_pub_address="tcp://127.0.0.1:12345",
|
|
load_balancer_factory=lambda: object(),
|
|
remote_replica_factory=remote_replica_factory,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_watch_replica_list_unregisters_disappeared_replicas(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
hub = FakeHub(
|
|
snapshots=[
|
|
_snapshot(_replica(0, "tcp://gone")),
|
|
_snapshot(),
|
|
]
|
|
)
|
|
controller = _controller(monkeypatch, pool, hub)
|
|
controller.WATCH_INTERVAL_S = 0
|
|
unregistered = []
|
|
|
|
async def _handle_unregister(stage_id, input_addr, *args, **kwargs):
|
|
unregistered.append((stage_id, input_addr))
|
|
controller._shutdown_event.set()
|
|
|
|
controller.handle_unregister = _handle_unregister # type: ignore[method-assign]
|
|
|
|
await asyncio.wait_for(controller._watch_replica_list(), timeout=1)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert unregistered == [(0, "tcp://gone")]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_closes_hub_then_cancels_watcher(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
hub = FakeHub()
|
|
controller = _controller(monkeypatch, pool, hub)
|
|
watcher = asyncio.create_task(asyncio.sleep(10))
|
|
controller._watcher_task = watcher
|
|
|
|
controller.shutdown()
|
|
await asyncio.sleep(0)
|
|
|
|
assert hub.closed is True
|
|
assert watcher.cancelled()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drain_tasks_waits_for_membership_tasks(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
controller = _controller(monkeypatch, pool, FakeHub())
|
|
completed = []
|
|
|
|
async def _task():
|
|
await asyncio.sleep(0)
|
|
completed.append(True)
|
|
|
|
controller._spawn_task(_task(), label="unit")
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert completed == [True]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_do_register_offloads_remote_factory(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
factory_thread_ids = []
|
|
|
|
def _factory(stage_id, replica_id):
|
|
factory_thread_ids.append(threading.get_ident())
|
|
return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"})
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
|
|
|
await controller._do_register(0, 1)
|
|
|
|
assert factory_thread_ids
|
|
assert factory_thread_ids[0] != threading.get_ident()
|
|
assert pool.added[0] == ("tcp://stage-0-replica-1", pool.clients[0], 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", ["missing-address", "add-client"])
|
|
async def test_do_register_shuts_down_rejected_client(monkeypatch, failure):
|
|
shutdown_calls = 0
|
|
|
|
class Client:
|
|
client_addresses = {} if failure == "missing-address" else {"input_address": "tcp://rejected"}
|
|
|
|
def shutdown(self):
|
|
nonlocal shutdown_calls
|
|
shutdown_calls += 1
|
|
|
|
pool = FakePool(stage_id=0)
|
|
if failure == "add-client":
|
|
monkeypatch.setattr(pool, "add_client", lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("reject")))
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=lambda *_args: Client())
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await controller._do_register(0, 1)
|
|
|
|
assert shutdown_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_then_register_restores_coordinator_replica_slot(monkeypatch):
|
|
class Client:
|
|
def __init__(self, input_addr: str) -> None:
|
|
self.client_addresses = {"input_address": input_addr}
|
|
self.shutdown_calls = 0
|
|
|
|
def shutdown(self) -> None:
|
|
self.shutdown_calls += 1
|
|
|
|
old_clients = [Client("tcp://old-0"), Client("tcp://old-1")]
|
|
replacement = Client("tcp://new-1")
|
|
pool = StagePool(0, old_clients)
|
|
|
|
def factory(stage_id: int, replica_id: int):
|
|
assert (stage_id, replica_id) == (0, 1)
|
|
return replacement
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory)
|
|
removed_replicas = []
|
|
|
|
async def cleanup(_request_ids):
|
|
return None
|
|
|
|
controller.install_unregister_handlers(
|
|
output_queue=asyncio.Queue(),
|
|
cleanup_callback=cleanup,
|
|
replica_removed_callback=lambda stage_id, replica_id: removed_replicas.append((stage_id, replica_id)),
|
|
)
|
|
await controller.handle_unregister(0, "tcp://old-0")
|
|
await controller.handle_unregister(0, "tcp://old-1")
|
|
|
|
await controller.handle_register(0, 1)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert pool.clients == [None, replacement]
|
|
assert pool.get_replica_id_by_addr("tcp://new-1") == 1
|
|
assert pool.get_replica_id_by_addr("tcp://old-1") is None
|
|
assert pool.available_replica_ids() == [1]
|
|
assert removed_replicas == [(0, 0), (0, 1)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_register_deduplicates_inflight_and_attached_replica(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
factory_started = threading.Event()
|
|
release_factory = threading.Event()
|
|
factory_calls = []
|
|
|
|
def _factory(stage_id, replica_id):
|
|
factory_calls.append((stage_id, replica_id))
|
|
factory_started.set()
|
|
assert release_factory.wait(timeout=1)
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
|
|
shutdown=lambda: None,
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
|
|
|
await asyncio.gather(
|
|
controller.handle_register(0, 2),
|
|
controller.handle_register(0, 2),
|
|
)
|
|
assert await asyncio.to_thread(factory_started.wait, 1)
|
|
|
|
await controller.handle_register(0, 2)
|
|
release_factory.set()
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
await controller.handle_register(0, 2)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert factory_calls == [(0, 2)]
|
|
assert len(pool.added) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_register_allows_retry_after_factory_failure(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
factory_calls = 0
|
|
|
|
def _factory(stage_id, replica_id):
|
|
nonlocal factory_calls
|
|
factory_calls += 1
|
|
if factory_calls == 1:
|
|
raise RuntimeError("injected registration failure")
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
|
|
shutdown=lambda: None,
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
|
|
|
await controller.handle_register(0, 3)
|
|
await controller.drain_tasks(timeout=1)
|
|
await controller.handle_register(0, 3)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert factory_calls == 2
|
|
assert len(pool.added) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_allows_same_replica_to_register_again(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
factory_calls = 0
|
|
|
|
def _factory(stage_id, replica_id):
|
|
nonlocal factory_calls
|
|
factory_calls += 1
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
|
|
shutdown=lambda: None,
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
|
|
|
await controller.handle_register(0, 4)
|
|
await controller.drain_tasks(timeout=1)
|
|
await controller.handle_unregister(0, "tcp://stage-0-replica-4")
|
|
await controller.handle_register(0, 4)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert factory_calls == 2
|
|
assert len(pool.added) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_register_ignores_registration_after_shutdown(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
factory_calls = 0
|
|
|
|
def _factory(stage_id, replica_id):
|
|
nonlocal factory_calls
|
|
factory_calls += 1
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
|
|
shutdown=lambda: None,
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
|
controller.shutdown()
|
|
|
|
await controller.handle_register(0, 5)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert factory_calls == 0
|
|
assert pool.added == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_discards_client_from_inflight_registration(monkeypatch):
|
|
pool = FakePool(stage_id=0)
|
|
factory_started = threading.Event()
|
|
release_factory = threading.Event()
|
|
client_shutdown = threading.Event()
|
|
|
|
def _factory(stage_id, replica_id):
|
|
factory_started.set()
|
|
assert release_factory.wait(timeout=1)
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
|
|
shutdown=client_shutdown.set,
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
|
|
|
await controller.handle_register(0, 6)
|
|
assert await asyncio.to_thread(factory_started.wait, 1)
|
|
controller.shutdown()
|
|
release_factory.set()
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert client_shutdown.is_set()
|
|
assert pool.added == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_watcher_reconciles_replica_that_disappears_during_registration(monkeypatch):
|
|
input_addr = "tcp://stage-0-replica-7"
|
|
pool = StagePool(0, [])
|
|
hub = FakeHub(
|
|
snapshots=[
|
|
_snapshot(_replica(0, input_addr)),
|
|
_snapshot(),
|
|
]
|
|
)
|
|
factory_started = threading.Event()
|
|
release_factory = threading.Event()
|
|
client_shutdown = threading.Event()
|
|
|
|
def _factory(stage_id, replica_id):
|
|
assert (stage_id, replica_id) == (0, 7)
|
|
factory_started.set()
|
|
assert release_factory.wait(timeout=1)
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": input_addr},
|
|
shutdown=client_shutdown.set,
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
|
|
controller.WATCH_INTERVAL_S = 0
|
|
original_handle_unregister = controller.handle_unregister
|
|
first_unregister = asyncio.Event()
|
|
second_unregister = asyncio.Event()
|
|
unregister_calls = 0
|
|
|
|
async def _track_unregister(stage_id, addr, *args, **kwargs):
|
|
nonlocal unregister_calls
|
|
await original_handle_unregister(stage_id, addr, *args, **kwargs)
|
|
unregister_calls += 1
|
|
if unregister_calls == 1:
|
|
first_unregister.set()
|
|
elif unregister_calls == 2:
|
|
second_unregister.set()
|
|
|
|
controller.handle_unregister = _track_unregister # type: ignore[method-assign]
|
|
|
|
await controller.handle_register(0, 7)
|
|
assert await asyncio.to_thread(factory_started.wait, 1)
|
|
watcher = controller.start()
|
|
|
|
try:
|
|
await asyncio.wait_for(first_unregister.wait(), timeout=1)
|
|
assert pool.live_num_replicas == 0
|
|
|
|
release_factory.set()
|
|
await asyncio.wait_for(second_unregister.wait(), timeout=1)
|
|
|
|
assert client_shutdown.is_set()
|
|
assert pool.live_num_replicas == 0
|
|
assert controller._attached_remote_replicas == set()
|
|
assert controller._remote_replica_keys_by_input_addr == {}
|
|
finally:
|
|
release_factory.set()
|
|
controller.shutdown()
|
|
try:
|
|
await watcher
|
|
except asyncio.CancelledError:
|
|
pass
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_attached_replica_survives_until_hub_observes_it_up(monkeypatch):
|
|
class EmptyHub:
|
|
def __init__(self):
|
|
self.polls = 0
|
|
|
|
def get_replica_list(self):
|
|
self.polls += 1
|
|
return SimpleNamespace(replicas=[], timestamp=0.0)
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
input_addr = "tcp://stage-0-replica-8"
|
|
pool = StagePool(0, [])
|
|
shutdown_calls = []
|
|
hub = EmptyHub()
|
|
|
|
def _factory(stage_id, replica_id):
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": input_addr},
|
|
shutdown=lambda: shutdown_calls.append((stage_id, replica_id)),
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
|
|
controller.WATCH_INTERVAL_S = 0.001
|
|
|
|
await controller.handle_register(0, 8)
|
|
await controller.drain_tasks(timeout=1)
|
|
assert pool.live_num_replicas == 1
|
|
|
|
watcher = controller.start()
|
|
try:
|
|
await _wait_until(lambda: hub.polls >= 3)
|
|
|
|
assert shutdown_calls == []
|
|
assert pool.live_num_replicas == 1
|
|
finally:
|
|
controller.shutdown()
|
|
try:
|
|
await watcher
|
|
except asyncio.CancelledError:
|
|
pass
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_up_observed_before_attach_does_not_qualify_attachment(monkeypatch):
|
|
import vllm_omni.engine.membership_controller as membership_mod
|
|
|
|
input_addr = "tcp://stage-0-replica-10"
|
|
pool = StagePool(0, [])
|
|
hub = FakeHub(
|
|
snapshots=[
|
|
_snapshot(_replica(0, input_addr)),
|
|
_snapshot(),
|
|
]
|
|
)
|
|
shutdown_calls = []
|
|
sleep_gates = asyncio.Queue()
|
|
|
|
async def _controlled_sleep(_delay):
|
|
gate = asyncio.Event()
|
|
await sleep_gates.put(gate)
|
|
await gate.wait()
|
|
|
|
monkeypatch.setattr(membership_mod.asyncio, "sleep", _controlled_sleep)
|
|
|
|
def _factory(stage_id, replica_id):
|
|
return SimpleNamespace(
|
|
client_addresses={"input_address": input_addr},
|
|
shutdown=lambda: shutdown_calls.append((stage_id, replica_id)),
|
|
)
|
|
|
|
controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
|
|
watcher = controller.start()
|
|
|
|
try:
|
|
first_sleep = await asyncio.wait_for(sleep_gates.get(), timeout=1)
|
|
await controller.handle_register(0, 10)
|
|
await controller.drain_tasks(timeout=1)
|
|
assert pool.live_num_replicas == 1
|
|
|
|
first_sleep.set()
|
|
await asyncio.wait_for(sleep_gates.get(), timeout=1)
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
assert shutdown_calls == []
|
|
assert pool.live_num_replicas == 1
|
|
finally:
|
|
controller.shutdown()
|
|
try:
|
|
await watcher
|
|
except asyncio.CancelledError:
|
|
pass
|
|
await controller.drain_tasks(timeout=1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reattached_address_requires_a_fresh_up_observation(monkeypatch):
|
|
class MutableHub:
|
|
def __init__(self, input_addr):
|
|
self.input_addr = input_addr
|
|
self.up = True
|
|
self.polls = 0
|
|
|
|
def get_replica_list(self):
|
|
self.polls += 1
|
|
replicas = [_replica(0, self.input_addr)] if self.up else []
|
|
return _snapshot(*replicas)
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
class Client:
|
|
def __init__(self, input_addr):
|
|
self.client_addresses = {"input_address": input_addr}
|
|
self.shutdown_calls = 0
|
|
|
|
def shutdown(self):
|
|
self.shutdown_calls += 1
|
|
|
|
input_addr = "tcp://stage-0-replica-9"
|
|
pool = StagePool(0, [])
|
|
hub = MutableHub(input_addr)
|
|
clients = []
|
|
|
|
def _factory(_stage_id, _replica_id):
|
|
client = Client(input_addr)
|
|
clients.append(client)
|
|
return client
|
|
|
|
controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
|
|
controller.WATCH_INTERVAL_S = 0.001
|
|
|
|
await controller.handle_register(0, 9)
|
|
await controller.drain_tasks(timeout=1)
|
|
watcher = controller.start()
|
|
|
|
try:
|
|
await _wait_until(lambda: hub.polls >= 2)
|
|
hub.up = False
|
|
await _wait_until(lambda: clients[0].shutdown_calls == 1)
|
|
assert pool.live_num_replicas == 0
|
|
|
|
await controller.handle_register(0, 9)
|
|
await controller.drain_tasks(timeout=1)
|
|
polls_after_reattach = hub.polls
|
|
await _wait_until(lambda: hub.polls >= polls_after_reattach + 3)
|
|
|
|
assert len(clients) == 2
|
|
assert clients[1].shutdown_calls == 0
|
|
assert pool.live_num_replicas == 1
|
|
finally:
|
|
controller.shutdown()
|
|
try:
|
|
await watcher
|
|
except asyncio.CancelledError:
|
|
pass
|
|
await controller.drain_tasks(timeout=1)
|