vllm-omni/vllm_omni/entrypoints/openai/api_server.py

import json
import multiprocessing
import multiprocessing.forkserver as forkserver
import os

# Image generation API imports
import time
from argparse import Namespace
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import fields
from http import HTTPStatus
from typing import Any

import vllm.envs as envs
from fastapi import Depends, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.datastructures import State
from vllm.config import VllmConfig
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import load_chat_template, resolve_hf_chat_template, resolve_mistral_chat_template
from vllm.entrypoints.launcher import serve_http
from vllm.entrypoints.logger import RequestLogger
from vllm.entrypoints.openai.api_server import (
    base,
    build_app,
    load_log_config,
    maybe_register_tokenizer_info_endpoint,
    router,
    setup_server,
    validate_json_request,
)
from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse
from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels
from vllm.entrypoints.openai.tool_parsers import ToolParserManager

# yapf conflicts with isort for this block
# yapf: disable
# yapf: enable
from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
from vllm.entrypoints.utils import load_aware_call, with_cancellation
from vllm.logger import init_logger
from vllm.tokenizers import MistralTokenizer
from vllm.utils.system_utils import decorate_logs

from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
from vllm_omni.entrypoints.async_diffusion import AsyncOmniDiffusion
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.openai.image_api_utils import (
    encode_image_base64,
    parse_size,
)
from vllm_omni.entrypoints.openai.protocol.images import (
    ImageData,
    ImageGenerationRequest,
    ImageGenerationResponse,
)
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat

logger = init_logger(__name__)


# Server entry points


async def omni_run_server(args, **uvicorn_kwargs) -> None:
    """Run a single-worker API server.

    Automatically detects if the model is a diffusion model and routes
    to the appropriate server implementation.
    """

    # Add process-specific prefix to stdout and stderr.
    decorate_logs("APIServer")

    listen_address, sock = setup_server(args)

    # Check if model is a diffusion model
    if is_diffusion_model(args.model):
        logger.info("Detected diffusion model, starting diffusion API server")
        await omni_run_diffusion_server_worker(listen_address, sock, args, **uvicorn_kwargs)
    else:
        await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)


async def omni_run_diffusion_server(args, **uvicorn_kwargs) -> None:
    """Run a diffusion model API server."""

    # Add process-specific prefix to stdout and stderr.
    decorate_logs("DiffusionAPIServer")

    listen_address, sock = setup_server(args)
    await omni_run_diffusion_server_worker(listen_address, sock, args, **uvicorn_kwargs)


async def omni_run_diffusion_server_worker(listen_address, sock, args, **uvicorn_kwargs) -> None:
    """Run a diffusion model API server worker."""

    # Load logging config for uvicorn if specified
    log_config = load_log_config(args.log_config_file)
    if log_config is not None:
        uvicorn_kwargs["log_config"] = log_config

    async with build_async_diffusion(args) as diffusion_engine:
        app = build_app(args)

        await omni_diffusion_init_app_state(diffusion_engine, app.state, args)

        logger.info("Starting vLLM Diffusion API server on %s", listen_address)

        shutdown_task = await serve_http(
            app,
            sock=sock,
            enable_ssl_refresh=getattr(args, "enable_ssl_refresh", False),
            host=args.host,
            port=args.port,
            log_level=args.uvicorn_log_level,
            access_log=not getattr(args, "disable_uvicorn_access_log", False),
            timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
            ssl_keyfile=getattr(args, "ssl_keyfile", None),
            ssl_certfile=getattr(args, "ssl_certfile", None),
            ssl_ca_certs=getattr(args, "ssl_ca_certs", None),
            ssl_cert_reqs=getattr(args, "ssl_cert_reqs", 0),
            h11_max_incomplete_event_size=getattr(args, "h11_max_incomplete_event_size", None),
            h11_max_header_count=getattr(args, "h11_max_header_count", None),
            **uvicorn_kwargs,
        )

    # NB: Await server shutdown only after the backend context is exited
    try:
        await shutdown_task
    finally:
        sock.close()


