""" Orchestrator for vLLM-Omni multi-stage runtime. Runs inside a background thread with its own asyncio event loop. Owns all StageEngineCoreClient instances, input/output processors, and handles stage-to-stage transfer logic. """ from __future__ import annotations import asyncio import copy import time as _time from dataclasses import dataclass, field from typing import Any import janus import torch from vllm.config import ModelConfig from vllm.logger import init_logger from vllm.outputs import RequestOutput from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreOutputs from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length from vllm_omni.engine import ( OmniEngineCoreRequest, ) from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics from vllm_omni.metrics.stats import StageStats from vllm_omni.metrics.utils import count_tokens_from_outputs logger = init_logger(__name__) def build_engine_core_request_from_tokens( request_id: str, prompt: dict[str, Any], params: SamplingParams | PoolingParams, arrival_time: float | None = None, model_config: ModelConfig | None = None, ) -> OmniEngineCoreRequest: """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt. Lightweight alternative to the full InputProcessor pipeline - skips tokenization, multimodal preprocessing, LoRA validation, and platform validation. Intended for stage 1+ where the upstream stage has already produced token IDs and optional embeddings. """ if arrival_time is None: arrival_time = _time.time() prompt_token_ids = prompt["prompt_token_ids"] # Clone params and set max_tokens if needed sampling_params = None pooling_params = None if isinstance(params, SamplingParams): sampling_params = params.clone() if sampling_params.max_tokens is None and model_config is not None: sampling_params.max_tokens = model_config.max_model_len - len(prompt_token_ids) else: pooling_params = params.clone() prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds") # Serialize additional_information if present additional_info_payload = serialize_additional_information( prompt.get("additional_information"), log_prefix=f"build_engine_core_request_from_tokens req={request_id}", ) return OmniEngineCoreRequest( request_id=request_id, prompt_token_ids=prompt_token_ids, mm_features=None, sampling_params=sampling_params, pooling_params=pooling_params, arrival_time=arrival_time, lora_request=None, cache_salt=None, data_parallel_rank=None, prompt_embeds=prompt_embeds, additional_information=additional_info_payload, ) # ============================================================ # Orchestrator internals (run inside the background thread) # ============================================================ @dataclass class OrchestratorRequestState: """Per-request bookkeeping inside the Orchestrator.""" request_id: str prompt: Any = None sampling_params_list: list[Any] = field(default_factory=list) final_stage_id: int = -1 # Metrics: timestamp when request was submitted to each stage stage_submit_ts: dict[int, float] = field(default_factory=dict) class Orchestrator: """Runs inside a background thread's asyncio event loop. Owns all StageEngineCoreClient instances, input/output processors, and handles stage-to-stage transfer logic. """ def __init__( self, request_async_queue: janus.AsyncQueue[dict[str, Any]], output_async_queue: janus.AsyncQueue[dict[str, Any]], rpc_async_queue: janus.AsyncQueue[dict[str, Any]], stage_clients: list[Any], output_processors: list[Any], stage_vllm_configs: list[Any], *, async_chunk: bool = False, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue self.rpc_async_queue = rpc_async_queue self.num_stages = len(stage_clients) self.async_chunk = bool(async_chunk) self.stage_clients: list[Any] = stage_clients self.output_processors: list[Any] = output_processors self.stage_vllm_configs: list[Any] = stage_vllm_configs # Per-request state self.request_states: dict[str, OrchestratorRequestState] = {} # CFG companion tracking self._companion_map: dict[str, dict[str, str]] = {} self._companion_to_parent: dict[str, str] = {} self._companion_ids: set[str] = set() self._companion_done: dict[str, set[str]] = {} self._deferred_parents: dict[str, dict[str, Any]] = {} # Per-stage metrics accumulators. self._batch_seq: list[int] = [0] * self.num_stages self._agg_total_tokens: list[int] = [0] * self.num_stages self._agg_total_gen_time_ms: list[float] = [0.0] * self.num_stages # Shutdown coordination self._shutdown_event = asyncio.Event() self._stages_shutdown = False async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" logger.info("[Orchestrator] Starting event loop") request_task = asyncio.create_task(self._request_handler(), name="orchestrator-request-handler") output_task = asyncio.create_task( self._orchestration_output_handler(), name="orchestrator-stage-output-handler", ) try: # Run both tasks concurrently; if either fails the other is cancelled. await asyncio.gather(request_task, output_task) except asyncio.CancelledError: raise except Exception: logger.exception("[Orchestrator] Fatal error in orchestrator tasks") raise finally: self._shutdown_event.set() for t in (request_task, output_task): if not t.done(): t.cancel() try: await asyncio.gather(request_task, output_task, return_exceptions=True) except Exception: pass self._shutdown_stages() # Cancel any remaining tasks spawned by wait_for / gather so # the event loop can close cleanly without "pending task" errors. loop = asyncio.get_running_loop() pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()] for t in pending: t.cancel() if pending: await asyncio.gather(*pending, return_exceptions=True) async def _request_handler(self) -> None: """Read messages from the main thread via request_async_queue.""" while True: msg = await self.request_async_queue.get() msg_type = msg.get("type") if msg_type == "add_request": await self._handle_add_request(msg) elif msg_type == "add_companion_request": await self._handle_add_companion(msg) elif msg_type == "abort": await self._handle_abort(msg) elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) elif msg_type == "shutdown": logger.info("[Orchestrator] Received shutdown signal") self._shutdown_event.set() self._shutdown_stages() break else: logger.warning(f"[Orchestrator] Unknown message type: {msg_type}") async def _orchestration_output_handler(self) -> None: """Poll all stages, handle transfers, send final outputs to main.""" try: await self._orchestration_loop() except asyncio.CancelledError: logger.debug("[Orchestrator] _orchestration_output_handler cancelled") return async def _orchestration_loop(self) -> None: """Inner loop for _orchestration_output_handler (clean cancellation). Control flow: poll raw → process through output processor → route. """ while not self._shutdown_event.is_set(): idle = True for stage_id in range(self.num_stages): if self._shutdown_event.is_set(): return # 1) Diffusion stage: poll non-blocking queue # TODO (Peiqi): the output of diffusion stage is OmniRequestOutput, # which is different from EngineCoreOutputs (LLM stages). We may want to unify # the output format in the future to simplify the processing logic in Orchestrator. stage_client = self.stage_clients[stage_id] if stage_client.stage_type == "diffusion": output = stage_client.get_diffusion_output_async() if output is not None: idle = False if isinstance(output, dict) and output.get("type") == "error": req_id = output.get("request_id") req_state = self.request_states.get(req_id) await self.output_async_queue.put( { "type": "error", "request_id": req_id, "stage_id": stage_id, "status_code": output.get("status_code", 500), "error": output.get("error", "Unknown diffusion stage error"), "error_type": output.get("error_type"), "detail": output.get("detail") or {}, "finished": True, "stage_submit_ts": req_state.stage_submit_ts.get(stage_id) if req_state else None, } ) self._cleanup_companion_state(req_id) self.request_states.pop(req_id, None) continue req_state = self.request_states.get(output.request_id) if req_state is not None: stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state) await self._route_output(stage_id, output, req_state, stage_metrics) continue # 1) Poll raw outputs from the stage try: raw_outputs = await asyncio.wait_for(self._poll_stage_raw(stage_id), timeout=0.001) except asyncio.TimeoutError: continue except asyncio.CancelledError: raise except Exception: if self._shutdown_event.is_set(): return logger.exception( "[Orchestrator] _poll_stage_raw failed for stage-%s", stage_id, ) raise if raw_outputs is None: continue idle = False # 2) Process raw outputs through the output processor request_outputs = await self._process_stage_outputs(stage_id, raw_outputs) # 3) Route each processed output for output in request_outputs: req_state = self.request_states.get(output.request_id) if req_state is None: logger.warning( "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)", output.request_id, stage_id, list(self.request_states.keys()), ) continue stage_metrics = None if output.finished: stage_metrics = self._build_stage_metrics( stage_id, output.request_id, [output], req_state, ) await self._route_output(stage_id, output, req_state, stage_metrics) if idle: await asyncio.sleep(0.001) else: await asyncio.sleep(0) async def _route_output( self, stage_id: int, output: Any, req_state: OrchestratorRequestState, stage_metrics: Any, ) -> None: """Route a processed output: send to main thread and/or forward to next stage.""" req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) stage_client = self.stage_clients[stage_id] # CFG companion handling: companions don't produce user-visible output # and don't forward to the next stage directly. if finished and req_id in self._companion_ids: parent_id = self._companion_to_parent.get(req_id) if parent_id is not None: self._companion_done.setdefault(parent_id, set()).add(req_id) logger.debug( "[Orchestrator] CFG companion %s done (parent=%s)", req_id, parent_id, ) # Check if parent is waiting and all companions are done if parent_id in self._deferred_parents and self._all_companions_done(parent_id): deferred = self._deferred_parents.pop(parent_id) parent_state = self.request_states.get(parent_id) if parent_state is not None: await self._forward_to_next_stage( parent_id, deferred["stage_id"], deferred["output"], parent_state, ) self.request_states.pop(req_id, None) return if stage_client.final_output: await self.output_async_queue.put( { "type": "output", "request_id": req_id, "stage_id": stage_id, "engine_outputs": output, "metrics": stage_metrics, "finished": finished and stage_id == req_state.final_stage_id, "stage_submit_ts": submit_ts, } ) elif stage_metrics is not None: await self.output_async_queue.put( { "type": "stage_metrics", "request_id": req_id, "stage_id": stage_id, "metrics": stage_metrics, "stage_submit_ts": submit_ts, } ) if finished and stage_id < req_state.final_stage_id and not self.async_chunk: # If this parent has CFG companions, defer forwarding until all done if req_id in self._companion_map and not self._all_companions_done(req_id): self._deferred_parents[req_id] = { "stage_id": stage_id, "output": output, } logger.debug( "[Orchestrator] Parent %s deferred, waiting for CFG companions", req_id, ) else: await self._forward_to_next_stage(req_id, stage_id, output, req_state) if finished and stage_id == req_state.final_stage_id: self._cleanup_companion_state(req_id) self.request_states.pop(req_id, None) def _cleanup_companion_state(self, parent_id: str) -> None: """Remove all companion tracking state for a completed parent.""" role_map = self._companion_map.pop(parent_id, {}) for cid in role_map.values(): self._companion_ids.discard(cid) self._companion_to_parent.pop(cid, None) self._companion_done.pop(parent_id, None) self._deferred_parents.pop(parent_id, None) def _all_companions_done(self, parent_id: str) -> bool: """Check whether all CFG companions for a parent request have finished.""" role_map = self._companion_map.get(parent_id, {}) if not role_map: return True done_set = self._companion_done.get(parent_id, set()) return all(cid in done_set for cid in role_map.values()) def _build_stage_metrics( self, stage_id: int, req_id: str, request_outputs: list[RequestOutput], req_state: OrchestratorRequestState, ) -> StageRequestMetrics: """Build StageRequestMetrics for a finished request at a stage. Reuses StageRequestMetrics so OrchestratorMetrics and downstream metric handlers can consume a stable schema. """ now = _time.time() submit_ts = req_state.stage_submit_ts.get(stage_id, now) stage_gen_time_ms = (now - submit_ts) * 1000.0 num_tokens_out = count_tokens_from_outputs(request_outputs) num_tokens_in = 0 if stage_id == 0: for ro in request_outputs: ptids = getattr(ro, "prompt_token_ids", None) if ptids is not None: num_tokens_in += len(ptids) # Monotonic batch counter per stage. self._batch_seq[stage_id] += 1 batch_id = self._batch_seq[stage_id] # Accumulate for running-average stage_stats self._agg_total_tokens[stage_id] += num_tokens_out self._agg_total_gen_time_ms[stage_id] += stage_gen_time_ms return StageRequestMetrics( num_tokens_in=num_tokens_in, num_tokens_out=num_tokens_out, stage_gen_time_ms=stage_gen_time_ms, batch_id=batch_id, batch_size=1, rx_decode_time_ms=0.0, rx_transfer_bytes=0, rx_in_flight_time_ms=0.0, stage_stats=StageStats( total_token=self._agg_total_tokens[stage_id], total_gen_time_ms=self._agg_total_gen_time_ms[stage_id], ), ) async def _forward_to_next_stage( self, req_id: str, stage_id: int, output: Any, req_state: OrchestratorRequestState, ) -> None: """Forward output from current stage to the next stage. Handles the full pipeline: set outputs on current stage, compute next-stage inputs, build lightweight requests, and submit them. """ next_stage_id = stage_id + 1 next_client = self.stage_clients[next_stage_id] params = req_state.sampling_params_list[next_stage_id] if next_client.stage_type == "diffusion": self.stage_clients[stage_id].set_engine_outputs([output]) if next_client.custom_process_input_func is not None: diffusion_prompt = next_client.custom_process_input_func( self.stage_clients, next_client.engine_input_source, req_state.prompt, False, ) if isinstance(diffusion_prompt, list): diffusion_prompt = diffusion_prompt[0] else: diffusion_prompt = req_state.prompt # Attach CFG companion KV request IDs so the diffusion model # runner can fetch companion KV caches alongside the primary one. cfg_ids = self._companion_map.get(req_id) if cfg_ids: from vllm_omni.inputs.data import OmniDiffusionSamplingParams if isinstance(params, OmniDiffusionSamplingParams): params = copy.deepcopy(params) params.cfg_kv_request_ids = cfg_ids logger.info( "[Orchestrator] Attaching cfg_kv_request_ids=%s to req %s", cfg_ids, req_id, ) if isinstance(diffusion_prompt, list): await next_client.add_batch_request_async( req_id, diffusion_prompt, params, ) else: await next_client.add_request_async(req_id, diffusion_prompt, params) req_state.stage_submit_ts[next_stage_id] = _time.time() return self.stage_clients[stage_id].set_engine_outputs([output]) # Process inputs for next stage try: next_inputs = next_client.process_engine_inputs( stage_list=self.stage_clients, prompt=req_state.prompt, ) except Exception: logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", req_id, next_stage_id, ) raise # Build and submit requests for each input for next_input in next_inputs: request = build_engine_core_request_from_tokens( request_id=req_id, prompt=next_input, params=params, model_config=self.stage_vllm_configs[next_stage_id].model_config, ) # TODO: Here we directly use the req id to assign. request.external_req_id = request.request_id self.output_processors[next_stage_id].add_request( request=request, prompt=None, parent_req=None, request_index=0, queue=None, ) await next_client.add_request_async(request) # Record submit timestamp for the next stage req_state.stage_submit_ts[next_stage_id] = _time.time() async def _poll_stage_raw(self, stage_id: int) -> EngineCoreOutputs | None: """Pull raw EngineCoreOutputs from a stage client without processing. Returns the raw outputs object, or None when there is nothing to consume. """ outputs = await self.stage_clients[stage_id].get_output_async() if not outputs.outputs: return None return outputs async def _process_stage_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> list[RequestOutput]: """Run the output processor on raw outputs, returning RequestOutputs. Also handles abort forwarding and scheduler stats updates. """ processor = self.output_processors[stage_id] processed = processor.process_outputs( raw_outputs.outputs, raw_outputs.timestamp, None, ) if processed.reqs_to_abort: await self.stage_clients[stage_id].abort_requests_async(processed.reqs_to_abort) if raw_outputs.scheduler_stats is not None: processor.update_scheduler_stats(raw_outputs.scheduler_stats) return processed.request_outputs async def _handle_add_request(self, msg: dict[str, Any]) -> None: """Handle an add_request message from the main thread.""" stage_id = 0 request_id = msg["request_id"] prompt = msg["prompt"] original_prompt = msg.get("original_prompt", prompt) sampling_params_list = msg["sampling_params_list"] if not sampling_params_list: raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") params = sampling_params_list[0] final_stage_id = msg["final_stage_id"] logger.info( "[Orchestrator] _handle_add_request: stage=%s req=%s " "prompt_type=%s original_prompt_type=%s final_stage=%s " "num_sampling_params=%d", stage_id, request_id, type(prompt).__name__, type(original_prompt).__name__, final_stage_id, len(sampling_params_list), ) # Track request state - use original_prompt so downstream stages # (e.g. thinker2talker) can access the raw dict with multi_modal_data. req_state = OrchestratorRequestState( request_id=request_id, prompt=original_prompt, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, ) req_state.stage_submit_ts[stage_id] = _time.time() self.request_states[request_id] = req_state # Stage-0 prompt is already a fully-formed OmniEngineCoreRequest # (pre-processed by AsyncOmniEngine.add_request, output processor # already registered there) - submit directly. request = prompt stage_client = self.stage_clients[stage_id] if stage_client.stage_type == "diffusion": if isinstance(prompt, list): await stage_client.add_batch_request_async( request_id, prompt, params, ) else: await stage_client.add_request_async(request_id, prompt, params) else: await stage_client.add_request_async(request) if self.async_chunk and stage_id == 0 and final_stage_id > 0: await self._prewarm_async_chunk_stages(request_id, request, req_state) async def _prewarm_async_chunk_stages( self, request_id: str, stage0_request: Any, req_state: OrchestratorRequestState, ) -> None: """Pre-submit downstream stages for async-chunk mode. In async-chunk mode, stages exchange data through connectors/chunk adapters, so downstream stages should be armed once at request start instead of waiting for stage-finished forwarding. """ if req_state.final_stage_id <= 0: return prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None) if prompt_token_ids is None: logger.warning( "[Orchestrator] async_chunk prewarm skipped for req=%s: stage0 prompt_token_ids missing", request_id, ) return # Pre-arm stage-1+ with placeholder prompt IDs. try: next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) except Exception: next_prompt_len = max(1, len(prompt_token_ids)) original_prompt = req_state.prompt if isinstance(original_prompt, dict): base_input = copy.deepcopy(original_prompt) else: base_input = {} base_input["prompt_token_ids"] = [0] * next_prompt_len base_input["multi_modal_data"] = None base_input["mm_processor_kwargs"] = None for next_stage_id in range(1, req_state.final_stage_id + 1): next_client = self.stage_clients[next_stage_id] params = req_state.sampling_params_list[next_stage_id] if next_client.stage_type == "diffusion": await next_client.add_request_async(request_id, req_state.prompt, params) req_state.stage_submit_ts[next_stage_id] = _time.time() continue request = build_engine_core_request_from_tokens( request_id=request_id, prompt=base_input, params=params, model_config=self.stage_vllm_configs[next_stage_id].model_config, ) request.external_req_id = request.request_id self.output_processors[next_stage_id].add_request( request=request, prompt=None, parent_req=None, request_index=0, queue=None, ) await next_client.add_request_async(request) req_state.stage_submit_ts[next_stage_id] = _time.time() async def _handle_add_companion(self, msg: dict[str, Any]) -> None: """Handle an add_companion_request message: submit companion to stage 0.""" companion_id = msg["companion_id"] parent_id = msg["parent_id"] role = msg["role"] companion_prompt = msg["prompt"] sampling_params_list = msg["sampling_params_list"] # Register companion mapping if parent_id not in self._companion_map: self._companion_map[parent_id] = {} self._companion_map[parent_id][role] = companion_id self._companion_ids.add(companion_id) self._companion_to_parent[companion_id] = parent_id self._companion_done.setdefault(parent_id, set()) companion_state = OrchestratorRequestState( request_id=companion_id, prompt=companion_prompt, sampling_params_list=sampling_params_list, final_stage_id=0, ) companion_state.stage_submit_ts[0] = _time.time() self.request_states[companion_id] = companion_state request = companion_prompt # Already a processed OmniEngineCoreRequest stage_client = self.stage_clients[0] await stage_client.add_request_async(request) logger.info( "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s)", companion_id, role, parent_id, ) async def _handle_abort(self, msg: dict[str, Any]) -> None: """Handle an abort message from the main thread.""" request_ids = msg["request_ids"] # Also abort any CFG companions for aborted parents companion_ids_to_abort: list[str] = [] for req_id in request_ids: role_map = self._companion_map.pop(req_id, {}) for cid in role_map.values(): companion_ids_to_abort.append(cid) self._companion_ids.discard(cid) self._companion_to_parent.pop(cid, None) self.request_states.pop(cid, None) self._companion_done.pop(req_id, None) self._deferred_parents.pop(req_id, None) all_ids_to_abort = list(request_ids) + companion_ids_to_abort for stage_id in range(self.num_stages): await self.stage_clients[stage_id].abort_requests_async(all_ids_to_abort) for req_id in request_ids: self.request_states.pop(req_id, None) logger.info("[Orchestrator] Aborted request(s) %s", request_ids) async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None: """Handle a control-plane RPC request from the main thread. TODO(AsyncOmni): parallelize stage dispatch if control latency becomes noticeable. The current sequential fanout keeps the first version simple and deterministic. """ rpc_id = msg["rpc_id"] method = msg["method"] timeout = msg.get("timeout") args = tuple(msg.get("args", ())) kwargs = dict(msg.get("kwargs") or {}) requested_stage_ids = msg.get("stage_ids") stage_ids = list(range(self.num_stages)) if requested_stage_ids is None else list(requested_stage_ids) results: list[Any] = [] for stage_id in stage_ids: if stage_id < 0 or stage_id >= self.num_stages: results.append( { "supported": False, "todo": True, "error": f"Invalid stage id {stage_id}", } ) continue stage_client = self.stage_clients[stage_id] try: if hasattr(stage_client, "collective_rpc_async"): stage_result = await stage_client.collective_rpc_async( method=method, timeout=timeout, args=args, kwargs=kwargs, ) else: stage_result = { "supported": False, "todo": True, "reason": (f"{stage_client.__class__.__name__}.collective_rpc_async is not implemented yet"), } except Exception as exc: logger.exception( "[Orchestrator] collective_rpc failed: stage=%s method=%s", stage_id, method, ) stage_result = { "supported": False, "error": str(exc), } results.append(stage_result) await self.rpc_async_queue.put( { "type": "collective_rpc_result", "rpc_id": rpc_id, "method": method, "stage_ids": stage_ids, "results": results, } ) def _shutdown_stages(self) -> None: """Shutdown all stage clients.""" if self._stages_shutdown: return self._stages_shutdown = True logger.info("[Orchestrator] Shutting down all stages") for stage_id, stage_client in enumerate(self.stage_clients): try: stage_client.shutdown() logger.info(f"[Orchestrator] Stage {stage_id} shut down") except Exception as e: logger.warning(f"[Orchestrator] Failed to shutdown stage {stage_id}: {e}")