""" 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.plugins.io_processors import get_io_processor 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.utils import random_uuid from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.engine.messages import ErrorMessage, OutputMessage from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni_base import ( OmniBase, OmniEngineDeadError, ) from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.inputs.preprocess import InputPreprocessor from vllm.tokenizers import TokenizerLike from vllm.v1.engine import PauseMode from vllm_omni.inputs.data import OmniPromptType logger = init_logger(__name__) _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 class AsyncEventResolver: """ A generic signal aggregator designed for synchronized handshakes in distributed or multi-stage environments. Supports waiting for a specified number (expected_count) of worker signals in both inline and multiprocess modes. """ def __init__(self, orchestrator=None): self._pending_tasks: dict[str, dict] = {} self.orchestrator = orchestrator self._lock = asyncio.Lock() def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: loop = asyncio.get_running_loop() fut = loop.create_future() self._pending_tasks[task_id] = { "future": fut, "expected_count": expected_count, "received": [], "start_time": time.time(), } return fut async def resolve(self, ack: OmniACK): tid = getattr(ack, "task_id", None) if tid is None and isinstance(ack, dict): tid = ack.get("task_id") async with self._lock: task_info = self._pending_tasks.get(tid) if task_info is None: logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.") return task_info["received"].append(ack) current_count = len(task_info["received"]) expected = task_info["expected_count"] orchestrator = self.orchestrator if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics: freed = getattr(ack, "freed_bytes", 0) if freed == 0 and isinstance(ack, dict): freed = ack.get("freed_bytes", 0) orchestrator.metrics.record_vram_reclaimed(freed) logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.") if current_count >= expected: self._pending_tasks.pop(tid) fut = task_info["future"] if not fut.done(): elapsed = time.time() - task_info["start_time"] logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.") fut.set_result(task_info["received"]) class AsyncOmni(EngineClient, OmniBase): """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. - stage_configs_path: Optional path to YAML file containing stage configurations. If None, configurations are resolved from 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 __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: OmniBase.__init__(self, model=model, **kwargs) self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False self._is_sleeping: bool = False self.final_output_task: asyncio.Task | None = None self.event_resolver = AsyncEventResolver(orchestrator=self) self.config_path = self.engine.config_path self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) self.input_processor = self.engine.input_processor stage_index = self._get_comprehension_stage_index() if stage_index is None: self.io_processor = None else: vllm_config = self.engine.stage_vllm_configs[stage_index] io_processor_plugin = vllm_config.model_config.io_processor_plugin renderer = self.renderer if renderer is None: from vllm.renderers import renderer_from_config renderer = renderer_from_config(vllm_config) self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None: """Look up the sticky-routed replica for (stage_id, request_id). Used as the ``replica_resolver`` callback by ``OrchestratorAggregator`` to label transfer_* metrics without plumbing replica ids through ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters. Returns None when stage_id is out of range or the request hasn't been bound to a replica yet — the metric emit then defensive-skips. """ pools = getattr(self.engine, "stage_pools", None) if pools is None or not (0 <= stage_id < len(pools)): return None return pools[stage_id].get_bound_replica_id(request_id) def _get_comprehension_stage_index(self) -> int | None: fallback_idx: int | None = None for idx, stage_client in enumerate(self.engine.stage_clients): stage_vllm_config = self.engine.stage_vllm_configs[idx] if stage_vllm_config is None: continue if fallback_idx is None: fallback_idx = idx if stage_client.is_comprehension: return idx return fallback_idx @property def renderer(self): """Return the renderer from the engine input processor when available.""" if self.input_processor is None: return None return self.input_processor.renderer @property def vllm_config(self): """Return the vLLM config for the comprehension stage when present.""" stage_index = self._get_comprehension_stage_index() if stage_index is None: return None return self.engine.stage_vllm_configs[stage_index] async def get_vllm_config(self) -> Any: """Compatibility helper for call sites expecting async vllm config access.""" return self.vllm_config def get_diffusion_od_config(self) -> Any | None: """Return the diffusion-stage config when the pipeline has one.""" for stage_client in self.engine.stage_clients: if getattr(stage_client, "stage_type", None) != "diffusion": continue od_config = getattr(stage_client, "od_config", None) if od_config is not None: return od_config inner_engine = getattr(stage_client, "_engine", None) od_config = getattr(inner_engine, "od_config", None) if od_config is not None: return od_config return None @property def model_config(self): """Return the model config for the comprehension stage when present.""" vllm_config = self.vllm_config if vllm_config is None: return None return vllm_config.model_config @staticmethod def _get_unique_request_id(external_request_id: str): """Get a random new request ID for this request; at the server level, this is usually set by the calling entrypoint, but in direct calls, we need to set it explicitly since we do not allow empty IDs. NOTE: in the upstream vLLM, this is done in the InputProcessor's `assign_request_id`. """ uuid = random_uuid() prefix = "" if not external_request_id else f"{external_request_id}-" return f"{prefix}{uuid:.8}" # ==================== 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. **Batch mode (diffusion only):** When *prompt* is a ``list``, all prompts are dispatched in a single ``DiffusionEngine.step()`` call at the diffusion stage. The combined result is yielded as one ``OmniRequestOutput`` with all generated images. Only a single *request_id* is used for the whole batch. Args: prompt: A single prompt **or** a list of prompts. A list triggers batch mode when the diffusion stage is reached. 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. In batch mode the diffusion stage yields one output containing all generated images. Raises: ValueError: If sampling_params_list has incorrect length. """ # 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 async with self._pause_cond: await self._pause_cond.wait_for(lambda: not self._paused) logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}") input_stream_task: asyncio.Task | None = None try: # 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, ) 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. if isinstance(prompt, AsyncGenerator): 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, ) 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, ) submit_ts = time.time() req_state.metrics.stage_first_ts[0] = submit_ts req_start_ts[request_id] = submit_ts # 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") self._log_summary_and_cleanup(request_id) except (asyncio.CancelledError, GeneratorExit): if input_stream_task is not None and not input_stream_task.done(): input_stream_task.cancel() self._fire_failure_counter_if_alive(request_id) await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: self._fire_failure_counter_if_alive(request_id) await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise 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, ) -> 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 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.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, resumable=True, ) has_submitted_first_chunk = True else: await 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, resumable=True, ) except (asyncio.CancelledError, GeneratorExit): cancelled = True except Exception as error: await req_state.queue.put( ErrorMessage( request_id=request_id, error=str(error), ) ) finally: 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.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, resumable=False, ) else: await 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, resumable=False, ) 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.") # ==================== Processing Methods ==================== async def _process_orchestrator_results( self, request_id: str, metrics: OrchestratorMetrics, final_stage_id_for_e2e: int, req_start_ts: dict[str, float], wall_start_ts: float, ) -> AsyncGenerator[OmniRequestOutput, None]: """Read results from the Orchestrator (via the request's asyncio.Queue) and yield OmniRequestOutput objects. The Orchestrator handles all stage-to-stage transfers. This method only processes final outputs that arrive on the per-request queue. """ req_state = self.request_states.get(request_id) if req_state is None: return while True: result = await req_state.queue.get() if isinstance(result, ErrorMessage): logger.error( "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s", request_id, result.stage_id, result.error, ) if result.fatal: raise OmniEngineDeadError( result.error, error_stage_id=result.stage_id, ) raise RuntimeError(result.error) if not isinstance(result, OutputMessage): logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) continue stage_id = result.stage_id self._check_engine_output_error(result, request_id, stage_id) # Process the result (constructs OmniRequestOutput) output_to_yield = self._process_single_result( result, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e, ) if output_to_yield: # Set the external request ID back to the user yielded input output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id logger.debug( "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s", request_id, stage_id, getattr(output_to_yield, "final_output_type", None), ) yield output_to_yield # The Orchestrator sets "finished" when the final stage is done if result.finished: break # ==================== Output Handler ==================== def _final_output_handler(self) -> None: """Start the final output handler if not already running. This handler reads messages from the Orchestrator output queue and routes them to per-request asyncio.Queues. """ if self.final_output_task is not None: return engine = self.engine async def _final_output_loop(): """Background coroutine that dispatches final outputs to request queues.""" try: while True: msg = await engine.try_get_output_async() if msg is None: await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) continue if isinstance(msg, dict) and msg.get("type") == "ack": ack_data = msg.get("ack") tid = getattr(ack_data, "task_id", "unknown") logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}") await self.event_resolver.resolve(ack_data) continue if isinstance(msg, OmniACK): logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}") await self.event_resolver.resolve(msg) continue if hasattr(msg, "task_id"): tid = getattr(msg, "task_id") logger.info(f"[{self._name}] Intercepted task-ID object: {tid}") await self.event_resolver.resolve(msg) continue should_continue, _, stage_id, req_state = self._handle_output_message(msg) if should_continue: continue req_state.stage_id = stage_id # Route to the per-request queue await req_state.queue.put(msg) except asyncio.CancelledError: raise except OmniEngineDeadError as e: logger.error("[AsyncOmni] Engine dead: %s", e) for req_state in list(self.request_states.values()): error_msg = ErrorMessage( error=str(e), fatal=True, request_id=req_state.request_id, stage_id=e.error_stage_id, ) await req_state.queue.put(error_msg) except EngineDeadError as e: logger.error("[AsyncOmni] Engine dead: %s", e) for req_state in list(self.request_states.values()): error_msg = ErrorMessage( error=str(e), fatal=True, request_id=req_state.request_id, ) await req_state.queue.put(error_msg) except Exception as e: logger.exception("[AsyncOmni] final_output_loop failed.") for req_state in list(self.request_states.values()): error_msg = ErrorMessage( request_id=req_state.request_id, error=str(e), ) await req_state.queue.put(error_msg) self.final_output_task = None self.final_output_task = asyncio.create_task(_final_output_loop()) logger.debug("[AsyncOmni] Final output handler started") # ==================== 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 @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]) -> 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) async def _abort_internal_requests(self, request_id: str | Iterable[str]): """Abort request(s) via the Orchestrator given internal request IDs, which take the format -. """ request_ids = [request_id] if isinstance(request_id, str) else list(request_id) # Request IDs are already internal, so we just need to get the matching states. internal_req_ids = [rid for rid in request_ids if rid in self.request_states] await self._abort(internal_req_ids) async def _abort(self, request_ids: list[str]) -> None: """Submit request IDs to be aborted to the engine.""" await self.engine.abort_async(request_ids) for rid in request_ids: self.request_states.pop(rid, None) if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) async def pause_generation( self, *, mode: PauseMode = "abort", wait_for_inflight_requests: bool = False, clear_cache: bool = True, ) -> None: """Pause generation.""" async with self._pause_cond: if self._paused: return self._paused = True # TODO: Implement request draining if wait_for_inflight_requests if clear_cache: # Clear caches for all stages. 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) -> None: """Resume generation.""" async with self._pause_cond: self._paused = False self._pause_cond.notify_all() async def is_paused(self) -> bool: """Check if 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 multi-modal cache for all stages. TODO: Forward to Orchestrator process via message. """ logger.warning("[AsyncOmni] reset_mm_cache not yet supported with Orchestrator process") 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]: self._final_output_handler() if stage_ids is None: stage_ids = list(range(len(self.engine.stage_clients))) total_workers = 0 for sid in stage_ids: client = self.engine.stage_clients[sid] # During the Diffusion phase, regardless of the TP amount, # currently only a summary ACK is reported at Rank 0. if getattr(client, "stage_type", "") == "diffusion": total_workers += 1 else: config = self.engine.stage_vllm_configs[sid] actual_tp = config.parallel_config.tensor_parallel_size if config else 1 total_workers += actual_tp task_id = str(uuid.uuid4()) self.event_resolver.watch_task(task_id, expected_count=total_workers) logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") 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 = [] 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) self._is_sleeping = True return final_acks async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]: self._final_output_handler() if stage_ids is None: stage_ids = list(range(len(self.engine.stage_clients))) total_workers = 0 for sid in stage_ids: client = self.engine.stage_clients[sid] if getattr(client, "stage_type", "") == "diffusion": total_workers += 1 else: config = self.engine.stage_vllm_configs[sid] total_workers += config.parallel_config.tensor_parallel_size if config else 1 task_id = str(uuid.uuid4()) self.event_resolver.watch_task(task_id, expected_count=total_workers) logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") task = OmniWakeTask(tags=tags, task_id=task_id) rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids) final_acks = [] 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) current_omni_platform.synchronize() await asyncio.sleep(0.1) self._is_sleeping = False logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.") 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 self._is_sleeping 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 ==================== @property def is_running(self) -> bool: """Check if the engine is running.""" orchestrator_alive = self.engine.is_alive() task_alive = self.final_output_task is not None and not self.final_output_task.done() return orchestrator_alive and task_alive @property def errored(self) -> bool: """Whether the engine is in a non-recoverable error state. Delegates to ``OmniBase.errored`` which checks the orchestrator thread and all stage clients. Redeclared here to satisfy the ``EngineClient`` abstract-property requirement (Python's ABC mechanism does not resolve abstract methods from sibling MRO entries). """ return OmniBase.errored.fget(self) # type: ignore[union-attr] @property def _name(self) -> str: return "AsyncOrchestrator" @property def is_stopped(self) -> bool: """EngineClient abstract property implementation.""" return self.errored @property def dead_error(self) -> BaseException: """EngineClient abstract property implementation.""" return OmniEngineDeadError() # ==================== EngineClient Interface ==================== async def get_input_preprocessor(self) -> InputPreprocessor: """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) -> None: """Finish the current weight update. Omni does not currently support weight transfer, so this is a no-op. """ 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) async def check_health(self) -> None: """Check engine health by verifying the Orchestrator process is alive.""" OmniBase.check_health(self) # ==================== Shutdown ==================== def shutdown(self, timeout: float | None = None) -> None: """Shutdown the engine.""" if self.final_output_task is not None: self.final_output_task.cancel() self.final_output_task = None OmniBase.shutdown(self)