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