diff --git a/docs/api/README.md b/docs/api/README.md index d9c4face4..fcd29d210 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -4,8 +4,8 @@ Main entry points for vLLM-Omni inference and serving. -- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][] -- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniStageLLM][] +- [vllm_omni.entrypoints.async_omni.AsyncOmni][] +- [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][] - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][] - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][] - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][] @@ -102,10 +102,6 @@ Configuration classes. Worker classes and model runners for distributed inference. -- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][] -- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][] -- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][] -- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][] - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][] - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][] - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][] diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 2ed46e93d..d582df801 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -44,7 +44,7 @@ stage_args: trust_remote_code: true # Needed by huggingface config parsing engine_output_type: latent # It claims that the stage will input latent hiddenstates besides token ids enable_prefix_caching: false # For request with hiddenstates output, the prefix caching is not supported now - is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmniLLM will use its tokenizer as default + is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmni will use its tokenizer as default final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role. final_output_type: text # What is the final output type. It can be text and audio now. default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM. @@ -206,7 +206,7 @@ Default: `false` ### `is_comprehension` -Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmniLLM` will use its tokenizer as the default tokenizer for the entire pipeline. +Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline. Default: `true` diff --git a/docs/design/vllm_omni_design.md b/docs/design/vllm_omni_design.md index 8fe328878..840106203 100644 --- a/docs/design/vllm_omni_design.md +++ b/docs/design/vllm_omni_design.md @@ -14,7 +14,7 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv ## Key Data Flow -API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core +API Server --> OmniLLM/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT) --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output) @@ -40,7 +40,7 @@ graph TD B --> C{Detect --omni flag} C -->|Yes| D[Parse OmniConfig] C -->|No| E[Forward to vLLM CLI] - D --> F[Initialize AsyncOmniLLM] + D --> F[Initialize AsyncOmni] F --> G[Start omni Server] G --> H[Multi-stage Processing] E --> I[Standard vLLM Pipeline] @@ -209,7 +209,7 @@ Similar to OmniLLM in offline inference, add some asynchronous processing, refer from vllm.v1.engine.async_llm import AsyncLLM -class AsyncOmniLLM(AsyncLLM): +class AsyncOmni(AsyncLLM): """Extended AsyncLLM supporting multiple engines and stage-based processing""" def __init__(self, stage_configs: List[StageConfig]): diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index eb7bbfdbe..535de983b 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/docs/mkdocs/hooks/generate_api_readme.py @@ -188,6 +188,11 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]: # Add classes (filter out internal ones) for class_name in classes: class_short_name = class_name.split(".")[-1] + + # TODO: Implement worker API reference + if "diffusion.worker.gpu_worker" in class_name: + continue + # Skip if it matches internal patterns (unless it's a main model class) if any(pattern in class_short_name for pattern in internal_patterns): # But include main model classes diff --git a/docs/user_guide/examples/online_serving/qwen2_5_omni.md b/docs/user_guide/examples/online_serving/qwen2_5_omni.md index 7918c5f84..bd51b88a9 100644 --- a/docs/user_guide/examples/online_serving/qwen2_5_omni.md +++ b/docs/user_guide/examples/online_serving/qwen2_5_omni.md @@ -48,7 +48,7 @@ This Web UI demo allows users to interact with the model through a web browser. ### Running Gradio Demo -Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by +Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by ```bash python gradio_demo.py --model Qwen/Qwen2.5-Omni-7B --port 7861 diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 34d82d17f..c593202c2 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -48,7 +48,7 @@ This Web UI demo allows users to interact with the model through a web browser. ### Running Gradio Demo -Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by +Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by ```bash python gradio_demo.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861 @@ -72,7 +72,7 @@ python gradio_demo.py \ ``` - `--model`: Model name -- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally. +- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally. - `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1) - `--ip`: Host/IP for Gradio server (default: 127.0.0.1) - `--port`: Port for Gradio server (default: 7861) diff --git a/examples/online_serving/qwen2_5_omni/README.md b/examples/online_serving/qwen2_5_omni/README.md index a351e3b53..2673e8fb0 100644 --- a/examples/online_serving/qwen2_5_omni/README.md +++ b/examples/online_serving/qwen2_5_omni/README.md @@ -45,7 +45,7 @@ This Web UI demo allows users to interact with the model through a web browser. ### Running Gradio Demo -Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by +Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by ```bash python gradio_demo.py --model Qwen/Qwen2.5-Omni-7B --port 7861 diff --git a/examples/online_serving/qwen2_5_omni/gradio_demo.py b/examples/online_serving/qwen2_5_omni/gradio_demo.py index b1f0c3309..ff77ea2d8 100644 --- a/examples/online_serving/qwen2_5_omni/gradio_demo.py +++ b/examples/online_serving/qwen2_5_omni/gradio_demo.py @@ -15,7 +15,7 @@ from PIL import Image from vllm.assets.video import video_get_metadata, video_to_ndarrays from vllm.sampling_params import SamplingParams -from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM +from vllm_omni.entrypoints.async_omni import AsyncOmni # Import utils from offline inference example sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni")) @@ -236,8 +236,8 @@ def process_video_file( return frames, metadata, audio_tuple -async def run_inference_async_omni_llm( - omni_llm: AsyncOmniLLM, +async def run_inference_async_omni( + omni: AsyncOmni, sampling_params: list[SamplingParams], prompt_args_template: SimpleNamespace, user_prompt: str, @@ -246,7 +246,7 @@ async def run_inference_async_omni_llm( video_file: Optional[str] = None, use_audio_in_video: bool = False, ): - """Run inference using AsyncOmniLLM directly with multimodal support.""" + """Run inference using AsyncOmni directly with multimodal support.""" if not user_prompt.strip() and not audio_file and not image_file and not video_file: return "Please provide at least a text prompt or multimodal input.", None @@ -328,7 +328,7 @@ async def run_inference_async_omni_llm( text_outputs: list[str] = [] audio_output = None - async for stage_outputs in omni_llm.generate( + async for stage_outputs in omni.generate( prompt=omni_prompt, request_id=request_id, sampling_params_list=sampling_params, @@ -362,12 +362,12 @@ async def run_inference_async_omni_llm( def build_interface( - omni_llm: AsyncOmniLLM, + omni: AsyncOmni, sampling_params: list[SamplingParams], prompt_args_template: SimpleNamespace, model: str, ): - """Build Gradio interface for AsyncOmniLLM mode.""" + """Build Gradio interface for AsyncOmni mode.""" async def run_inference( user_prompt: str, @@ -376,8 +376,8 @@ def build_interface( video_file: Optional[str], use_audio_in_video: bool, ): - return await run_inference_async_omni_llm( - omni_llm, + return await run_inference_async_omni( + omni, sampling_params, prompt_args_template, user_prompt, @@ -475,7 +475,7 @@ def build_interface( def main(): args = parse_args() - omni_llm = None + omni = None model_name = "/".join(args.model.split("/")[-2:]) assert model_name in SUPPORTED_MODELS, ( @@ -485,9 +485,9 @@ def main(): # Register signal handlers for graceful shutdown def signal_handler(sig, frame): print("\nReceived interrupt signal, shutting down...") - if omni_llm is not None: + if omni is not None: try: - omni_llm.shutdown() + omni.shutdown() except Exception as e: print(f"Error during shutdown: {e}") sys.exit(0) @@ -495,21 +495,21 @@ def main(): signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) - print(f"Initializing AsyncOmniLLM with model: {args.model}") + print(f"Initializing AsyncOmni with model: {args.model}") if args.stage_configs_path: print(f"Using custom stage configs: {args.stage_configs_path}") sampling_params = build_sampling_params(SEED, model_name) - omni_llm = AsyncOmniLLM( + omni = AsyncOmni( model=args.model, stage_configs_path=args.stage_configs_path, init_timeout=ASYNC_INIT_TIMEOUT, ) - print("✓ AsyncOmniLLM initialized successfully") + print("✓ AsyncOmni initialized successfully") prompt_args_template = create_prompt_args(args) demo = build_interface( - omni_llm, + omni, sampling_params, prompt_args_template, args.model, @@ -524,9 +524,9 @@ def main(): print("\nShutting down...") finally: # Cleanup - if omni_llm is not None: + if omni is not None: try: - omni_llm.shutdown() + omni.shutdown() except Exception as e: print(f"Error during cleanup: {e}") diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index ec80cedc0..06edb5471 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -45,7 +45,7 @@ This Web UI demo allows users to interact with the model through a web browser. ### Running Gradio Demo -Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by +Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by ```bash python gradio_demo.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861 diff --git a/examples/online_serving/qwen3_omni/gradio_demo.py b/examples/online_serving/qwen3_omni/gradio_demo.py index b6d802f85..d73c3ab5e 100644 --- a/examples/online_serving/qwen3_omni/gradio_demo.py +++ b/examples/online_serving/qwen3_omni/gradio_demo.py @@ -15,7 +15,7 @@ from PIL import Image from vllm.assets.video import video_get_metadata, video_to_ndarrays from vllm.sampling_params import SamplingParams -from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM +from vllm_omni.entrypoints.async_omni import AsyncOmni # Import utils from offline inference example sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni")) @@ -239,8 +239,8 @@ def process_video_file( return frames, metadata, audio_tuple -async def run_inference_async_omni_llm( - omni_llm: AsyncOmniLLM, +async def run_inference_async_omni( + omni: AsyncOmni, sampling_params: list[SamplingParams], prompt_args_template: SimpleNamespace, user_prompt: str, @@ -249,7 +249,7 @@ async def run_inference_async_omni_llm( video_file: Optional[str] = None, use_audio_in_video: bool = False, ): - """Run inference using AsyncOmniLLM directly with multimodal support.""" + """Run inference using AsyncOmni directly with multimodal support.""" if not user_prompt.strip() and not audio_file and not image_file and not video_file: return "Please provide at least a text prompt or multimodal input.", None @@ -331,7 +331,7 @@ async def run_inference_async_omni_llm( text_outputs: list[str] = [] audio_output = None - async for stage_outputs in omni_llm.generate( + async for stage_outputs in omni.generate( prompt=omni_prompt, request_id=request_id, sampling_params_list=sampling_params, @@ -366,12 +366,12 @@ async def run_inference_async_omni_llm( def build_interface( - omni_llm: AsyncOmniLLM, + omni: AsyncOmni, sampling_params: list[SamplingParams], prompt_args_template: SimpleNamespace, model: str, ): - """Build Gradio interface for AsyncOmniLLM mode.""" + """Build Gradio interface for AsyncOmni mode.""" async def run_inference( user_prompt: str, @@ -380,8 +380,8 @@ def build_interface( video_file: Optional[str], use_audio_in_video: bool, ): - return await run_inference_async_omni_llm( - omni_llm, + return await run_inference_async_omni( + omni, sampling_params, prompt_args_template, user_prompt, @@ -479,7 +479,7 @@ def build_interface( def main(): args = parse_args() - omni_llm = None + omni = None model_name = "/".join(args.model.split("/")[-2:]) assert model_name in SUPPORTED_MODELS, ( @@ -489,9 +489,9 @@ def main(): # Register signal handlers for graceful shutdown def signal_handler(sig, frame): print("\nReceived interrupt signal, shutting down...") - if omni_llm is not None: + if omni is not None: try: - omni_llm.shutdown() + omni.shutdown() except Exception as e: print(f"Error during shutdown: {e}") sys.exit(0) @@ -499,21 +499,21 @@ def main(): signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) - print(f"Initializing AsyncOmniLLM with model: {args.model}") + print(f"Initializing AsyncOmni with model: {args.model}") if args.stage_configs_path: print(f"Using custom stage configs: {args.stage_configs_path}") sampling_params = build_sampling_params(SEED, model_name) - omni_llm = AsyncOmniLLM( + omni = AsyncOmni( model=args.model, stage_configs_path=args.stage_configs_path, init_timeout=ASYNC_INIT_TIMEOUT, ) - print("✓ AsyncOmniLLM initialized successfully") + print("✓ AsyncOmni initialized successfully") prompt_args_template = create_prompt_args(args) demo = build_interface( - omni_llm, + omni, sampling_params, prompt_args_template, args.model, @@ -528,9 +528,9 @@ def main(): print("\nShutting down...") finally: # Cleanup - if omni_llm is not None: + if omni is not None: try: - omni_llm.shutdown() + omni.shutdown() except Exception as e: print(f"Error during cleanup: {e}") diff --git a/mkdocs.yml b/mkdocs.yml index 31ca350fa..067e766e8 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -86,6 +86,7 @@ plugins: api_root_uri: "api" nav_item_prefix: "" # No prefix in navigation tree (clean names) show_full_namespace: false # Show only module name, not full path + on_implicit_namespace_package: skip # Skip implicit namespace packages without warning exclude: - "re:vllm_omni\\._.*" # Internal modules - "vllm_omni.diffusion.models.qwen_image" # avoid importing vllm in mkdocs building diff --git a/tests/test_omni_llm.py b/tests/test_omni_llm.py index 13153d6eb..567208b27 100644 --- a/tests/test_omni_llm.py +++ b/tests/test_omni_llm.py @@ -342,14 +342,14 @@ def _setup_log_mocks(monkeypatch): def mock_get_config(monkeypatch): """Auto-mock get_config and related model loading functions to avoid model path validation.""" # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports - # This prevents ImportError when async_omni_llm is imported (which happens via omni_stage) + # This prevents ImportError when async_omni is imported (which happens via omni_stage) import sys fake_tokenizer = MagicMock() fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3]) fake_tokenizer.decode = MagicMock(return_value="test") - # Mock init_tokenizer_from_configs (used in async_omni_llm) + # Mock init_tokenizer_from_configs (used in async_omni) def _mock_init_tokenizer_from_configs(model_config=None, **kwargs): return fake_tokenizer @@ -367,8 +367,9 @@ def mock_get_config(monkeypatch): tokenizer_module = sys.modules[tokenizer_module_path] setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) - # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni_llm - # This is because async_omni_llm imports processor.py, which imports this function at module level + # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni + + # This is because async_omni imports processor.py, which imports this function at module level # Mock length_from_prompt_token_ids_or_embeds (used in processor.py) def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None): # Return a reasonable default length @@ -402,19 +403,19 @@ def mock_get_config(monkeypatch): processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds ) - # Strategy 3: Now mock async_omni_llm AFTER length_from_prompt_token_ids_or_embeds is mocked - # This prevents ImportError when async_omni_llm imports processor.py + # Strategy 3: Now mock async_omni AFTER length_from_prompt_token_ids_or_embeds is mocked + # This prevents ImportError when async_omni imports processor.py monkeypatch.setattr( - "vllm_omni.entrypoints.async_omni_llm.init_tokenizer_from_configs", + "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False, ) - # Strategy 4: If async_omni_llm is already imported, patch it directly - async_omni_llm_path = "vllm_omni.entrypoints.async_omni_llm" - if async_omni_llm_path in sys.modules: - async_omni_llm_module = sys.modules[async_omni_llm_path] - setattr(async_omni_llm_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) + # Strategy 4: If async_omni is already imported, patch it directly + async_omni_path = "vllm_omni.entrypoints.async_omni" + if async_omni_path in sys.modules: + async_omni_module = sys.modules[async_omni_path] + setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) # Now mock get_config and other functions fake_hf_config = MagicMock() diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index fd7f9174f..5e37834c2 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -21,7 +21,7 @@ except ModuleNotFoundError as exc: # pragma: no cover - optional dependency patch = None # type: ignore from .config import OmniModelConfig -from .entrypoints.async_omni_llm import AsyncOmniLLM +from .entrypoints.async_omni import AsyncOmni # Main entry points from .entrypoints.omni import Omni @@ -34,7 +34,7 @@ __all__ = [ "__version_tuple__", # Main components "Omni", - "AsyncOmniLLM", + "AsyncOmni", # Configuration "OmniModelConfig", # All other components are available through their respective modules diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni.py similarity index 99% rename from vllm_omni/entrypoints/async_omni_llm.py rename to vllm_omni/entrypoints/async_omni.py index c45acdc44..4785978f5 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -51,8 +51,8 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -class AsyncOmniLLM(EngineClient): - """Async entry point for vLLM-Omni inference. +class AsyncOmni(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 @@ -75,7 +75,7 @@ class AsyncOmniLLM(EngineClient): **kwargs: Additional keyword arguments passed to stage engines Example: - >>> async_llm = AsyncOmniLLM(model="Qwen/Qwen2.5-Omni-7B") + >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B") >>> async for output in async_llm.generate( ... prompt="Hello", ... request_id="req-1", @@ -175,7 +175,7 @@ class AsyncOmniLLM(EngineClient): 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. + when done using the AsyncOmni instance. """ for q in self._stage_in_queues: try: @@ -611,13 +611,13 @@ class AsyncOmniLLM(EngineClient): **kwargs, ): """Generate outputs for a request from a pooling model.""" - raise NotImplementedError("encode() is not implemented for AsyncOmniLLM") + raise NotImplementedError("encode() is not implemented for AsyncOmni") async def start_profile(self) -> None: - raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM") + raise NotImplementedError("start_profile() is not implemented for AsyncOmni") async def stop_profile(self) -> None: - raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM") + raise NotImplementedError("stop_profile() is not implemented for AsyncOmni") class AsyncOmniStageLLM(AsyncLLM): diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 90bb55276..5150ea27e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -671,7 +671,7 @@ async def _stage_worker_async( import logging as _logging import time as _time - from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM + from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM from vllm_omni.entrypoints.log_utils import ( compute_and_log_stage_request_stats, count_tokens_from_outputs, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 61cec05b5..f40e579bc 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -38,7 +38,7 @@ from vllm.logger import init_logger from vllm.transformers_utils.tokenizer import MistralTokenizer from vllm.utils import decorate_logs -from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM +from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat logger = init_logger(__name__) @@ -65,7 +65,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, if log_config is not None: uvicorn_kwargs["log_config"] = log_config - async with build_async_omni_llm( + async with build_async_omni( args, client_config=client_config, ) as engine_client: @@ -108,15 +108,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, @asynccontextmanager -async def build_async_omni_llm( +async def build_async_omni( args: Namespace, *, disable_frontend_multiprocessing: Optional[bool] = None, client_config: Optional[dict[str, Any]] = None, ) -> AsyncIterator[EngineClient]: - """Build an AsyncOmniLLM instance from command-line arguments. + """Build an AsyncOmni instance from command-line arguments. - Creates an async context manager that yields an AsyncOmniLLM instance + 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. @@ -127,7 +127,7 @@ async def build_async_omni_llm( client_config: Optional client configuration dictionary Yields: - EngineClient instance (AsyncOmniLLM) ready for use + EngineClient instance (AsyncOmni) ready for use """ if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver": # The executor is expected to be mp. @@ -138,25 +138,25 @@ async def build_async_omni_llm( forkserver.ensure_running() logger.debug("Forkserver setup complete!") - # Context manager to handle async_omni_llm lifecycle + # Context manager to handle async_omni lifecycle # Ensures everything is shutdown and cleaned up on error/exit - async with build_async_omni_llm_from_stage_config( + async with build_async_omni_from_stage_config( args, disable_frontend_multiprocessing=disable_frontend_multiprocessing, - ) as async_omni_llm: - yield async_omni_llm + ) as async_omni: + yield async_omni @asynccontextmanager -async def build_async_omni_llm_from_stage_config( +async def build_async_omni_from_stage_config( args: Namespace, *, disable_frontend_multiprocessing: bool = False, client_config: Optional[dict[str, Any]] = None, ) -> AsyncIterator[EngineClient]: - """Create AsyncOmniLLM from stage configuration. + """Create AsyncOmni from stage configuration. - Creates an AsyncOmniLLM instance either in-process or using multiprocess + Creates an AsyncOmni instance either in-process or using multiprocess RPC. Loads stage configurations from the model or from a specified path. Args: @@ -166,7 +166,7 @@ async def build_async_omni_llm_from_stage_config( client_config: Optional client configuration dictionary Yields: - EngineClient instance (AsyncOmniLLM) ready for use + EngineClient instance (AsyncOmni) ready for use Note: Stage configurations are loaded from args.stage_configs_path if provided, @@ -182,21 +182,21 @@ async def build_async_omni_llm_from_stage_config( "To disable frontend multiprocessing, set VLLM_USE_V1=0." ) - async_omni_llm: Optional[EngineClient] = None + async_omni: Optional[EngineClient] = None try: if getattr(args, "stage_configs_path", None): - async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path) + async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path) else: - async_omni_llm = AsyncOmniLLM(model=args.model) + async_omni = AsyncOmni(model=args.model) # # Don't keep the dummy data in memory # await async_llm.reset_mm_cache() - yield async_omni_llm + yield async_omni finally: - if async_omni_llm: - async_omni_llm.shutdown() + if async_omni: + async_omni.shutdown() async def omni_init_app_state( @@ -211,7 +211,7 @@ async def omni_init_app_state( and other server configuration needed for handling API requests. Args: - engine_client: Engine client instance (AsyncOmniLLM) + engine_client: Engine client instance (AsyncOmni) vllm_config: vLLM configuration object state: FastAPI application state object to initialize args: Parsed command-line arguments