[Bugfix] Fix remote replica membership lifecycle races (#5277)
Signed-off-by: hbhflw2000 <417911774@qq.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Samit <285365963@qq.com>
This commit is contained in:
2 files changed
+503
-8
No files matched your search
@@ -30,6 +30,14 @@ class FakeHub:
|
||||
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
|
||||
@@ -159,6 +167,29 @@ async def test_do_register_offloads_remote_factory(monkeypatch):
|
||||
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:
|
||||
@@ -188,3 +219,367 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
|
||||
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]
|
||||
|
||||
|
||||
@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)
|
||||
Reference in new issue
Block a user