diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index dbf05a26b..8356adac0 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -1,16 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import multiprocessing as mp -import os import time from collections.abc import AsyncGenerator, Iterable -from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import asdict from pprint import pformat from typing import Any -from omegaconf import OmegaConf from vllm.config import VllmConfig from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger @@ -23,50 +19,24 @@ from vllm.v1.engine.exceptions import EngineDeadError # Internal imports (our code) from vllm_omni.config import OmniModelConfig from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.distributed.omni_connectors import ( - get_stage_connector_config, - initialize_orchestrator_connectors, -) from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector -from vllm_omni.distributed.ray_utils.utils import ( - create_placement_group, - get_ray_queue_class, - try_close_ray, -) from vllm_omni.engine.input_processor import OmniInputProcessor from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.log_utils import ( OrchestratorMetrics, ) +from vllm_omni.entrypoints.omni import OmniBase from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load from vllm_omni.entrypoints.utils import ( get_final_stage_id_for_e2e, - load_stage_configs_from_model, - load_stage_configs_from_yaml, - resolve_model_config_path, ) from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -def _dummy_snapshot_download(model_id): - return model_id - - -def omni_snapshot_download(model_id) -> str: - # TODO: this is just a workaround for quickly use modelscope, we should support - # modelscope in weight loading feature instead of using `snapshot_download` - if os.environ.get("VLLM_USE_MODELSCOPE", False): - from modelscope.hub.snapshot_download import snapshot_download - - return snapshot_download(model_id) - else: - return _dummy_snapshot_download(model_id) - - -class AsyncOmni: +class AsyncOmni(OmniBase): """Asynchronous unified entry point supporting multi-stage pipelines for LLM and Diffusion models. Similar to the Omni class, but provides an asynchronous interface supporting @@ -102,23 +72,6 @@ class AsyncOmni: """ def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None: - model = args[0] if args else kwargs.get("model", "") - assert model != "", "Null model id detected, please specify a model id." - model = omni_snapshot_download(model) - if args: - args[0] = model - elif kwargs.get("model", "") != "": - kwargs["model"] = model - - # Stage management attributes - self.stage_list: list[OmniStage] = [] - self._stage_in_queues: list[mp.Queue] = [] - self._stage_out_queues: list[mp.Queue] = [] - self._stages_ready: set[int] = set() - self._ray_pg = None - self._queue_cls = None - self._ctx = None - # Pause/resume control attributes self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False @@ -127,352 +80,107 @@ class AsyncOmni: self.request_states: dict[str, ClientRequestState] = {} self.output_handler: asyncio.Task | None = None - # Initialize stages - each stage will create appropriate instances based on stage_type - # Stage workers will automatically create AsyncOmniLLM or AsyncOmniDiffusion instances - # based on stage_type in YAML config (handled in omni_stage.py) - logger.info(f"Initializing async stages for model: {model}") - # Use kwargs-based initialization logic to avoid conflicts with old _initialize_stages signature - self._initialize_stages_from_kwargs(model, kwargs) + super().__init__(*args, **kwargs) - def _initialize_stages_from_kwargs(self, model: str, kwargs: dict[str, Any]) -> None: - """Initialize stage list management. - - Each stage will create appropriate instances (AsyncOmniLLM or AsyncOmniDiffusion) - based on stage_type in YAML config. - """ - init_sleep_seconds = kwargs.get("init_sleep_seconds", 20) - shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536) - init_timeout = kwargs.get("init_timeout", 300) - worker_backend = kwargs.get("worker_backend", "multi_process") - ray_address = kwargs.get("ray_address", None) - batch_timeout = kwargs.get("batch_timeout", 10) - stage_configs_path = kwargs.get("stage_configs_path", None) - log_stats = kwargs.get("log_stats", False) - - # Load stage configs from YAML - if stage_configs_path is None: - self.config_path = resolve_model_config_path(model) - self.stage_configs = load_stage_configs_from_model(model) - if not self.stage_configs: - default_stage_cfg = [ - { - "stage_id": 0, - "stage_type": "diffusion", - "runtime": { - "process": True, - "devices": "0", - "max_batch_size": 1, - }, - "engine_args": { - "parallel_config": DiffusionParallelConfig( - pipeline_parallel_size=1, - data_parallel_size=1, - tensor_parallel_size=1, - sequence_parallel_size=1, - ulysses_degree=1, - ring_degree=1, - cfg_parallel_size=1, - ), - "vae_use_slicing": kwargs.get("vae_use_slicing", False), - "vae_use_tiling": kwargs.get("vae_use_tiling", False), - "cache_backend": kwargs.get("cache_backend", "none"), - "cache_config": kwargs.get("cache_config", None), - }, - "final_output": True, - "final_output_type": "image", - } - ] - default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" - self.stage_configs = OmegaConf.create(default_stage_cfg) + def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]: + """Create default diffusion stage configuration.""" + # TODO: here is different from the Omni class. We should merge the two in the future. + devices = "0" + if "parallel_config" in kwargs: + parallel_config = kwargs["parallel_config"] + num_devices = kwargs["parallel_config"].world_size + for i in range(1, num_devices): + devices += f",{i}" else: - self.config_path = stage_configs_path - self.stage_configs = load_stage_configs_from_yaml(stage_configs_path) - - # Initialize connectors - self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors( - self.config_path, worker_backend=worker_backend, shm_threshold_bytes=shm_threshold_bytes - ) - - # Initialize stats paths - self._enable_stats: bool = bool(log_stats) - - self.worker_backend = worker_backend - self.ray_address = ray_address - self.batch_timeout = batch_timeout - - # Build OmniStage instances in parallel, preserving original order - def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]: - idx, cfg = idx_cfg - return idx, OmniStage(cfg) - - with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor: - futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)] - results: list[tuple[int, OmniStage]] = [] - for fut in as_completed(futures): - results.append(fut.result()) - results.sort(key=lambda x: x[0]) - self.stage_list = [st for _, st in results] - self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list] - - self.output_modalities = [st.final_output_type for st in self.stage_list] - logger.debug("[AsyncOrchestrator] Loaded %d stages", len(self.stage_list)) - - if self.worker_backend == "ray": - self._queue_cls = get_ray_queue_class() - else: - self._ctx = mp.get_context("spawn") - self._queue_cls = lambda: self._ctx.Queue(maxsize=0) - - self._init_sleep_seconds = max(0, int(init_sleep_seconds)) - self._shm_threshold_bytes = max(0, int(shm_threshold_bytes)) - self._start_stages(model) - # Wait for all stages to report readiness before seeding - self._wait_for_stages_ready(timeout=init_timeout) - - def _start_stages(self, model: str) -> None: - """Start all stage processes.""" - if self.worker_backend == "ray": - # Initialize Ray cluster - self._ray_pg = create_placement_group( - number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK" + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + cfg_parallel_size=1, ) + default_stage_cfg = [ + { + "stage_id": 0, + "stage_type": "diffusion", + "runtime": { + "process": True, + "devices": devices, + "max_batch_size": 1, + }, + "engine_args": { + "parallel_config": parallel_config, + "vae_use_slicing": kwargs.get("vae_use_slicing", False), + "vae_use_tiling": kwargs.get("vae_use_tiling", False), + "cache_backend": kwargs.get("cache_backend", "none"), + "cache_config": kwargs.get("cache_config", None), + }, + "final_output": True, + "final_output_type": "image", + } + ] + default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" + return default_stage_cfg - for stage_id, stage in enumerate(self.stage_list): - in_q = self._queue_cls() - out_q = self._queue_cls() - self._stage_in_queues.append(in_q) - self._stage_out_queues.append(out_q) - stage.attach_queues(in_q, out_q) - - stage_connectors_config = get_stage_connector_config( - self.omni_transfer_config, - stage_id, - ) - - stage.init_stage_worker( - model, - shm_threshold_bytes=self._shm_threshold_bytes, - ctx=self._ctx if self.worker_backend != "ray" else None, - batch_timeout=self.batch_timeout, - connectors_config=stage_connectors_config, - worker_backend=self.worker_backend, - ray_placement_group=self._ray_pg, - ) - - logger.debug("[AsyncOrchestrator] Stage-%s process started", stage_id) - time.sleep(self._init_sleep_seconds) + def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str, Any]) -> None: + # Store vllm_config received from worker process (may be None for diffusion stages) + vllm_config = result.get("vllm_config") + if vllm_config is not None: + stage.set_vllm_config(vllm_config) + tokenizer = result.get("tokenizer") + if tokenizer is not None: + stage.set_tokenizer(tokenizer) + is_tracing_enabled = result.get("is_tracing_enabled") + if is_tracing_enabled is not None: + stage.set_is_tracing_enabled(is_tracing_enabled) + super()._process_stage_ready(stage, stage_id, result) def _wait_for_stages_ready(self, timeout: int = 120) -> None: """Wait for all stages to report readiness.""" - deadline = time.time() + max(0, int(timeout)) - num_stages = len(self.stage_list) - while len(self._stages_ready) < num_stages and time.time() < deadline: - progressed = False - for stage_id, stage in enumerate(self.stage_list): - if stage_id in self._stages_ready: - continue - result = stage.try_collect() - if result is None: - continue - progressed = True - if result.get("type") == "stage_ready": - self._stages_ready.add(stage_id) - # Store vllm_config received from worker process (may be None for diffusion stages) - vllm_config = result.get("vllm_config") - if vllm_config is not None: - stage.set_vllm_config(vllm_config) - tokenizer = result.get("tokenizer") - if tokenizer is not None: - stage.set_tokenizer(tokenizer) - is_tracing_enabled = result.get("is_tracing_enabled") - if is_tracing_enabled is not None: - stage.set_is_tracing_enabled(is_tracing_enabled) - logger.info("[AsyncOrchestrator] Stage-%s reported ready", stage_id) - else: - # No user data should arrive before seeding; ignore other messages - pass - if not progressed: - time.sleep(0.01) - if len(self._stages_ready) < num_stages: - not_ready = sorted(set(range(num_stages)) - set(self._stages_ready)) + super()._wait_for_stages_ready(timeout) + for stage in self.stage_list: + if stage.vllm_config is not None and stage.tokenizer is not None: + try: + vllm_config = stage.vllm_config + tokenizer = stage.tokenizer + # Initialize input_processor + self.input_processor = OmniInputProcessor( + vllm_config=vllm_config, + tokenizer=tokenizer, + ) + # Initialize model_config + self.model_config = vllm_config.model_config + # Initialize io_processor + io_processor_plugin = self.model_config.io_processor_plugin + self.io_processor = get_io_processor(vllm_config, io_processor_plugin) + + logger.info( + f"[{self._name}] Initialized input_processor, " + f"io_processor, and model_config from stage-{stage.stage_id}", + ) + break + except Exception as e: + logger.warning( + f"[{self._name}] Failed to initialize processors from stage-{stage.stage_id}: {e}", + ) + # If no LLM stage found, set processors to None + if not hasattr(self, "input_processor") or self.input_processor is None: logger.warning( - "[AsyncOrchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s", - len(self._stages_ready), - num_stages, - not_ready, + f"[{self._name}] No LLM stage found, processors will not be available. " + "This may cause issues with OpenAIServingModels." ) - # Provide actionable suggestions before shutdown - try: - suggestions = [ - "Verify GPU/device assignment in config (runtime.devices) is correct.", - "Check GPU/host memory availability; reduce model or batch size if needed.", - "Check model weights path and network reachability (if loading remotely).", - "Increase initialization wait time (init_sleep_seconds or call-site timeout).", - ] - logger.error( - "[AsyncOrchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s", - "\n- ".join(suggestions), - ) - except Exception: - # Best-effort logging of suggestions - logger.error( - "[AsyncOrchestrator] Stage initialization failed and an error occurred while logging suggestions", - ) - elif len(self._stages_ready) == num_stages: - logger.info("[AsyncOrchestrator] All stages initialized successfully") - # Initialize input_processor, io_processor, and model_config for API server compatibility - # Find the first LLM stage (with vllm_config) to get vllm_config and tokenizer - for stage in self.stage_list: - if stage.vllm_config is not None and stage.tokenizer is not None: - try: - vllm_config = stage.vllm_config - tokenizer = stage.tokenizer - - # Initialize input_processor - self.input_processor = OmniInputProcessor( - vllm_config=vllm_config, - tokenizer=tokenizer, - ) - - # Initialize model_config - self.model_config = vllm_config.model_config - - # Initialize io_processor - io_processor_plugin = self.model_config.io_processor_plugin - self.io_processor = get_io_processor(vllm_config, io_processor_plugin) - - logger.info( - "[AsyncOrchestrator] Initialized input_processor, " - "io_processor, and model_config from stage-%s", - stage.stage_id, - ) - break - except Exception as e: - logger.warning( - "[AsyncOrchestrator] Failed to initialize processors from stage-%s: %s", - stage.stage_id, - e, - ) - # If no LLM stage found, set processors to None - if not hasattr(self, "input_processor") or self.input_processor is None: - logger.warning( - "[AsyncOrchestrator] No LLM stage found, processors will not be available. " - "This may cause issues with OpenAIServingModels." - ) - self.input_processor = None - self.io_processor = None - self.model_config = None - - def _initialize_stages( - self, - model: str, - init_sleep_seconds: int, - shm_threshold_bytes: int, - init_timeout: int, - ) -> None: - self.stage_list: list[OmniStage] = [] - - # Build OmniStage instances in parallel, preserve original order - def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]: - idx, cfg = idx_cfg - return idx, OmniStage(cfg) - - with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor: - futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)] - results: list[tuple[int, OmniStage]] = [] - for fut in as_completed(futures): - results.append(fut.result()) - results.sort(key=lambda x: x[0]) - self.stage_list = [st for _, st in results] - self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list] - self.output_modalities = [st.final_output_type for st in self.stage_list] - logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) - - if self.worker_backend == "ray": - self._queue_cls = get_ray_queue_class() - else: - self._ctx = mp.get_context("spawn") - self._queue_cls = lambda: self._ctx.Queue(maxsize=0) - - self._stage_in_queues: list[mp.Queue] = [] - self._stage_out_queues: list[mp.Queue] = [] - self._init_sleep_seconds = max(0, int(init_sleep_seconds)) - self._shm_threshold_bytes = max(0, int(shm_threshold_bytes)) - self._start_stages(model) - # Wait for all stages to report readiness before seeding - self._stages_ready: set[int] = set() - self._wait_for_stages_ready(timeout=init_timeout) - - def _start_stages(self, model: str) -> None: - if self.worker_backend == "ray": - # Initialize Ray Cluster - self._ray_pg = create_placement_group( - number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK" - ) - - for stage_id, stage in enumerate(self.stage_list): - # Use unbounded queues to avoid deadlock when seeding many requests - in_q = self._queue_cls() - out_q = self._queue_cls() - self._stage_in_queues.append(in_q) - self._stage_out_queues.append(out_q) - - # Attach queues and start Stage-owned worker process - stage.attach_queues(in_q, out_q) - - # Build connectors config for this stage - stage_connectors_config = get_stage_connector_config( - self.omni_transfer_config, - stage_id, - ) - - stage.init_stage_worker( - model, - is_async=True, - shm_threshold_bytes=self._shm_threshold_bytes, - ctx=self._ctx if self.worker_backend != "ray" else None, - batch_timeout=self.batch_timeout, - connectors_config=stage_connectors_config, - worker_backend=self.worker_backend, - ray_placement_group=self._ray_pg, - ) - - logger.debug("[Orchestrator] Stage-%s process started", stage_id) - time.sleep(self._init_sleep_seconds) + self.input_processor = None + self.io_processor = None + self.model_config = None def close(self) -> None: - """Close all stage processes and clean up resources. - - Sends shutdown signals to all stage input queues and stops - all stage worker processes. This method should be called - when done using the AsyncOmni instance. - """ - for q in self._stage_in_queues: - try: - q.put_nowait(None) - except Exception as e: - logger.warning( - "[Orchestrator] Failed to send shutdown signal to \ - stage input queue: %s", - e, - ) - for stage in self.stage_list: - try: - stage.stop_stage_worker() - except Exception as e: - logger.warning("[Orchestrator] Failed to stop stage worker: %s", e) - + """Close all stage processes and clean up resources.""" if self.output_handler is not None: self.output_handler.cancel() self.output_handler = None - - try_close_ray(self._ray_pg) - - def __del__(self) -> None: # best-effort - try: - self.close() - except Exception as e: - logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) + super().close() def shutdown(self): """Shutdown, cleaning up the background proc and IPC. @@ -523,7 +231,7 @@ class AsyncOmni: async with self._pause_cond: await self._pause_cond.wait_for(lambda: not self._paused) - logger.debug("[AsyncOrchestrator] generate() called") + logger.debug(f"[{self._name}] generate() called") # Start output handler on the first call to generate() self._run_output_handler() @@ -577,7 +285,7 @@ class AsyncOmni: _wall_start_ts, ) # Seed stage-0 queue with all requests - logger.debug("[AsyncOrchestrator] Seeding request into stage-0") + logger.debug(f"[{self._name}] Seeding request into stage-0") req_state = ClientRequestState(request_id) self.request_states[request_id] = req_state @@ -592,9 +300,9 @@ class AsyncOmni: } self.stage_list[0].submit(task) _req_start_ts[request_id] = time.time() - logger.debug("[AsyncOrchestrator] Enqueued request %s to stage-0", request_id) + logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0") - logger.debug("[AsyncOrchestrator] Entering scheduling loop: stages=%d", num_stages) + logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}") for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): result = await req_state.queue.get() assert stage_id == req_state.stage_id @@ -602,10 +310,7 @@ class AsyncOmni: req_id = result.get("request_id") if "error" in result: logger.error( - "Stage %s error on request %s: %s", - stage_id, - req_id, - result["error"], + f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}", ) raise RuntimeError(result) # Request Finished due to error @@ -618,23 +323,16 @@ class AsyncOmni: metrics.on_stage_metrics(stage_id, req_id, _m) except Exception as e: logger.exception( - "[AsyncOrchestrator] Failed to process metrics for stage %s, req %s: %s", - stage_id, - req_id, - e, + f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}", ) logger.debug( - "[AsyncOrchestrator] Stage-%s completed request %s; forwarding or finalizing", - stage_id, - req_id, + f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing", ) stage.set_engine_outputs(engine_outputs) if getattr(stage, "final_output", False): logger.debug( - "[AsyncOrchestrator] Request %s finalized at stage-%s", - req_id, - stage_id, + f"[{self._name}] Request {req_id} finalized at stage-{stage_id}", ) # End-to-end timing and time-per-token for final output @@ -650,10 +348,7 @@ class AsyncOmni: ) except Exception as e: logger.exception( - "[AsyncOrchestrator] Finalize request handling error for req %s at stage %s: %s", - req_id, - stage_id, - e, + f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}", ) if isinstance(engine_outputs, list): @@ -709,93 +404,26 @@ class AsyncOmni: # because continuing would cause the request to be silently dropped # and the orchestrator to hang waiting for completion. error_msg = ( - f"[AsyncOrchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. " + f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. " "Configure a connector for this edge or inspect connector logs for details." ) logger.error(error_msg) raise RuntimeError(error_msg) - logger.debug( - "[AsyncOrchestrator] Forwarded request %s to stage-%s", - req_id, - next_stage_id, - ) + logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}") else: - logger.debug("[AsyncOrchestrator] Request %s fully completed", req_id) + logger.debug(f"[{self._name}] Request {req_id} fully completed") - logger.debug("[AsyncOrchestrator] All requests completed") + logger.debug(f"[{self._name}] All requests completed") # Summarize and print stats try: summary = metrics.build_and_log_summary(final_stage_id_for_e2e) logger.info("[Summary] %s", pformat(summary, sort_dicts=False)) except Exception as e: - logger.exception("[AsyncOrchestrator] Failed to build/log summary: %s", e) + logger.exception(f"[{self._name}] Failed to build/log summary: {e}") finally: self.request_states.pop(request_id, None) - def _wait_for_stages_ready(self, timeout: int = 120) -> None: - num_stages = len(self.stage_list) - while len(self._stages_ready) < num_stages: - progressed = False - for stage_id, stage in enumerate(self.stage_list): - if stage_id in self._stages_ready: - continue - result = stage.try_collect() - if result is None: - continue - progressed = True - if result.get("type") == "stage_ready": - self._stages_ready.add(stage_id) - # Store vllm_config received from worker process - vllm_config = result.get("vllm_config") - if vllm_config is not None: - stage.set_vllm_config(vllm_config) - tokenizer = result.get("tokenizer") - if tokenizer is not None: - stage.set_tokenizer(tokenizer) - # input_preprocessor = result.get("input_preprocessor") - # if input_preprocessor is not None: - # stage.set_input_preprocessor(input_preprocessor) - is_tracing_enabled = result.get("is_tracing_enabled") - if is_tracing_enabled is not None: - stage.set_is_tracing_enabled(is_tracing_enabled) - logger.debug("[Orchestrator] Stage-%s reported ready", stage_id) - else: - # No user data should arrive before seeding; ignore other messages - pass - if not progressed: - time.sleep(0.01) - if len(self._stages_ready) < num_stages: - not_ready = sorted(set(range(num_stages)) - set(self._stages_ready)) - logger.warning( - "[Orchestrator] Initialization timeout: only %s/%s stages are \ - ready; not ready: %s", - len(self._stages_ready), - num_stages, - not_ready, - ) - # Provide actionable suggestions before shutdown - try: - suggestions = [ - "Verify GPU/device assignment in config (runtime.devices) is \ - correct.", - "Check GPU/host memory availability; reduce model or batch size if needed.", # noqa: E501 - "Check model weights path and network reachability (if loading remotely).", # noqa: E501 - "Increase initialization wait time (init_sleep_seconds or \ - call-site timeout).", - ] - logger.error( - "[Orchestrator] Stage initialization failed, shutting down. \ - Suggestions:\n- %s", - "\n- ".join(suggestions), - ) - except Exception: - # Best-effort logging of suggestions - logger.error( - "[Orchestrator] Stage initialization failed and an error \ - occurred while logging suggestions", - ) - def _run_output_handler(self) -> None: if self.output_handler is not None: return @@ -821,10 +449,8 @@ class AsyncOmni: req_state = request_states.get(req_id) if req_state is None: logger.debug( - "[Orchestrator] Request may have been aborted; \ - dropping output for req %s at stage-%s ", - req_id, - stage_id, + f"[{self._name}] Request may have been aborted; \ + dropping output for req {req_id} at stage-{stage_id}" ) continue await req_state.queue.put(result) @@ -854,6 +480,14 @@ class AsyncOmni: def errored(self) -> bool: return not self.is_running + @property + def _name(self) -> str: + return "AsyncOrchestrator" + + @property + def is_async(self) -> bool: + return True + @property def dead_error(self) -> BaseException: return EngineDeadError() diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index b57b063c1..fc6646984 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -54,8 +54,8 @@ def omni_snapshot_download(model_id) -> str: return _dummy_snapshot_download(model_id) -class Omni: - """Unified entrypoint for both LLM and Diffusion models for better usability. +class OmniBase: + """Base class for serving Omni models. Args: *args: Variable length argument list. @@ -75,11 +75,6 @@ class Omni: - batch_timeout: Timeout in seconds for batching requests within a stage - init_timeout: Timeout in seconds for waiting for all stages to initialize - Additional keyword arguments passed to stage engines. - - Example: - >>> omni = Omni(model="Qwen/Qwen2.5-Omni-7B") - >>> outputs = omni.generate(prompts="Hello, world!", sampling_params_list=[SamplingParams()]) - >>> print(outputs) """ def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None: @@ -106,12 +101,38 @@ class Omni: logger.info(f"Initializing stages for model: {model}") self._initialize_stages(model, kwargs) - def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None: - """Initialize stage list management. + def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]: + """Create default diffusion stage configuration.""" + # We temporally create a default config for diffusion stage. + # In the future, we should merge the default config with the user-provided config. + # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. + if "dtype" in kwargs: + kwargs["dtype"] = str(kwargs["dtype"]) + # TODO: hack, calculate devices based on parallel config. + devices = "0" + if "parallel_config" in kwargs: + num_devices = kwargs["parallel_config"].world_size + for i in range(1, num_devices): + devices += f",{i}" + default_stage_cfg = [ + { + "stage_id": 0, + "stage_type": "diffusion", + "runtime": { + "process": True, + "devices": devices, + "max_batch_size": 1, + }, + "engine_args": OmegaConf.create(kwargs), + "final_output": True, + "final_output_type": "image", + } + ] + default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" + return default_stage_cfg - Each stage will create appropriate instance (OmniLLM or OmniDiffusion) - based on stage_type in YAML config (handled in omni_stage.py). - """ + def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None: + """Initialize stage list management.""" init_sleep_seconds = kwargs.get("init_sleep_seconds", 20) shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536) init_timeout = kwargs.get("init_timeout", 300) @@ -126,31 +147,7 @@ class Omni: self.config_path = resolve_model_config_path(model) self.stage_configs = load_stage_configs_from_model(model) if not self.stage_configs: - # TODO: hack here, convert dtype to string to avoid non-premitive omegaconf create error. - if "dtype" in kwargs: - kwargs["dtype"] = str(kwargs["dtype"]) - # TODO: hack, calculate devices based on parallel config. - devices = "0" - if "parallel_config" in kwargs: - num_devices = kwargs["parallel_config"].world_size - for i in range(1, num_devices): - devices += f",{i}" - logger.info(f"model: {model}, kwargs: {kwargs}") - default_stage_cfg = [ - { - "stage_id": 0, - "stage_type": "diffusion", - "runtime": { - "process": True, - "devices": devices, - "max_batch_size": 1, - }, - "engine_args": OmegaConf.create(kwargs), - "final_output": True, - "final_output_type": "image", - } - ] - default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" + default_stage_cfg = self._create_default_diffusion_stage_cfg(kwargs) self.stage_configs = OmegaConf.create(default_stage_cfg) else: self.config_path = stage_configs_path @@ -182,7 +179,7 @@ class Omni: self.stage_list = [st for _, st in results] self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list] self.output_modalities = [st.final_output_type for st in self.stage_list] - logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) + logger.debug(f"[{self._name}] Loaded {len(self.stage_list)} stages") if self.worker_backend == "ray": self._queue_cls = get_ray_queue_class() @@ -218,6 +215,7 @@ class Omni: stage.init_stage_worker( model, + is_async=self.is_async, shm_threshold_bytes=self._shm_threshold_bytes, ctx=self._ctx if self.worker_backend != "ray" else None, batch_timeout=self.batch_timeout, @@ -226,9 +224,13 @@ class Omni: ray_placement_group=self._ray_pg, ) - logger.debug("[Orchestrator] Stage-%s process started", stage_id) + logger.debug(f"[{self._name}] Stage-{stage_id} process started") time.sleep(self._init_sleep_seconds) + def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str, Any]) -> None: + self._stages_ready.add(stage_id) + logger.info(f"[{self._name}] Stage-{stage_id} reported ready") + def _wait_for_stages_ready(self, timeout: int = 120) -> None: """Wait for all stages to report readiness.""" deadline = time.time() + max(0, int(timeout)) @@ -243,40 +245,100 @@ class Omni: continue progressed = True if result.get("type") == "stage_ready": - self._stages_ready.add(stage_id) - logger.info("[Orchestrator] Stage-%s reported ready", stage_id) - else: - # No user data should arrive before seeding; ignore other messages - pass + self._process_stage_ready(stage, stage_id, result) if not progressed: time.sleep(0.01) if len(self._stages_ready) < num_stages: not_ready = sorted(set(range(num_stages)) - set(self._stages_ready)) logger.warning( - "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s", - len(self._stages_ready), - num_stages, - not_ready, + f"[{self._name}] Initialization timeout: only {len(self._stages_ready)}/{num_stages}" + f" stages are ready; not ready: {not_ready}.", ) # Provide actionable suggestions before shutdown try: - suggestions = [ - "Verify GPU/device assignment in config (runtime.devices) is correct.", - "Check GPU/host memory availability; reduce model or batch size if needed.", - "Check model weights path and network reachability (if loading remotely).", - "Increase initialization wait time (init_sleep_seconds or call-site timeout).", - ] + suggestions = "".join( + [ + "Verify GPU/device assignment in config (runtime.devices) is correct.", + "Check GPU/host memory availability; reduce model or batch size if needed.", + "Check model weights path and network reachability (if loading remotely).", + "Increase initialization wait time (init_sleep_seconds or call-site timeout).", + ] + ) logger.error( - "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s", - "\n- ".join(suggestions), + f"[{self._name}] Stage initialization failed, shutting down. Suggestions:\n- {suggestions}", ) except Exception: # Best-effort logging of suggestions logger.error( - "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions", + f"[{self._name}] Stage initialization failed and an error occurred while logging suggestions", ) elif len(self._stages_ready) == num_stages: - logger.info("[Orchestrator] All stages initialized successfully") + logger.info(f"[{self._name}] All stages initialized successfully") + + def close(self) -> None: + """Close all stage processes and clean up resources.""" + # Close stages if they exist (for LLM models) + if self.stage_list: + for q in self._stage_in_queues: + try: + q.put_nowait(None) + except Exception as e: + logger.warning( + f"[{self._name}] Failed to send shutdown signal to stage input queue: {e}", + ) + for stage in self.stage_list: + try: + stage.stop_stage_worker() + except Exception as e: + logger.warning(f"[{self._name}] Failed to stop stage worker: {e}") + + try_close_ray(self._ray_pg) + + def __del__(self): # pragma: no cover - best effort cleanup + try: + self.close() + except Exception: + logger.debug(f"[{self._name}] __del__ close() raised", exc_info=True) + + @property + def _name(self) -> str: + return "OmniBase" + + @property + def is_async(self) -> bool: + return False + + +class Omni(OmniBase): + """Unified entrypoint for both LLM and Diffusion models for better usability. + + Args: + *args: Variable length argument list. + - args[0]: Model name or path to load. + **kwargs: Arbitrary keyword arguments. + - model: Model name or path to load (if not in args). + - stage_configs_path: Optional path to YAML file containing stage + configurations. If None, configurations are loaded from the model. + - log_stats: Whether to enable statistics logging + be written to files with stage-specific suffixes. + - init_sleep_seconds: Number of seconds to sleep between starting + each stage process during initialization + - shm_threshold_bytes: Threshold in bytes for using shared memory + for IPC. Objects larger than this threshold will use shared memory. + - worker_backend: Backend for worker processes. Default is "multi_process". + - ray_address: Address of Ray cluster for Ray backend, if using Ray backend. + - batch_timeout: Timeout in seconds for batching requests within a stage + - init_timeout: Timeout in seconds for waiting for all stages to initialize + - Additional keyword arguments passed to stage engines. + + Example: + >>> omni = Omni(model="Qwen/Qwen2.5-Omni-7B") + >>> outputs = omni.generate(prompts="Hello, world!", sampling_params_list=[SamplingParams()]) + >>> print(outputs) + """ + + def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None: + super().__init__(*args, **kwargs) def generate(self, *args: Any, **kwargs: dict[str, Any]) -> list[OmniRequestOutput]: """Generate outputs for the given prompts. @@ -336,7 +398,7 @@ class Omni: sampling_params_list: Any | Sequence[Any] | None = None, ) -> list[OmniRequestOutput]: """Run generation through all stages in the pipeline.""" - logger.debug("[Orchestrator] generate() called") + logger.debug(f"[{self._name}] generate() called") if sampling_params_list is None: raise ValueError("sampling_params_list is required for pipelined generation") @@ -388,7 +450,7 @@ class Omni: ) # Seed stage-0 queue with all requests - logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)) + logger.debug(f"[{self._name}] Seeding {len(request_prompts)} requests into stage-0") # Mark first input time for stage-0 metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() @@ -401,7 +463,7 @@ class Omni: } self.stage_list[0].submit(task) _req_start_ts[req_id] = time.time() - logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id) + logger.debug(f"[{self._name}] Enqueued request {req_id} to stage-0") # For each stage, forward results to next stage; collect finals at the end # We pipeline by continually polling output queues in stage order @@ -410,9 +472,7 @@ class Omni: total_requests = len(request_prompts) logger.debug( - "[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", - total_requests, - num_stages, + f"[{self._name}] Entering scheduling loop: total_requests={total_requests}, stages={num_stages}", ) while completed_requests < total_requests: made_progress = False @@ -425,10 +485,7 @@ class Omni: req_id = result.get("request_id") if "error" in result: logger.error( - "Stage %s error on request %s: %s", - stage_id, - req_id, - result["error"], + f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}", ) continue @@ -447,15 +504,10 @@ class Omni: metrics.on_stage_metrics(stage_id, req_id, _m) except Exception as e: logger.exception( - "[Orchestrator] Failed to process metrics for stage %s, req %s: %s", - stage_id, - req_id, - e, + f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}", ) logger.debug( - "[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", - stage_id, - req_id, + f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing", ) stage.set_engine_outputs(engine_outputs) @@ -468,9 +520,7 @@ class Omni: ) ) logger.debug( - "[Orchestrator] Request %s finalized at stage-%s", - req_id, - stage_id, + f"[{self._name}] Request {req_id} finalized at stage-{stage_id}", ) # End-to-end timing and time-per-token for final output @@ -486,10 +536,7 @@ class Omni: ) except Exception as e: logger.exception( - "[Orchestrator] Finalize request handling error for req %s at stage %s: %s", - req_id, - stage_id, - e, + f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}", ) next_stage_id = stage_id + 1 @@ -499,10 +546,8 @@ class Omni: next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) except Exception as e: logger.exception( - "[Orchestrator] Process engine inputs error for req %s at stage %s: %s", - req_id, - next_stage_id, - e, + f"[{self._name}] Process engine inputs error for req {req_id}" + f" at stage {next_stage_id}: {e}", ) continue sp_next = sampling_params_list[next_stage_id] # type: ignore[index] @@ -526,59 +571,32 @@ class Omni: if not sent_via_connector: raise RuntimeError( - f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. " + f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. " "Configure a connector for this edge or inspect connector logs for details." ) logger.debug( - "[Orchestrator] Forwarded request %s to stage-%s", - req_id, - next_stage_id, + f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}", ) remaining_by_stage[next_stage_id] += 1 else: completed_requests += 1 logger.debug( - "[Orchestrator] Request %s fully completed (%d/%d)", - req_id, - completed_requests, - total_requests, + f"[{self._name}] Request {req_id} fully completed ({completed_requests}/{total_requests})", ) if not made_progress: time.sleep(0.005) - logger.debug("[Orchestrator] All requests completed") + logger.debug(f"[{self._name}] All requests completed") # Summarize and print stats try: summary = metrics.build_and_log_summary(final_stage_id_to_prompt) logger.info("[Summary] %s", pformat(summary, sort_dicts=False)) except Exception as e: - logger.exception("[Orchestrator] Failed to build/log summary: %s", e) + logger.exception(f"[{self._name}] Failed to build/log summary: {e}") return final_outputs - def close(self) -> None: - """Close all stage processes and clean up resources.""" - # Close stages if they exist (for LLM models) - if self.stage_list: - for q in self._stage_in_queues: - try: - q.put_nowait(None) - except Exception as e: - logger.warning( - "[Orchestrator] Failed to send shutdown signal to stage input queue: %s", - e, - ) - for stage in self.stage_list: - try: - stage.stop_stage_worker() - except Exception as e: - logger.warning("[Orchestrator] Failed to stop stage worker: %s", e) - - try_close_ray(self._ray_pg) - - def __del__(self): # pragma: no cover - best effort cleanup - try: - self.close() - except Exception: - logger.debug("[Orchestrator] __del__ close() raised", exc_info=True) + @property + def _name(self) -> str: + return "Orchestrator"