async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None:
    """Run a single API server worker."""

    if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
        ToolParserManager.import_tool_parser(args.tool_parser_plugin)

    # Load logging config for uvicorn if specified
    log_config = load_log_config(args.log_config_file)
    if log_config is not None:
        uvicorn_kwargs["log_config"] = log_config

    async with build_async_omni(
        args,
        client_config=client_config,
    ) as engine_client:
        maybe_register_tokenizer_info_endpoint(args)
        app = build_app(args)

        vllm_config = await engine_client.get_vllm_config()
        await omni_init_app_state(engine_client, vllm_config, app.state, args)

        logger.info(
            "Starting vLLM API server %d on %s",
            vllm_config.parallel_config._api_process_rank,
            listen_address,
        )
        shutdown_task = await serve_http(
            app,
            sock=sock,
            enable_ssl_refresh=args.enable_ssl_refresh,
            host=args.host,
            port=args.port,
            log_level=args.uvicorn_log_level,
            # NOTE: When the 'disable_uvicorn_access_log' value is True,
            # no access log will be output.
            access_log=not args.disable_uvicorn_access_log,
            timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
            ssl_keyfile=args.ssl_keyfile,
            ssl_certfile=args.ssl_certfile,
            ssl_ca_certs=args.ssl_ca_certs,
            ssl_cert_reqs=args.ssl_cert_reqs,
            h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
            h11_max_header_count=args.h11_max_header_count,
            **uvicorn_kwargs,
        )

    # NB: Await server shutdown only after the backend context is exited
    try:
        await shutdown_task
    finally:
        sock.close()


