vllm-omni/vllm_omni/engine/membership_controller.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""MembershipController: distributed replica lifecycle management.

Extracted from Orchestrator to keep request-flow code free of distributed
concerns. Owns the OmniCoordClientForHub, watches for replica
disappearances, and handles register/unregister by building head-side
clients via an injected factory and mutating StagePool membership.
"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from typing import Any

from vllm.logger import init_logger

from vllm_omni.distributed.omni_coordinator import (
    LoadBalancer,
)
from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus
from vllm_omni.distributed.omni_coordinator.omni_coord_client_for_hub import OmniCoordClientForHub
from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage
from vllm_omni.engine.stage_pool import StagePool

logger = init_logger(__name__)

RemoteReplicaFactory = Callable[[int, int], Any]


class MembershipController:
    """Manages dynamic replica attach/detach for distributed mode.

    Constructed by DistStageRuntime and passed to Orchestrator.
    The Orchestrator delegates register/unregister messages here.
    """

    WATCH_INTERVAL_S: float = 0.5

    def __init__(
        self,
        stage_pools: list[StagePool],
        coordinator_pub_address: str,
        load_balancer_factory: Callable[[], LoadBalancer],
        remote_replica_factory: RemoteReplicaFactory,
    ) -> None:
        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
        self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None
        self._replica_removed_callback: Callable[[int, int], None] | None = None

        self._hub = OmniCoordClientForHub(coordinator_pub_address)
        factory = load_balancer_factory
        for pool in self._stage_pools:
            pool.attach_hub(self._hub)
            pool.attach_load_balancer(factory())

    def start(self) -> asyncio.Task[None]:
        """Start the replica watcher as a background task. Returns the task."""
        self._watcher_task = asyncio.create_task(self._watch_replica_list(), name="membership-watcher")
        return self._watcher_task

    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._register_once(stage_id, replica_id),
            label=f"register-s{stage_id}-r{replica_id}",
        )

    async def handle_unregister(
        self,
        stage_id: int,
        input_addr: str,
        output_queue: asyncio.Queue[EngineQueueMessage] | None = None,
        cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None,
    ) -> None:
        """Handle an unregister_remote_replica message."""
        pool = self._pool_for_stage_id(stage_id)
        if pool is None:
            return
        effective_output_queue = output_queue if output_queue is not None else self._output_queue
        effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback
        replica_id = pool.get_replica_id_by_addr(input_addr)
        client = pool.get_client_by_addr(input_addr)
        affected = pool.invalidate_addr(input_addr)
        self._detach_replica(stage_id, input_addr)
        if client is not None and replica_id is not None and self._replica_removed_callback is not None:
            self._replica_removed_callback(stage_id, replica_id)
        if affected and effective_cleanup_callback is not None:
            await effective_cleanup_callback(affected)
        if affected and effective_output_queue is not None:
            for req_id in affected:
                await effective_output_queue.put(
                    ErrorMessage(error="stage replica disappeared", request_id=req_id, stage_id=stage_id)
                )

    def shutdown(self) -> None:
        """Signal stop and close the hub."""
        self._shutdown_event.set()
        if self._hub is not None:
            self._hub.close()
            self._hub = None
        if self._watcher_task is not None and not self._watcher_task.done():
            self._watcher_task.cancel()

    async def drain_tasks(self, timeout: float = 10.0) -> None:
        """Wait for in-flight membership tasks to complete."""
        if self._membership_tasks:
            await asyncio.wait(self._membership_tasks, timeout=timeout)

    # ---- Internal ----

    def _spawn_task(self, coro: Awaitable[None], *, label: str) -> None:
        task = asyncio.create_task(coro, name=f"membership-{label}")
        self._membership_tasks.add(task)

        def _on_done(t: asyncio.Task[None]) -> None:
            self._membership_tasks.discard(t)
            if t.cancelled():
                return
            exc = t.exception()
            if exc is not None:
                logger.error("[MembershipController] %s task crashed", label, exc_info=exc)

        task.add_done_callback(_on_done)

    async def _watch_replica_list(self) -> None:
        """Convert hub replica disappearances into unregister actions."""
        last_up: set[tuple[int, str]] = set()
        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}
                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}",
                    )
                last_up = current
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("[MembershipController] _watch_replica_list iteration failed")

            try:
                await asyncio.sleep(self.WATCH_INTERVAL_S)
            except asyncio.CancelledError:
                raise

    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 None
        client = await asyncio.to_thread(self._remote_replica_factory, stage_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
        client = pool.remove_client(input_addr)
        if client is None:
            return
        try:
            client.shutdown()
        except Exception:
            logger.exception("[MembershipController] failed to shutdown client stage=%d addr=%s", stage_id, input_addr)
        logger.info("[MembershipController] detached replica stage=%d addr=%s", stage_id, input_addr)

    def install_unregister_handlers(
        self,
        *,
        output_queue: asyncio.Queue[EngineQueueMessage],
        cleanup_callback: Callable[[list[str]], Awaitable[None]],
        replica_removed_callback: Callable[[int, int], None] | None = None,
    ) -> None:
        """Install shared cleanup sinks for watcher-driven unregister events."""
        self._output_queue = output_queue
        self._cleanup_callback = cleanup_callback
        self._replica_removed_callback = replica_removed_callback

    def _pool_for_stage_id(self, stage_id: int) -> StagePool | None:
        if not (0 <= stage_id < len(self._stage_pools)):
            return None
        return self._stage_pools[stage_id]