vllm-omni/vllm_omni/entrypoints/async_omni.py

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

"""
AsyncOmni - Refactored async orchestrator using AsyncOmniEngine.

This is the new implementation that uses AsyncOmniEngine (which manages
StageEngineCoreClient instances) instead of OmniStage with worker processes.
"""

from __future__ import annotations

import asyncio
import time
import uuid
from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, cast

import anyio
from vllm import TokensPrompt
from vllm.engine.protocol import EngineClient, StreamingInput
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
from vllm.outputs import PoolingRequestOutput
from vllm.pooling_params import PoolingParams
from vllm.renderers.inputs.preprocess import extract_prompt_components
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.tasks import SupportedTask

from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.messages import ErrorMessage
from vllm_omni.entrypoints.async_omni_base import ABORT_TIMEOUT_S, AsyncOmniBase
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.errors import client_error_metadata
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
from vllm_omni.outputs import OmniRequestOutput

if TYPE_CHECKING:
    from vllm.tokenizers import TokenizerLike
    from vllm.v1.engine import PauseMode
    from vllm.v1.engine.input_processor import InputProcessor

    from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType

logger = init_logger(__name__)

CACHE_RESET_TIMEOUT_S = 60.0


class AsyncOmni(AsyncOmniBase, EngineClient):
    """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine.

    This is the refactored version that uses AsyncOmniEngine instead of
    OmniStage workers. It provides the same interface as AsyncOmni but with
    a cleaner architecture.

    Args:
        model: Model name or path to load.
        **kwargs: Additional keyword arguments.
            - deploy_config: Optional path to a deploy YAML. If None,
              configurations are resolved from the model pipeline factory.
            - log_stats: Whether to enable statistics logging.
            - stage_init_timeout: Timeout for per-stage initialization.
            - init_timeout: Total timeout for orchestrator startup.
            - async_chunk: Whether to enable async chunk mode.
            - output_modalities: Requested output modalities.
            - Additional keyword arguments passed to stage engines.

    Example:
        >>> async_omni = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B")
        >>> async for output in async_omni.generate(
        ...     prompt="Hello",
        ...     request_id="req-1",
        ...     sampling_params_list=[SamplingParams(), SamplingParams()]
        ... ):
        ...     print(output)
    """

    engine: AsyncOmniEngine

    def _create_engine(self, **engine_kwargs: Any) -> AsyncOmniEngine:
        return AsyncOmniEngine(**engine_kwargs)

    def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None:
        tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None)
        AsyncOmniBase.__init__(self, *args, model=model, **kwargs)
        self.tts_max_instructions_length = tts_max_instructions_length
        self._pause_cond: asyncio.Condition = asyncio.Condition()
        self._paused: bool = False
        # In-flight EngineCore submits (non-streaming add_request, or each
        # streaming ADD/update/final marker). sleep() waits for this to hit
        # zero so a pipelined request cannot race into EngineCore during
        # drain/offload. Streaming generate does not hold a slot while
        # waiting for the next client chunk.
        self._admitting: int = 0
        # True after pause_generation() or AR EngineCore sleep; wake_up must
        # not reopen generate() until resume_generation(). Diffusion-only
        # sleep uses _paused as a temporary admission gate and clears it
        # on wake so sleep → wake → generate keeps working.
        self._hold_admission_until_resume: bool = False
        # Stages whose scheduler pause_generation closed (AR in any mode,
        # diffusion with mode="keep"); admission stays closed until
        # resume_generation has reopened all of them.
        self._paused_stage_ids: set[int] = set()
        self._sleeping_tags: set[str] = set()
        self._stage_sleeping_tags: dict[int, set[str]] = {}
        self._level2_sleeping: bool = False

    # ==================== Generate Method ====================

    async def generate(
        self,
        prompt: OmniPromptType | AsyncGenerator[StreamingInput, None] | list[OmniPromptType],
        sampling_params: Any = None,
        request_id: str = "",
        *,
        prompt_text: str | None = None,
        lora_request: Any = None,
        tokenization_kwargs: dict[str, Any] | None = None,
        sampling_params_list: Sequence[OmniSamplingParams] | None = None,
        output_modalities: list[str] | None = None,
        trace_headers: Mapping[str, str] | None = None,
        priority: int = 0,
        data_parallel_rank: int | None = None,
        session_id: str | None = None,
        reasoning_ended: bool | None = None,
        reasoning_parser_kwargs: dict[str, Any] | None = None,
        arrival_time: float | None = None,
    ) -> AsyncGenerator[OmniRequestOutput, None]:
        """Generate outputs for the given prompt(s) asynchronously.

        Coordinates multi-stage pipeline execution. Processes the prompt
        through all stages in the pipeline and yields outputs as they become
        available.

        ``session_id`` is accepted for EngineClient protocol compatibility
        and is not duplex-session plumbing.

        **Diffusion batching:**
        Diffusion stages accept only a single prompt per request.  Passing a
        ``list`` of prompts to a diffusion stage will raise ``ValueError``.
        To batch multiple diffusion prompts, submit each as an independent
        request; the scheduler will automatically co-batch compatible requests.

        Args:
            prompt: A single prompt **or** a list of prompts.  For diffusion
                stages, only a single prompt is accepted; a list will be
                rejected with an error.
            request_id: Unique identifier for this request. If one is not provided,
                a random one will be generated.
            sampling_params_list: List of SamplingParams, one per stage.
                Must have the same length as the number of stages.
                If *None*, uses default sampling params for each stage.
            output_modalities: Optional list of output modalities.

        Yields:
            OmniRequestOutput objects as they are produced by each stage.

        Raises:
            ValueError: If sampling_params_list has incorrect length, or
                if a list prompt is submitted to a diffusion stage.
        """
        # Append a random UUID suffix to the request_id to ensure it is unique
        # and non-empty, similar to vLLM's input processor. The suffix is used
        # only for internal tracking throughout the request's life.
        external_request_id = request_id
        request_id = self._get_unique_request_id(external_request_id)

        # Wait until generation is resumed if the engine is paused. Non-streaming
        # generate holds an admission slot until add_request completes so sleep()
        # cannot race a just-unblocked generate into EngineCore during offload.
        # Streaming generate does **not** hold the slot while waiting for the
        # first client chunk; the input pump acquires it immediately before each
        # EngineCore ADD/update.
        streaming_input = isinstance(prompt, AsyncGenerator)
        async with self._pause_cond:
            await self._pause_cond.wait_for(lambda: not self._paused)
            if not streaming_input:
                self._admitting = getattr(self, "_admitting", 0) + 1
        admitting = not streaming_input

        logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")

        input_stream_task: asyncio.Task | None = None
        try:
            _sleeping_tags = getattr(self, "_sleeping_tags", None)
            if _sleeping_tags:
                raise RuntimeError(
                    f"Generation rejected: Engine is partially or fully asleep. "
                    f"Currently sleeping tags: {list(_sleeping_tags)}. "
                    f"Please perform a full wake_up before generating."
                )

            # Reject diffusion list-prompt early with a clear API error.
            if isinstance(prompt, list) and any(
                stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs
            ):
                raise ValueError(
                    "Diffusion stages accept only a single prompt per request. "
                    "Submit multiple independent requests to use scheduler batching."
                )

            # Start final output dispatcher on the first call to generate()
            self._final_output_handler()

            # Forward bare sampling_params (e.g. from /v1/completions) as the stage-0 entry.
            if sampling_params_list is None and sampling_params is not None:
                if self.num_stages == 1:
                    sampling_params_list = [sampling_params]
                else:
                    default = list(self.default_sampling_params_list)
                    default[0] = sampling_params
                    sampling_params_list = default

            # Expand sampling params for PD disaggregation (user may provide N-1 params)
            if (
                sampling_params_list is not None
                and isinstance(sampling_params_list, Sequence)
                and not isinstance(sampling_params_list, (str, bytes))
            ):
                sampling_params_list = self._maybe_expand_sampling_params(list(sampling_params_list))

            # Set the output kind to delta output if sampling params were omitted,
            # since AsyncOmni is typically used for streaming.
            sampling_params_list = self.resolve_sampling_params_list(
                sampling_params_list,
                allow_delta_coercion=True,
            )

            # Track per-request metrics
            wall_start_ts = float(arrival_time) if arrival_time is not None else time.time()
            req_start_ts: dict[str, float] = {}

            # Determine the final stage for E2E stats
            final_stage_id_for_e2e = self._compute_final_stage_id(output_modalities)
            final_output_stage_ids = self._compute_final_output_stage_ids(output_modalities) or [final_stage_id_for_e2e]

            metrics = OrchestratorMetrics(
                self.num_stages,
                self.log_stats,
                wall_start_ts,
                final_stage_id_for_e2e,
                transfer_emitter=getattr(self, "transfer_metrics", None),
                replica_resolver=self._resolve_transfer_replica,
            )

            req_state = ClientRequestState(
                request_id=request_id,
                external_request_id=external_request_id,
                final_stage_id=final_stage_id_for_e2e,
            )
            req_state.metrics = metrics
            req_state.request_arrival_ts = wall_start_ts
            self.request_states[request_id] = req_state

            # PD disaggregation: modify prefill-stage sampling params per request
            req_sp_list = list(sampling_params_list)
            pd_pair = self._get_pd_separation_pair()
            if pd_pair is not None:
                p_id = pd_pair[0]
                req_sp_list[p_id] = self._prepare_prefill_sampling_params(request_id, req_sp_list[p_id])

            # Add request(s) to stage 0. For streaming inputs, submit
            # chunks incrementally through streaming_update. The helper
            # returns as soon as the pump task is created; each ADD/update
            # takes its own admission slot inside the pump.
            if streaming_input:
                first_chunk_submitted = asyncio.get_running_loop().create_future()
                input_stream_task = await self._add_streaming_input_request(
                    request_id=request_id,
                    input_stream=cast(AsyncGenerator, prompt),
                    sampling_params_list=req_sp_list,
                    final_stage_id=final_stage_id_for_e2e,
                    final_output_stage_ids=final_output_stage_ids,
                    arrival_time=wall_start_ts,
                    lora_request=lora_request,
                    first_chunk_submitted=first_chunk_submitted,
                )
                await first_chunk_submitted
            else:
                await self.engine.add_request_async(
                    request_id=request_id,
                    prompt=prompt,
                    sampling_params_list=req_sp_list,
                    final_stage_id=final_stage_id_for_e2e,
                    final_output_stage_ids=final_output_stage_ids,
                    arrival_time=wall_start_ts,
                    lora_request=lora_request,
                )
            submit_ts = time.time()
            stage_first_ts = cast(list[float | None], req_state.metrics.stage_first_ts)
            stage_first_ts[0] = submit_ts
            req_start_ts[request_id] = submit_ts
            if admitting:
                await self._release_generate_admission()
                admitting = False
            # Refresh gauges on arrival.
            self._publish_request_gauges(len(self.request_states))

            # Process results based on mode
            # Both sequential and async_chunk modes read the same message stream
            # from Orchestrator; stage-transfer behavior differs inside
            # Orchestrator._route_output().
            async for output in self._process_orchestrator_results(
                request_id,
                metrics,
                final_stage_id_for_e2e,
                req_start_ts,
                wall_start_ts,
            ):
                yield output

            logger.debug(f"[AsyncOmni] Request {request_id} completed")

        except (asyncio.CancelledError, GeneratorExit):
            self._record_request_failure_once(request_id, reason="client_disconnect")
            # ASGI cancellation also affects subsequent awaits. Shield the
            # bounded engine abort before removing the local request state.
            with anyio.CancelScope(shield=True):
                await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S)
            logger.info(f"[AsyncOmni] Request {request_id} aborted.")
            raise
        except Exception as e:
            self._record_request_failure_once(request_id, reason="stage_error")
            await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S)
            logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}")
            raise
        finally:
            if input_stream_task is not None and not input_stream_task.done():
                input_stream_task.cancel()
            if admitting:
                await self._release_generate_admission()
            self._log_summary_and_cleanup(request_id)

    async def _release_generate_admission(self) -> None:
        """Drop one in-flight generate admission slot held across add_request."""
        async with self._pause_cond:
            self._admitting = max(getattr(self, "_admitting", 1) - 1, 0)
            self._pause_cond.notify_all()

    async def _submit_with_admission(self, awaitable):
        """Wait for resume, hold one admission slot for a single EngineCore submit."""
        async with self._pause_cond:
            await self._pause_cond.wait_for(lambda: not self._paused)
            self._admitting = getattr(self, "_admitting", 0) + 1
        try:
            return await awaitable
        finally:
            await self._release_generate_admission()

    async def _add_streaming_input_request(
        self,
        *,
        request_id: str,
        input_stream: AsyncGenerator[StreamingInput, None],
        sampling_params_list: Sequence[OmniSamplingParams],
        final_stage_id: int,
        final_output_stage_ids: Sequence[int],
        arrival_time: float,
        lora_request: Any = None,
        first_chunk_submitted: asyncio.Future[None] | None = None,
    ) -> asyncio.Task:
        """Submit a streaming input generator as incremental stage-0 updates."""
        if not sampling_params_list:
            raise ValueError("sampling_params_list cannot be empty for streaming input")
        # only check thinker's sampling params now
        stage0_params = sampling_params_list[0]
        self._validate_streaming_input_sampling_params(stage0_params)
        stage0_params = cast(SamplingParams, stage0_params)
        req_state = self.request_states[request_id]
        has_submitted_first_chunk = False
        input_error_reported = False

        # NOTE: InputProcessor in vLLM should generally do this too, but for
        # now we do it defensively. TODO (Alex) ensure clones/copying are optimized
        if not stage0_params.skip_clone:
            stage0_params = stage0_params.clone()
            stage0_params.skip_clone = True

        def _mark_first_chunk_submitted() -> None:
            if first_chunk_submitted is not None and not first_chunk_submitted.done():
                first_chunk_submitted.set_result(None)

        async def _report_input_error(error: Exception) -> None:
            nonlocal input_error_reported
            if input_error_reported:
                return
            input_error_reported = True
            status_code, error_type = client_error_metadata(error)
            await req_state.queue.put(
                ErrorMessage(
                    request_id=request_id,
                    error=str(error),
                    status_code=status_code,
                    error_type=error_type,
                )
            )

        async def handle_inputs() -> None:
            nonlocal has_submitted_first_chunk
            cancelled = False
            try:
                async for chunk in input_stream:
                    chunk_params = getattr(chunk, "sampling_params", None) or stage0_params
                    self._validate_streaming_input_sampling_params(chunk_params)
                    chunk_sampling_params_list = list(sampling_params_list)
                    chunk_sampling_params_list[0] = chunk_params
                    chunk_prompt = chunk.prompt
                    prompt_text, _, _ = extract_prompt_components(self.model_config, chunk_prompt)

                    if not has_submitted_first_chunk:
                        await self._submit_with_admission(
                            self.engine.add_request_async(
                                request_id=request_id,
                                prompt=chunk_prompt,
                                prompt_text=prompt_text,
                                sampling_params_list=chunk_sampling_params_list,
                                final_stage_id=final_stage_id,
                                final_output_stage_ids=final_output_stage_ids,
                                arrival_time=arrival_time,
                                lora_request=lora_request,
                                resumable=True,
                            )
                        )
                        has_submitted_first_chunk = True
                        _mark_first_chunk_submitted()
                    else:
                        await self._submit_with_admission(
                            self.engine.add_streaming_update_async(
                                request_id=request_id,
                                prompt=chunk_prompt,
                                prompt_text=prompt_text,
                                sampling_params_list=chunk_sampling_params_list,
                                final_stage_id=final_stage_id,
                                final_output_stage_ids=final_output_stage_ids,
                                arrival_time=arrival_time,
                                lora_request=lora_request,
                                resumable=True,
                            )
                        )
            except (asyncio.CancelledError, GeneratorExit):
                cancelled = True
            except Exception as error:
                await _report_input_error(error)
            finally:
                try:
                    if not cancelled:
                        # Send empty final request to indicate that inputs have
                        # finished. Don't send if canceled (session was aborted).
                        final_sampling_params_list = list(sampling_params_list)
                        final_sampling_params_list[0] = stage0_params
                        final_prompt = TokensPrompt(prompt_token_ids=[0])

                        if has_submitted_first_chunk:
                            await self._submit_with_admission(
                                self.engine.add_streaming_update_async(
                                    request_id=request_id,
                                    prompt=final_prompt,
                                    prompt_text=None,
                                    sampling_params_list=final_sampling_params_list,
                                    final_stage_id=final_stage_id,
                                    final_output_stage_ids=final_output_stage_ids,
                                    arrival_time=arrival_time,
                                    lora_request=lora_request,
                                    resumable=False,
                                )
                            )
                        else:
                            await self._submit_with_admission(
                                self.engine.add_request_async(
                                    request_id=request_id,
                                    prompt=final_prompt,
                                    prompt_text=None,
                                    sampling_params_list=final_sampling_params_list,
                                    final_stage_id=final_stage_id,
                                    final_output_stage_ids=final_output_stage_ids,
                                    arrival_time=arrival_time,
                                    lora_request=lora_request,
                                    resumable=False,
                                )
                            )
                            has_submitted_first_chunk = True
                except Exception as error:
                    await _report_input_error(error)
                finally:
                    # Unblock generate() even on cancel / empty stream / submit
                    # failure so it can observe the queued terminal result or error.
                    _mark_first_chunk_submitted()

        input_stream_task = asyncio.create_task(handle_inputs())
        req_state.input_stream_task = input_stream_task
        return input_stream_task

    @staticmethod
    def _validate_streaming_input_sampling_params(params: OmniSamplingParams) -> None:
        if (
            not isinstance(params, SamplingParams)
            or params.n > 1
            or params.output_kind == RequestOutputKind.FINAL_ONLY
            or params.stop
        ):
            raise ValueError(
                "Input streaming is currently supported only for SamplingParams "
                "with n == 1, output_kind != FINAL_ONLY, and without stop strings."
            )

    async def encode(
        self,
        prompt: Any,
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: LoRARequest | None = None,
        trace_headers: dict[str, str] | None = None,
        priority: int = 0,
        tokenization_kwargs: dict[str, Any] | None = None,
        reasoning_ended: bool | None = None,
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
        """EngineClient.encode() stub.

        Omni pipeline currently exposes only generate() API at orchestrator level.
        """
        raise NotImplementedError("AsyncOmni.encode is not implemented.")

    # ==================== Control Methods ====================

    async def collective_rpc(
        self,
        method: str,
        timeout: float | None = None,
        args: tuple[Any, ...] = (),
        kwargs: dict[str, Any] | None = None,
        stage_ids: list[int] | None = None,
    ) -> list[Any]:
        """Execute a best-effort control RPC on selected stages.

        Unsupported stages currently return a TODO-style result dict instead of
        failing the entire call. This keeps AsyncOmni usable while the orchestrator
        control plane is still being filled out.
        """
        results = await self.engine.collective_rpc_async(
            method=method,
            timeout=timeout,
            args=args,
            kwargs=kwargs,
            stage_ids=stage_ids,
        )

        unsupported_results = [
            index for index, result in enumerate(results) if isinstance(result, dict) and result.get("todo")
        ]
        if unsupported_results:
            logger.warning(
                "[AsyncOmni] collective_rpc(%s) has TODO support in replica result(s) %s (requested stages: %s)",
                method,
                unsupported_results,
                stage_ids,
            )

        return results

    async def _engine_core_rpc(
        self,
        method: str,
        *,
        stage_ids: list[int],
        args: tuple[Any, ...] = (),
        kwargs: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> list[Any]:
        """Call EngineCore helpers and reject failures from any replica."""
        timeout_kwargs = {"timeout": timeout} if timeout is not None else {}
        results = await self.collective_rpc(
            method=method,
            args=args,
            kwargs=kwargs,
            stage_ids=stage_ids,
            **timeout_kwargs,
        )

        def check_result(result: Any) -> None:
            if isinstance(result, (list, tuple)):
                for item in result:
                    check_result(item)
            elif isinstance(result, dict):
                if result.get("error") or result.get("todo") or result.get("supported") is False:
                    raise RuntimeError(f"{method} failed: {result}")

        for result in results:
            check_result(result)
        return results

    @staticmethod
    def _coerce_stage_bool(result: Any) -> bool:
        """Reduce a stage RPC result to a boolean.

        Some stage RPCs may return worker-level lists like ``[True]``;
        diffusion wrappers usually return a plain bool.
        """
        if isinstance(result, list):
            return all(bool(item) for item in result)
        return bool(result)

    async def abort(self, request_id: str | Iterable[str], *, timeout: float | None = None) -> None:
        """Abort request(s) via the Orchestrator."""
        request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
        # Map the external user request IDs to internal IDs used by the Orchestrator.
        # NOTE: If the user request_id matches multiple requests, all of them will be
        # aborted. This is also what happens in this case in vLLM's output processor.
        internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id in request_ids]
        await self._abort(internal_ids, timeout=timeout)

    async def submit_interaction_async(
        self,
        request_id: str,
        *,
        interaction: OmniInteractionPrompt,
    ) -> None:
        """Apply a midway interaction to an active streaming diffusion request.

        ``request_id`` is the external id created by the server-side session,
        matching the value passed to :meth:`generate`.
        """
        event = interaction.get("event")
        prompt = event.get("prompt") if isinstance(event, dict) else None
        if isinstance(event, dict) and "prompt" in event and (not isinstance(prompt, str) or not prompt):
            raise ValueError("prompt must be non-empty")
        transition_chunks = interaction.get("transition_chunks")
        if transition_chunks is not None and transition_chunks < 0:
            raise ValueError("transition_chunks must be >= 0")

        if self.num_stages != 1:
            raise ValueError("interaction requires single-stage diffusion")
        stage_meta = self.engine.get_stage_metadata(0)
        if stage_meta.stage_type != "diffusion":
            raise ValueError("interaction requires a diffusion stage")

        internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id == request_id]
        if not internal_ids:
            raise ValueError(f"No active request for interaction: {request_id!r}")
        if len(internal_ids) > 1:
            raise ValueError(
                f"interaction requires exactly one active request for {request_id!r}, found {len(internal_ids)}"
            )

        await self.engine.submit_interaction_async(
            internal_ids[0],
            interaction=interaction,
        )
        if self.log_stats:
            logger.info("[AsyncOmni] Queued interaction for request %s", request_id)

    def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]:
        """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC)."""
        n_stages = len(self.engine.stage_configs)
        if stage_ids is None:
            stage_ids = list(range(n_stages))
        else:
            invalid = [sid for sid in stage_ids if not isinstance(sid, int) or sid < 0 or sid >= n_stages]
            if invalid:
                raise ValueError(
                    f"Invalid stage_ids {invalid}; valid range is 0..{n_stages - 1}"
                    if n_stages
                    else f"Invalid stage_ids {invalid}; this engine has no stages"
                )
        ar_stage_ids: list[int] = []
        diffusion_stage_ids: list[int] = []
        for sid in stage_ids:
            stage_config = self.engine.stage_configs[sid]
            if stage_config.stage_type == "diffusion":
                diffusion_stage_ids.append(sid)
            else:
                ar_stage_ids.append(sid)
        return ar_stage_ids, diffusion_stage_ids

    def _sleeping_tags_for_stages(self, stage_ids: list[int]) -> set[str]:
        """Union of sleeping tags recorded for the given stages."""
        per_stage = getattr(self, "_stage_sleeping_tags", None) or {}
        tags: set[str] = set()
        for sid in stage_ids:
            tags.update(per_stage.get(sid, set()))
        return tags

    def _refresh_union_sleeping_tags(self) -> None:
        """Keep ``_sleeping_tags`` as the engine-wide union of per-stage tags."""
        per_stage = getattr(self, "_stage_sleeping_tags", None) or {}
        union: set[str] = set()
        for stage_tags in per_stage.values():
            union.update(stage_tags)
        self._sleeping_tags = union

    def _record_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None:
        per_stage = getattr(self, "_stage_sleeping_tags", None)
        if per_stage is None:
            per_stage = {}
            self._stage_sleeping_tags = per_stage
        tag_set = set(tags)
        for sid in stage_ids:
            per_stage.setdefault(sid, set()).update(tag_set)
        self._refresh_union_sleeping_tags()

    def _clear_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None:
        per_stage = getattr(self, "_stage_sleeping_tags", None)
        if per_stage is None:
            self._sleeping_tags = set()
            return
        tag_set = set(tags)
        for sid in stage_ids:
            remaining = per_stage.get(sid)
            if remaining is None:
                continue
            remaining.difference_update(tag_set)
            if not remaining:
                per_stage.pop(sid, None)
        self._refresh_union_sleeping_tags()

    async def pause_generation(
        self,
        *,
        mode: PauseMode = "abort",
        wait_for_inflight_requests: bool = False,
        clear_cache: bool = True,
        stage_ids: list[int] | None = None,
    ) -> None:
        """Pause generation, mirroring vLLM AsyncLLM.pause_generation.

        1. Stop frontend admission (``_paused``).
        2. For AR/LLM stages, call EngineCore.pause_scheduler via the
           Orchestrator loop (abort/wait/keep + optional cache clear).
        3. For diffusion stages, ``mode="keep"`` pauses the DiffusionEngine
           scheduler and returns once the batch that was running has finished
           on every worker; that batch is delivered before any control RPC
           issued after this call runs, so the documented pause -> sleep order
           is safe. Queued requests stay queued until
           :meth:`resume_generation`. Other modes pause frontend admission
           only.

        Note: ``sleep()`` already pauses the AR scheduler internally (same as
        vLLM EngineCore.sleep). Call this API when you need pause *without*
        freeing GPU memory (e.g. weight sync).
        """
        if wait_for_inflight_requests:
            mode = "wait"

        async with self._pause_cond:
            # Keep running EngineCore pause + cache clear even when frontend
            # admission is already paused (sleep or a prior pause_generation).
            self._paused = True
            self._hold_admission_until_resume = True

        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
        if mode != "keep":
            diffusion_stage_ids = []
        # Recorded before the RPCs so a failed or cancelled pause can still
        # be undone with an explicit resume_generation.
        self._paused_stage_ids.update(ar_stage_ids, diffusion_stage_ids)
        if ar_stage_ids:
            logger.info(
                "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)",
                self._name,
                ar_stage_ids,
                mode,
            )
            # Same API name as vLLM AsyncMPClient.pause_scheduler_async; routed
            # through collective_rpc so it runs on the orchestrator event loop.
            await self._engine_core_rpc(
                "pause_scheduler",
                stage_ids=ar_stage_ids,
                kwargs={"mode": mode, "clear_cache": clear_cache},
            )
        if diffusion_stage_ids:
            logger.info(
                "[%s] Pausing diffusion stage(s) %s via DiffusionEngine pause_scheduler(mode=keep)",
                self._name,
                diffusion_stage_ids,
            )
            await self._engine_core_rpc(
                "pause_scheduler",
                stage_ids=diffusion_stage_ids,
                kwargs={"mode": "keep"},
            )

        # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler
        # already clears AR-side caches when clear_cache=True.
        if clear_cache and 0 in ar_stage_ids:
            await self._clear_frontend_mm_cache()

    async def resume_generation(self, stage_ids: list[int] | None = None) -> None:
        """Resume generation after :meth:`pause_generation`."""
        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
        if ar_stage_ids:
            logger.info("[%s] Resuming AR stage(s) %s via EngineCore", self._name, ar_stage_ids)
            await self._engine_core_rpc("resume_scheduler", stage_ids=ar_stage_ids)
            self._paused_stage_ids.difference_update(ar_stage_ids)
        diffusion_stage_ids = [sid for sid in diffusion_stage_ids if sid in self._paused_stage_ids]
        if diffusion_stage_ids:
            logger.info("[%s] Resuming diffusion stage(s) %s via DiffusionEngine", self._name, diffusion_stage_ids)
            await self._engine_core_rpc("resume_scheduler", stage_ids=diffusion_stage_ids)
            self._paused_stage_ids.difference_update(diffusion_stage_ids)

        if self._paused_stage_ids:
            # Reopening admission now would let new requests queue on a stage
            # whose scheduler is still closed.
            logger.info(
                "[%s] Admission stays paused: stage(s) %s are still paused",
                self._name,
                sorted(self._paused_stage_ids),
            )
            return

        async with self._pause_cond:
            self._paused = False
            self._hold_admission_until_resume = False
            self._pause_cond.notify_all()

    async def is_paused(self) -> bool:
        """Check if frontend admission is paused."""
        async with self._pause_cond:
            return self._paused

    # EngineClient exposes async operations; OmniBase implements the sync API.
    async def start_profile(  # type: ignore[override]
        self,
        profile_prefix: str | None = None,
        stages: list[int] | None = None,
    ) -> list[Any]:
        """Start profiling specified stages.

        Uses vLLM-compatible profile(is_start=True, profile_prefix) interface.

        Args:
            profile_prefix: Optional prefix for the trace file names.
            stages: List of stage IDs to profile. If None, profiles all stages.
        """
        return await self.collective_rpc(method="profile", args=(True, profile_prefix), stage_ids=stages)

    async def stop_profile(self, stages: list[int] | None = None) -> list[Any]:  # type: ignore[override]
        """Stop profiling specified stages.

        Uses vLLM-compatible profile(is_start=False) interface.

        Args:
            stages: List of stage IDs to profile. If None, stops all stages.
        """
        return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages)

    async def _clear_frontend_mm_cache(self) -> None:
        """Clear P0 through the renderer's serialized multimodal executor."""
        renderer = self.renderer
        if renderer is not None:
            await renderer.clear_mm_cache_async()

    async def reset_mm_cache(
        self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S
    ) -> None:
        """Reset sender and receiver MM caches on selected AR stages.

        By default all AR stages are selected; diffusion stages are skipped.
        Stage 0's sender is the frontend renderer. Downstream senders are
        cleared by the orchestrator before resetting their engine cores.
        Call while generation is paused to avoid racing new inputs.
        """
        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
        if 0 in ar_stage_ids:
            await asyncio.wait_for(self._clear_frontend_mm_cache(), timeout=timeout)
        if ar_stage_ids:
            await self._engine_core_rpc("reset_mm_cache", stage_ids=ar_stage_ids, timeout=timeout)

    async def reset_encoder_cache(
        self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S
    ) -> None:
        """Reset encoder caches on selected AR stages (all by default).

        Diffusion stages are skipped. RPC failures are raised to the caller.
        """
        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
        if ar_stage_ids:
            await self._engine_core_rpc("reset_encoder_cache", stage_ids=ar_stage_ids, timeout=timeout)

    async def reset_prefix_cache(
        self,
        reset_running_requests: bool = False,
        reset_connector: bool = False,
        *,
        stage_ids: list[int] | None = None,
        timeout: float = CACHE_RESET_TIMEOUT_S,
    ) -> bool:
        """Reset prefix caches on selected AR stages (all by default).

        Diffusion stages are skipped. Return False if a stage cannot reset
        its cache; unsupported operations and RPC failures raise instead of
        silently retaining KV computed under previous model weights.
        """
        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
        if not ar_stage_ids:
            return True
        results = await self._engine_core_rpc(
            "reset_prefix_cache",
            stage_ids=ar_stage_ids,
            args=(reset_running_requests, reset_connector),
            timeout=timeout,
        )
        return all(self._coerce_stage_bool(result) for result in results)

    async def sleep(
        self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort"
    ) -> list[OmniACK]:
        """Put stages to sleep.

        AR/LLM stages use EngineCore.sleep (pause scheduler, wait idle, then
        offload/discard memory) — matching vLLM AsyncLLM.sleep.

        Diffusion stages keep the worker-level handle_sleep_task RPC, which
        does not stop the DiffusionEngine scheduler; quiesce a busy diffusion
        stage first with ``pause_generation(mode="keep")`` (or abort it).

        Frontend admission is blocked at the start of this call (``_paused``)
        so pipelined :meth:`generate` cannot race into stages while sleep is
        in flight. This does **not** invoke EngineCore.pause_scheduler again
        (sleep already pauses the AR scheduler).

        For AR / mixed engines, ``wake_up`` does **not** clear ``_paused``;
        callers must :meth:`resume_generation` when ready (typical trainer
        order: pause → abort → sleep → train → wake → resume). Diffusion-only
        engines have no EngineCore pause to hold, so ``wake_up`` restores
        admission and ``sleep → wake → generate`` keeps working.
        """
        # Block admission before any sleep RPC so generate() waits on
        # _pause_cond during the drain/offload window. Wait until generate()
        # coroutines that already passed the pause check have submitted (or
        # failed) so EngineCore does not see ADD frames while sleeping.
        async with self._pause_cond:
            self._paused = True
            await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0)

        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
        # EngineCore.sleep resets receiver caches itself; only clear P0 here.
        if 0 in ar_stage_ids:
            await self._clear_frontend_mm_cache()

        self._final_output_handler()
        final_acks: list[OmniACK] = []
        if ar_stage_ids:
            self._hold_admission_until_resume = True
            logger.info(
                "[%s] Sleeping AR stage(s) %s via EngineCore.sleep(level=%s, mode=%s)",
                self._name,
                ar_stage_ids,
                level,
                mode,
            )
            await self._engine_core_rpc(
                "sleep",
                stage_ids=ar_stage_ids,
                args=(level, mode),
            )
            # EngineCore.sleep has no OmniACK handshake; emit stage-level SUCCESS
            # markers so callers/tests that count ACKs keep a stable API.
            task_id = f"engine_core-sleep-{uuid.uuid4().hex[:8]}"
            final_acks.extend(
                OmniACK(
                    task_id=task_id,
                    status="SUCCESS",
                    stage_id=sid,
                    rank=0,
                    metadata={"path": "engine_core", "level": level, "mode": mode},
                )
                for sid in ar_stage_ids
            )

        if diffusion_stage_ids:
            final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level))

        self._record_stage_sleep(
            ar_stage_ids + diffusion_stage_ids,
            [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value],
        )
        if level == 2:
            self._level2_sleeping = True
        return final_acks

    async def _sleep_diffusion(self, stage_ids: list[int], level: int) -> list[OmniACK]:
        """Worker-level sleep RPC for diffusion stages only."""
        # Diffusion reports one summary ACK at rank 0 regardless of TP.
        total_workers = len(stage_ids)
        task_id = str(uuid.uuid4())
        self.event_resolver.watch_task(task_id, expected_count=total_workers)
        logger.info("[%s] Sleep (diffusion) initiated (Task: %s).", self._name, task_id)
        task = OmniSleepTask(level=level, task_id=task_id)
        rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
        final_acks: list[OmniACK] = []
        for stage_res in rpc_results:
            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
            for ack in worker_acks:
                if ack is not None:
                    await self.event_resolver.resolve(ack)
                    final_acks.append(ack)
        return final_acks

    async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
        """Wake stages after sleep.

        AR/LLM stages use EngineCore.wake_up (restore memory, auto-resume
        scheduler). Diffusion stages keep the worker-level wake RPC.

        Does **not** clear the frontend ``_paused`` admission gate when
        :meth:`pause_generation` ran or AR stages were slept — call
        :meth:`resume_generation` when the trainer is ready to admit new
        requests. Diffusion-only ``sleep`` uses ``_paused`` only as a race
        guard; this method restores admission after a successful wake.
        """
        self._final_output_handler()

        if getattr(self, "_level2_sleeping", False):
            raise NotImplementedError(
                "wake_up() after sleep(level=2) is not yet implemented: weights were "
                "discarded from GPU and reloading from disk is not yet supported. "
                "Use sleep(level=1) instead, which offloads weights to CPU RAM "
                "and supports fast DMA restore."
            )
        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
        target_stage_ids = ar_stage_ids + diffusion_stage_ids
        _current_tags = self._sleeping_tags_for_stages(target_stage_ids)
        per_stage = getattr(self, "_stage_sleeping_tags", None) or {}
        if not _current_tags and not per_stage:
            _current_tags = set(getattr(self, "_sleeping_tags", set()))
        if tags is None:
            requested_tags = list(_current_tags)
        else:
            requested_tags = [t for t in tags if t in _current_tags]
        if not requested_tags:
            logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
            return []

        final_acks: list[OmniACK] = []
        if ar_stage_ids:
            logger.info("[%s] Waking AR stage(s) %s via EngineCore", self._name, ar_stage_ids)
            await self._engine_core_rpc(
                "wake_up",
                stage_ids=ar_stage_ids,
                kwargs={"tags": requested_tags},
            )
            task_id = f"engine_core-wake-{uuid.uuid4().hex[:8]}"
            final_acks.extend(
                OmniACK(
                    task_id=task_id,
                    status="SUCCESS",
                    stage_id=sid,
                    rank=0,
                    metadata={"path": "engine_core", "tags": list(requested_tags)},
                )
                for sid in ar_stage_ids
            )

        if diffusion_stage_ids:
            final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags))

        self._clear_stage_sleep(target_stage_ids, requested_tags)
        # Only clear the level-2 flag once all tags are warm, in case partial
        # wake support (e.g. tags=["kv_cache"] only) is added in the future.
        if not getattr(self, "_sleeping_tags", None):
            self._level2_sleeping = False
        logger.info(
            "[%s] Wake-up complete for stage(s) %s.",
            self._name,
            ar_stage_ids + diffusion_stage_ids,
        )
        # Diffusion-only sleep uses `_paused` as a race guard. Restore
        # generate() admission after memory is back. AR/mixed sleep and
        # pause_generation keep the trainer hold until resume_generation.
        if not getattr(self, "_hold_admission_until_resume", False):
            async with self._pause_cond:
                self._paused = False
                self._pause_cond.notify_all()
        return final_acks

    async def _wake_diffusion(self, stage_ids: list[int], requested_tags: list[str]) -> list[OmniACK]:
        """Worker-level wake RPC for diffusion stages only."""
        total_workers = len(stage_ids)
        task_id = str(uuid.uuid4())
        self.event_resolver.watch_task(task_id, expected_count=total_workers)
        logger.info("[%s] Wake-up (diffusion) initiated (Task: %s).", self._name, task_id)
        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
        rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
        final_acks: list[OmniACK] = []
        for stage_res in rpc_results:
            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
            for ack in worker_acks:
                if ack is not None:
                    await self.event_resolver.resolve(ack)
                    final_acks.append(ack)
        await asyncio.sleep(0.1)
        return final_acks

    async def is_sleeping(self) -> bool:
        """Return whether all stages are sleeping.

        TODO(AsyncOmni): query the orchestrator once all stage backends expose
        a real sleeping-state RPC. For now we track the requested state locally.
        """
        return bool(getattr(self, "_sleeping_tags", None))

    async def add_lora(self, lora_request: LoRARequest) -> bool:
        """Load a new LoRA adapter into all stages.

        Returns True only if all concretely-implemented stages report success.
        """
        results = await self.collective_rpc(method="add_lora", args=(lora_request,))
        concrete_results = [r for r in results if not (isinstance(r, dict) and r.get("todo"))]
        return all(self._coerce_stage_bool(r) for r in concrete_results) if concrete_results else False

    async def remove_lora(self, adapter_id: int) -> bool:
        """Remove a LoRA adapter from all stages.

        TODO(AsyncOmni): add richer per-stage error reporting to the public API.
        """
        results = await self.collective_rpc(method="remove_lora", args=(adapter_id,))
        concrete_results = [r for r in results if not (isinstance(r, dict) and r.get("todo"))]
        return all(self._coerce_stage_bool(r) for r in concrete_results) if concrete_results else False

    async def list_loras(self) -> list[int]:
        """List all loaded LoRA adapter IDs across stages."""
        results = await self.collective_rpc(method="list_loras")
        merged: set[int] = set()
        for result in results:
            if isinstance(result, dict) and result.get("todo"):
                continue
            if isinstance(result, (list, set)):
                for item in result:
                    if isinstance(item, (list, set)):
                        merged.update(item)
                    elif isinstance(item, int):
                        merged.add(item)
            elif isinstance(result, int):
                merged.add(result)
        return sorted(merged)

    async def pin_lora(self, adapter_id: int) -> bool:
        """Pin a LoRA adapter across stages."""
        results = await self.collective_rpc(method="pin_lora", args=(adapter_id,))
        concrete_results = [r for r in results if not (isinstance(r, dict) and r.get("todo"))]
        return all(self._coerce_stage_bool(r) for r in concrete_results) if concrete_results else False

    # ==================== Properties ====================

    async def get_input_preprocessor(self) -> InputProcessor:
        """Get input preprocessor."""
        return self.input_processor

    async def get_tokenizer(self) -> TokenizerLike:
        """Get tokenizer for the comprehension stage."""
        stage_index = self._get_comprehension_stage_index()
        if stage_index is not None:
            processor = self.engine.output_processors[stage_index]
            assert processor is not None
            tokenizer = processor.tokenizer
            if tokenizer is not None:
                return tokenizer
        processor = self.input_processor
        assert processor is not None
        return processor.tokenizer

    async def is_tracing_enabled(self) -> bool:
        """Check if tracing is enabled."""
        return False

    async def notify_kv_transfer_request_rejected(
        self,
        request_id: str,
        kv_transfer_params: dict[str, Any],
        *,
        data_parallel_rank: int | None = None,
    ) -> None:
        """Notify engine that a KV-transfer request was rejected before admission.

        Omni does not currently use KV-transfer pre-admission resources,
        so this is a no-op.
        """
        logger.debug(
            "KV-transfer request rejected (no-op in omni): request_id=%s",
            request_id,
        )

    async def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
        """Start a new weight update.

        Omni does not currently support weight transfer, so this is a no-op.
        """
        logger.debug("Weight update start requested (no-op in omni)")

    async def finish_weight_update(self, weight_version: str | None = None) -> None:
        """Finish the current weight update.

        Omni does not currently support weight transfer, so this is a no-op.
        ``weight_version`` is accepted for upstream ``EngineClient`` protocol
        compatibility (RLHF weight-transfer routers pass it positionally).
        """
        logger.debug("Weight update finish requested (no-op in omni)")

    async def do_log_stats(self) -> None:
        """Log statistics.

        TODO: Forward to Orchestrator process via message.
        """
        pass

    async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
        """Return the task set exposed by the orchestrator-backed engine."""
        return tuple(self.engine.supported_tasks)