@asynccontextmanager
async def build_async_omni(
    args: Namespace,
    *,
    disable_frontend_multiprocessing: bool | None = None,
    client_config: dict[str, Any] | None = None,
) -> AsyncIterator[EngineClient]:
    """Build an AsyncOmni instance from command-line arguments.

    Creates an async context manager that yields an AsyncOmni 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 (AsyncOmni) 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
        logger.debug("Setup forkserver with pre-imports")
        multiprocessing.set_start_method("forkserver")
        multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
        forkserver.ensure_running()
        logger.debug("Forkserver setup complete!")

    # Context manager to handle async_omni lifecycle
    # Ensures everything is shutdown and cleaned up on error/exit
    async with build_async_omni_from_stage_config(
        args,
        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
    ) as async_omni:
        yield async_omni


@asynccontextmanager
async def build_async_diffusion(
    args: Namespace,
    **kwargs: Any,
) -> AsyncIterator[AsyncOmniDiffusion]:
    """Build an AsyncOmniDiffusion instance from command-line arguments.

    Creates an async context manager that yields an AsyncOmniDiffusion
    instance configured from the provided arguments.

    Args:
        args: Parsed command-line arguments containing model and configuration
        **kwargs: Additional keyword arguments passed to AsyncOmniDiffusion

    Yields:
        AsyncOmniDiffusion instance ready for use
    """
    diffusion_engine: AsyncOmniDiffusion | None = None

    try:
        # Build diffusion kwargs by extracting matching OmniDiffusionConfig fields from args
        config_field_names = {f.name for f in fields(OmniDiffusionConfig)}
        diffusion_kwargs: dict[str, Any] = {"model": args.model}

        # Diffusion parallelism configuration (e.g. `--usp 2`).
        parallel_config_kwargs: dict[str, Any] = {}
        for field in fields(DiffusionParallelConfig):
            if not hasattr(args, field.name):
                continue
            value = getattr(args, field.name)
            if value is None:
                continue
            parallel_config_kwargs[field.name] = value
        if parallel_config_kwargs:
            diffusion_kwargs["parallel_config"] = DiffusionParallelConfig(**parallel_config_kwargs)

        for field_name in config_field_names:
            if not hasattr(args, field_name):
                continue
            value = getattr(args, field_name)
            if value is None:
                continue
            # Special handling for cache_config JSON string
            if field_name == "cache_config" and isinstance(value, str):
                try:
                    value = json.loads(value)
                except json.JSONDecodeError as e:
                    logger.warning(f"Failed to parse cache_config JSON: {e}")
                    continue
            diffusion_kwargs[field_name] = value

        diffusion_kwargs.update(kwargs)
        logger.info(f"diffusion_kwargs: {diffusion_kwargs}")
        logger.info(
            "Building AsyncOmniDiffusion with model=%s, num_gpus=%s",
            args.model,
            diffusion_kwargs.get("num_gpus", 1),
        )
        diffusion_engine = AsyncOmniDiffusion(**diffusion_kwargs)

        yield diffusion_engine
    finally:
        if diffusion_engine:
            diffusion_engine.shutdown()


@asynccontextmanager
async def build_async_omni_from_stage_config(
    args: Namespace,
    *,
    disable_frontend_multiprocessing: bool = False,
    client_config: dict[str, Any] | None = None,
) -> AsyncIterator[EngineClient]:
    """Create AsyncOmni from stage configuration.

    Creates an AsyncOmni 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 (AsyncOmni) ready for use

    Note:
        Stage configurations are loaded from args.stage_configs_path if provided,
        otherwise from the model's default configuration.
    """

    # V1 AsyncLLM.
    if disable_frontend_multiprocessing:
        logger.warning("V1 is enabled, but got --disable-frontend-multiprocessing.")

    async_omni: EngineClient | None = None

    try:
        async_omni = AsyncOmni(model=args.model, cli_args=args)

        # # Don't keep the dummy data in memory
        # await async_llm.reset_mm_cache()

        yield async_omni
    finally:
        if async_omni:
            async_omni.shutdown()


async def omni_init_app_state(
    engine_client: EngineClient,
    vllm_config: VllmConfig,
    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 (AsyncOmni)
        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:
        served_model_names = [args.model]

    if args.enable_log_requests:
        request_logger = RequestLogger(max_log_len=args.max_log_len)
    else:
        request_logger = None

    base_model_paths = [BaseModelPath(name=name, model_path=args.model) for name in served_model_names]
    state.engine_client = engine_client
    state.log_stats = not args.disable_log_stats
    state.vllm_config = vllm_config
    _model_config = vllm_config.model_config
    state.log_stats = not args.disable_log_stats

    # For omni models
    state.stage_configs = engine_client.stage_configs

    resolved_chat_template = load_chat_template(args.chat_template)
    if resolved_chat_template is not None:
        # Get the tokenizer to check official template
        tokenizer = await engine_client.get_tokenizer()

        if isinstance(tokenizer, MistralTokenizer):
            # The warning is logged in resolve_mistral_chat_template.
            resolved_chat_template = resolve_mistral_chat_template(chat_template=resolved_chat_template)
        else:
            hf_chat_template = resolve_hf_chat_template(
                tokenizer=tokenizer,
                chat_template=None,
                tools=None,
                model_config=vllm_config.model_config,
            )

            if hf_chat_template != resolved_chat_template:
                logger.warning(
                    "Using supplied chat template: %s\nIt is different from official chat template '%s'. This discrepancy may lead to performance degradation.",  # noqa: E501
                    resolved_chat_template,
                    args.model,
                )

    if args.tool_server == "demo":
        tool_server: ToolServer | None = DemoToolServer()
        assert isinstance(tool_server, DemoToolServer)
        await tool_server.init_and_validate()
    elif args.tool_server:
        tool_server = MCPToolServer()
        await tool_server.add_tool_server(args.tool_server)
    else:
        tool_server = None

    # Merge default_mm_loras into the static lora_modules
    default_mm_loras = vllm_config.lora_config.default_mm_loras if vllm_config.lora_config is not None else {}

    lora_modules = args.lora_modules
    if default_mm_loras:
        default_mm_lora_paths = [
            LoRAModulePath(
                name=modality,
                path=lora_path,
            )
            for modality, lora_path in default_mm_loras.items()
        ]
        if args.lora_modules is None:
            lora_modules = default_mm_lora_paths
        else:
            lora_modules += default_mm_lora_paths

    state.openai_serving_models = OpenAIServingModels(
        engine_client=engine_client,
        base_model_paths=base_model_paths,
        lora_modules=lora_modules,
    )
    await state.openai_serving_models.init_static_loras()
    state.openai_serving_chat = OmniOpenAIServingChat(
        engine_client,
        state.openai_serving_models,
        args.response_role,
        request_logger=request_logger,
        chat_template=resolved_chat_template,
        chat_template_content_format=args.chat_template_content_format,
        trust_request_chat_template=args.trust_request_chat_template,
        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
        enable_auto_tools=args.enable_auto_tool_choice,
        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
        tool_parser=args.tool_call_parser,
        reasoning_parser=args.structured_outputs_config.reasoning_parser,
        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
        enable_force_include_usage=args.enable_force_include_usage,
        enable_log_outputs=args.enable_log_outputs,
        log_error_stack=args.log_error_stack,
    )

    state.enable_server_load_tracking = args.enable_server_load_tracking
    state.server_load_metrics = 0


def Omnichat(request: Request) -> OmniOpenAIServingChat | None:
    return request.app.state.openai_serving_chat


async def omni_diffusion_init_app_state(
    diffusion_engine: AsyncOmniDiffusion,
    state: State,
    args: Namespace,
) -> None:
    """Initialize the FastAPI application state for diffusion model API server.

    Sets up the application state with diffusion model information and
    chat completion handler for image generation via /v1/chat/completions.

    Args:
        diffusion_engine: AsyncOmniDiffusion engine instance
        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:
        served_model_names = [args.model]

    model_name = served_model_names[0] if served_model_names else args.model

    state.diffusion_engine = diffusion_engine
    state.diffusion_model_name = model_name  # Store for image endpoints
    state.log_stats = not getattr(args, "disable_log_stats", False)

    # Initialize chat handler with diffusion engine (uses /v1/chat/completions endpoint)
    # Note: Request-level parameters (num_inference_steps, guidance_scale, seed, height, width, etc.)
    # are passed per-request via the API, not as server defaults
    state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion(
        diffusion_engine=diffusion_engine,
        model_name=model_name,
    )

    # Set other handlers to None for diffusion-only mode
    state.engine_client = None
    state.vllm_config = None

    state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False)
    state.server_load_metrics = 0

    logger.info("Diffusion API server initialized for model: %s", model_name)


