From f0808c97939f56d6a8e4fa8b910e2605b0073b49 Mon Sep 17 00:00:00 2001 From: Wei <48665416+hbhflw2000@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:34:59 +0800 Subject: [PATCH] [Bugfix] Fix remote replica membership lifecycle races (#5277) Signed-off-by: hbhflw2000 <417911774@qq.com> Co-authored-by: Gao Han Co-authored-by: Samit <285365963@qq.com> --- tests/engine/test_membership_controller.py | 395 +++++++++++++++++++++ vllm_omni/engine/membership_controller.py | 116 +++++- 2 files changed, 503 insertions(+), 8 deletions(-) diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py index 1b8bf2adb..410bb3030 100644 --- a/tests/engine/test_membership_controller.py +++ b/tests/engine/test_membership_controller.py @@ -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) diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py index a7676de54..2053f6160 100644 --- a/vllm_omni/engine/membership_controller.py +++ b/vllm_omni/engine/membership_controller.py @@ -49,6 +49,13 @@ class MembershipController: self._stage_pools = stage_pools self._remote_replica_factory = remote_replica_factory self._membership_tasks: set[asyncio.Task[None]] = set() + self._attaching_remote_replicas: set[tuple[int, int]] = set() + self._attached_remote_replicas: set[tuple[int, int]] = set() + self._remote_replica_keys_by_input_addr: dict[str, tuple[int, int]] = {} + self._watch_generation = 0 + self._attach_start_generation: dict[tuple[int, int], int] = {} + self._up_generation_by_input_addr: dict[tuple[int, str], int] = {} + self._observed_up_attached_addrs: set[tuple[int, str]] = set() self._shutdown_event = asyncio.Event() self._watcher_task: asyncio.Task[None] | None = None self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None @@ -67,8 +74,34 @@ class MembershipController: async def handle_register(self, stage_id: int, replica_id: int) -> None: """Handle a register_remote_replica message (fire-and-forget).""" + if self._shutdown_event.is_set(): + logger.info( + "[MembershipController] ignoring register during shutdown stage=%d replica=%d", + stage_id, + replica_id, + ) + return + + key = (stage_id, replica_id) + if key in self._attached_remote_replicas: + logger.info( + "[MembershipController] ignoring register for attached replica stage=%d replica=%d", + stage_id, + replica_id, + ) + return + if key in self._attaching_remote_replicas: + logger.info( + "[MembershipController] ignoring register while attach is in flight stage=%d replica=%d", + stage_id, + replica_id, + ) + return + + self._attaching_remote_replicas.add(key) + self._attach_start_generation[key] = self._watch_generation self._spawn_task( - self._do_register(stage_id, replica_id), + self._register_once(stage_id, replica_id), label=f"register-s{stage_id}-r{replica_id}", ) @@ -131,8 +164,24 @@ class MembershipController: while not self._shutdown_event.is_set(): try: snap = self._hub.get_replica_list() + self._watch_generation += 1 current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} - for stage_id, addr in last_up - current: + attached = { + (stage_id, input_addr) + for input_addr, (stage_id, _replica_id) in self._remote_replica_keys_by_input_addr.items() + } + self._observed_up_attached_addrs.update(attached & current) + attaching_stages = {stage_id for stage_id, _replica_id in self._attaching_remote_replicas} + for addr_key in current: + if addr_key[0] in attaching_stages: + self._up_generation_by_input_addr[addr_key] = self._watch_generation + + # Preserve watcher cleanup for untracked/static clients. A + # controller-attached client is eligible only after this + # attachment has been observed UP at least once. + untracked_disappeared = (last_up - current) - attached + confirmed_attached_disappeared = self._observed_up_attached_addrs - current + for stage_id, addr in untracked_disappeared | confirmed_attached_disappeared: self._spawn_task( self.handle_unregister(stage_id, addr), label=f"unregister-s{stage_id}", @@ -148,24 +197,75 @@ class MembershipController: except asyncio.CancelledError: raise - async def _do_register(self, stage_id: int, replica_id: int) -> None: + async def _register_once(self, stage_id: int, replica_id: int) -> None: + key = (stage_id, replica_id) + attach_start_generation = self._attach_start_generation[key] + try: + input_addr = await self._do_register(stage_id, replica_id) + if input_addr is None: + return + self._attached_remote_replicas.add(key) + self._remote_replica_keys_by_input_addr[input_addr] = key + addr_key = (stage_id, input_addr) + if self._up_generation_by_input_addr.get(addr_key, -1) > attach_start_generation: + self._observed_up_attached_addrs.add(addr_key) + finally: + self._attaching_remote_replicas.discard(key) + self._attach_start_generation.pop(key, None) + if not any(attaching_stage_id == stage_id for attaching_stage_id, _ in self._attaching_remote_replicas): + self._up_generation_by_input_addr = { + addr_key: generation + for addr_key, generation in self._up_generation_by_input_addr.items() + if addr_key[0] != stage_id + } + + async def _do_register(self, stage_id: int, replica_id: int) -> str | None: pool = self._pool_for_stage_id(stage_id) if pool is None: logger.warning("[MembershipController] register: stage_id %d out of range", stage_id) - return + return None client = await asyncio.to_thread(self._remote_replica_factory, stage_id, replica_id) - input_addr = StagePool._client_input_addr(client) - if input_addr is None: - raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address") - pool.add_client(input_addr, client, replica_id=replica_id) + if self._shutdown_event.is_set(): + try: + client.shutdown() + except Exception: + logger.exception( + "[MembershipController] failed to shutdown late replica client stage=%d replica=%d", + stage_id, + replica_id, + ) + return None + try: + input_addr = StagePool._client_input_addr(client) + if input_addr is None: + raise RuntimeError( + f"remote replica factory for stage {stage_id} produced a client without input address" + ) + pool.add_client(input_addr, client, replica_id=replica_id) + except Exception: + try: + client.shutdown() + except Exception: + logger.exception( + "[MembershipController] failed to shutdown rejected replica client stage=%d replica=%d", + stage_id, + replica_id, + ) + raise logger.info( "[MembershipController] attached remote replica stage=%d replica=%d addr=%s", stage_id, replica_id, input_addr, ) + return input_addr def _detach_replica(self, stage_id: int, input_addr: str) -> None: + key = self._remote_replica_keys_by_input_addr.pop(input_addr, None) + if key is not None: + self._attached_remote_replicas.discard(key) + self._attaching_remote_replicas.discard(key) + self._observed_up_attached_addrs.discard((stage_id, input_addr)) pool = self._pool_for_stage_id(stage_id) if pool is None: return