vllm-omni/vllm_omni/engine/membership_controller.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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._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._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)."""
        self._spawn_task(
            self._do_register(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
        affected = pool.invalidate_addr(input_addr)
        self._detach_replica(stage_id, input_addr)
        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()
                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:
                    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 _do_register(self, stage_id: int, replica_id: int) -> 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
        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)
        logger.info(
            "[MembershipController] attached remote replica stage=%d replica=%d addr=%s",
            stage_id,
            replica_id,
            input_addr,
        )

    def _detach_replica(self, stage_id: int, input_addr: str) -> None:
        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]],
    ) -> None:
        """Install shared cleanup sinks for watcher-driven unregister events."""
        self._output_queue = output_queue
        self._cleanup_callback = cleanup_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]