feat: enhance serving functionality and add comprehensive examples

- Improve API server with better error handling and response formatting
- Enhance CLI with additional options for DiT stages and configuration
- Add comprehensive examples in examples/basic/ including:
  - API client with health checks and text generation
  - Docker setup and usage examples
  - Simple usage patterns for different scenarios
- Add utility scripts for model downloading and Docker setup
- Update documentation with implementation details and testing guidelines
- Fix configuration validation issues in OmniLLM
- Improve stage configuration handling for AR and DiT stages
- Add proper error handling and fallback mechanisms

Tested with Qwen3-0.6B model:
- Server starts successfully on port 8000
- Health and info endpoints working correctly
- Text generation with various parameters functioning
- API client examples working as expected
- CLI help and configuration options working properly
This commit is contained in:
hsliu committed 2025-09-26 15:47:55 +08:00
1 parent 51d2570464
commit afe5c2692c
16 files changed
+1051 -134

No files matched your search

+76
View File
@@ -0,0 +1,76 @@
# Basic vLLM-omni Examples
This directory contains simple examples demonstrating how to use vLLM-omni for text generation.
## Examples
### 1. `simple_usage.py` - Direct Library Usage
Demonstrates how to use vLLM-omni directly in your Python code:
```bash
# Run the example
python simple_usage.py
```
**Features:**
- Synchronous text generation
- Asynchronous text generation
- Multi-prompt processing
- Shows both text output and token information
### 2. `api_client.py` - API Server Usage
Demonstrates how to interact with the vLLM-omni API server:
```bash
# First, start the server
vllm serve Qwen/Qwen3-0.6B --omni --port 8000
# Then run the client (in another terminal)
python api_client.py
```
**Features:**
- Health check
- Text generation via HTTP API
- Multiple prompt types
- Full response structure inspection
## Quick Start
1. **Install vLLM-omni:**
```bash
pip install -e .
```
2. **Run direct usage example:**
```bash
cd examples/basic
python simple_usage.py
```
3. **Run API server example:**
```bash
# Terminal 1: Start server
vllm serve Qwen/Qwen3-0.6B --omni --port 8000
# Terminal 2: Run client
cd examples/basic
python api_client.py
```
## Expected Output
The examples will show:
- ✅ Generated text responses
- ✅ Token counts and completion status
- ✅ Stage processing information
- ✅ Error handling
## Requirements
- Python 3.8+
- vLLM-omni installed
- For API examples: HTTP server running
+128
View File
@@ -0,0 +1,128 @@
#!/usr/bin/env python3
"""
Simple API client example for vLLM-omni server.
This example shows how to interact with the vLLM-omni API server
using HTTP requests.
"""
import requests
import json
import time
def test_health(host="localhost", port=8000):
"""Test if the server is healthy."""
url = f"http://{host}:{port}/health"
try:
response = requests.get(url, timeout=5)
if response.status_code == 200:
print("✅ Server is healthy")
return True
else:
print(f"❌ Server health check failed: {response.status_code}")
return False
except requests.exceptions.RequestException as e:
print(f"❌ Cannot connect to server: {e}")
return False
def generate_text(prompt, host="localhost", port=8000, **kwargs):
"""Generate text using the API server."""
url = f"http://{host}:{port}/generate"
# Default parameters
params = {
"prompts": [prompt],
"max_tokens": kwargs.get("max_tokens", 100),
"temperature": kwargs.get("temperature", 0.7),
"top_p": kwargs.get("top_p", 1.0),
"frequency_penalty": kwargs.get("frequency_penalty", 0.0),
"presence_penalty": kwargs.get("presence_penalty", 0.0),
}
# Add optional parameters
if "stop" in kwargs:
params["stop"] = kwargs["stop"]
try:
response = requests.post(url, json=params, timeout=30)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
print(f"❌ API request failed: {e}")
return None
def main():
"""Main example function."""
print("vLLM-omni API Client Example")
print("=" * 40)
# Test server health
if not test_health():
print("Please start the server first:")
print(" vllm serve Qwen/Qwen3-0.6B --omni --port 8000")
return
# Example 1: Simple generation
print("\n1. Simple text generation:")
result = generate_text("Hello, how are you?", max_tokens=50)
if result:
text = result["outputs"][0]["text"]
print(f" Input: Hello, how are you?")
print(f" Output: {text}")
# Example 2: Creative writing
print("\n2. Creative writing:")
result = generate_text(
"Write a short story about a robot learning to paint",
max_tokens=150,
temperature=0.9
)
if result:
text = result["outputs"][0]["text"]
print(f" Input: Write a short story about a robot learning to paint")
print(f" Output: {text}")
# Example 3: Question answering
print("\n3. Question answering:")
result = generate_text(
"What is the capital of France?",
max_tokens=30,
temperature=0.3
)
if result:
text = result["outputs"][0]["text"]
print(f" Input: What is the capital of France?")
print(f" Output: {text}")
# Example 4: Multiple prompts
print("\n4. Multiple prompts:")
prompts = [
"Tell me a joke",
"What's 2+2?",
"Describe the color blue"
]
for prompt in prompts:
result = generate_text(prompt, max_tokens=30)
if result:
text = result["outputs"][0]["text"]
print(f" Q: {prompt}")
print(f" A: {text}")
# Example 5: Show full response structure
print("\n5. Full response structure:")
result = generate_text("Hello world", max_tokens=20)
if result:
print(" Full API response:")
print(json.dumps(result, indent=2))
print("\n" + "=" * 40)
print("API client examples completed!")
if __name__ == "__main__":
main()
+178
View File
@@ -0,0 +1,178 @@
#!/usr/bin/env python3
"""
Docker example for vLLM-omni.
This example shows how to test the vLLM-omni API server running in Docker.
"""
import requests
import json
import time
import subprocess
import sys
def check_docker_running():
"""Check if Docker is running."""
try:
result = subprocess.run(['docker', 'info'],
capture_output=True, text=True, timeout=5)
return result.returncode == 0
except (subprocess.TimeoutExpired, FileNotFoundError):
return False
def start_docker_server():
"""Start the vLLM-omni server in Docker."""
print("🐳 Starting vLLM-omni server in Docker...")
# Check if container is already running
result = subprocess.run(['docker', 'ps', '--filter', 'name=vllm-omni-cpu', '--format', '{{.Names}}'],
capture_output=True, text=True)
if 'vllm-omni-cpu' in result.stdout:
print("✅ Server already running")
return True
# Start the server
cmd = [
'docker', 'run', '-d', '--name', 'vllm-omni-cpu',
'-p', '8000:8000',
'vllm-omni-cpu:latest',
'vllm', 'serve', 'Qwen/Qwen3-0.6B', '--omni', '--port', '8000'
]
try:
subprocess.run(cmd, check=True)
print("✅ Server started successfully")
# Wait for server to be ready
print("⏳ Waiting for server to be ready...")
for i in range(30): # Wait up to 30 seconds
if test_server_health():
print("✅ Server is ready!")
return True
time.sleep(1)
print("❌ Server failed to start within 30 seconds")
return False
except subprocess.CalledProcessError as e:
print(f"❌ Failed to start server: {e}")
return False
def test_server_health():
"""Test if the server is healthy."""
try:
response = requests.get("http://localhost:8000/health", timeout=2)
return response.status_code == 200
except requests.exceptions.RequestException:
return False
def test_generation():
"""Test text generation via API."""
print("\n🧪 Testing text generation...")
test_cases = [
{
"prompt": "Hello, how are you?",
"max_tokens": 50,
"temperature": 0.7
},
{
"prompt": "Tell me a joke",
"max_tokens": 100,
"temperature": 0.9
},
{
"prompt": "What is artificial intelligence?",
"max_tokens": 150,
"temperature": 0.5
}
]
for i, test_case in enumerate(test_cases, 1):
print(f"\nTest {i}: {test_case['prompt']}")
try:
response = requests.post(
"http://localhost:8000/generate",
json={
"prompts": [test_case["prompt"]],
"max_tokens": test_case["max_tokens"],
"temperature": test_case["temperature"]
},
timeout=30
)
if response.status_code == 200:
result = response.json()
text = result["outputs"][0]["text"]
print(f"✅ Response: {text}")
else:
print(f"❌ Error: {response.status_code} - {response.text}")
except requests.exceptions.RequestException as e:
print(f"❌ Request failed: {e}")
def cleanup():
"""Clean up Docker container."""
print("\n🧹 Cleaning up...")
try:
subprocess.run(['docker', 'stop', 'vllm-omni-cpu'],
capture_output=True, text=True)
subprocess.run(['docker', 'rm', 'vllm-omni-cpu'],
capture_output=True, text=True)
print("✅ Cleanup completed")
except Exception as e:
print(f"⚠️ Cleanup warning: {e}")
def main():
"""Main example function."""
print("🐳 vLLM-omni Docker Example")
print("=" * 40)
# Check Docker
if not check_docker_running():
print("❌ Docker is not running. Please start Docker Desktop.")
sys.exit(1)
print("✅ Docker is running")
# Check if image exists
result = subprocess.run(['docker', 'images', 'vllm-omni-cpu:latest'],
capture_output=True, text=True)
if 'vllm-omni-cpu' not in result.stdout:
print("❌ Docker image not found. Please build it first:")
print(" docker build -f docker/Dockerfile.cpu -t vllm-omni-cpu .")
sys.exit(1)
print("✅ Docker image found")
try:
# Start server
if not start_docker_server():
sys.exit(1)
# Test generation
test_generation()
print("\n🎉 Docker example completed successfully!")
except KeyboardInterrupt:
print("\n⏹️ Example interrupted by user")
finally:
# Ask if user wants to keep server running
keep_running = input("\n🤔 Keep the server running? (y/N): ").lower().strip()
if keep_running != 'y':
cleanup()
else:
print("✅ Server kept running. Access it at http://localhost:8000")
if __name__ == "__main__":
main()
+137
View File
@@ -0,0 +1,137 @@
#!/usr/bin/env python3
"""
Simple vLLM-omni usage example.
This example demonstrates how to use vLLM-omni for basic text generation
with a single AR (Autoregressive) stage.
"""
import asyncio
from vllm_omni.core.omni_llm import OmniLLM, AsyncOmniLLM
from vllm_omni.config import create_ar_stage_config
def sync_example():
"""Synchronous usage example."""
print("=== Synchronous vLLM-omni Example ===")
# Create a simple AR stage configuration
stage_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"]
)
# Initialize OmniLLM
omni_llm = OmniLLM([stage_config])
# Prepare stage arguments
stage_args = [{
"prompt": "Hello, how are you today?",
"max_tokens": 50,
"temperature": 0.7
}]
# Generate text
print(f"Input: {stage_args[0]['prompt']}")
outputs = omni_llm.generate(stage_args)
# Display results
for i, output in enumerate(outputs):
print(f"Output {i}:")
if hasattr(output, 'outputs') and output.outputs:
for completion in output.outputs:
print(f" Text: {completion.text}")
print(f" Finished: {completion.finish_reason != 'length'}")
print(f" Tokens: {len(completion.token_ids)} tokens")
async def async_example():
"""Asynchronous usage example."""
print("\n=== Asynchronous vLLM-omni Example ===")
# Create a simple AR stage configuration
stage_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"]
)
# Initialize AsyncOmniLLM
omni_llm = AsyncOmniLLM([stage_config])
# Prepare stage arguments
stage_args = [{
"prompt": "What is artificial intelligence?",
"max_tokens": 100,
"temperature": 0.8
}]
# Generate text asynchronously
print(f"Input: {stage_args[0]['prompt']}")
outputs = await omni_llm.generate_async(stage_args)
# Display results
for i, output in enumerate(outputs):
print(f"Output {i}:")
if hasattr(output, 'outputs') and output.outputs:
for completion in output.outputs:
print(f" Text: {completion.text}")
print(f" Finished: {completion.finish_reason != 'length'}")
print(f" Tokens: {len(completion.token_ids)} tokens")
def multi_prompt_example():
"""Example with multiple prompts."""
print("\n=== Multi-Prompt Example ===")
# Create stage configuration
stage_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"]
)
# Initialize OmniLLM
omni_llm = OmniLLM([stage_config])
# Multiple prompts
prompts = [
"Tell me a joke",
"What's the weather like?",
"Explain quantum computing"
]
for prompt in prompts:
print(f"\nInput: {prompt}")
stage_args = [{
"prompt": prompt,
"max_tokens": 30,
"temperature": 0.9
}]
outputs = omni_llm.generate(stage_args)
if outputs and hasattr(outputs[0], 'outputs') and outputs[0].outputs:
completion = outputs[0].outputs[0]
print(f"Output: {completion.text}")
if __name__ == "__main__":
print("vLLM-omni Simple Usage Examples")
print("=" * 40)
# Run synchronous example
sync_example()
# Run asynchronous example
asyncio.run(async_example())
# Run multi-prompt example
multi_prompt_example()
print("\n" + "=" * 40)
print("Examples completed!")