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

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__)


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)
    """

    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
        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,
        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.

        **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(
                getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
            ):
                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=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()
            req_state.metrics.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")
            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)
        req_state = self.request_states[request_id]
        has_submitted_first_chunk = 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 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:
                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,
                    )
                )
            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
                finally:
                    # Unblock generate() even on cancel / empty stream / submit
                    # failure so it can observe a terminal abort or empty result.
                    _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_stage_ids: list[int] = []
        effective_stage_ids = stage_ids or list(range(len(results)))
        for index, result in enumerate(results):
            if isinstance(result, dict) and result.get("todo"):
                unsupported_stage_ids.append(effective_stage_ids[index])

        if unsupported_stage_ids:
            logger.warning(
                "[AsyncOmni] collective_rpc(%s) has TODO support on stage(s): %s",
                method,
                unsupported_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,
    ) -> list[Any]:
        """Call an AR EngineCore helper via collective_rpc (orchestrator loop).

        StagePool resolves ``{method}_async`` on the AR client when present
        (vLLM AsyncMPClient convention). Raises if any replica reports failure.
        """
        results = await self.collective_rpc(
            method=method,
            args=args,
            kwargs=kwargs,
            stage_ids=stage_ids,
        )
        for result in results:
            if isinstance(result, dict) and result.get("error"):
                raise RuntimeError(f"{method} failed: {result['error']}")
        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_clients)
        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:
            client = self.engine.stage_clients[sid]
            if getattr(client, "stage_type", "llm") == "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. Diffusion stages have no EngineCore scheduler — only frontend
           admission is paused for them.

        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 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},
            )

        # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler
        # already clears AR-side caches when clear_cache=True.
        if clear_cache:
            await self.reset_prefix_cache(
                reset_running_requests=not wait_for_inflight_requests,
                reset_connector=True,
            )
            await self.reset_mm_cache()
            await self.reset_encoder_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)

        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

    async def start_profile(
        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]:
        """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 reset_mm_cache(self) -> None:
        """Reset the frontend (P0) multimodal processor cache.

        ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache.
        Clearing P0 avoids hash-only follow-up requests after that reset.
        """
        renderer = self.renderer
        if renderer is not None:
            await renderer.clear_mm_cache_async()

    async def reset_encoder_cache(self) -> None:
        """Reset the encoder cache for all stages.

        TODO: Forward to Orchestrator process via message.
        """
        logger.warning("[AsyncOmni] reset_encoder_cache not yet supported with Orchestrator process")

    async def reset_prefix_cache(
        self,
        reset_running_requests: bool = False,
        reset_connector: bool = False,
    ) -> bool:
        """Reset the prefix cache for all stages.

        TODO: Forward to Orchestrator process via message.
        """
        logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process")
        return True

    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 existing worker-level handle_sleep_task RPC
        because StageDiffusionProc does not expose EngineCore.pause_scheduler.

        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)

        # P0 sender cache must drop hashes before EngineCore.sleep clears P1.
        await self.reset_mm_cache()

        self._final_output_handler()
        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
        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:
            tokenizer = self.engine.output_processors[stage_index].tokenizer
            if tokenizer is not None:
                return tokenizer
        return self.input_processor.tokenizer  # type: ignore[return-value]

    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)