diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml deleted file mode 100644 index b79b37033..000000000 --- a/.github/workflows/docs.yml +++ /dev/null @@ -1,51 +0,0 @@ -name: Deploy Documentation - -on: - push: - branches: - - main - - master - paths: - - 'docs/**' - - 'mkdocs.yml' - - 'vllm_omni/**' - - '.github/workflows/docs.yml' - workflow_dispatch: - -permissions: - contents: write - pages: write - id-token: write - -jobs: - deploy: - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - cache: 'pip' - - - name: Install dependencies - run: | - pip install -e ".[docs]" - - - name: Build documentation - run: mkdocs build - - - name: Setup Pages - uses: actions/configure-pages@v4 - - - name: Upload artifact - uses: actions/upload-pages-artifact@v3 - with: - path: ./site - - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 - diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85a7d1527..585150b27 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,4 +26,4 @@ repos: rev: 7.1.1 hooks: - id: flake8 - args: ["--max-line-length=88", "--extend-ignore=E203,W503"] + args: ["--max-line-length=160", "--extend-ignore=E203,W503"] diff --git a/README.md b/README.md index 83c769f34..e8a345c5e 100644 --- a/README.md +++ b/README.md @@ -100,15 +100,7 @@ uv pip install -e . ## Run examples (Qwen2.5-omni) -Get into the example folder -```bash -cd examples/offline_inference/qwen_2_5_omni -``` -Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run. -```bash -bash run.sh -``` -The output audio is saved in ./output_audio +Please check the folder of [examples](examples) ## Further details diff --git a/docs/DOCS_GUIDE.md b/docs/DOCS_GUIDE.md index d52ad1c01..ccba0663f 100644 --- a/docs/DOCS_GUIDE.md +++ b/docs/DOCS_GUIDE.md @@ -40,15 +40,15 @@ Example docstring: ```python class OmniLLM: """Main entry point for vLLM-omni inference. - + This class provides a high-level interface for running multi-modal inference with non-autoregressive models. - + Args: model: Model name or path stage_configs: Optional stage configurations **kwargs: Additional arguments passed to the engine - + Example: >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni") >>> outputs = llm.generate(prompts="Hello") diff --git a/docs/api/overview.md b/docs/api/overview.md index 8e386044b..8a33a9095 100644 --- a/docs/api/overview.md +++ b/docs/api/overview.md @@ -66,4 +66,3 @@ Worker classes and model runners for distributed inference. - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][] - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][] - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][] - diff --git a/docs/contributing/design_documents/api_design_doc.md b/docs/contributing/design_documents/api_design_doc.md index 4ca442645..f67f325fe 100644 --- a/docs/contributing/design_documents/api_design_doc.md +++ b/docs/contributing/design_documents/api_design_doc.md @@ -2,7 +2,7 @@ This document contains comprehensive API design documentation for all core modules in vLLM-omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules. -## 📋 Module API +## 📋 Module API ### Core Module API **Core module** provides fundamental scheduling, caching, and resource management functionality. diff --git a/docs/contributing/design_documents/architecture_overview.md b/docs/contributing/design_documents/architecture_overview.md index 99b21e4b4..54319bce0 100644 --- a/docs/contributing/design_documents/architecture_overview.md +++ b/docs/contributing/design_documents/architecture_overview.md @@ -1,4 +1,4 @@ -# Architecture Overview +# Architecture Overview # Introduction diff --git a/docs/contributing/design_documents/index.md b/docs/contributing/design_documents/index.md index 5598b46ca..42e360ed4 100644 --- a/docs/contributing/design_documents/index.md +++ b/docs/contributing/design_documents/index.md @@ -10,4 +10,3 @@ This section contains design documents and architecture specifications for vLLM- ## API Design Documentation - [vLLM-omni API Documentation](api_design_doc.md) - diff --git a/docs/index.md b/docs/index.md index 733b3e4e8..0db4f2908 100644 --- a/docs/index.md +++ b/docs/index.md @@ -28,6 +28,5 @@ ## 📚 Documentation Navigation -- To run open-source models on vLLM-Omni, we recommend starting with the [:material-code-tags: User Quide](user_guide/getting_started/quickstart.md) +- To run open-source models on vLLM-Omni, we recommend starting with the [:material-code-tags: User Quide](user_guide/getting_started/quickstart.md) - To develop and contribute to vLLM-Omni, we recommend starting with the [:material-tools: Developer Guide](contributing/README.md) - diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 518a66f7a..82b1ba06f 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -103,11 +103,20 @@ class Example: if self.path.is_file(): return [] # Text file extensions that can be safely included - text_extensions = {".py", ".md", ".sh", ".yaml", ".yml", ".json", ".txt", ".toml", ".cfg", ".ini"} - is_other_file = lambda file: ( - file.is_file() - and file != self.main_file - and file.suffix in text_extensions + text_extensions = { + ".py", + ".md", + ".sh", + ".yaml", + ".yml", + ".json", + ".txt", + ".toml", + ".cfg", + ".ini", + } + is_other_file = lambda file: ( # noqa: E731 + file.is_file() and file != self.main_file and file.suffix in text_extensions ) return [file for file in self.path.rglob("*") if is_other_file(file)] @@ -172,7 +181,7 @@ class Example: f"{code_fence}\n" ) else: - with open(self.main_file) as f: + with open(self.main_file, encoding="utf-8") as f: # Skip the title from md snippets as it's been included above main_content = f.readlines()[1:] content += self.fix_relative_links("".join(main_content)) diff --git a/docs/mkdocs/stylesheets/extra.css b/docs/mkdocs/stylesheets/extra.css index e306f7f12..efc636b53 100644 --- a/docs/mkdocs/stylesheets/extra.css +++ b/docs/mkdocs/stylesheets/extra.css @@ -53,4 +53,3 @@ body[data-md-color-scheme="slate"] .md-nav__item--section > label.md-nav__link . .md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav > .md-nav__list > .md-nav__item--nested.md-nav__item--active > .md-nav { display: block; } - diff --git a/docs/user_guide/examples/index.md b/docs/user_guide/examples/index.md index 3f3e5fd00..d6e084006 100644 --- a/docs/user_guide/examples/index.md +++ b/docs/user_guide/examples/index.md @@ -6,4 +6,3 @@ vLLM-omni's examples are split into two categories: - If you are using vLLM-omni from an HTTP application or client, see the *Online Serving* section. For detailed example documentation, check the [examples directory](https://github.com/vllm-project/vllm-omni/tree/main/examples) in the repository. - diff --git a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md new file mode 100644 index 000000000..576338003 --- /dev/null +++ b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md @@ -0,0 +1,64 @@ +# Offline Example of vLLM-omni for Qwen2.5-omni + +Source . + + +## 🛠️ Installation + +Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md) + +## Run examples (Qwen2.5-omni) +### Multiple Prompts +Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can: +```bash +tar -xf /seedtts_testset.tar +cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst +python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \ + --input examples/offline_inference/qwen2_5_omni/meta.lst \ + --output examples/offline_inference/qwen2_5_omni/top100.txt \ + --topk 100 +``` +Get into the example folder +```bash +cd examples/offline_inference/qwen2_5_omni +``` +Then run the command below. +```bash +bash run_multiple_prompts.sh +``` +### Single Prompts +Get into the example folder +```bash +cd examples/offline_inference/qwen2_5_omni +``` +Then run the command below. +```bash +bash run_single_prompt.sh +``` + +## Example materials + +??? abstract "end2end.py" + ``````py + --8<-- "examples\offline_inference\qwen2_5_omni\end2end.py" + `````` +??? abstract "extract_prompts.py" + ``````py + --8<-- "examples\offline_inference\qwen2_5_omni\extract_prompts.py" + `````` +??? abstract "processing_omni.py" + ``````py + --8<-- "examples\offline_inference\qwen2_5_omni\processing_omni.py" + `````` +??? abstract "run_multiple_prompts.sh" + ``````sh + --8<-- "examples\offline_inference\qwen2_5_omni\run_multiple_prompts.sh" + `````` +??? abstract "run_single_prompt.sh" + ``````sh + --8<-- "examples\offline_inference\qwen2_5_omni\run_single_prompt.sh" + `````` +??? abstract "utils.py" + ``````py + --8<-- "examples\offline_inference\qwen2_5_omni\utils.py" + `````` diff --git a/docs/user_guide/examples/online_serving/README.md b/docs/user_guide/examples/online_serving/README.md new file mode 100644 index 000000000..7b3341255 --- /dev/null +++ b/docs/user_guide/examples/online_serving/README.md @@ -0,0 +1,35 @@ +# Online serving Example of vLLM-omni for Qwen2.5-omni + +Source . + + +## 🛠️ Installation + +Please refer to [README.md](https://github.com/vllm-project/vllm-omni/blob/main/README.md) + +## Run examples (Qwen2.5-omni) + +Launch the server +```bash +vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 +``` + +If you have custom stage configs file, launch the server with command below +```bash +vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file +``` + +Get into the example folder +```bash +cd examples/online_serving +``` + +Send request via python +```bash +python openai_chat_completion_client_for_multimodal_generation.py +``` + +Send request via curl +```bash +bash run_curl_multimodal_generation.sh +``` diff --git a/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md b/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md new file mode 100644 index 000000000..066b27624 --- /dev/null +++ b/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md @@ -0,0 +1,7 @@ +# OpenAI Chat Completion Client For Multimodal Generation + +Source . + +``````py +--8<-- "examples\online_serving\openai_chat_completion_client_for_multimodal_generation.py" +`````` diff --git a/docs/user_guide/getting_started/quickstart.md b/docs/user_guide/getting_started/quickstart.md index a28696754..8207b23e8 100644 --- a/docs/user_guide/getting_started/quickstart.md +++ b/docs/user_guide/getting_started/quickstart.md @@ -53,4 +53,3 @@ The output audio is saved in ./output_audio - Read the [architecture documentation](../../contributing/design_documents/vllm_omni_design.md) - Check out the [API reference](../../api/overview.md) - Explore the [examples](../examples/index.md) - diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md new file mode 100644 index 000000000..112ef5335 --- /dev/null +++ b/examples/offline_inference/qwen2_5_omni/README.md @@ -0,0 +1,34 @@ +# Offline Example of vLLM-omni for Qwen2.5-omni + +## 🛠️ Installation + +Please refer to [README.md](../../../README.md) + +## Run examples (Qwen2.5-omni) +### Multiple Prompts +Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can: +```bash +tar -xf /seedtts_testset.tar +cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst +python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \ + --input examples/offline_inference/qwen2_5_omni/meta.lst \ + --output examples/offline_inference/qwen2_5_omni/top100.txt \ + --topk 100 +``` +Get into the example folder +```bash +cd examples/offline_inference/qwen2_5_omni +``` +Then run the command below. +```bash +bash run_multiple_prompts.sh +``` +### Single Prompts +Get into the example folder +```bash +cd examples/offline_inference/qwen2_5_omni +``` +Then run the command below. +```bash +bash run_single_prompt.sh +``` diff --git a/examples/offline_inference/qwen_2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py similarity index 94% rename from examples/offline_inference/qwen_2_5_omni/end2end.py rename to examples/offline_inference/qwen2_5_omni/end2end.py index bbd4e0f53..f205afcbb 100644 --- a/examples/offline_inference/qwen_2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -7,8 +7,8 @@ import numpy as np import soundfile as sf import torch from utils import make_omni_prompt -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import SamplingParams from vllm_omni.entrypoints.omni_llm import OmniLLM _os_env_toggle.environ["VLLM_USE_V1"] = "1" @@ -109,10 +109,14 @@ def parse_args(): parser.add_argument("--use-torchvision", action="store_true") parser.add_argument("--tokenize", action="store_true") parser.add_argument( - "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)." + "--output-wav", + default="output.wav", + help="[Deprecated] Output wav directory (use --output-dir).", ) parser.add_argument( - "--output-dir", default="outputs", help="Output directory to save text and wav files together." + "--output-dir", + default="outputs", + help="Output directory to save text and wav files together.", ) parser.add_argument( "--thinker-hidden-states-dir", @@ -168,7 +172,9 @@ def main(): raise if args.prompts is None: - raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts (with --prompt_type text)") + raise ValueError( + "No prompts provided. Use --prompts ... or --txt-prompts (with --prompt_type text)" + ) omni_llm = OmniLLM( model=model_name, log_stats=args.enable_stats, @@ -217,7 +223,9 @@ def main(): omni_outputs = omni_llm.generate(prompt, sampling_params_list) # Determine output directory: prefer --output-dir; fallback to --output-wav - output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav + output_dir = ( + args.output_dir if getattr(args, "output_dir", None) else args.output_wav + ) os.makedirs(output_dir, exist_ok=True) for stage_outputs in omni_outputs: if stage_outputs.final_output_type == "text": diff --git a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py b/examples/offline_inference/qwen2_5_omni/extract_prompts.py similarity index 79% rename from examples/offline_inference/qwen_2_5_omni/extract_prompts.py rename to examples/offline_inference/qwen2_5_omni/extract_prompts.py index 7380e701b..cf39f46a2 100644 --- a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py +++ b/examples/offline_inference/qwen2_5_omni/extract_prompts.py @@ -1,26 +1,30 @@ #!/usr/bin/env python3 import argparse from typing import Optional -import os -import torch def extract_prompt(line: str) -> Optional[str]: # 提取第一个 '|' 与第二个 '|' 之间的内容 - i = line.find('|') + i = line.find("|") if i == -1: return None - j = line.find('|', i + 1) + j = line.find("|", i + 1) if j == -1: return None - return line[i + 1:j].strip() + return line[i + 1 : j].strip() def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--input", "-i", required=True, help="Input .lst file path") parser.add_argument("--output", "-o", required=True, help="Output file path") - parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)") + parser.add_argument( + "--topk", + "-k", + type=int, + default=100, + help="Extract the top K prompts (default: 100)", + ) args = parser.parse_args() prompts = [] @@ -39,5 +43,3 @@ def main() -> None: if __name__ == "__main__": main() - - diff --git a/examples/offline_inference/qwen_2_5_omni/processing_omni.py b/examples/offline_inference/qwen2_5_omni/processing_omni.py similarity index 100% rename from examples/offline_inference/qwen_2_5_omni/processing_omni.py rename to examples/offline_inference/qwen2_5_omni/processing_omni.py diff --git a/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh b/examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh similarity index 89% rename from examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh rename to examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh index ba07616ce..78d3dd54f 100644 --- a/examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh +++ b/examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh @@ -1,4 +1,3 @@ -export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ --voice-type "m02" \ --dit-ckpt none \ diff --git a/examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh b/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh similarity index 91% rename from examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh rename to examples/offline_inference/qwen2_5_omni/run_single_prompt.sh index 8dc2c74e5..739902b25 100644 --- a/examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh +++ b/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh @@ -1,4 +1,3 @@ -export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH python end2end.py --model Qwen/Qwen2.5-Omni-7B \ --voice-type "m02" \ --dit-ckpt none \ diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen2_5_omni/utils.py similarity index 99% rename from examples/offline_inference/qwen_2_5_omni/utils.py rename to examples/offline_inference/qwen2_5_omni/utils.py index 42894c4a7..c25b0734f 100644 --- a/examples/offline_inference/qwen_2_5_omni/utils.py +++ b/examples/offline_inference/qwen2_5_omni/utils.py @@ -9,6 +9,7 @@ import soundfile as sf import torch import torchvision.io from processing_omni import fetch_image, fetch_video +from transformers import AutoConfig, AutoProcessor from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt diff --git a/examples/offline_inference/qwen_2_5_omni/README.md b/examples/offline_inference/qwen_2_5_omni/README.md deleted file mode 100644 index d07e115ee..000000000 --- a/examples/offline_inference/qwen_2_5_omni/README.md +++ /dev/null @@ -1,92 +0,0 @@ -# Offline Example of vLLM-omni for Qwen2.5-omni - -## 🛠️ Installation - -### Environment setup - -Use Docker to keep consistent basic environment (Optional, Recommended) -```bash -docker run --gpus all --ipc=host --network=host -v $source_dir:$container_dir --name $container_name -it nvcr.io/nvidia/pytorch:25.01-py3 bash -``` - -Set up basic uv environment -```bash -pip install uv -uv venv --python 3.12 --seed -source .venv/bin/activate -``` - -### Installation of vLLM (for users) - -Now we build it based on vLLM v0.11.0. Please install it with command below. -```bash -uv pip install vllm==0.11.0 --torch-backend=auto -``` - -### Installation of vLLM (for developers) - -If you want to check or debug with source code of vLLM, install stable release version of vllm with 0.11.0 from source with pre-built wheel file. - -```bash -git clone https://github.com/vllm-project/vllm.git -cd vllm -git checkout v0.11.0 -``` -Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set VLLM_PRECOMPILED_WHEEL_LOCATION as your local absolute path of whl file. -```bash -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl -``` -Install vllm with command below (If you have no existing PyTorch). -```bash -uv pip install --editable . -``` - -Install vllm with command below (If you already have PyTorch). -```bash -python use_existing_torch.py -uv pip install -r requirements/build.txt -uv pip install --no-build-isolation --editable . -``` - -### Verification for successful installation of vLLM -Just run the command below. If no error, it demonstrates that the installation is successfull. -```bash -python -c "import vllm._C" -``` - -### Installation of vLLM-omni -Install additional requirements for vllm-omni -```bash -git clone https://github.com/vllm-project/vllm-omni.git -cd vllm_omni -uv pip install -e . -``` - -## Run examples (Qwen2.5-omni) -### Multiple Prompts -Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can: -```bash -tar -xf /seedtts_testset.tar -cp seedtts_testset/en/meta.lst examples/offline_inference/qwen_2_5_omni/meta.lst -python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \ - --input examples/offline_inference/qwen_2_5_omni/meta.lst \ - --output examples/offline_inference/qwen_2_5_omni/top100.txt \ - --topk 100 -``` -Get into the example folder -```bash -cd examples/offline_inference/qwen_2_5_omni -``` -Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run. -```bash -bash run_multiple_prompts.sh -``` -### Single Prompts -Get into the example folder -```bash -cd examples/offline_inference/qwen_2_5_omni -``` -Modify PYTHONPATH in run_single_prompt.sh as your path of vllm_omni. Then run. -```bash -bash run_single_prompt.sh -``` diff --git a/examples/online_serving/README.md b/examples/online_serving/README.md new file mode 100644 index 000000000..b64989326 --- /dev/null +++ b/examples/online_serving/README.md @@ -0,0 +1,32 @@ +# Online serving Example of vLLM-omni for Qwen2.5-omni + +## 🛠️ Installation + +Please refer to [README.md](../../README.md) + +## Run examples (Qwen2.5-omni) + +Launch the server +```bash +vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 +``` + +If you have custom stage configs file, launch the server with command below +```bash +vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file +``` + +Get into the example folder +```bash +cd examples/online_serving +``` + +Send request via python +```bash +python openai_chat_completion_client_for_multimodal_generation.py +``` + +Send request via curl +```bash +bash run_curl_multimodal_generation.sh +``` diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py new file mode 100644 index 000000000..3bc864852 --- /dev/null +++ b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py @@ -0,0 +1,104 @@ +import base64 + +from openai import OpenAI + +# Modify OpenAI's API key and API base to use vLLM's API server. +openai_api_key = "EMPTY" +openai_api_base = "http://localhost:8091/v1" + +client = OpenAI( + # defaults to os.environ.get("OPENAI_API_KEY") + api_key=openai_api_key, + base_url=openai_api_base, +) + +SEED = 42 + + +def get_system_prompt(): + + return { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual inputs, " + "as well as generating text and speech." + ), + } + ], + } + + +def run_text_to_audio(model: str) -> None: + thinker_sampling_params = { + "temperature": 0.0, # Deterministic - no randomness + "top_p": 1.0, # Disable nucleus sampling + "top_k": -1, # Disable top-k sampling + "max_tokens": 2048, + "seed": SEED, # Fixed seed for sampling + "detokenize": True, + "repetition_penalty": 1.1, + } + talker_sampling_params = { + "temperature": 0.9, + "top_p": 0.8, + "top_k": 40, + "max_tokens": 2048, + "seed": SEED, # Fixed seed for sampling + "detokenize": True, + "repetition_penalty": 1.05, + "stop_token_ids": [8294], + } + code2wav_sampling_params = { + "temperature": 0.0, # Deterministic - no randomness + "top_p": 1.0, # Disable nucleus sampling + "top_k": -1, # Disable top-k sampling + "max_tokens": 2048, + "seed": SEED, # Fixed seed for sampling + "detokenize": True, + "repetition_penalty": 1.1, + } + + sampling_params_list = [ + thinker_sampling_params, + talker_sampling_params, + code2wav_sampling_params, + ] + + chat_completion = client.chat.completions.create( + messages=[ + get_system_prompt(), + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.", + }, + ], + }, + ], + model=model, + extra_body={ + "sampling_params_list": sampling_params_list + }, # Optional, it has a default setting in stage_configs of the corresponding model. + ) + + count = 0 + for choice in chat_completion.choices: + if choice.message.audio: + audio_data = base64.b64decode(choice.message.audio.data) + audio_file_path = f"audio_{count}.wav" + with open(audio_file_path, "wb") as f: + f.write(audio_data) + print(f"Audio saved to {audio_file_path}") + count += 1 + elif choice.message.content: + print("Chat completion output from text:", choice.message.content) + + +if __name__ == "__main__": + run_text_to_audio("Qwen/Qwen2.5-Omni-7B") diff --git a/examples/online_serving/run_curl_multimodal_generation.sh b/examples/online_serving/run_curl_multimodal_generation.sh new file mode 100644 index 000000000..6d2afd4a8 --- /dev/null +++ b/examples/online_serving/run_curl_multimodal_generation.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +set -euo pipefail + +SEED=42 + +thinker_sampling_params='{ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 2048, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.1 +}' + +talker_sampling_params='{ + "temperature": 0.9, + "top_p": 0.8, + "top_k": 40, + "max_tokens": 2048, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05, + "stop_token_ids": [8294] +}' + +code2wav_sampling_params='{ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 2048, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.1 +}' +# Above is optional, it has a default setting in stage_configs of the corresponding model. + +output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d @- < OmniModelConfig: + # First, get the base ModelConfig from the parent class + base_config = super().create_model_config() + + # Create OmniModelConfig by copying all base config attributes + # and adding the new omni-specific fields + config_dict = base_config.__dict__.copy() + + # Add the new omni-specific fields + config_dict["stage_id"] = self.stage_id + config_dict["model_stage"] = self.model_stage + config_dict["model_arch"] = self.model_arch + config_dict["engine_output_type"] = self.engine_output_type + + # Create and return the OmniModelConfig instance + omni_config = OmniModelConfig(**config_dict) + omni_config.hf_config.architectures = omni_config.architectures + + return omni_config diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py index d9fb00c13..d74022535 100644 --- a/vllm_omni/engine/processor.py +++ b/vllm_omni/engine/processor.py @@ -7,6 +7,7 @@ import torch from vllm.config import VllmConfig from vllm.inputs import ProcessorInputs, PromptType from vllm.inputs.parse import split_enc_dec_inputs +from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.inputs import MultiModalFeatureSpec @@ -25,6 +26,8 @@ from vllm_omni.engine import ( ) from vllm_omni.inputs.preprocess import OmniInputPreprocessor +logger = init_logger(__name__) + class OmniProcessor(Processor): @staticmethod diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py new file mode 100644 index 000000000..cc1114afc --- /dev/null +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -0,0 +1,778 @@ +import asyncio +import multiprocessing as mp +import os +import socket +import sys +import time +from collections.abc import AsyncGenerator, Iterable +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Mapping, Optional, Sequence, Union + +import torch + +# External library imports (vLLM) +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.engine.protocol import EngineClient +from vllm.inputs import PromptType +from vllm.inputs.preprocess import InputPreprocessor +from vllm.logger import init_logger +from vllm.lora.request import LoRARequest +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.sampling_params import SamplingParams +from vllm.tracing import init_tracer +from vllm.transformers_utils.config import maybe_register_config_serialize_by_value +from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs +from vllm.usage.usage_lib import UsageContext +from vllm.utils import Device, deprecate_kwargs +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.engine.core_client import EngineCoreClient +from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.executor.abstract import Executor +from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager + +# Internal imports (our code) +from vllm_omni.config import OmniModelConfig +from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs +from vllm_omni.engine.output_processor import MultimodalOutputProcessor +from vllm_omni.engine.processor import OmniProcessor +from vllm_omni.entrypoints.log_utils import ( + OrchestratorMetrics, + configure_orchestrator_logger, + init_stats_paths, + remove_old_logs, +) +from vllm_omni.entrypoints.omni_stage import OmniStage +from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode +from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load +from vllm_omni.entrypoints.stage_utils import serialize_obj as _ser +from vllm_omni.entrypoints.utils import ( + load_stage_configs_from_model, + load_stage_configs_from_yaml, +) +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + + +class AsyncOmniLLM(EngineClient): + def __init__( + self, + model: str, + stage_configs_path: Optional[str] = None, + log_stats: bool = False, + log_file: Optional[str] = None, + init_sleep_seconds: int = 20, + shm_threshold_bytes: int = 65536, + batch_timeout: int = 10, + init_timeout: int = 300, + **kwargs, + ): + + self.batch_timeout = batch_timeout + self._enable_stats: bool = bool(log_stats) + + if stage_configs_path is None: + self.stage_configs = load_stage_configs_from_model(model) + else: + self.stage_configs = load_stage_configs_from_yaml(stage_configs_path) + + self.stage_list: list[OmniStage] = [] + self.default_sampling_params_list: list[SamplingParams] = [] + # Optional file handler for orchestrator + self._log_file = log_file + if self._log_file: + remove_old_logs(self._log_file, len(self.stage_list)) + configure_orchestrator_logger(logger, self._log_file) + + self._stats_file, self._overall_stats_file = init_stats_paths( + self._enable_stats, self._log_file + ) + self._initialize_stages( + model, init_sleep_seconds, shm_threshold_bytes, init_timeout + ) + + 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 + ] + logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) + + self._ctx = mp.get_context("spawn") + 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_stage_processes(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_stage_processes(self, model: str) -> None: + for stage_id, stage in enumerate(self.stage_list): + # Use unbounded queues to avoid deadlock when seeding many requests + in_q: mp.Queue = self._ctx.Queue(maxsize=0) + out_q: mp.Queue = self._ctx.Queue(maxsize=0) + 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) + stage.init_stage_worker( + model, + is_async=True, + log_file=self._log_file, + shm_threshold_bytes=self._shm_threshold_bytes, + ctx=self._ctx, + batch_timeout=self.batch_timeout, + ) + logger.debug("[Orchestrator] Stage-%s process started", stage_id) + time.sleep(self._init_sleep_seconds) + + def close(self) -> None: + 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) + + def __del__(self) -> None: # best-effort + try: + self.close() + except Exception as e: + logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) + + def shutdown(self): + """Shutdown, cleaning up the background proc and IPC.""" + try: + self.close() + except Exception as e: + logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) + + async def generate( + self, + prompt: PromptType, + request_id: str, + sampling_params_list: Optional[ + Union[SamplingParams, Sequence[SamplingParams]] + ] = None, + lora_request: Optional[LoRARequest] = None, + trace_headers: Optional[Mapping[str, str]] = None, + priority: int = 0, + data_parallel_rank: Optional[int] = None, + ) -> AsyncGenerator[OmniRequestOutput, None]: + logger.debug("[Orchestrator] generate() called") + if sampling_params_list is None: + sampling_params_list = self.default_sampling_params_list + if len(sampling_params_list) != len(self.stage_list): + raise ValueError( + f"Expected {len(self.stage_list)} sampling params, \ + got {len(sampling_params_list)}" + ) + + # Orchestrator keeps stage objects for input derivation + num_stages = len(self.stage_list) + # Track per-request start time for end-to-end timing + _req_start_ts: dict[int, float] = {} + _wall_start_ts: float = time.time() + # _last_finish_ts: float = _wall_start_ts + + # Determine the final stage for E2E stats (highest stage_id with + # final_output=True; fallback to last stage) + final_stage_id_for_e2e = -1 + try: + for _sid, _st in enumerate(self.stage_list): + if getattr(_st, "final_output", False): + final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid) + if final_stage_id_for_e2e < 0: + final_stage_id_for_e2e = len(self.stage_list) - 1 + except Exception as e: + logger.debug( + "[Orchestrator] Failed to determine final stage for E2E; \ + falling back to last: %s", + e, + exc_info=True, + ) + final_stage_id_for_e2e = len(self.stage_list) - 1 + # Metrics/aggregation helper + metrics = OrchestratorMetrics( + num_stages, + self._enable_stats, + self._stats_file, + self._overall_stats_file, + _wall_start_ts, + ) + # Seed stage-0 queue with all requests + logger.debug("[Orchestrator] Seeding request into stage-0") + # Mark first input time for stage-0 + metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() + + sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] + task = { + "request_id": request_id, + "engine_inputs": prompt, + "sampling_params": sp0, + } + self.stage_list[0].submit(task) + _req_start_ts[request_id] = time.time() + logger.debug("[Orchestrator] Enqueued request %s to stage-0", request_id) + + logger.debug("[Orchestrator] Entering scheduling loop: stages=%d", num_stages) + finished = False + while not finished: + made_progress = False + for stage_id, stage in enumerate(self.stage_list): + result = stage.try_collect() + if result is None: + continue + + made_progress = True + 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"], + ) + continue + + if result.get("type") == "stage_ready": + # Only happens when stage is initialized slower than expected, + # so we wait for a short time and try again + time.sleep(0.05) + continue + + engine_outputs = _load( + result, obj_key="engine_outputs", shm_key="engine_outputs_shm" + ) + # Mark last output time for this stage whenever we receive outputs + metrics.stage_last_ts[stage_id] = max( + metrics.stage_last_ts[stage_id] or 0.0, time.time() + ) + try: + _m = result.get("metrics") + if _m is not None: + 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, + ) + logger.debug( + "[Orchestrator] Stage-%s completed request %s; \ + forwarding or finalizing", + stage_id, + req_id, + ) + stage.set_engine_outputs(engine_outputs) + + if getattr(stage, "final_output", False): + logger.debug( + "[Orchestrator] Request %s finalized at stage-%s", + req_id, + stage_id, + ) + + # End-to-end timing and time-per-token for final output + # (only once per request at the designated final stage) + try: + rid_int = ( + int(req_id) + if isinstance(req_id, (int, str)) and str(req_id).isdigit() + else req_id + ) + if ( + stage_id == final_stage_id_for_e2e + and rid_int not in metrics.e2e_done + ): + metrics.on_finalize_request( + stage_id, + req_id, + engine_outputs, + _req_start_ts.get(req_id, _wall_start_ts), + ) + except Exception as e: + logger.exception( + "[Orchestrator] Finalize request handling error for \ + req %s at stage %s: %s", + req_id, + stage_id, + e, + ) + + if isinstance(engine_outputs, list): + engine_outputs = engine_outputs[0] + yield OmniRequestOutput( + stage_id=stage_id, + final_output_type=stage.final_output_type, + request_output=engine_outputs, + ) + + next_stage_id = stage_id + 1 + if next_stage_id < num_stages: + next_stage: OmniStage = self.stage_list[next_stage_id] + next_inputs = next_stage.process_engine_inputs( + self.stage_list, prompt + ) + sp_next: SamplingParams = sampling_params_list[next_stage_id] + try: + # Measure transfer size and time (encode + enqueue) + size_bytes = 0 + try: + size_bytes = len(_ser(next_inputs)) + except Exception: + size_bytes = 0 + t0 = time.time() + ipc_payload = _encode( + next_inputs, + getattr(self, "_shm_threshold_bytes", 65536), + obj_key="engine_inputs", + shm_key="engine_inputs_shm", + ) + ipc_payload.update( + { + "request_id": req_id, + "sampling_params": sp_next, + "sent_ts": time.time(), + } + ) + self.stage_list[next_stage_id].submit(ipc_payload) + t1 = time.time() + tx_ms = (t1 - t0) * 1000.0 + metrics.on_forward( + stage_id, + next_stage_id, + req_id, + int(size_bytes), + float(tx_ms), + bool("engine_inputs_shm" in ipc_payload), + ) + except Exception as e: + logger.warning( + "[Orchestrator] IPC encode failed for req %s: %s; \ + falling back to inline payload", + req_id, + e, + ) + self.stage_list[next_stage_id].submit( + { + "request_id": req_id, + "engine_inputs": next_inputs, + "sampling_params": sp_next, + } + ) + logger.debug( + "[Orchestrator] Forwarded request %s to stage-%s", + req_id, + next_stage_id, + ) + else: + finished = True + logger.debug("[Orchestrator] Request %s fully completed", req_id) + + if not made_progress: + time.sleep(0.005) + + logger.debug("[Orchestrator] All requests completed") + + # Summarize and print stats + try: + summary = metrics.build_and_log_summary(final_stage_id_for_e2e) + logger.info("[Summary] %s", summary) + except Exception as e: + logger.exception("[Orchestrator] Failed to build/log summary: %s", e) + + def _wait_for_stages_ready(self, timeout: int = 120) -> None: + 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 + 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).", + ] + if getattr(self, "_log_file", None): + suggestions.append( + f"Inspect per-stage log files for details: \ + {self._log_file}.stage.log" + ) + 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", + ) + + # Attempt graceful shutdown of all stages before exiting + try: + self.close() + except Exception: + pass + + # Terminate the current process with non-zero exit code + try: + sys.exit(1) + except SystemExit: + raise + except Exception: + os._exit(1) + + @property + def is_running(self) -> bool: + # Is None before the loop is started. + return len(self._stage_in_queues) > 0 + + @property + def is_stopped(self) -> bool: + return self.errored + + @property + def errored(self) -> bool: + return not self.is_running + + @property + def dead_error(self) -> BaseException: + return EngineDeadError() + + async def abort(self, request_id: Union[str, Iterable[str]]) -> None: + pass + + async def get_vllm_config(self) -> VllmConfig: + for stage in self.stage_list: + if stage.is_comprehension: + # Use the vllm_config received from worker process + if stage.vllm_config is not None: + return stage.vllm_config + return None + + async def get_model_config(self) -> OmniModelConfig: + for stage in self.stage_list: + if stage.is_comprehension: + # Use the vllm_config received from worker process + if stage.vllm_config is not None: + return stage.vllm_config.model_config + 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: + for stage in self.stage_list: + if stage.is_comprehension: + return stage.tokenizer + return None + + async def is_tracing_enabled(self) -> bool: + for stage in self.stage_list: + if stage.is_comprehension: + return stage.is_tracing_enabled + return False + + async def do_log_stats(self) -> None: + pass + + async def check_health(self) -> None: + pass + + async def reset_mm_cache(self) -> None: + pass + + async def reset_prefix_cache(self, device: Optional[Device] = None) -> None: + pass + + async def sleep(self, level: int = 1) -> None: + pass + + async def wake_up(self, tags: Optional[list[str]] = None) -> None: + pass + + async def is_sleeping(self) -> bool: + """Check whether the engine is sleeping""" + return False + + async def add_lora(self, lora_request: LoRARequest) -> bool: + """Load a new LoRA adapter into the engine for future requests.""" + return False + + async def encode( + self, + *args, + **kwargs, + ): + """Generate outputs for a request from a pooling model.""" + raise NotImplementedError("encode() is not implemented for AsyncOmniLLM") + + async def start_profile(self) -> None: + raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM") + + async def stop_profile(self) -> None: + raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM") + + +class AsyncOmniStageLLM(AsyncLLM): + def __init__( + self, + engine_args: AsyncOmniEngineArgs, + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, + use_cached_outputs: bool = False, + log_requests: bool = True, + start_engine_loop: bool = True, + stat_loggers: Optional[list[StatLoggerFactory]] = None, + client_addresses: Optional[dict[str, str]] = None, + client_count: int = 1, + client_index: int = 0, + ) -> None: + """ + Create an AsyncLLM. + + Args: + vllm_config: global configuration. + executor_class: an Executor impl, e.g. MultiprocExecutor. + log_stats: Whether to log stats. + usage_context: Usage context of the LLM. + mm_registry: Multi-modal registry. + use_cached_outputs: Whether to use cached outputs. + log_requests: Whether to log requests. + start_engine_loop: Whether to start the engine loop. + stat_loggers: customized stat loggers for the engine. + If not provided, default stat loggers will be used. + PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE + IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE. + + Returns: + None + """ + if not envs.VLLM_USE_V1: + raise ValueError( + "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. " + "This should not happen. As a workaround, try using " + "AsyncLLMEngine.from_vllm_config(...) or explicitly set " + "VLLM_USE_V1=0 or 1 and report this issue on Github." + ) + + # Ensure we can serialize custom transformer configs + maybe_register_config_serialize_by_value() + + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.observability_config = vllm_config.observability_config + self.log_requests = log_requests + + self.log_stats = log_stats or (stat_loggers is not None) + if not log_stats and stat_loggers is not None: + logger.info( + "AsyncLLM created with log_stats=False and non-empty custom " + "logger list; enabling logging without default stat loggers" + ) + + if self.model_config.skip_tokenizer_init: + self.tokenizer = None + else: + # Tokenizer (+ ensure liveness if running in another process). + self.tokenizer = init_tokenizer_from_configs( + model_config=vllm_config.model_config + ) + + # Processor (converts Inputs --> EngineCoreRequests). + self.processor = OmniProcessor( + vllm_config=vllm_config, + tokenizer=self.tokenizer, + mm_registry=mm_registry, + ) + + # OutputProcessor (converts EngineCoreOutputs --> RequestOutput). + self.output_processor = MultimodalOutputProcessor( + tokenizer=self.tokenizer, + log_stats=self.log_stats, + engine_core_output_type=engine_args.engine_output_type, + ) + if self.observability_config.otlp_traces_endpoint is not None: + tracer = init_tracer( + "vllm.llm_engine", self.observability_config.otlp_traces_endpoint + ) + self.output_processor.tracer = tracer + + # EngineCore (starts the engine in background process). + self.engine_core = EngineCoreClient.make_async_mp_client( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=self.log_stats, + client_addresses=client_addresses, + client_count=client_count, + client_index=client_index, + ) + + # Loggers. + self.logger_manager: Optional[StatLoggerManager] = None + if self.log_stats: + self.logger_manager = StatLoggerManager( + vllm_config=vllm_config, + engine_idxs=self.engine_core.engine_ranks_managed, + custom_stat_loggers=stat_loggers, + enable_default_loggers=log_stats, + client_count=client_count, + ) + self.logger_manager.log_engine_initialized() + + self.output_handler: Optional[asyncio.Task] = None + try: + # Start output handler eagerly if we are in the asyncio eventloop. + asyncio.get_running_loop() + self._run_output_handler() + except RuntimeError: + pass + + if envs.VLLM_TORCH_PROFILER_DIR: + logger.info( + "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501 + envs.VLLM_TORCH_PROFILER_DIR, + ) + worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm" + self.profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + ], + with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK, + on_trace_ready=torch.profiler.tensorboard_trace_handler( + envs.VLLM_TORCH_PROFILER_DIR, worker_name=worker_name, use_gzip=True + ), + ) + else: + self.profiler = None + + @classmethod + @deprecate_kwargs( + "disable_log_requests", + additional_message=( + "This argument will have no effect. " "Use `enable_log_requests` instead." + ), + ) + def from_vllm_config( + cls, + vllm_config: VllmConfig, + engine_args: AsyncOmniEngineArgs, + start_engine_loop: bool = True, + usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, + stat_loggers: Optional[list[StatLoggerFactory]] = None, + enable_log_requests: bool = False, + disable_log_stats: bool = False, + client_addresses: Optional[dict[str, str]] = None, + client_count: int = 1, + client_index: int = 0, + disable_log_requests: bool = True, # Deprecated, will be removed + ) -> "AsyncLLM": + if not envs.VLLM_USE_V1: + raise ValueError( + "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. " + "This should not happen. As a workaround, try using " + "AsyncLLMEngine.from_vllm_config(...) or explicitly set " + "VLLM_USE_V1=0 or 1 and report this issue on Github." + ) + + # Create the LLMEngine. + return cls( + vllm_config=vllm_config, + executor_class=Executor.get_class(vllm_config), + start_engine_loop=start_engine_loop, + stat_loggers=stat_loggers, + log_requests=enable_log_requests, + log_stats=not disable_log_stats, + usage_context=usage_context, + client_addresses=client_addresses, + client_count=client_count, + client_index=client_index, + engine_args=engine_args, + ) diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 18d5ce9b6..3fb8bf188 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -2,24 +2,54 @@ CLI entry point for vLLM-omni that intercepts vLLM commands. """ +import importlib.metadata import sys -from vllm_omni.entrypoints.cli.serve import OmniServeCommand - def main(): """Main CLI entry point that intercepts vLLM commands.""" # Check if --omni flag is present - if "--omni" in sys.argv: - # Remove --omni flag and "serve" command, process with vLLM-omni - omni_args = [arg for arg in sys.argv[1:] if arg not in ["--omni", "serve"]] - omni_serve = OmniServeCommand() - omni_serve.run(omni_args) - else: - # Forward to original vLLM CLI + if "--omni" not in sys.argv: from vllm.entrypoints.cli.main import main as vllm_main vllm_main() + return + else: + import vllm_omni.entrypoints.cli.serve + from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup + from vllm.utils import FlexibleArgumentParser + + CMD_MODULES = [ + vllm_omni.entrypoints.cli.serve, + ] + + cli_env_setup() + + parser = FlexibleArgumentParser( + description="vLLM OMNI CLI", + epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"), + ) + parser.add_argument( + "-v", + "--version", + action="version", + version=importlib.metadata.version("vllm_omni"), + ) + subparsers = parser.add_subparsers(required=False, dest="subparser") + cmds = {} + for cmd_module in CMD_MODULES: + new_cmds = cmd_module.cmd_init() + for cmd in new_cmds: + cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) + cmds[cmd.name] = cmd + args = parser.parse_args() + if args.subparser in cmds: + cmds[args.subparser].validate(args) + + if hasattr(args, "dispatch_function"): + args.dispatch_function(args) + else: + parser.print_help() if __name__ == "__main__": diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 81645732e..5fda170fa 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -3,93 +3,67 @@ Omni serve command for vLLM-omni. """ import argparse -from typing import List + +import uvloop + +from vllm.entrypoints.cli.types import CLISubcommand +from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args +from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG +from vllm.logger import init_logger +from vllm.utils import FlexibleArgumentParser +from vllm_omni.entrypoints.openai.api_server import omni_run_server + +logger = init_logger(__name__) + +DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni LLM +completions via HTTP. Defaults to Qwen/Qwen2.5-Omni-7B if no model is specified. + +Search by using: `--help=` to explore options by section (e.g., +--help=ModelConfig, --help=Frontend) + Use `--help=all` to show all available flags at once. +""" -class OmniServeCommand: - """Command handler for vLLM-omni serve command.""" +class OmniServeCommand(CLISubcommand): + """The `serve` subcommand for the vLLM CLI.""" - def __init__(self): - self.parser = self._create_parser() + name = "serve" - def _create_parser(self) -> argparse.ArgumentParser: - """Create argument parser for omni serve command.""" - parser = argparse.ArgumentParser( - description="vLLM-omni: Multi-modality models inference and serving" + @staticmethod + def cmd(args: argparse.Namespace) -> None: + # If model is specified in CLI (as positional arg), it takes precedence + if hasattr(args, "model_tag") and args.model_tag is not None: + args.model = args.model_tag + + uvloop.run(omni_run_server(args)) + + def validate(self, args: argparse.Namespace) -> None: + validate_parsed_serve_args(args) + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + serve_parser = subparsers.add_parser( + self.name, + description=DESCRIPTION, + usage="vllm serve [model_tag] --omni [options]", ) - # Model arguments - make it optional with default - parser.add_argument( - "model", - nargs="?", - default="Qwen/Qwen3-0.6B", - help="Path to the model or model name (default: Qwen/Qwen3-0.6B)", - ) - - # Server arguments - parser.add_argument( - "--port", - type=int, - default=8000, - help="Port to run the server on", - ) - - parser.add_argument( - "--host", - type=str, - default="0.0.0.0", - help="Host to run the server on", - ) - - # Stage configuration arguments - parser.add_argument("--ar-stage", type=str, help="AR stage model path") - - parser.add_argument("--dit-stage", type=str, help="DiT stage model path") - - parser.add_argument( - "--dit-steps", - type=int, - default=50, - help="Number of DiT inference steps", - ) - - parser.add_argument( - "--dit-guidance-scale", - type=float, - default=7.5, - help="DiT guidance scale", - ) - - parser.add_argument( - "--use-diffusers", + serve_parser = make_arg_parser(serve_parser) + serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name) + serve_parser.add_argument( + "--omni", action="store_true", - help="Use diffusers pipeline for DiT stage", + help="Enable vLLM-omni mode", ) - - # Other arguments - parser.add_argument( - "--log-stats", action="store_true", help="Enable logging statistics" + serve_parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.", ) + return serve_parser - return parser - def run(self, args: List[str]) -> None: - """Run the omni serve command.""" - pass - - def _create_stage_configs(self, args) -> List: - """Create stage configurations based on arguments.""" - pass - - async def _start_server(self, omni_llm, args) -> None: - """Start the API server.""" - try: - # Import here to avoid circular imports - from vllm_omni.entrypoints.api_server import run_server - - await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port) - except KeyboardInterrupt: - print("\nShutting down server...") - except Exception as e: - print(f"Error starting server: {e}") - raise +def cmd_init() -> list[CLISubcommand]: + return [OmniServeCommand()] diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py index 11b98ed81..8f616db0f 100644 --- a/vllm_omni/entrypoints/log_utils.py +++ b/vllm_omni/entrypoints/log_utils.py @@ -1,10 +1,9 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional - import logging import os import time +from typing import Any, Dict, List, Optional from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl @@ -49,22 +48,32 @@ def remove_old_logs(log_file: Optional[str], num_stages: int) -> None: pass -def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None: +def configure_orchestrator_logger( + logger: logging.Logger, log_file: Optional[str] +) -> None: try: if not log_file: return - has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers) + has_file_handler = any( + isinstance(h, logging.FileHandler) for h in logger.handlers + ) if not has_file_handler: fh = logging.FileHandler(log_file) fh.setLevel(logging.DEBUG) - fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s")) + fh.setFormatter( + logging.Formatter( + "%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s" + ) + ) logger.addHandler(fh) logger.setLevel(logging.DEBUG) except Exception: pass -def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]: +def init_stats_paths( + enable_stats: bool, log_file: Optional[str] +) -> tuple[Optional[str], Optional[str]]: stats_file: Optional[str] = None overall_file: Optional[str] = None try: @@ -86,122 +95,217 @@ def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None: pass -def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any, - size_bytes: int, tx_time_ms: float, used_shm: bool) -> None: - _safe_append_jsonl(stats_file, { - "type": "transfer_stats", - "from_stage": from_stage, - "to_stage": to_stage, - "request_id": request_id, - "size_bytes": int(size_bytes), - "tx_time_ms": float(tx_time_ms), - "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0), - "used_shm": bool(used_shm), - }) +def log_transfer_tx( + stats_file: Optional[str], + from_stage: int, + to_stage: int, + request_id: Any, + size_bytes: int, + tx_time_ms: float, + used_shm: bool, +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "transfer_stats", + "from_stage": from_stage, + "to_stage": to_stage, + "request_id": request_id, + "size_bytes": int(size_bytes), + "tx_time_ms": float(tx_time_ms), + "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0), + "used_shm": bool(used_shm), + }, + ) -def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any, - rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None: - _safe_append_jsonl(stats_file, { - "type": "transfer_rx_stats", - "from_stage": from_stage, - "to_stage": to_stage, - "request_id": request_id, - "rx_bytes": int(rx_bytes), - "rx_decode_time_ms": float(rx_decode_time_ms), - "in_flight_time_ms": float(in_flight_time_ms), - "rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0, - }) +def log_transfer_rx( + stats_file: Optional[str], + from_stage: int, + to_stage: int, + request_id: Any, + rx_bytes: int, + rx_decode_time_ms: float, + in_flight_time_ms: float, +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "transfer_rx_stats", + "from_stage": from_stage, + "to_stage": to_stage, + "request_id": request_id, + "rx_bytes": int(rx_bytes), + "rx_decode_time_ms": float(rx_decode_time_ms), + "in_flight_time_ms": float(in_flight_time_ms), + "rx_time_per_kb_ms": ( + (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) + if rx_bytes > 0 + else 0.0 + ), + }, + ) -def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any, - size_bytes: int, tx_time_ms: float, in_flight_time_ms: float, - rx_decode_time_ms: float, total_time_ms: float) -> None: - _safe_append_jsonl(stats_file, { - "type": "transfer_total_stats", - "from_stage": from_stage, - "to_stage": to_stage, - "request_id": request_id, - "size_bytes": int(size_bytes), - "tx_time_ms": float(tx_time_ms), - "in_flight_time_ms": float(in_flight_time_ms), - "rx_decode_time_ms": float(rx_decode_time_ms), - "total_time_ms": float(total_time_ms), - "total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0, - }) +def log_transfer_total( + stats_file: Optional[str], + from_stage: int, + to_stage: int, + request_id: Any, + size_bytes: int, + tx_time_ms: float, + in_flight_time_ms: float, + rx_decode_time_ms: float, + total_time_ms: float, +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "transfer_total_stats", + "from_stage": from_stage, + "to_stage": to_stage, + "request_id": request_id, + "size_bytes": int(size_bytes), + "tx_time_ms": float(tx_time_ms), + "in_flight_time_ms": float(in_flight_time_ms), + "rx_decode_time_ms": float(rx_decode_time_ms), + "total_time_ms": float(total_time_ms), + "total_time_per_kb_ms": ( + float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) + if size_bytes > 0 + else 0.0 + ), + }, + ) -def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int, - e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None: - _safe_append_jsonl(stats_file, { - "type": "orchestrator_request_e2e", - "request_id": request_id, - "final_stage_id": final_stage_id, - "e2e_time_ms": float(e2e_time_ms), - "num_tokens_out": int(num_tokens_out), - "e2e_time_per_token_ms": float(e2e_time_per_token_ms), - }) +def log_orchestrator_e2e( + stats_file: Optional[str], + request_id: Any, + final_stage_id: int, + e2e_time_ms: float, + num_tokens_out: int, + e2e_time_per_token_ms: float, +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "orchestrator_request_e2e", + "request_id": request_id, + "final_stage_id": final_stage_id, + "e2e_time_ms": float(e2e_time_ms), + "num_tokens_out": int(num_tokens_out), + "e2e_time_per_token_ms": float(e2e_time_per_token_ms), + }, + ) -def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None: +def log_orchestrator_summary( + stats_file: Optional[str], summary: Dict[str, Any] +) -> None: _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary}) -def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None: +def log_overall_summary( + overall_stats_file: Optional[str], summary: Dict[str, Any] +) -> None: _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary}) -def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None: +def log_overall_record( + overall_stats_file: Optional[str], record: Dict[str, Any] +) -> None: _safe_append_jsonl(overall_stats_file, record) -def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int, - num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float, - rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None: - _safe_append_jsonl(stats_file, { - "type": "stage_request_stats", - "stage_id": stage_id, - "request_id": request_id, - "batch_size": int(batch_size), - "num_tokens_out": int(num_tokens_out), - "stage_gen_time_ms": float(stage_gen_time_ms), - "tokens_per_s": float(tokens_per_s), - "rx_transfer_bytes": int(rx_transfer_bytes), - "rx_decode_time_ms": float(rx_decode_time_ms), - "rx_mbps": float(rx_mbps), - }) +def log_stage_request_stats( + stats_file: Optional[str], + stage_id: int, + request_id: Any, + batch_size: int, + num_tokens_out: int, + stage_gen_time_ms: float, + tokens_per_s: float, + rx_transfer_bytes: int, + rx_decode_time_ms: float, + rx_mbps: float, +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "stage_request_stats", + "stage_id": stage_id, + "request_id": request_id, + "batch_size": int(batch_size), + "num_tokens_out": int(num_tokens_out), + "stage_gen_time_ms": float(stage_gen_time_ms), + "tokens_per_s": float(tokens_per_s), + "rx_transfer_bytes": int(rx_transfer_bytes), + "rx_decode_time_ms": float(rx_decode_time_ms), + "rx_mbps": float(rx_mbps), + }, + ) -def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int, - total_gen_time_ms: float, avg_tokens_per_s: float) -> None: - _safe_append_jsonl(stats_file, { - "type": "stage_running_avg", - "stage_id": stage_id, - "total_tokens": int(total_tokens), - "total_gen_time_ms": float(total_gen_time_ms), - "avg_tokens_per_s": float(avg_tokens_per_s), - }) +def log_stage_running_avg( + stats_file: Optional[str], + stage_id: int, + total_tokens: int, + total_gen_time_ms: float, + avg_tokens_per_s: float, +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "stage_running_avg", + "stage_id": stage_id, + "total_tokens": int(total_tokens), + "total_gen_time_ms": float(total_gen_time_ms), + "avg_tokens_per_s": float(avg_tokens_per_s), + }, + ) -def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int, - batch_gen_time_ms: float, request_ids: List[Any]) -> None: - _safe_append_jsonl(stats_file, { - "type": "stage_batch_stats", - "stage_id": stage_id, - "batch_size": int(batch_size), - "batch_gen_time_ms": float(batch_gen_time_ms), - "request_ids": list(request_ids), - }) +def log_stage_batch_stats( + stats_file: Optional[str], + stage_id: int, + batch_size: int, + batch_gen_time_ms: float, + request_ids: List[Any], +) -> None: + _safe_append_jsonl( + stats_file, + { + "type": "stage_batch_stats", + "stage_id": stage_id, + "batch_size": int(batch_size), + "batch_gen_time_ms": float(batch_gen_time_ms), + "request_ids": list(request_ids), + }, + ) -def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, - batch_size: int, engine_outputs: List[Any], - stage_gen_time_ms: float, rx_transfer_bytes: int, - rx_decode_time_ms: float) -> None: +def compute_and_log_stage_request_stats( + stats_file: Optional[str], + stage_id: int, + request_id: Any, + batch_size: int, + engine_outputs: List[Any], + stage_gen_time_ms: float, + rx_transfer_bytes: int, + rx_decode_time_ms: float, +) -> None: """Compute per-request metrics and log them in one call.""" num_tokens = count_tokens_from_outputs(engine_outputs) - tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0 - rx_mbps = (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) if rx_transfer_bytes > 0 else 0.0 + tokens_per_s = ( + (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0 + ) + rx_mbps = ( + (float(rx_transfer_bytes) * 8.0) + / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) + if rx_transfer_bytes > 0 + else 0.0 + ) log_stage_request_stats( stats_file, stage_id, @@ -218,18 +322,23 @@ def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int # ----------------- Aggregation helpers for orchestrator ----------------- -def record_stage_metrics(per_request: Dict[int, Dict[str, Any]], - stage_req_counts: List[int], - stage_total_time_ms: List[float], - stage_total_tokens: List[int], - stage_id: int, - req_id: Any, - metrics: Dict[str, Any]) -> None: + +def record_stage_metrics( + per_request: Dict[int, Dict[str, Any]], + stage_req_counts: List[int], + stage_total_time_ms: List[float], + stage_total_tokens: List[int], + stage_id: int, + req_id: Any, + metrics: Dict[str, Any], +) -> None: try: stage_req_counts[stage_id] += 1 stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0)) rid_int = int(req_id) - pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr = per_request.setdefault( + rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0} + ) pr_stages = pr["stages"] # type: ignore[index] pr_stages[stage_id] = { "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)), @@ -239,23 +348,32 @@ def record_stage_metrics(per_request: Dict[int, Dict[str, Any]], pass -def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]], - transfer_agg: Dict[tuple[int, int], Dict[str, float]], - per_request: Dict[int, Dict[str, Any]], - stage_id: int, - req_id: Any, - rx_bytes: float, - rx_ms: float, - in_flight_ms: float) -> Optional[tuple[int, float, float]]: +def aggregate_rx_and_maybe_total( + transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]], + transfer_agg: Dict[tuple[int, int], Dict[str, float]], + per_request: Dict[int, Dict[str, Any]], + stage_id: int, + req_id: Any, + rx_bytes: float, + rx_ms: float, + in_flight_ms: float, +) -> Optional[tuple[int, float, float]]: try: # Update RX aggregates for (stage_id-1 -> stage_id) if stage_id > 0: key = (stage_id - 1, stage_id) agg = transfer_agg.get(key) if agg is None: - agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, - "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, - "sum_total_ms": 0.0, "total_count": 0.0} + agg = { + "sum_bytes": 0.0, + "sum_ms": 0.0, + "count": 0.0, + "sum_rx_bytes": 0.0, + "sum_rx_ms": 0.0, + "rx_count": 0.0, + "sum_total_ms": 0.0, + "total_count": 0.0, + } transfer_agg[key] = agg agg["sum_rx_bytes"] += float(rx_bytes) agg["sum_rx_ms"] += float(rx_ms) @@ -276,7 +394,9 @@ def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], D agg["total_count"] += 1.0 # accumulate per-request transfer totals try: - pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr = per_request.setdefault( + rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0} + ) pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index] pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index] except Exception: @@ -287,20 +407,29 @@ def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], D return None -def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]], - transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]], - from_stage: int, - to_stage: int, - req_id: Any, - size_bytes: int, - tx_ms: float) -> None: +def record_sender_transfer_agg( + transfer_agg: Dict[tuple[int, int], Dict[str, float]], + transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]], + from_stage: int, + to_stage: int, + req_id: Any, + size_bytes: int, + tx_ms: float, +) -> None: try: key = (from_stage, to_stage) agg = transfer_agg.get(key) if agg is None: - agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, - "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, - "sum_total_ms": 0.0, "total_count": 0.0} + agg = { + "sum_bytes": 0.0, + "sum_ms": 0.0, + "count": 0.0, + "sum_rx_bytes": 0.0, + "sum_rx_ms": 0.0, + "rx_count": 0.0, + "sum_total_ms": 0.0, + "total_count": 0.0, + } transfer_agg[key] = agg agg["sum_bytes"] += float(size_bytes) agg["sum_ms"] += float(tx_ms) @@ -332,9 +461,11 @@ def count_tokens_from_outputs(engine_outputs: List[Any]) -> int: return total -def build_stage_summary(stage_req_counts: List[int], - stage_total_tokens: List[int], - stage_total_time_ms: List[float]) -> List[Dict[str, Any]]: +def build_stage_summary( + stage_req_counts: List[int], + stage_total_tokens: List[int], + stage_total_time_ms: List[float], +) -> List[Dict[str, Any]]: summary: List[Dict[str, Any]] = [] for sid in range(len(stage_req_counts)): reqs = stage_req_counts[sid] @@ -342,53 +473,74 @@ def build_stage_summary(stage_req_counts: List[int], total_ms = float(stage_total_time_ms[sid]) avg_req = (total_ms / reqs) if reqs > 0 else 0.0 avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 - summary.append({ - "stage_id": sid, - "requests": int(reqs), - "tokens": int(tokens), - "total_time_ms": total_ms, - "avg_time_per_request_ms": avg_req, - "avg_tokens_per_s": avg_tok, - }) + summary.append( + { + "stage_id": sid, + "requests": int(reqs), + "tokens": int(tokens), + "total_time_ms": total_ms, + "avg_time_per_request_ms": avg_req, + "avg_tokens_per_s": avg_tok, + } + ) return summary -def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]: +def build_transfer_summary( + transfer_agg: Dict[tuple[int, int], Dict[str, float]] +) -> List[Dict[str, Any]]: summary: List[Dict[str, Any]] = [] for (src, dst), agg in transfer_agg.items(): sum_bytes = float(agg.get("sum_bytes", 0.0)) sum_ms = float(agg.get("sum_ms", 0.0)) samples = int(agg.get("count", 0.0)) - tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + tx_mbps = ( + (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + ) sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0)) sum_rx_ms = float(agg.get("sum_rx_ms", 0.0)) samples_rx = int(agg.get("rx_count", 0.0)) - rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0 + rx_mbps = ( + (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) + if sum_rx_bytes > 0 + else 0.0 + ) sum_total_ms = float(agg.get("sum_total_ms", 0.0)) samples_total = int(agg.get("total_count", 0.0)) - total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 - summary.append({ - "from_stage": src, - "to_stage": dst, - "samples": samples, - "total_bytes": int(sum_bytes), - "total_time_ms": sum_ms, - "tx_mbps": tx_mbps, - "rx_samples": samples_rx, - "rx_total_bytes": int(sum_rx_bytes), - "rx_total_time_ms": sum_rx_ms, - "rx_mbps": rx_mbps, - "total_samples": samples_total, - "total_transfer_time_ms": sum_total_ms, - "total_mbps": total_mbps, - }) + total_mbps = ( + (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) + if sum_bytes > 0 + else 0.0 + ) + summary.append( + { + "from_stage": src, + "to_stage": dst, + "samples": samples, + "total_bytes": int(sum_bytes), + "total_time_ms": sum_ms, + "tx_mbps": tx_mbps, + "rx_samples": samples_rx, + "rx_total_bytes": int(sum_rx_bytes), + "rx_total_time_ms": sum_rx_ms, + "rx_mbps": rx_mbps, + "total_samples": samples_total, + "total_transfer_time_ms": sum_total_ms, + "total_mbps": total_mbps, + } + ) return summary class OrchestratorMetrics: - def __init__(self, num_stages: int, enable_stats: bool, - stats_file: Optional[str], overall_stats_file: Optional[str], - wall_start_ts: float) -> None: + def __init__( + self, + num_stages: int, + enable_stats: bool, + stats_file: Optional[str], + overall_stats_file: Optional[str], + wall_start_ts: float, + ) -> None: self.num_stages = int(num_stages) self.enable_stats = bool(enable_stats) self.stats_file = stats_file @@ -406,11 +558,19 @@ class OrchestratorMetrics: self.sum_per_request_transfer_ms: float = 0.0 self.wall_start_ts: float = float(wall_start_ts) self.last_finish_ts: float = float(wall_start_ts) - self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)} - self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)] - self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)] + self.stage_seen_batches: Dict[int, set] = { + sid: set() for sid in range(self.num_stages) + } + self.stage_first_ts: List[Optional[float]] = [ + None for _ in range(self.num_stages) + ] + self.stage_last_ts: List[Optional[float]] = [ + None for _ in range(self.num_stages) + ] - def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None: + def on_stage_metrics( + self, stage_id: int, req_id: Any, metrics: Dict[str, Any] + ) -> None: record_stage_metrics( self.per_request, self.stage_req_counts, @@ -425,7 +585,9 @@ class OrchestratorMetrics: if batch_id_raw is not None: batch_id = int(batch_id_raw) if batch_id not in self.stage_seen_batches[stage_id]: - self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0)) + self.stage_total_time_ms[stage_id] += float( + metrics.get("stage_gen_time_ms", 0.0) + ) self.stage_seen_batches[stage_id].add(batch_id) except Exception: pass @@ -466,8 +628,15 @@ class OrchestratorMetrics: float(total_ms_c), ) - def on_forward(self, from_stage: int, to_stage: int, req_id: Any, - size_bytes: int, tx_ms: float, used_shm: bool) -> None: + def on_forward( + self, + from_stage: int, + to_stage: int, + req_id: Any, + size_bytes: int, + tx_ms: float, + used_shm: bool, + ) -> None: # Mark first input time for the destination stage if not set if self.stage_first_ts[to_stage] is None: self.stage_first_ts[to_stage] = time.time() @@ -491,8 +660,9 @@ class OrchestratorMetrics: float(tx_ms), ) - def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any], - req_start_ts: float) -> None: + def on_finalize_request( + self, stage_id: int, req_id: Any, engine_outputs: List[Any], req_start_ts: float + ) -> None: try: rid_int = int(req_id) except Exception: @@ -509,7 +679,9 @@ class OrchestratorMetrics: self.e2e_total_tokens += int(num_tokens) self.e2e_count += 1 self.e2e_done.add(rid_int) - pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr = self.per_request.setdefault( + rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0} + ) per_req_record = { "type": "overall_request", "request_id": rid_int, @@ -524,7 +696,9 @@ class OrchestratorMetrics: log_overall_record(self.overall_stats_file, per_req_record) if self.enable_stats and self.stats_file: e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0 - log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt) + log_orchestrator_e2e( + self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt + ) def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]: # Compute stage summary using wall time between first input and last output per stage @@ -532,22 +706,34 @@ class OrchestratorMetrics: for sid in range(self.num_stages): first_ts = self.stage_first_ts[sid] last_ts = self.stage_last_ts[sid] - total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0 + total_ms = ( + (max(0.0, (last_ts - first_ts)) * 1000.0) + if (first_ts is not None and last_ts is not None) + else 0.0 + ) reqs = self.stage_req_counts[sid] tokens = self.stage_total_tokens[sid] avg_req = (total_ms / reqs) if reqs > 0 else 0.0 avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 - stage_summary.append({ - "stage_id": sid, - "requests": int(reqs), - "tokens": int(tokens), - "total_time_ms": float(total_ms), - "avg_time_per_request_ms": float(avg_req), - "avg_tokens_per_s": float(avg_tok), - }) + stage_summary.append( + { + "stage_id": sid, + "requests": int(reqs), + "tokens": int(tokens), + "total_time_ms": float(total_ms), + "avg_time_per_request_ms": float(avg_req), + "avg_tokens_per_s": float(avg_tok), + } + ) transfer_summary = build_transfer_summary(self.transfer_agg) - e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 - e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0 + e2e_avg_req = ( + (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 + ) + e2e_avg_tok = ( + (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) + if self.e2e_total_ms > 0 + else 0.0 + ) wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0) summary: Dict[str, Any] = { "e2e_requests": int(self.e2e_count), @@ -566,5 +752,3 @@ class OrchestratorMetrics: if self.enable_stats and self.overall_stats_file: log_overall_summary(self.overall_stats_file, summary) return summary - - diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index c1b12a1b2..2eb493e80 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,11 +1,9 @@ -from typing import Any, Optional, Sequence, Union -import logging import multiprocessing as mp -from concurrent.futures import ThreadPoolExecutor, as_completed import os -import queue import sys import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Optional, Sequence, Union import cloudpickle from pydantic import ValidationError @@ -26,18 +24,19 @@ from vllm.v1.engine.llm_engine import LLMEngine from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.processor import OmniProcessor -from vllm_omni.entrypoints.omni_stage import OmniStage -from vllm_omni.entrypoints.utils import load_stage_configs_from_model from vllm_omni.entrypoints.log_utils import ( - remove_old_logs, + OrchestratorMetrics, configure_orchestrator_logger, init_stats_paths, - OrchestratorMetrics, + remove_old_logs, ) -from vllm_omni.entrypoints.stage_utils import ( - maybe_load_from_ipc as _load, - encode_for_ipc as _encode, - serialize_obj as _ser, +from vllm_omni.entrypoints.omni_stage import OmniStage +from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode +from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load +from vllm_omni.entrypoints.stage_utils import serialize_obj as _ser +from vllm_omni.entrypoints.utils import ( + load_stage_configs_from_model, + load_stage_configs_from_yaml, ) from vllm_omni.outputs import OmniRequestOutput @@ -52,23 +51,25 @@ class OmniLLM: - max_inflight=1 per stage (serial within a stage), but pipeline across stages """ - def __init__(self, model: str, - stage_configs=None, - log_stats: bool = False, - log_file: Optional[str] = None, - init_sleep_seconds: int = 20, - shm_threshold_bytes: int = 65536, - batch_timeout: int = 10, - init_timeout: int = 300, - **kwargs): + def __init__( + self, + model: str, + stage_configs_path: Optional[str] = None, + log_stats: bool = False, + log_file: Optional[str] = None, + init_sleep_seconds: int = 20, + shm_threshold_bytes: int = 65536, + batch_timeout: int = 10, + init_timeout: int = 300, + **kwargs, + ): self.batch_timeout = batch_timeout self._enable_stats: bool = bool(log_stats) # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent. - if stage_configs is None: + if stage_configs_path is None: self.stage_configs = load_stage_configs_from_model(model) else: - self.stage_configs = stage_configs - + self.stage_configs = load_stage_configs_from_yaml(stage_configs_path) # Optional file handler for orchestrator self._log_file = log_file @@ -76,18 +77,34 @@ class OmniLLM: remove_old_logs(self._log_file, len(self.stage_list)) configure_orchestrator_logger(logger, self._log_file) - self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) - self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout) + self._stats_file, self._overall_stats_file = init_stats_paths( + self._enable_stats, self._log_file + ) + self._initialize_stages( + model, init_sleep_seconds, shm_threshold_bytes, init_timeout + ) - def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> 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)] + 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()) @@ -104,7 +121,6 @@ class OmniLLM: # 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_stage_processes(self, model: str) -> None: for stage_id, stage in enumerate(self.stage_list): @@ -131,7 +147,10 @@ class OmniLLM: try: q.put_nowait(None) except Exception as e: - logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", 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() @@ -150,10 +169,27 @@ class OmniLLM: sampling_params_list: Optional[ Union[SamplingParams, Sequence[SamplingParams]] ] = None, + ) -> list[OmniRequestOutput]: + try: + return self._run_generation(prompts, sampling_params_list) + except Exception as e: + logger.exception("[Orchestrator] Failed to run generation: %s", e) + raise e + finally: + self.close() + + def _run_generation( + self, + prompts: Union[PromptType, Sequence[PromptType]], + sampling_params_list: Optional[ + Union[SamplingParams, Sequence[SamplingParams]] + ] = None, ) -> list[OmniRequestOutput]: logger.debug("[Orchestrator] generate() called") if sampling_params_list is None: - raise ValueError("sampling_params_list is required for pipelined generation") + raise ValueError( + "sampling_params_list is required for pipelined generation" + ) if len(sampling_params_list) != len(self.stage_list): raise ValueError( f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}" @@ -171,12 +207,13 @@ class OmniLLM: num_stages = len(self.stage_list) # Map from request_id to original prompt - request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)} + request_id_to_prompt: dict[int, PromptType] = { + i: p for i, p in enumerate(request_prompts) + } # Track per-request start time for end-to-end timing _req_start_ts: dict[int, float] = {} _wall_start_ts: float = time.time() - _last_finish_ts: float = _wall_start_ts # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage) final_stage_id_for_e2e = -1 @@ -187,13 +224,25 @@ class OmniLLM: if final_stage_id_for_e2e < 0: final_stage_id_for_e2e = len(self.stage_list) - 1 except Exception as e: - logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True) + logger.debug( + "[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", + e, + exc_info=True, + ) final_stage_id_for_e2e = len(self.stage_list) - 1 # Metrics/aggregation helper - metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts) + metrics = OrchestratorMetrics( + num_stages, + self._enable_stats, + self._stats_file, + self._overall_stats_file, + _wall_start_ts, + ) # Seed stage-0 queue with all requests - logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)) + logger.debug( + "[Orchestrator] Seeding %d requests into stage-0", len(request_prompts) + ) # Mark first input time for stage-0 metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() @@ -214,7 +263,11 @@ class OmniLLM: completed_requests = 0 total_requests = len(request_prompts) - logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages) + logger.debug( + "[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", + total_requests, + num_stages, + ) while completed_requests < total_requests: made_progress = False for stage_id, stage in enumerate(self.stage_list): @@ -225,24 +278,42 @@ class OmniLLM: made_progress = True 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"]) + logger.error( + "Stage %s error on request %s: %s", + stage_id, + req_id, + result["error"], + ) continue - + if result.get("type") == "stage_ready": - #Only happens when stage is initialized slower than expected, so we wait for a short time and try again + # Only happens when stage is initialized slower than expected, so we wait for a short time and try again time.sleep(0.05) continue - engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") + engine_outputs = _load( + result, obj_key="engine_outputs", shm_key="engine_outputs_shm" + ) # Mark last output time for this stage whenever we receive outputs - metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) + metrics.stage_last_ts[stage_id] = max( + metrics.stage_last_ts[stage_id] or 0.0, time.time() + ) try: _m = result.get("metrics") if _m is not None: 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) - logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id) + logger.exception( + "[Orchestrator] Failed to process metrics for stage %s, req %s: %s", + stage_id, + req_id, + e, + ) + logger.debug( + "[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", + stage_id, + req_id, + ) stage.set_engine_outputs(engine_outputs) if getattr(stage, "final_output", False): @@ -253,20 +324,43 @@ class OmniLLM: request_output=engine_outputs, ) ) - logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id) + logger.debug( + "[Orchestrator] Request %s finalized at stage-%s", + req_id, + stage_id, + ) # End-to-end timing and time-per-token for final output (only once per request at the designated final stage) try: - rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id - if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done: - metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts)) + rid_int = ( + int(req_id) + if isinstance(req_id, (int, str)) and str(req_id).isdigit() + else req_id + ) + if ( + stage_id == final_stage_id_for_e2e + and rid_int not in metrics.e2e_done + ): + metrics.on_finalize_request( + stage_id, + req_id, + engine_outputs, + _req_start_ts.get(req_id, _wall_start_ts), + ) except Exception as e: - logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e) + logger.exception( + "[Orchestrator] Finalize request handling error for req %s at stage %s: %s", + req_id, + stage_id, + e, + ) next_stage_id = stage_id + 1 if next_stage_id < num_stages: next_stage: OmniStage = self.stage_list[next_stage_id] - next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) + next_inputs = next_stage.process_engine_inputs( + self.stage_list, [request_id_to_prompt[req_id]] + ) sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index] try: # Measure transfer size and time (encode + enqueue) @@ -282,27 +376,51 @@ class OmniLLM: obj_key="engine_inputs", shm_key="engine_inputs_shm", ) - ipc_payload.update({ - "request_id": req_id, - "sampling_params": sp_next, - "sent_ts": time.time(), - }) + ipc_payload.update( + { + "request_id": req_id, + "sampling_params": sp_next, + "sent_ts": time.time(), + } + ) self.stage_list[next_stage_id].submit(ipc_payload) t1 = time.time() tx_ms = (t1 - t0) * 1000.0 - metrics.on_forward(stage_id, next_stage_id, req_id, int(size_bytes), float(tx_ms), bool("engine_inputs_shm" in ipc_payload)) + metrics.on_forward( + stage_id, + next_stage_id, + req_id, + int(size_bytes), + float(tx_ms), + bool("engine_inputs_shm" in ipc_payload), + ) except Exception as e: - logger.warning("[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", req_id, e) - self.stage_list[next_stage_id].submit({ - "request_id": req_id, - "engine_inputs": next_inputs, - "sampling_params": sp_next, - }) - logger.debug("[Orchestrator] Forwarded request %s to stage-%s", req_id, next_stage_id) + logger.warning( + "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", + req_id, + e, + ) + self.stage_list[next_stage_id].submit( + { + "request_id": req_id, + "engine_inputs": next_inputs, + "sampling_params": sp_next, + } + ) + logger.debug( + "[Orchestrator] Forwarded request %s to stage-%s", + req_id, + 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) + logger.debug( + "[Orchestrator] Request %s fully completed (%d/%d)", + req_id, + completed_requests, + total_requests, + ) if not made_progress: time.sleep(0.005) @@ -341,7 +459,9 @@ class OmniLLM: 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, + len(self._stages_ready), + num_stages, + not_ready, ) # Provide actionable suggestions before shutdown try: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 8f95cf62e..bbb33c41d 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -1,44 +1,58 @@ """ Stage manager for orchestrating multiple engines in vLLM-omni. +Enhanced to encapsulate per-stage process lifecycle and worker logic +(device setup, LLM init, batching, shared-memory IPC), while preserving +the original input processing utilities for cross-stage data wiring. + Enhanced to encapsulate per-stage process lifecycle and worker logic (device setup, LLM init, batching, shared-memory IPC), while preserving the original input processing utilities for cross-stage data wiring. """ +import asyncio import importlib import logging import multiprocessing as mp from typing import Any, Dict, List, Optional, Union -from vllm_omni.entrypoints.stage_utils import ( - _to_dict, - set_stage_gpu_devices, - maybe_load_from_ipc_with_metrics, - maybe_dump_to_shm, -) - from vllm.inputs import TextPrompt +from vllm.inputs.preprocess import InputPreprocessor +from vllm.logger import init_logger +from vllm.sampling_params import SamplingParams +from vllm.transformers_utils.tokenizer import AnyTokenizer +from vllm.usage.usage_lib import UsageContext from vllm.v1.engine import EngineCoreOutput from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.llm_engine import LLMEngine +from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs +from vllm_omni.entrypoints.stage_utils import ( + _to_dict, + maybe_dump_to_shm, + maybe_load_from_ipc_with_metrics, + set_stage_gpu_devices, +) from vllm_omni.inputs.data import OmniTokensPrompt +logger = init_logger(__name__) + class OmniStage: def __init__(self, stage_config): self.stage_config = stage_config self.engine = None self.async_engine = None + self.vllm_config = None + self.tokenizer = None + self.input_preprocessor = None + self.is_tracing_enabled = False self.stage_id = stage_config.stage_id self.engine_args = stage_config.engine_args self.model_stage = stage_config.engine_args.model_stage - if hasattr(stage_config, "engine_input_source"): - self.engine_input_source = stage_config.engine_input_source - else: - self.engine_input_source = [] + self.engine_input_source = getattr(stage_config, "engine_input_source", []) self.engine_output_type = stage_config.engine_args.engine_output_type self.engine_outputs = None + self.is_comprehension = getattr(stage_config, "is_comprehension", False) if hasattr(stage_config, "custom_process_input_func"): # Import the module specified in the config (already a full module path) module_path, func_name = stage_config.custom_process_input_func.rsplit( @@ -49,16 +63,12 @@ class OmniStage: else: self.custom_process_input_func = None - if hasattr(stage_config, "final_output"): - self.final_output = stage_config.final_output - else: - self.final_output = False - - if hasattr(stage_config, "final_output_type"): - self.final_output_type = stage_config.final_output_type - else: - self.final_output_type = None - + self.final_output = getattr(stage_config, "final_output", False) + self.final_output_type = getattr(stage_config, "final_output_type", None) + default_sampling_params = getattr(stage_config, "default_sampling_params", {}) + self.default_sampling_params = SamplingParams( + **_to_dict(default_sampling_params) + ) # Runtime orchestration state (added) self._in_q: Optional[mp.Queue] = None self._out_q: Optional[mp.Queue] = None @@ -75,6 +85,22 @@ class OmniStage: """Initialize the async engine for the 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).""" + self.vllm_config = vllm_config + + def set_tokenizer(self, tokenizer: AnyTokenizer) -> None: + """Set the tokenizer for the stage (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).""" + 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).""" + self.is_tracing_enabled = is_tracing_enabled + def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None: """Set the engine output for the stage.""" self.engine_outputs = engine_outputs @@ -88,12 +114,15 @@ class OmniStage: self, model: str, *, + is_async: bool = False, log_file: Optional[str] = None, shm_threshold_bytes: int = 65536, ctx: Optional[mp.context.BaseContext] = None, batch_timeout: int = 10, ) -> None: - assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process" + 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 ctx = ctx or mp.get_context("spawn") @@ -106,10 +135,31 @@ class OmniStage: "runtime": runtime_cfg, "shm_threshold_bytes": self._shm_threshold_bytes, } - self._proc = ctx.Process( - target=_stage_worker, - args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout), - ) + if is_async: + self._proc = ctx.Process( + target=_stage_worker_async_entry, + args=( + self, + model, + stage_payload, + self._in_q, + self._out_q, + self._log_file, + batch_timeout, + ), + ) + else: + self._proc = ctx.Process( + target=_stage_worker, + args=( + model, + stage_payload, + self._in_q, + self._out_q, + self._log_file, + batch_timeout, + ), + ) self._proc.start() def stop_stage_worker(self) -> None: @@ -117,17 +167,23 @@ class OmniStage: try: self._in_q.put_nowait(None) except Exception as e: - self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e) + self._logger.warning( + "[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e + ) if self._proc is not None: try: self._proc.join(timeout=5) except Exception as e: - self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True) + self._logger.debug( + "[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True + ) if self._proc.is_alive(): try: self._proc.terminate() except Exception as e: - self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e) + self._logger.warning( + "[Stage-%s] terminate() failed: %s", self.stage_id, e + ) def submit(self, payload: Dict[str, Any]) -> None: assert self._in_q is not None @@ -150,9 +206,11 @@ class OmniStage: raise ValueError("engine_input_source is empty") source_stage_id = self.engine_input_source[0] source_outputs = stage_list[source_stage_id].engine_outputs + if not isinstance(prompt, list): + prompt = [prompt] multi_modal_data = { - source_output.request_id: prompt.get("multi_modal_data", None) - for source_output, prompt in zip(source_outputs, prompt) + source_output.request_id: p.get("multi_modal_data", None) + for source_output, p in zip(source_outputs, prompt) } for source_output in source_outputs: @@ -184,15 +242,16 @@ def _stage_worker( ) -> None: """Stage worker entry: device setup, LLM init, batching, SHM IPC.""" import logging as _logging - from vllm_omni.entrypoints.omni_llm import OmniStageLLM # noqa: WPS433 + import time as _time + from vllm_omni.entrypoints.log_utils import ( # noqa: WPS433 - log_stage_running_avg, - log_stage_batch_stats, compute_and_log_stage_request_stats, count_tokens_from_outputs, + log_stage_batch_stats, + log_stage_running_avg, ) - import queue as _queue - import time as _time + from vllm_omni.entrypoints.omni_llm import OmniStageLLM # noqa: WPS433 + # no inline JSONL/serialization imports; logging handled by utilities stage_id = stage_payload["stage_id"] @@ -207,11 +266,17 @@ def _stage_worker( stage_log.setLevel(_logging.DEBUG) fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log") fh.setLevel(_logging.DEBUG) - fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")) + fh.setFormatter( + _logging.Formatter( + "%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s" + ) + ) + class _StageFilter(_logging.Filter): def filter(self, record: _logging.LogRecord) -> bool: setattr(record, "stage", stage_id) return True + fh.addFilter(_StageFilter()) stage_log.addHandler(fh) except Exception: @@ -230,10 +295,16 @@ def _stage_worker( try: set_stage_gpu_devices(stage_id, runtime_cfg.get("devices")) except Exception as e: - _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e) + _logging.getLogger(__name__).warning( + "[Stage-%s] Device setup failed: %s", stage_id, e + ) # Init LLM - _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys())) + _logging.getLogger(__name__).debug( + "[Stage-%s] Initializing engine with args keys=%s", + stage_id, + list(engine_args.keys()), + ) stage_engine = OmniStageLLM(model=model, **engine_args) _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id) # Signal readiness to orchestrator @@ -247,7 +318,9 @@ def _stage_worker( task = in_q.get() _recv_dequeue_ts = _time.time() if task is None: - _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id) + _logging.getLogger(__name__).debug( + "[Stage-%s] Received shutdown signal", stage_id + ) break max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) @@ -289,23 +362,37 @@ def _stage_worker( elif isinstance(ein, dict): batch_engine_inputs.append(ein) else: - _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein)) + _logging.getLogger(__name__).exception( + "[Stage-%s] Invalid engine input type: %s", stage_id, type(ein) + ) sampling_params = batch_tasks[0]["sampling_params"] - _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids) + _logging.getLogger(__name__).debug( + "[Stage-%s] Received batch size=%d, request_ids=%s", + stage_id, + len(batch_tasks), + batch_request_ids, + ) print("--------------------------------", flush=True) - print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True) + print( + f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", + flush=True, + ) print("--------------------------------", flush=True) try: _batch_seq += 1 gen_outputs: List[Any] = [] _gen_t0 = _time.time() - for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False): + for ro in stage_engine.generate( + batch_engine_inputs, sampling_params, use_tqdm=False + ): gen_outputs.append(ro) _gen_t1 = _time.time() _gen_ms = (_gen_t1 - _gen_t0) * 1000.0 # Group outputs per request id with fallback - req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids} + req_to_outputs: Dict[Any, List[Any]] = { + rid: [] for rid in batch_request_ids + } unmapped: List[Any] = [] for ro in gen_outputs: rid = getattr(ro, "request_id", None) @@ -326,12 +413,27 @@ def _stage_worker( _num_tokens = count_tokens_from_outputs(_r_outputs) _agg_total_tokens += _num_tokens _agg_total_gen_time_ms += _gen_ms - _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0 if _stats_file: - _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0 - log_stage_running_avg(_stats_file, stage_id, int(_agg_total_tokens), float(_agg_total_gen_time_ms), float(_avg_tokens_per_s)) - log_stage_batch_stats(_stats_file, stage_id, len(batch_tasks), float(_gen_ms), list(batch_request_ids)) + _avg_tokens_per_s = ( + (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) + if _agg_total_gen_time_ms > 0 + else 0.0 + ) + log_stage_running_avg( + _stats_file, + stage_id, + int(_agg_total_tokens), + float(_agg_total_gen_time_ms), + float(_avg_tokens_per_s), + ) + log_stage_batch_stats( + _stats_file, + stage_id, + len(batch_tasks), + float(_gen_ms), + list(batch_request_ids), + ) # Emit per-request results for rid in batch_request_ids: @@ -344,7 +446,9 @@ def _stage_worker( "batch_id": int(_batch_seq), "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), - "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + "rx_in_flight_time_ms": float( + _in_flight_ms_by_rid.get(rid, 0.0) + ), } if _stats_file: compute_and_log_stage_request_stats( @@ -354,42 +458,339 @@ def _stage_worker( len(batch_tasks), r_outputs, float(_gen_ms), - int(_metrics["rx_transfer_bytes"]), # type: ignore[index] - float(_metrics["rx_decode_time_ms"]), # type: ignore[index] + int(_metrics["rx_transfer_bytes"]), # type: ignore[index] + float(_metrics["rx_decode_time_ms"]), # type: ignore[index] ) if use_shm: - out_q.put({ + out_q.put( + { + "request_id": rid, + "stage_id": stage_id, + "engine_outputs_shm": payload, + "metrics": _metrics, + } + ) + else: + out_q.put( + { + "request_id": rid, + "stage_id": stage_id, + "engine_outputs": payload, + "metrics": _metrics, + } + ) + except Exception: + out_q.put( + { + "request_id": rid, + "stage_id": stage_id, + "engine_outputs": r_outputs, + "metrics": { + "num_tokens_out": int( + count_tokens_from_outputs(r_outputs) + ), + "stage_gen_time_ms": _gen_ms, + "rx_decode_time_ms": float( + _rx_decode_ms_by_rid.get(rid, 0.0) + ), + "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), + "rx_in_flight_time_ms": float( + _in_flight_ms_by_rid.get(rid, 0.0) + ), + }, + } + ) + _logging.getLogger(__name__).debug( + "[Stage-%s] Enqueued result for request %s to downstream", + stage_id, + rid, + ) + except Exception as e: + _logging.getLogger(__name__).exception( + "[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e + ) + for rid in batch_request_ids: + out_q.put( + { + "request_id": rid, + "stage_id": stage_id, + "error": str(e), + } + ) + + +def _stage_worker_async_entry( + omni_stage: OmniStage, + model: str, + stage_payload: Dict[str, Any], + in_q: mp.Queue, + out_q: mp.Queue, + log_file: Optional[str] = None, + batch_timeout: int = 10, +) -> None: + asyncio.run( + _stage_worker_async( + omni_stage, model, stage_payload, in_q, out_q, log_file, batch_timeout + ) + ) + + +async def _stage_worker_async( + omni_stage: OmniStage, + model: str, + stage_payload: Dict[str, Any], + in_q: mp.Queue, + out_q: mp.Queue, + log_file: Optional[str] = None, + batch_timeout: int = 10, +) -> None: + """Stage worker entry: device setup, LLM init, batching, SHM IPC.""" + import logging as _logging + import time as _time + + from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM # noqa: WPS433 + from vllm_omni.entrypoints.log_utils import ( # noqa: WPS433 + compute_and_log_stage_request_stats, + count_tokens_from_outputs, + log_stage_batch_stats, + log_stage_running_avg, + ) + + # no inline JSONL/serialization imports; logging handled by utilities + + stage_id = stage_payload["stage_id"] + engine_args = stage_payload.get("engine_args", {}) + runtime_cfg = stage_payload.get("runtime", {}) + shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) + + # Per-stage file logger (optional) + try: + if log_file: + stage_log = _logging.getLogger(__name__) + stage_log.setLevel(_logging.DEBUG) + fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log") + fh.setLevel(_logging.DEBUG) + fh.setFormatter( + _logging.Formatter( + "%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s" # noqa: E501 + ) + ) + + class _StageFilter(_logging.Filter): + def filter(self, record: _logging.LogRecord) -> bool: + setattr(record, "stage", stage_id) + return True + + fh.addFilter(_StageFilter()) + stage_log.addHandler(fh) + except Exception: + pass + + # Stage stats JSONL file + _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None + + # Aggregates for running average + _agg_total_tokens = 0 + _agg_total_gen_time_ms = 0.0 + # Monotonic batch id per stage process for orchestrator dedup on time + # aggregation + _batch_seq = 0 + + # Device mapping + try: + set_stage_gpu_devices(stage_id, runtime_cfg.get("devices")) + except Exception as e: + _logging.getLogger(__name__).warning( + "[Stage-%s] Device setup failed: %s", stage_id, e + ) + + # Init LLM + _logging.getLogger(__name__).debug( + "[Stage-%s] Initializing engine with args keys=%s", + stage_id, + list(engine_args.keys()), + ) + omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args) + usage_context = UsageContext.OPENAI_API_SERVER + vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context) + stage_engine = AsyncOmniStageLLM.from_vllm_config( + vllm_config=vllm_config, + usage_context=usage_context, + engine_args=omni_engine_args, + ) + omni_stage.set_async_engine(stage_engine) + # Don't keep the dummy data in memory + await stage_engine.reset_mm_cache() + _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id) + # Signal readiness to orchestrator and send vllm_config back to main process + try: + # Send vllm_config back to main process so it can be accessed via + # get_vllm_config(). This is needed because async_engine is only available + # in the worker process + + # input_preprocessor = await stage_engine.get_input_preprocessor() + out_q.put( + { + "type": "stage_ready", + "stage_id": stage_id, + "vllm_config": vllm_config, + "tokenizer": getattr(stage_engine, "tokenizer", None), + "is_tracing_enabled": await stage_engine.is_tracing_enabled(), + # "input_preprocessor": input_preprocessor, + } + ) + except Exception as e: + _logging.getLogger(__name__).warning( + "[Stage-%s] Failed to send stage ready signal: %s", stage_id, e + ) + + # Batch processing loop + while True: + task = in_q.get() + _recv_dequeue_ts = _time.time() + if task is None: + _logging.getLogger(__name__).debug( + "[Stage-%s] Received shutdown signal", stage_id + ) + break + + _rx_bytes_by_rid: Dict[Any, int] = {} + _rx_decode_ms_by_rid: Dict[Any, float] = {} + _in_flight_ms_by_rid: Dict[Any, float] = {} + + rid = task["request_id"] + try: + sent_ts = ( + float(task.get("sent_ts", None)) if isinstance(task, dict) else None + ) + if sent_ts is not None: + _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0 + else: + _in_flight_ms_by_rid[rid] = 0.0 + except Exception: + _in_flight_ms_by_rid[rid] = 0.0 + ein, _rx_metrics = maybe_load_from_ipc_with_metrics( + task, obj_key="engine_inputs", shm_key="engine_inputs_shm" + ) + _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0)) + _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0)) + + sampling_params = task["sampling_params"] + _logging.getLogger(__name__).debug( + "[Stage-%s] Received batch size=1, request_ids=%d", stage_id, rid + ) + print("--------------------------------", flush=True) + print( + f"[Stage-{stage_id}] Received batch size=1, request_ids={rid}", flush=True + ) + print("--------------------------------", flush=True) + try: + _batch_seq += 1 + _gen_t0 = _time.time() + if isinstance(ein, list): + ein = ein[0] + + async for res in stage_engine.generate(ein, sampling_params, rid): + gen_output = res + _gen_t1 = _time.time() + _gen_ms = (_gen_t1 - _gen_t0) * 1000.0 + + r_outputs = [gen_output] + _num_tokens = count_tokens_from_outputs(r_outputs) + _agg_total_tokens += _num_tokens + _agg_total_gen_time_ms += _gen_ms + + if _stats_file: + _avg_tokens_per_s = ( + (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) + if _agg_total_gen_time_ms > 0 + else 0.0 + ) + log_stage_running_avg( + _stats_file, + stage_id, + int(_agg_total_tokens), + float(_agg_total_gen_time_ms), + float(_avg_tokens_per_s), + ) + log_stage_batch_stats(_stats_file, stage_id, 1, float(_gen_ms), [rid]) + + try: + use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes) + _metrics = { + "num_tokens_out": int(count_tokens_from_outputs(r_outputs)), + "stage_gen_time_ms": _gen_ms, + "batch_id": int(_batch_seq), + "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), + "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), + "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + } + if _stats_file: + compute_and_log_stage_request_stats( + _stats_file, + stage_id, + rid, + 1, + r_outputs, + float(_gen_ms), + int(_metrics["rx_transfer_bytes"]), # type: ignore[index] + float(_metrics["rx_decode_time_ms"]), # type: ignore[index] + ) + if use_shm: + out_q.put( + { "request_id": rid, "stage_id": stage_id, "engine_outputs_shm": payload, "metrics": _metrics, - }) - else: - out_q.put({ + } + ) + else: + out_q.put( + { "request_id": rid, "stage_id": stage_id, "engine_outputs": payload, "metrics": _metrics, - }) - except Exception: - out_q.put({ + } + ) + except Exception as e: + _logging.getLogger(__name__).exception( + "[Stage-%s] Failed to enqueue result for request %s: %s", + stage_id, + rid, + e, + ) + out_q.put( + { "request_id": rid, "stage_id": stage_id, "engine_outputs": r_outputs, "metrics": { "num_tokens_out": int(count_tokens_from_outputs(r_outputs)), "stage_gen_time_ms": _gen_ms, - "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), + "rx_decode_time_ms": float( + _rx_decode_ms_by_rid.get(rid, 0.0) + ), "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), - "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + "rx_in_flight_time_ms": float( + _in_flight_ms_by_rid.get(rid, 0.0) + ), }, - }) - _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid) + } + ) + _logging.getLogger(__name__).debug( + "[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid + ) + except Exception as e: - _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e) - for rid in batch_request_ids: - out_q.put({ + _logging.getLogger(__name__).exception( + "[Stage-%s] Failed on request %s: %s", stage_id, rid, e + ) + out_q.put( + { "request_id": rid, "stage_id": stage_id, "error": str(e), - }) + } + ) diff --git a/vllm_omni/entrypoints/openai/__init__.py b/vllm_omni/entrypoints/openai/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py new file mode 100644 index 000000000..4c3931466 --- /dev/null +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -0,0 +1,346 @@ +import multiprocessing +import multiprocessing.forkserver as forkserver +import os +from argparse import Namespace +from contextlib import asynccontextmanager +from http import HTTPStatus +from typing import Any, AsyncIterator, Optional + +from fastapi import Depends, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from starlette.datastructures import State + +import vllm.envs as envs +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.transformers_utils.tokenizer import MistralTokenizer +from vllm.utils import decorate_logs +from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + +logger = init_logger(__name__) + + +async def omni_run_server(args, **uvicorn_kwargs) -> None: + """Run a single-worker API server.""" + + # Add process-specific prefix to stdout and stderr. + decorate_logs("APIServer") + + listen_address, sock = setup_server(args) + await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) + + +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_llm( + 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_llm( + args: Namespace, + *, + disable_frontend_multiprocessing: Optional[bool] = None, + client_config: Optional[dict[str, Any]] = None, +) -> AsyncIterator[EngineClient]: + + 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_llm lifecycle + # Ensures everything is shutdown and cleaned up on error/exit + async with build_async_omni_llm_from_stage_config( + args, + disable_frontend_multiprocessing=disable_frontend_multiprocessing, + ) as async_omni_llm: + yield async_omni_llm + + +@asynccontextmanager +async def build_async_omni_llm_from_stage_config( + args: Namespace, + *, + 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 + + Returns the AsyncOmniLLM or None if the creation failed. + """ + + # V1 AsyncLLM. + assert envs.VLLM_USE_V1 + + if disable_frontend_multiprocessing: + logger.warning( + "V1 is enabled, but got --disable-frontend-multiprocessing. " + "To disable frontend multiprocessing, set VLLM_USE_V1=0." + ) + + async_omni_llm: 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 + ) + else: + async_omni_llm = AsyncOmniLLM(model=args.model) + + # # Don't keep the dummy data in memory + # await async_llm.reset_mm_cache() + + yield async_omni_llm + finally: + if async_omni_llm: + async_omni_llm.shutdown() + + +async def omni_init_app_state( + engine_client: EngineClient, + vllm_config: VllmConfig, + state: State, + args: Namespace, +) -> None: + + 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\n" + "It is different from official chat template '%s'. " + "This discrepancy may lead to performance degradation.", + resolved_chat_template, + args.model, + ) + + if args.tool_server == "demo": + tool_server: Optional[ToolServer] = 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, + model_config=model_config, + 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, + model_config, + 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) -> Optional[OmniOpenAIServingChat]: + return request.app.state.openai_serving_chat + + +@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: + 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.error.code + ) + + elif isinstance(generator, ChatCompletionResponse): + return JSONResponse(content=generator.model_dump()) + + return StreamingResponse(content=generator, media_type="text/event-stream") diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py new file mode 100644 index 000000000..b4f2bcb59 --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -0,0 +1,727 @@ +import asyncio +import base64 +import json +import time +import uuid +from collections.abc import AsyncGenerator, AsyncIterator +from datetime import datetime, timedelta, timezone +from io import BytesIO +from typing import Optional, Union + +import jinja2 +from fastapi import Request +from pydantic import TypeAdapter + +try: + import soundfile +except ImportError: + soundfile = None + +from openai.types.chat.chat_completion_audio import ( + ChatCompletionAudio as OpenAIChatCompletionAudio, +) + +from vllm.entrypoints.chat_utils import ( + ConversationMessage, + get_history_tool_calls_cnt, + make_tool_call_id, +) +from vllm.entrypoints.harmony_utils import parse_chat_output +from vllm.entrypoints.openai.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + ErrorResponse, + FunctionCall, + FunctionDefinition, + PromptTokenUsageInfo, + RequestResponseMetadata, + ToolCall, + UsageInfo, +) +from vllm.entrypoints.openai.serving_chat import OpenAIServingChat +from vllm.entrypoints.openai.serving_engine import RequestPrompt, clamp_prompt_logprobs +from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall +from vllm.inputs.data import PromptType +from vllm.logger import init_logger +from vllm.lora.request import LoRARequest +from vllm.outputs import RequestOutput +from vllm.sampling_params import SamplingParams +from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer +from vllm.transformers_utils.tokenizers import ( + maybe_serialize_tool_calls, + truncate_tool_call_ids, + validate_request_params, +) +from vllm.utils import as_list +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + + +class OmniOpenAIServingChat(OpenAIServingChat): + async def create_chat_completion( + self, + request: ChatCompletionRequest, + raw_request: Optional[Request] = None, + ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]: + """ + Chat Completion API similar to OpenAI's API. + + See https://platform.openai.com/docs/api-reference/chat/create + for the API specification. This API mimics the OpenAI + Chat Completion API. + """ + error_check_ret = await self._check_model(request) + if error_check_ret is not None: + logger.error("Error with model %s", error_check_ret) + return error_check_ret + + # If the engine is dead, raise the engine's DEAD_ERROR. + # This is required for the streaming case, where we return a + # success status before we actually start generating text :). + if self.engine_client.errored: + raise self.engine_client.dead_error + + try: + lora_request = self._maybe_get_adapters( + request, supports_default_mm_loras=True + ) + + model_name = self.models.model_name(lora_request) + + tokenizer = await self.engine_client.get_tokenizer() + + tool_parser = self.tool_parser + + if isinstance(tokenizer, MistralTokenizer): + # because of issues with pydantic we need to potentially + # re-serialize the tool_calls field of the request + # for more info: see comment in `maybe_serialize_tool_calls` + maybe_serialize_tool_calls(request) + truncate_tool_call_ids(request) + validate_request_params(request) + + if ( + request.tool_choice == "auto" + and not (self.enable_auto_tools and tool_parser is not None) + and not isinstance(tokenizer, MistralTokenizer) + and not self.use_harmony + ): + # for hf tokenizers, "auto" tools requires + # --enable-auto-tool-choice and --tool-call-parser + return self.create_error_response( + '"auto" tool choice requires ' + "--enable-auto-tool-choice and --tool-call-parser to be set" + ) + + if request.tools is None or ( + request.tool_choice == "none" + and self.exclude_tools_when_tool_choice_none + ): + tool_dicts = None + else: + tool_dicts = [tool.model_dump() for tool in request.tools] + + # Common case. + request_chat_template = request.chat_template + chat_template_kwargs = request.chat_template_kwargs + if not self.trust_request_chat_template and ( + request_chat_template is not None + or ( + chat_template_kwargs + and chat_template_kwargs.get("chat_template") is not None + ) + ): + return self.create_error_response( + "Chat template is passed with request, but " + "--trust-request-chat-template is not set. " + "Refused request with untrusted chat template." + ) + ( + conversation, + request_prompts, + engine_prompts, + ) = await self._preprocess_chat( + request, + tokenizer, + request.messages, + chat_template=request_chat_template or self.chat_template, + chat_template_content_format=self.chat_template_content_format, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + tool_dicts=tool_dicts, + documents=request.documents, + chat_template_kwargs=request.chat_template_kwargs, + tool_parser=tool_parser, + add_special_tokens=request.add_special_tokens, + ) + + except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: + logger.exception("Error in preprocessing prompt inputs") + return self.create_error_response(f"{e} {e.__cause__}") + + request_id = ( + "chatcmpl-" f"{self._base_request_id(raw_request, request.request_id)}" + ) + + request_metadata = RequestResponseMetadata(request_id=request_id) + if raw_request: + raw_request.state.request_metadata = request_metadata + + # Schedule the request and get the result generator. + generators: list[AsyncGenerator[RequestOutput, None]] = [] + try: + for i, engine_prompt in enumerate(engine_prompts): + + if hasattr(request, "sampling_params_list"): + sampling_params_list = self._to_sampling_params_list( + request.sampling_params_list + ) + else: + sampling_params_list = None + + self._log_inputs( + request_id, + request_prompts[i], + params_list=sampling_params_list, + lora_request=lora_request, + ) + + trace_headers = ( + None + if raw_request is None + else await self._get_trace_headers(raw_request.headers) + ) + + generator = self.engine_client.generate( + prompt=engine_prompt, + request_id=request_id, + sampling_params_list=sampling_params_list, + lora_request=lora_request, + trace_headers=trace_headers, + priority=request.priority, + ) + + generators.append(generator) + except ValueError as e: + # TODO: Use a vllm-specific Validation Error + return self.create_error_response(str(e)) + + assert len(generators) == 1 + (result_generator,) = generators + + # Streaming response + if request.stream: + return self.chat_completion_stream_generator( + request, + result_generator, + request_id, + model_name, + conversation, + tokenizer, + request_metadata, + enable_force_include_usage=self.enable_force_include_usage, + ) + + try: + return await self.chat_completion_full_generator( + request, + result_generator, + request_id, + model_name, + conversation, + tokenizer, + request_metadata, + ) + except ValueError as e: + # TODO: Use a vllm-specific Validation Error + return self.create_error_response(str(e)) + + def _to_sampling_params_list( + self, sampling_params_list: list[dict] + ) -> list[SamplingParams]: + + final_sampling_params_list = [] + for sampling_params in sampling_params_list: + if isinstance(sampling_params, dict): + final_sampling_params_list.append(SamplingParams(**sampling_params)) + elif isinstance(sampling_params, SamplingParams): + final_sampling_params_list.append(sampling_params) + else: + raise ValueError(f"Invalid sampling params: {sampling_params}") + return final_sampling_params_list + + def _log_inputs( + self, + request_id: str, + inputs: Union[RequestPrompt, PromptType], + params_list: Optional[list[SamplingParams]], + lora_request: Optional[LoRARequest], + ) -> None: + if self.request_logger is None: + return + prompt, prompt_token_ids, prompt_embeds = None, None, None + if isinstance(inputs, str): + prompt = inputs + elif isinstance(inputs, list): + prompt_token_ids = inputs + else: + prompt = getattr(inputs, "prompt", None) + prompt_token_ids = getattr(inputs, "prompt_token_ids", None) + + logger.info( + "Received request %s: prompt: %r, " + "params_list: %s, prompt_token_ids: %s, " + "prompt_embeds shape: %s, " + "lora_request: %s.", + request_id, + prompt, + params_list, + prompt_token_ids, + prompt_embeds.shape if prompt_embeds is not None else None, + lora_request, + ) + + async def chat_completion_full_generator( + self, + request: ChatCompletionRequest, + result_generator: AsyncIterator[RequestOutput], + request_id: str, + model_name: str, + conversation: list[ConversationMessage], + tokenizer: AnyTokenizer, + request_metadata: RequestResponseMetadata, + ) -> Union[ErrorResponse, ChatCompletionResponse]: + + created_time = int(time.time()) + final_res: Optional[RequestOutput] = None + + final_outputs: list[OmniRequestOutput] = [] + try: + async for res in result_generator: + final_outputs.append(res) + except asyncio.CancelledError: + return self.create_error_response("Client disconnected") + except ValueError as e: + # TODO: Use a vllm-specific Validation Error + return self.create_error_response(str(e)) + + assert final_outputs is not None + + choices: list[ChatCompletionResponseChoice] = [] + + usage = UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0) + role = self.get_chat_request_role(request) + prompt_logprobs = None + prompt_token_ids = None + kv_transfer_params = None + + for omni_outputs in final_outputs: + choices_data = [] + if omni_outputs.final_output_type == "text": + ( + choices_data, + usage, + prompt_logprobs, + prompt_token_ids, + kv_transfer_params, + ) = self._create_text_choice( + request, omni_outputs, tokenizer, conversation, role + ) + elif omni_outputs.final_output_type == "audio": + choices_data = self._create_audio_choice(omni_outputs, role) + else: + logger.warning( + f"Unsupported final output type: {omni_outputs.final_output_type}" + ) + continue + choices.extend(choices_data) + + response = ChatCompletionResponse( + id=request_id, + created=created_time, + model=model_name, + choices=choices, + usage=usage, + prompt_logprobs=prompt_logprobs, + prompt_token_ids=prompt_token_ids, + kv_transfer_params=kv_transfer_params, + ) + + # Log complete response if output logging is enabled + if self.enable_log_outputs and self.request_logger: + for choice in choices: + output_text = "" + if choice.message.content: + output_text = choice.message.content + elif choice.message.tool_calls: + # For tool calls, log the function name and arguments + tool_call_descriptions = [] + for tc in choice.message.tool_calls: + if hasattr(tc.function, "name") and hasattr( + tc.function, "arguments" + ): + tool_call_descriptions.append( + f"{tc.function.name}({tc.function.arguments})" + ) + tool_calls_str = ", ".join(tool_call_descriptions) + output_text = f"[tool_calls: {tool_calls_str}]" + + if output_text: + # Get the corresponding output token IDs + output_token_ids = None + if choice.index < len(final_res.outputs): + output_token_ids = final_res.outputs[choice.index].token_ids + + self.request_logger.log_outputs( + request_id=request_id, + outputs=output_text, + output_token_ids=output_token_ids, + finish_reason=choice.finish_reason, + is_streaming=False, + delta=False, + ) + + return response + + def _create_text_choice( + self, + request: ChatCompletionRequest, + omni_outputs: OmniRequestOutput, + tokenizer: AnyTokenizer, + conversation: list[ConversationMessage], + role: str, + ): + final_res = omni_outputs.request_output + if self.tool_call_id_type == "kimi_k2": + history_tool_call_cnt = get_history_tool_calls_cnt(conversation) + else: + history_tool_call_cnt = 0 + + choices: list[ChatCompletionResponseChoice] = [] + + for output in final_res.outputs: + token_ids = output.token_ids + out_logprobs = output.logprobs + tool_call_info = None + + if request.logprobs and request.top_logprobs is not None: + assert out_logprobs is not None, "Did not output logprobs" + logprobs = self._create_chat_logprobs( + token_ids=token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=request.top_logprobs, + tokenizer=tokenizer, + return_as_token_id=request.return_tokens_as_token_ids, + ) + else: + logprobs = None + + if self.use_harmony: + reasoning_content, content, _ = parse_chat_output(token_ids) + if not request.include_reasoning: + reasoning_content = None + + if self.tool_parser is not None: + tool_parser = self.tool_parser(tokenizer) + # NOTE: We use token_ids for openai tool parser + tool_call_info = tool_parser.extract_tool_calls( + "", + request=request, + token_ids=token_ids, # type: ignore + ) + content = tool_call_info.content + message = ChatMessage( + role=role, + reasoning_content=reasoning_content, + content=content, + tool_calls=tool_call_info.tool_calls, + ) + else: + message = ChatMessage( + role=role, + reasoning_content=reasoning_content, + content=content, + ) + + choice_data = ChatCompletionResponseChoice( + index=output.index, + message=message, + logprobs=logprobs, + finish_reason=( + "tool_calls" + if (tool_call_info is not None and tool_call_info.tools_called) + else (output.finish_reason if output.finish_reason else "stop") + ), + stop_reason=output.stop_reason, + ) + choices.append(choice_data) + continue + + if self.reasoning_parser: + try: + reasoning_parser = self.reasoning_parser(tokenizer) + except RuntimeError as e: + logger.exception("Error in reasoning parser creation.") + return self.create_error_response(str(e)) + # If the reasoning parser is enabled, + # tool calls are extracted exclusively from the content. + reasoning_content, content = reasoning_parser.extract_reasoning_content( + output.text, request=request + ) + if not request.include_reasoning: + reasoning_content = None + else: + reasoning_content = None + content = output.text + + auto_tools_called = False + # if auto tools are not enabled, and a named tool choice using + # outlines is not being used + if (not self.enable_auto_tools or not self.tool_parser) and ( + not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) + and request.tool_choice != "required" + ): + message = ChatMessage( + role=role, reasoning_content=reasoning_content, content=content + ) + + # if the request uses tools and specified a tool choice + elif ( + request.tool_choice + and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam + ): + + tool_call_class = ( + MistralToolCall + if isinstance(tokenizer, MistralTokenizer) + else ToolCall + ) + message = ChatMessage( + role=role, + reasoning_content=reasoning_content, + content="", + tool_calls=[ + tool_call_class( + function=FunctionCall( + name=request.tool_choice.function.name, + arguments=content, + ) + ) + ], + ) + + elif request.tool_choice and request.tool_choice == "required": + tool_call_class = ( + MistralToolCall + if isinstance(tokenizer, MistralTokenizer) + else ToolCall + ) + + # the fields of FunctionDefinition are a superset of the + # tool call outputs and can be used for parsing + assert content is not None + tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json( + content + ) + tool_call_ids = [] + for tool_call in tool_calls: + tool_call_ids.append( + make_tool_call_id( + id_type=self.tool_call_id_type, + func_name=tool_call.name, + idx=history_tool_call_cnt, + ) + ) + history_tool_call_cnt += 1 + message = ChatMessage( + role=role, + content="", + tool_calls=[ + tool_call_class( + id=tool_call_ids[i], + function=FunctionCall( + name=tool_call.name, + arguments=json.dumps( + tool_call.parameters, ensure_ascii=False + ), + ), + ) + for i, tool_call in enumerate(tool_calls) + ], + reasoning_content=reasoning_content, + ) + + # if the request doesn't use tool choice + # OR specifies to not use a tool + elif not request.tool_choice or request.tool_choice == "none": + + message = ChatMessage( + role=role, reasoning_content=reasoning_content, content=content + ) + + # handle when there are tools and tool choice is auto + elif ( + request.tools + and (request.tool_choice == "auto" or request.tool_choice is None) + and self.enable_auto_tools + and self.tool_parser + ): + + try: + tool_parser = self.tool_parser(tokenizer) + except RuntimeError as e: + logger.exception("Error in tool parser creation.") + return self.create_error_response(str(e)) + + tool_call_info = tool_parser.extract_tool_calls( + content if content is not None else "", request=request + ) + # In the OpenAI API the finish_reason is "tools_called" + # if the tool choice is auto and the model produced a tool + # call. The same is not true for named function calls + auto_tools_called = tool_call_info.tools_called + if tool_call_info.tools_called: + message = ChatMessage( + role=role, + reasoning_content=reasoning_content, + content=tool_call_info.content, + tool_calls=tool_call_info.tool_calls, + ) + + else: + # FOR NOW make it a chat message; we will have to detect + # the type to make it later. + ret_content = content + + # try to use content return from tool parser first, + # tool parser may do some modify for the content. + if tool_call_info.content and len(tool_call_info.content) > 0: + ret_content = tool_call_info.content + message = ChatMessage( + role=role, + reasoning_content=reasoning_content, + content=ret_content, + ) + + # undetermined case that is still important to handle + else: + logger.error( + "Error in chat_completion_full_generator - cannot determine" + " if tools should be extracted. Returning a standard chat " + "completion." + ) + message = ChatMessage( + role=role, reasoning_content=reasoning_content, content=content + ) + + choice_data = ChatCompletionResponseChoice( + index=output.index, + message=message, + logprobs=logprobs, + finish_reason=( + "tool_calls" + if auto_tools_called + else output.finish_reason if output.finish_reason else "stop" + ), + stop_reason=output.stop_reason, + token_ids=( + as_list(output.token_ids) if request.return_token_ids else None + ), + ) + choices.append(choice_data) + + if request.echo: + last_msg_content: Union[str, list[dict[str, str]]] = "" + if ( + conversation + and "content" in conversation[-1] + and conversation[-1].get("role") == role + ): + last_msg_content = conversation[-1]["content"] or "" + if isinstance(last_msg_content, list): + last_msg_content = "\n".join(msg["text"] for msg in last_msg_content) + + for choice in choices: + full_message = last_msg_content + (choice.message.content or "") + choice.message.content = full_message + + assert final_res.prompt_token_ids is not None + num_prompt_tokens = len(final_res.prompt_token_ids) + if final_res.encoder_prompt_token_ids is not None: + num_prompt_tokens += len(final_res.encoder_prompt_token_ids) + num_generated_tokens = sum( + len(output.token_ids) for output in final_res.outputs + ) + usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=num_generated_tokens, + total_tokens=num_prompt_tokens + num_generated_tokens, + ) + if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + usage.prompt_tokens_details = PromptTokenUsageInfo( + cached_tokens=final_res.num_cached_tokens + ) + + prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs) + prompt_token_ids = ( + final_res.prompt_token_ids if request.return_token_ids else None + ) + kv_transfer_params = final_res.kv_transfer_params + + return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params + + def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str): + choices: list[ChatCompletionResponseChoice] = [] + final_res = omni_outputs.request_output + audio_tensor = final_res.multimodal_output["audio"].detach().cpu().numpy() + + # Convert numpy array to WAV bytes and encode as base64 + if soundfile is None: + raise ImportError( + "soundfile is required for audio generation. " + "Please install it with: pip install soundfile" + ) + + # Default sample rate for TTS models (typically 24000 Hz) + # You may need to adjust this based on your model's configuration + sample_rate = 24000 + + # Ensure audio is 1D (flatten if needed) + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.flatten() + + # Convert to WAV format and encode as base64 + with BytesIO() as buffer: + soundfile.write(buffer, audio_tensor, sample_rate, format="WAV") + wav_bytes = buffer.getvalue() + + audio_base64 = base64.b64encode(wav_bytes).decode("utf-8") + + # Generate unique ID for the audio + audio_id = f"audio-{uuid.uuid4().hex[:16]}" + + # Set expiration time (e.g., 24 hours from now) as Unix timestamp + expires_at = int((datetime.now(timezone.utc) + timedelta(hours=24)).timestamp()) + + # Create OpenAIChatCompletionAudio object with all required fields + audio_obj = OpenAIChatCompletionAudio( + id=audio_id, + data=audio_base64, + expires_at=expires_at, + transcript="", # Empty transcript if not available + ) + + for output in final_res.outputs: + choice_data = ChatCompletionResponseChoice( + index=output.index, + message=ChatMessage(role=role, audio=audio_obj), + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + choices.append(choice_data) + return choices diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 16f8c6331..dc38e617a 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -1,14 +1,14 @@ from __future__ import annotations -from typing import Any, Dict, Optional, Tuple, Union - -import logging import json +import logging import os import pickle from multiprocessing import shared_memory as _shm +from typing import Any, Dict, Optional, Tuple, Union import cloudpickle +from omegaconf import OmegaConf logger = logging.getLogger(__name__) @@ -37,12 +37,18 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> selected_physical = int(toks[0]) logger.debug( "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s", - stage_id, devices, selected_physical, + stage_id, + devices, + selected_physical, ) except Exception as e: - logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e) + logger.debug( + "[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e + ) selected_physical = None - elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()): + elif isinstance(devices, (int, str)) and ( + isinstance(devices, int) or str(devices).isdigit() + ): logical_idx = max(0, int(devices)) vis = os.environ.get("CUDA_VISIBLE_DEVICES") if vis: @@ -51,43 +57,72 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> if 0 <= logical_idx < len(mapping): selected_physical = mapping[logical_idx] except Exception as e: - logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e) + logger.debug( + "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", + stage_id, + e, + ) selected_physical = None if selected_physical is None: selected_physical = int(logical_idx) os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) logger.debug( "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device", - stage_id, logical_idx + 1, selected_physical, + stage_id, + logical_idx + 1, + selected_physical, ) elif devices in (None, "cpu"): - logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices) + logger.debug( + "[Stage-%s] Using default device visibility (devices=%s)", + stage_id, + devices, + ) else: selected_physical = int(str(devices)) os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) - logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical) + logger.debug( + "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", + stage_id, + selected_physical, + ) try: import torch # noqa: WPS433 + if torch.cuda.is_available(): try: torch.cuda.set_device(0) except Exception as e: - logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True) + logger.debug( + "[Stage-%s] torch.cuda.set_device(0) failed: %s", + stage_id, + e, + exc_info=True, + ) num = torch.cuda.device_count() info = [] for i in range(num): total = torch.cuda.get_device_properties(i).total_memory free, _ = torch.cuda.mem_get_info(i) - info.append({ - "idx": i, - "name": torch.cuda.get_device_name(i), - "total": int(total), - "free": int(free), - }) - logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info) + info.append( + { + "idx": i, + "name": torch.cuda.get_device_name(i), + "total": int(total), + "free": int(free), + } + ) + logger.debug( + "[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info + ) except Exception as e: - logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True) + logger.debug( + "[Stage-%s] Failed to query CUDA devices: %s", + stage_id, + e, + exc_info=True, + ) except Exception as e: logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e) @@ -110,7 +145,7 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]: try: shm.close() except Exception as e: - logger.debug("Ensure parent dir failed for %s: %s", path, e) + logger.debug("Failed to close shared memory: %s", e) return meta @@ -180,7 +215,9 @@ def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) - return container[obj_key] -def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]: +def maybe_load_from_ipc_with_metrics( + container: Dict[str, Any], obj_key: str, shm_key: str +) -> tuple[Any, Dict[str, float]]: """Load object and return (object, metrics) with RX bytes and decode time. Metrics keys: @@ -188,6 +225,7 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh - rx_decode_time_ms: float """ import time as _time # local import to avoid overhead at module import + t0 = _time.time() if shm_key in container: meta = container[shm_key] # type: ignore[index] @@ -211,7 +249,9 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh } -def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]: +def encode_for_ipc( + obj: Any, threshold: int, obj_key: str, shm_key: str +) -> Dict[str, Any]: """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key). When serialized size exceeds threshold, returns {shm_key: {name,size}}; diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 383c9f47d..a0de525bf 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,3 +1,6 @@ +from __future__ import annotations + +import logging import os from pathlib import Path @@ -5,6 +8,8 @@ from omegaconf import OmegaConf from vllm.transformers_utils.config import get_config +logger = logging.getLogger(__name__) + # Get the project root directory (2 levels up from this file) PROJECT_ROOT = Path(__file__).parent.parent.parent diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 5bae25b8d..31c491ae3 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -4,10 +4,13 @@ from typing_extensions import assert_never from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt from vllm.inputs.preprocess import InputPreprocessor +from vllm.logger import init_logger from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni from vllm_omni.inputs.parse import parse_singleton_prompt_omni +logger = init_logger(__name__) + class OmniInputPreprocessor(InputPreprocessor): def _process_tokens( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index 1768c4395..f6e2ab7fc 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -279,7 +279,9 @@ class Qwen2_5OmniForConditionalGeneration( if input_ids is None and additional_information is None: input_ids = torch.zeros( - inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device + inputs_embeds.shape[0], + dtype=torch.long, + device=inputs_embeds.device, ) additional_information = {} self.thinker_reply_part = torch.zeros_like(inputs_embeds) @@ -293,13 +295,19 @@ class Qwen2_5OmniForConditionalGeneration( # ------- Request-scoped additional information (no cross-request concat) ------- request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered - request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans") + request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get( + "request_token_spans" + ) addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id") runtime_addi = kwargs.get("runtime_additional_information") # Normalize runtime_addi into a mapping by request_id for convenience runtime_addi_by_req: dict[str, dict] = {} - if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids): + if ( + isinstance(request_ids, list) + and isinstance(runtime_addi, list) + and len(runtime_addi) == len(request_ids) + ): for i, rid in enumerate(request_ids): if isinstance(rid, str) and isinstance(runtime_addi[i], dict): runtime_addi_by_req[rid] = runtime_addi[i] @@ -332,9 +340,19 @@ class Qwen2_5OmniForConditionalGeneration( otoks = info.get("thinker_output_token_ids") # list[int] if not isinstance(pe, torch.Tensor): - pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + pe = torch.zeros( + 0, + self.talker.config.hidden_size, + dtype=inputs_embeds.dtype, + device=self._module_device(self.model), + ) if not isinstance(tr, torch.Tensor): - tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + tr = torch.zeros( + 0, + self.talker.config.hidden_size, + dtype=inputs_embeds.dtype, + device=self._module_device(self.model), + ) if not isinstance(ptoks, (list, torch.Tensor)): ptoks = [] if not isinstance(otoks, (list, torch.Tensor)): @@ -342,22 +360,35 @@ class Qwen2_5OmniForConditionalGeneration( req_input_ids, req_embeds = self._thinker_to_talker_prefill( voice_type=voice_type, - output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_prompt_embeds=tr.to(inputs_embeds.dtype).to( + self._module_device(self.model) + ), output_token_ids=otoks, - thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)), + thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to( + self._module_device(self.model) + ), prompt_token_ids=ptoks, ) seg_len = min(span_len, req_embeds.shape[0]) inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] - if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: + if ( + isinstance(req_input_ids, torch.Tensor) + and req_input_ids.numel() == seg_len + ): input_ids[s : s + seg_len] = req_input_ids # Prepare per-request reply queue for subsequent decode: drop first row if tr.ndim == 2 and tr.shape[0] > 0: - update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous() + update_by_req_id.setdefault(rid, {})[ + "thinker_reply_part_per_request" + ] = (tr[1:].detach().to("cpu").contiguous()) # ------- Decode: span_len == 1 ------- - if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list): + if ( + not is_profile + and isinstance(request_ids, list) + and isinstance(request_token_spans, list) + ): for idx_req, rid in enumerate(request_ids): s, e = request_token_spans[idx_req] if (int(e) - int(s)) != 1: @@ -367,18 +398,34 @@ class Qwen2_5OmniForConditionalGeneration( # A) runtime queue q = None if isinstance(rid, str): - q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request") + q = runtime_addi_by_req.get(rid, {}).get( + "thinker_reply_part_per_request" + ) if isinstance(q, torch.Tensor) and q.numel() > 0: step_vec = q[0:1] new_q = q[1:].detach().to("cpu").contiguous() - update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q + update_by_req_id.setdefault(rid, {})[ + "thinker_reply_part_per_request" + ] = new_q else: # B) per-request provided decode vector (optional) - info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {} - dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None + info = ( + addi_by_req.get(rid, {}) + if isinstance(addi_by_req, dict) + else {} + ) + dv = ( + info.get("decode_output_prompt_embeds") + if isinstance(info, dict) + else None + ) if isinstance(dv, torch.Tensor) and dv.numel() > 0: step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1) - elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0: + elif ( + hasattr(self, "thinker_reply_part") + and isinstance(self.thinker_reply_part, torch.Tensor) + and self.thinker_reply_part.numel() > 0 + ): # C) fallback shared pool step_vec = self.thinker_reply_part[0:1] self.thinker_reply_part = self.thinker_reply_part[1:] @@ -386,7 +433,9 @@ class Qwen2_5OmniForConditionalGeneration( if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0: one_id = input_ids[s : s + 1] _, one_embed = self._thinker_to_talker_decode_one_step( - output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to( + self._module_device(self.model) + ), output_token_ids=one_id, ) inputs_embeds[s] = one_embed[0] @@ -400,11 +449,15 @@ class Qwen2_5OmniForConditionalGeneration( multimodal_outputs: dict = None # Return updates if any if update_by_req_id: - multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id} + multimodal_outputs = { + "additional_information_update_by_req_id": update_by_req_id + } if sampling_metadata is not None: # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size - sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448 + sampling_metadata.prompt_token_ids[ + sampling_metadata.prompt_token_ids == 152064 + ] = 8448 return OmniOutput( text_hidden_states=talker_hidden, @@ -421,13 +474,11 @@ class Qwen2_5OmniForConditionalGeneration( device=inputs_embeds.device, ) ) - + code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code - - audio_tensor = self.generate_audio( - code, voice_type - ) + + audio_tensor = self.generate_audio(code, voice_type) return OmniOutput( text_hidden_states=None, multimodal_outputs={"audio": audio_tensor} ) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py index 6e9527cdd..8944b693b 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py @@ -166,8 +166,16 @@ class Qwen2_5OmniTalkerForConditionalGeneration( ) return hidden_states + def bad_word_processor(self, logits: torch.Tensor) -> torch.Tensor: + if hasattr(self.config, "tts_codec_start_token_id"): + bos_id = int(getattr(self.config, "tts_codec_start_token_id")) + logits[..., bos_id] = -1e9 + return logits + def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: - return self.language_model.compute_logits(hidden_states) + logits = self.language_model.compute_logits(hidden_states) + logits = self.bad_word_processor(logits) + return logits def sample( self, diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index 620da3a89..54d2b2359 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -15,8 +15,17 @@ stage_args: trust_remote_code: true engine_output_type: latent # change the param name,such as pooling_output enable_prefix_caching: false + is_comprehension: true final_output: true final_output_type: text + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.1 - stage_id: 1 runtime: process: true @@ -34,6 +43,15 @@ stage_args: engine_output_type: latent engine_input_source: [0] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker + default_sampling_params: + temperature: 0.9 + top_p: 0.8 + top_k: 40 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.05 + stop_token_ids: [8294] - stage_id: 2 runtime: process: true @@ -52,6 +70,14 @@ stage_args: engine_input_source: [1] final_output: true final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.1 # Top-level runtime config (concise): default windows and stage edges runtime: diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 02da0bd20..16717cdfc 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -4,9 +4,12 @@ import torch from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt + TALKER_CODEC_PAD_TOKEN_ID = 8292 TALKER_CODEC_START_TOKEN_ID = 8293 TALKER_CODEC_END_TOKEN_ID = 8294 + + def thinker2talker( stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None ): @@ -19,9 +22,11 @@ def thinker2talker( raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") thinker_outputs = stage_list[source_stage_id].engine_outputs talker_inputs = [] + if not isinstance(prompt, list): + prompt = [prompt] multi_modal_data = { - thinker_output.request_id: prompt.get("multi_modal_data", None) - for thinker_output, prompt in zip(thinker_outputs, prompt) + thinker_output.request_id: p.get("multi_modal_data", None) + for thinker_output, p in zip(thinker_outputs, prompt) } for i, thinker_output in enumerate(thinker_outputs): @@ -31,21 +36,28 @@ def thinker2talker( prompt_token_ids_len = len(prompt_token_ids) latent = output.multimodal_output["latent"] thinker_hidden_states = latent.clone().detach().to(latent.device) + addtional_information = { + "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to( + torch.float32 + ), + "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to( + torch.float32 + ), + "prompt_token_ids": prompt_token_ids, + "thinker_output_token_ids": thinker_output_ids, + "thinker_result_shape": list( + thinker_hidden_states[prompt_token_ids_len:].shape + ), + "prompt_embeds_shape": list( + thinker_hidden_states[:prompt_token_ids_len].shape + ), + } talker_inputs.append( OmniTokensPrompt( - prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + [TALKER_CODEC_END_TOKEN_ID], - additional_information={ - "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to( - torch.float32 - ), - "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to( - torch.float32 - ), - "prompt_token_ids": prompt_token_ids, - "thinker_output_token_ids": thinker_output_ids, - "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape), - "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape), - }, + prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + + [TALKER_CODEC_END_TOKEN_ID], + additional_information=addtional_information, multi_modal_data=( multi_modal_data[thinker_output.request_id] if multi_modal_data is not None diff --git a/vllm_omni/version.py b/vllm_omni/version.py new file mode 100644 index 000000000..92007aa40 --- /dev/null +++ b/vllm_omni/version.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +try: + from ._version import __version__, __version_tuple__ +except Exception as e: + import warnings + + warnings.warn(f"Failed to read commit hash:\n{e}", RuntimeWarning, stacklevel=2) + + __version__ = "dev" + __version_tuple__ = (0, 0, __version__) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d91defd81..8f757bb6a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -138,9 +138,15 @@ class GPUARModelRunner(OmniGPUModelRunner): # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor if overlay_len > 0: try: - seg = v[ - num_computed_tokens : num_computed_tokens + overlay_len - ].detach().to("cpu").contiguous() + seg = ( + v[ + num_computed_tokens : num_computed_tokens + + overlay_len + ] + .detach() + .to("cpu") + .contiguous() + ) except Exception: seg = v.detach().to("cpu").contiguous() req_info[k] = seg @@ -380,7 +386,9 @@ class GPUARModelRunner(OmniGPUModelRunner): model_kwargs_extra = {} # Pass per-request additional information map for this step (no concat) if per_req_additional_information: - model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information + model_kwargs_extra["additional_information_by_req_id"] = ( + per_req_additional_information + ) # Always pass per-request runtime additional_information (persisted in request state) try: per_req_runtime_info = [] @@ -433,12 +441,9 @@ class GPUARModelRunner(OmniGPUModelRunner): # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. try: - if ( - isinstance(multimodal_outputs, dict) - and ( - "additional_information_update" in multimodal_outputs - or "additional_information_update_by_req_id" in multimodal_outputs - ) + if isinstance(multimodal_outputs, dict) and ( + "additional_information_update" in multimodal_outputs + or "additional_information_update_by_req_id" in multimodal_outputs ): # Option A: list[dict] in batch order updates_list = multimodal_outputs.get( @@ -446,7 +451,9 @@ class GPUARModelRunner(OmniGPUModelRunner): ) if isinstance(updates_list, list): for idx, upd in enumerate(updates_list): - if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): + if not isinstance(upd, dict) or idx >= len( + self.input_batch.req_ids + ): continue req_id = self.input_batch.req_ids[idx] self._merge_additional_information_update(req_id, upd) @@ -462,7 +469,10 @@ class GPUARModelRunner(OmniGPUModelRunner): continue self._merge_additional_information_update(req_id, upd) except Exception as e: - logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}") + logger.error( + f"Error merging for requests:{self.input_batch.req_ids} additional \ + information update: {e}, with the multimodal_outputs as {multimodal_outputs}" + ) if not self.broadcast_pp_output: # Common case. @@ -624,7 +634,7 @@ class GPUARModelRunner(OmniGPUModelRunner): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, ) - + def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: req_state = self.requests.get(req_id) if req_state is None: @@ -646,4 +656,4 @@ class GPUARModelRunner(OmniGPUModelRunner): merged[k] = new_list else: merged[k] = v - setattr(req_state, "additional_information_cpu", merged) \ No newline at end of file + setattr(req_state, "additional_information_cpu", merged) diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py index e8cd68491..c8b294c5d 100644 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/worker/gpu_diffusion_model_runner.py @@ -501,13 +501,16 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner): else: cudagraph_runtime_mode = _cg_mode - with self.maybe_randomize_inputs(input_ids), set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens_after_padding, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + with ( + self.maybe_randomize_inputs(input_ids), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_after_padding, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ), ): outputs = self.model( input_ids=input_ids, diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5369da899..67dfdf537 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -545,14 +545,17 @@ class OmniGPUModelRunner(GPUModelRunner): if num_tokens_across_dp is not None: num_tokens_across_dp[:] = num_tokens_after_padding - with self.maybe_randomize_inputs(input_ids), set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens_after_padding, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, - ubatch_slices=ubatch_slices, + with ( + self.maybe_randomize_inputs(input_ids), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_after_padding, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ubatch_slices=ubatch_slices, + ), ): outputs = self.model( input_ids=input_ids,