hsliu
afe5c2692c
- 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
129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
#!/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()
|
|
|