@router.post(
    "/v1/chat/completions",
    dependencies=[Depends(validate_json_request)],
    responses={
        HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
    },
)
@with_cancellation
@load_aware_call
async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
    handler = Omnichat(raw_request)
    if handler is None:
        return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
    try:
        generator = await handler.create_chat_completion(request, raw_request)
    except Exception as e:
        logger.exception("Chat completion failed: %s", e)
        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e

    if isinstance(generator, ErrorResponse):
        return JSONResponse(
            content=generator.model_dump(), status_code=generator.code if hasattr(generator, "code") else 400
        )

    elif isinstance(generator, ChatCompletionResponse):
        return JSONResponse(content=generator.model_dump())

    return StreamingResponse(content=generator, media_type="text/event-stream")


# Image generation API endpoints


@router.post(
    "/v1/images/generations",
    dependencies=[Depends(validate_json_request)],
    responses={
        HTTPStatus.OK.value: {"model": ImageGenerationResponse},
        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
        HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
    },
)
async def generate_images(request: ImageGenerationRequest, raw_request: Request) -> ImageGenerationResponse:
    """Generate images from text prompts using diffusion models.

    OpenAI DALL-E compatible endpoint for text-to-image generation.

    Args:
        request: Image generation request with prompt and parameters
        raw_request: Raw FastAPI request for accessing app state

    Returns:
        ImageGenerationResponse with generated images as base64 PNG

    Raises:
        HTTPException: For validation errors, missing engine, or generation failures
    """
    # Get diffusion engine from app state
    diffusion_engine: AsyncOmniDiffusion | None = getattr(raw_request.app.state, "diffusion_engine", None)
    if diffusion_engine is None:
        raise HTTPException(
            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
            detail="Diffusion engine not initialized. Start server with a diffusion model.",
        )

    # Get server's loaded model
    model_name = getattr(raw_request.app.state, "diffusion_model_name", "unknown")

    # Validate model field (warn if mismatch, don't error)
    if request.model is not None and request.model != model_name:
        logger.warning(
            f"Model mismatch: request specifies '{request.model}' but "
            f"server is running '{model_name}'. Using server model."
        )

    try:
        # Build params - pass through user values directly
        gen_params = {
            "prompt": request.prompt,
            "num_outputs_per_prompt": request.n,
        }

        # Parse and add size if provided
        if request.size:
            width, height = parse_size(request.size)
            gen_params["height"] = height
            gen_params["width"] = width
            size_str = f"{width}x{height}"
        else:
            size_str = "model default"

        # Add optional parameters ONLY if provided
        if request.num_inference_steps is not None:
            gen_params["num_inference_steps"] = request.num_inference_steps
        if request.negative_prompt is not None:
            gen_params["negative_prompt"] = request.negative_prompt
        if request.guidance_scale is not None:
            gen_params["guidance_scale"] = request.guidance_scale
        if request.true_cfg_scale is not None:
            gen_params["true_cfg_scale"] = request.true_cfg_scale
        if request.seed is not None:
            gen_params["seed"] = request.seed

        logger.info(f"Generating {request.n} image(s) {size_str}")

        # Generate images using AsyncOmniDiffusion
        result = await diffusion_engine.generate(**gen_params)

        # Extract images from result
        images = result.images if hasattr(result, "images") else []

        logger.info(f"Successfully generated {len(images)} image(s)")

        # Encode images to base64
        image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]

        return ImageGenerationResponse(
            created=int(time.time()),
            data=image_data,
        )

    except HTTPException:
        # Re-raise HTTPExceptions as-is
        raise
    except ValueError as e:
        logger.error(f"Validation error: {e}")
        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
    except Exception as e:
        logger.exception(f"Image generation failed: {e}")
        raise HTTPException(
            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
        )