[Docs] optimize and supplement docs system (#86)

Signed-off-by: Gaohan123 <gaohan19@huawei.com>

vllm_omni/entrypoints/async_omni_llm.py

@@ -53,6 +53,38 @@ logger = init_logger(__name__)
 
 
 class AsyncOmniLLM(EngineClient):
+    """Async entry point for vLLM-omni inference.
+
+    This class provides an asynchronous interface for running multi-modal
+    comprehension and generation models. It orchestrates multiple
+    stages in a pipeline, where each stage runs in a separate process.
+    Designed for use with async/await patterns and streaming generation.
+
+    Args:
+        model: Model name or path to load
+        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
+        log_file: Optional path prefix for log files. If provided, logs will
+            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.
+        batch_timeout: Timeout in seconds for batching requests within a stage
+        init_timeout: Timeout in seconds for waiting for all stages to initialize
+        **kwargs: Additional keyword arguments passed to stage engines
+
+    Example:
+        >>> async_llm = AsyncOmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+        >>> async for output in async_llm.generate(
+        ...     prompt="Hello",
+        ...     request_id="req-1",
+        ...     sampling_params_list=[SamplingParams(), SamplingParams()]
+        ... ):
+        ...     print(output)
+    """
+
     def __init__(
         self,
         model: str,
@@ -63,7 +95,7 @@ class AsyncOmniLLM(EngineClient):
         shm_threshold_bytes: int = 65536,
         batch_timeout: int = 10,
         init_timeout: int = 300,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
@@ -140,6 +172,12 @@ class AsyncOmniLLM(EngineClient):
             time.sleep(self._init_sleep_seconds)
 
     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 AsyncOmniLLM instance.
+        """
         for q in self._stage_in_queues:
             try:
                 q.put_nowait(None)
@@ -162,7 +200,11 @@ class AsyncOmniLLM(EngineClient):
             logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
 
     def shutdown(self):
-        """Shutdown, cleaning up the background proc and IPC."""
+        """Shutdown, cleaning up the background proc and IPC.
+
+        Alias for close() method. Cleans up all stage processes
+        and inter-process communication resources.
+        """
         try:
             self.close()
         except Exception as e:
@@ -178,6 +220,33 @@ class AsyncOmniLLM(EngineClient):
         priority: int = 0,
         data_parallel_rank: Optional[int] = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
+        """Generate outputs for the given prompt asynchronously.
+
+        Processes the prompt through all stages in the pipeline and yields
+        outputs as they become available. Each stage uses its corresponding
+        sampling parameters from the sampling_params_list.
+
+        Args:
+            prompt: Prompt to process. Can be a text string, token IDs,
+                or multimodal prompt.
+            request_id: Unique identifier for this request
+            sampling_params_list: List of SamplingParams, one for each stage.
+                Must have the same length as the number of stages.
+                If None, uses default sampling params for each stage.
+            lora_request: Optional LoRA adapter request for this generation
+            trace_headers: Optional tracing headers for observability
+            priority: Request priority (higher values processed first)
+            data_parallel_rank: Optional data parallel rank for distributed
+                inference
+
+        Yields:
+            OmniRequestOutput objects as they are produced by each stage.
+            Each output contains the stage_id, final_output_type, and
+            the request_output from that stage.
+
+        Raises:
+            ValueError: If sampling_params_list has incorrect length.
+        """
         logger.debug("[Orchestrator] generate() called")
         if sampling_params_list is None:
             sampling_params_list = self.default_sampling_params_list
@@ -510,9 +579,6 @@ class AsyncOmniLLM(EngineClient):
         return None
 
     async def get_input_preprocessor(self) -> InputPreprocessor:
-        # for stage in self.stage_list:
-        #     if stage.is_comprehension:
-        #         return stage.input_preprocessor
         return None
 
     async def get_tokenizer(self) -> AnyTokenizer:
@@ -569,6 +635,31 @@ class AsyncOmniLLM(EngineClient):
 
 
 class AsyncOmniStageLLM(AsyncLLM):
+    """Async single-stage LLM engine for use within a stage worker process.
+
+    This class extends the base vLLM AsyncLLM class with omni-specific
+    processors for handling multimodal inputs and outputs. It is used
+    internally by AsyncOmniStage workers and should not be instantiated
+    directly by users.
+
+    Args:
+        engine_args: AsyncOmniEngineArgs containing engine configuration
+        vllm_config: Global vLLM configuration
+        executor_class: Executor implementation class, e.g. MultiprocExecutor
+        log_stats: Whether to log statistics
+        usage_context: Usage context of the LLM (default: ENGINE_CONTEXT)
+        mm_registry: Multi-modal registry for processing multimodal inputs
+        use_cached_outputs: Whether to use cached outputs
+        log_requests: Whether to log requests
+        start_engine_loop: Whether to start the engine loop automatically
+        stat_loggers: Customized stat loggers for the engine.
+            If not provided, default stat loggers will be used.
+            Note: Stat logger interface may change in V1.
+        client_addresses: Optional dictionary mapping client names to addresses
+        client_count: Total number of clients (default: 1)
+        client_index: Index of this client (default: 0)
+    """
+
     def __init__(
         self,
         engine_args: AsyncOmniEngineArgs,

vllm_omni/entrypoints/omni_llm.py

@@ -42,11 +42,36 @@ logger = init_logger(__name__)
 
 
 class OmniLLM:
-    """Multi-process pipelined OmniLLM.
+    """Main entry point for vLLM-omni inference.
 
-    - Per-stage process with copy-based IPC (Queues)
-    - window=-1 across stages (downstream starts when upstream finishes)
-    - max_inflight=1 per stage (serial within a stage), but pipeline across stages
+    This class provides a high-level interface for running multi-modal
+    comprehension and generation models. It orchestrates multiple
+    stages in a pipeline, where each stage runs in a separate process
+    with copy-based IPC (Queues). Downstream stages start when upstream
+    stages finish (window=-1), and each stage processes requests serially
+    (max_inflight=1) but pipelines across stages.
+
+    Args:
+        model: Model name or path to load
+        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
+        log_file: Optional path prefix for log files. If provided, logs will
+            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.
+        batch_timeout: Timeout in seconds for batching requests within a stage
+        init_timeout: Timeout in seconds for waiting for all stages to initialize
+        **kwargs: Additional keyword arguments passed to stage engines
+
+    Example:
+        >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+        >>> outputs = llm.generate(
+        ...     prompts="Hello",
+        ...     sampling_params_list=[SamplingParams(), SamplingParams()]
+        ... )
     """
 
     def __init__(
@@ -59,7 +84,7 @@ class OmniLLM:
         shm_threshold_bytes: int = 65536,
         batch_timeout: int = 10,
         init_timeout: int = 300,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
@@ -132,6 +157,12 @@ class OmniLLM:
             time.sleep(self._init_sleep_seconds)
 
     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 OmniLLM instance.
+        """
         for q in self._stage_in_queues:
             try:
                 q.put_nowait(None)
@@ -157,6 +188,27 @@ class OmniLLM:
         prompts: Union[PromptType, Sequence[PromptType]],
         sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
     ) -> list[OmniRequestOutput]:
+        """Generate outputs for the given prompts.
+
+        Processes prompts through all stages in the pipeline and returns
+        the final outputs. Each stage uses its corresponding sampling
+        parameters from the sampling_params_list.
+
+        Args:
+            prompts: Single prompt or sequence of prompts to process.
+                Can be text strings, token IDs, or multimodal prompts.
+            sampling_params_list: List of SamplingParams, one for each stage.
+                Must have the same length as the number of stages.
+                Required for pipelined generation.
+
+        Returns:
+            List of OmniRequestOutput objects, one for each input prompt.
+            Each output contains the stage_id, final_output_type, and
+            the request_output from the final stage.
+
+        Raises:
+            ValueError: If sampling_params_list is None or has incorrect length.
+        """
         try:
             return self._run_generation(prompts, sampling_params_list)
         except Exception as e:
@@ -465,13 +517,31 @@ class OmniLLM:
 
 
 class OmniStageLLM(LLM):
+    """Single-stage LLM engine for use within a stage worker process.
+
+    This class extends the base vLLM LLM class with omni-specific
+    processors for handling multimodal inputs and outputs. It is used
+    internally by OmniStage workers and should not be instantiated directly
+    by users.
+
+    Args:
+        model: Model name or path to load
+        compilation_config: Optional compilation configuration. Can be an
+            integer (compilation level), dict, or CompilationConfig instance.
+        hf_overrides: Optional HuggingFace model configuration overrides
+        structured_outputs_config: Optional structured outputs configuration.
+            Can be a dict or StructuredOutputsConfig instance.
+        **kwargs: Additional keyword arguments passed to the base LLM class
+            and engine
+    """
+
     def __init__(
         self,
         model: str,
         compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
         hf_overrides: Optional[HfOverrides] = None,
         structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """LLM constructor."""
         if "disable_log_stats" not in kwargs:

vllm_omni/entrypoints/omni_stage.py

@@ -39,7 +39,18 @@ logger = init_logger(__name__)
 
 
 class OmniStage:
-    def __init__(self, stage_config):
+    """Stage manager for orchestrating a single stage in the omni pipeline.
+
+    Encapsulates per-stage process lifecycle and worker logic, including
+    device setup, LLM initialization, batching, and shared-memory IPC.
+    Preserves input processing utilities for cross-stage data wiring.
+
+    Args:
+        stage_config: Stage configuration object containing engine arguments,
+            runtime settings, and stage-specific parameters
+    """
+
+    def __init__(self, stage_config: Any):
         self.stage_config = stage_config
         self.engine = None
         self.async_engine = None
@@ -76,35 +87,69 @@ class OmniStage:
         self._logger = logging.getLogger(__name__)
 
     def set_engine(self, engine: LLMEngine) -> None:
-        """Initialize the engine for the stage."""
+        """Set the LLM engine for this stage.
+
+        Args:
+            engine: LLMEngine instance to use for this stage
+        """
         self.engine = engine
 
     def set_async_engine(self, async_engine: AsyncLLM) -> None:
-        """Initialize the async engine for the stage."""
+        """Set the async LLM engine for this stage.
+
+        Args:
+            async_engine: AsyncLLM instance to use for this stage
+        """
         self.async_engine = async_engine
 
-    def set_vllm_config(self, vllm_config) -> None:
-        """Set the vllm_config for the stage (received from worker process)."""
+    def set_vllm_config(self, vllm_config: Any) -> None:
+        """Set the vLLM configuration for this stage.
+
+        Args:
+            vllm_config: VllmConfig instance received from worker process
+        """
         self.vllm_config = vllm_config
 
     def set_tokenizer(self, tokenizer: AnyTokenizer) -> None:
-        """Set the tokenizer for the stage (received from worker process)."""
+        """Set the tokenizer for this stage.
+
+        Args:
+            tokenizer: Tokenizer instance received from worker process
+        """
         self.tokenizer = tokenizer
 
     def set_input_preprocessor(self, input_preprocessor: InputPreprocessor) -> None:
-        """Set the input preprocessor for the stage (received from worker process)."""
+        """Set the input preprocessor for this stage.
+
+        Args:
+            input_preprocessor: InputPreprocessor instance received from worker process
+        """
         self.input_preprocessor = input_preprocessor
 
     def set_is_tracing_enabled(self, is_tracing_enabled: bool) -> None:
-        """Set the is_tracing_enabled for the stage (received from worker process)."""
+        """Set whether tracing is enabled for this stage.
+
+        Args:
+            is_tracing_enabled: Boolean indicating if tracing is enabled
+        """
         self.is_tracing_enabled = is_tracing_enabled
 
     def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
-        """Set the engine output for the stage."""
+        """Set the engine outputs for this stage.
+
+        Args:
+            engine_outputs: EngineCoreOutput from this stage's processing
+        """
         self.engine_outputs = engine_outputs
 
     # ----------------- New Orchestration APIs -----------------
     def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None:
+        """Attach input and output queues for IPC communication.
+
+        Args:
+            in_q: Input queue for receiving tasks from orchestrator
+            out_q: Output queue for sending results to orchestrator
+        """
         self._in_q = in_q
         self._out_q = out_q
 
@@ -118,6 +163,22 @@ class OmniStage:
         ctx: Optional[mp.context.BaseContext] = None,
         batch_timeout: int = 10,
     ) -> None:
+        """Initialize and start the stage worker process.
+
+        Creates a worker process that runs the LLM engine for this stage.
+        The worker handles batching, generation, and IPC communication.
+
+        Args:
+            model: Model name or path to load
+            is_async: Whether to use async engine (default: False)
+            log_file: Optional log file path prefix for stage-specific logs
+            shm_threshold_bytes: Threshold for using shared memory for IPC
+            ctx: Optional multiprocessing context (default: spawn)
+            batch_timeout: Timeout in seconds for batching requests
+
+        Raises:
+            AssertionError: If queues are not attached before calling this method
+        """
         assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
         self._log_file = log_file
         self._shm_threshold_bytes = shm_threshold_bytes
@@ -156,6 +217,11 @@ class OmniStage:
         self._proc.start()
 
     def stop_stage_worker(self) -> None:
+        """Stop the stage worker process gracefully.
+
+        Sends shutdown signal to the worker and waits for it to terminate.
+        If graceful shutdown fails, forcefully terminates the process.
+        """
         if self._in_q is not None:
             try:
                 self._in_q.put_nowait(None)
@@ -173,10 +239,22 @@ class OmniStage:
                     self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
 
     def submit(self, payload: dict[str, Any]) -> None:
+        """Submit a task to the stage worker.
+
+        Args:
+            payload: Dictionary containing task data (request_id, engine_inputs,
+                sampling_params, etc.)
+        """
         assert self._in_q is not None
         self._in_q.put(payload)
 
     def try_collect(self) -> Optional[dict[str, Any]]:
+        """Try to collect a result from the stage worker without blocking.
+
+        Returns:
+            Result dictionary if available, None otherwise. Result contains
+            request_id, engine_outputs (or engine_outputs_shm), and metrics.
+        """
         assert self._out_q is not None
         try:
             return self._out_q.get_nowait()
@@ -184,9 +262,24 @@ class OmniStage:
             return None
 
     def process_engine_inputs(
-        self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
+        self, stage_list: list[Any], prompt: Union[OmniTokensPrompt, TextPrompt] = None
     ) -> list[Union[OmniTokensPrompt, TextPrompt]]:
-        """Process the engine input for the stage."""
+        """Process engine inputs for this stage from upstream stage outputs.
+
+        Derives inputs for this stage from outputs of upstream stages.
+        Uses engine_input_source configuration to determine which upstream
+        stage outputs to use. Supports custom processing functions.
+
+        Args:
+            stage_list: List of all stages in the pipeline
+            prompt: Optional original prompt (for multimodal data preservation)
+
+        Returns:
+            List of processed engine inputs ready for this stage
+
+        Raises:
+            ValueError: If engine_input_source is empty or invalid
+        """
         if self.custom_process_input_func is None:
             engine_inputs = []
             if len(self.engine_input_source) == 0:

vllm_omni/entrypoints/openai/api_server.py

@@ -114,6 +114,21 @@ async def build_async_omni_llm(
     disable_frontend_multiprocessing: Optional[bool] = None,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
+    """Build an AsyncOmniLLM instance from command-line arguments.
+
+    Creates an async context manager that yields an AsyncOmniLLM instance
+    configured from the provided arguments. Handles forkserver setup if
+    needed and ensures proper cleanup on exit.
+
+    Args:
+        args: Parsed command-line arguments containing model and configuration
+        disable_frontend_multiprocessing: Optional flag to disable frontend
+            multiprocessing (deprecated in V1)
+        client_config: Optional client configuration dictionary
+
+    Yields:
+        EngineClient instance (AsyncOmniLLM) ready for use
+    """
     if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
         # The executor is expected to be mp.
         # Pre-import heavy modules in the forkserver process
@@ -139,12 +154,23 @@ async def build_async_omni_llm_from_stage_config(
     disable_frontend_multiprocessing: bool = False,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """
-    Create AsyncOmniLLM, either:
-        - in-process using the AsyncOmniLLM Directly
-        - multiprocess using AsyncOmniLLM RPC
+    """Create AsyncOmniLLM from stage configuration.
 
-    Returns the AsyncOmniLLM or None if the creation failed.
+    Creates an AsyncOmniLLM instance either in-process or using multiprocess
+    RPC. Loads stage configurations from the model or from a specified path.
+
+    Args:
+        args: Parsed command-line arguments containing model and stage configs
+        disable_frontend_multiprocessing: Flag to disable frontend multiprocessing
+            (deprecated in V1)
+        client_config: Optional client configuration dictionary
+
+    Yields:
+        EngineClient instance (AsyncOmniLLM) ready for use
+
+    Note:
+        Stage configurations are loaded from args.stage_configs_path if provided,
+        otherwise from the model's default configuration.
     """
 
     # V1 AsyncLLM.
@@ -179,6 +205,17 @@ async def omni_init_app_state(
     state: State,
     args: Namespace,
 ) -> None:
+    """Initialize the FastAPI application state for omni API server.
+
+    Sets up the application state with model information, request logger,
+    and other server configuration needed for handling API requests.
+
+    Args:
+        engine_client: Engine client instance (AsyncOmniLLM)
+        vllm_config: vLLM configuration object
+        state: FastAPI application state object to initialize
+        args: Parsed command-line arguments
+    """
     if args.served_model_name is not None:
         served_model_names = args.served_model_name
     else:

vllm_omni/entrypoints/utils.py

@@ -13,8 +13,22 @@ logger = logging.getLogger(__name__)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 
 
-def load_stage_configs_from_model(model: str):
-    """Load stage configs from model."""
+def load_stage_configs_from_model(model: str) -> list:
+    """Load stage configurations from model's default config file.
+
+    Loads stage configurations based on the model type. Looks for a
+    YAML configuration file in the stage_configs directory matching
+    the model's model_type.
+
+    Args:
+        model: Model name or path (used to determine model_type)
+
+    Returns:
+        List of stage configuration dictionaries
+
+    Raises:
+        FileNotFoundError: If no stage config file exists for the model type
+    """
     hf_config = get_config(model, trust_remote_code=True)
     model_type = hf_config.model_type
     stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
@@ -25,7 +39,14 @@ def load_stage_configs_from_model(model: str):
     return stage_configs
 
 
-def load_stage_configs_from_yaml(config_path: str):
-    """Load stage configs from yaml file."""
+def load_stage_configs_from_yaml(config_path: str) -> list:
+    """Load stage configurations from a YAML file.
+
+    Args:
+        config_path: Path to the YAML configuration file
+
+    Returns:
+        List of stage configuration dictionaries from the file's stage_args
+    """
     config_data = OmegaConf.load(config_path)
     return config_data.stage_args