update ar_dit_test
This commit is contained in:
hsliu committed
1 parent
ff9c7d4536
commit
81f04d3cf9
32 files changed
+1192
-1908
No files matched your search
@@ -1,83 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# vLLM-omni Docker Setup Script for macOS
|
||||
# This script helps set up Docker for vLLM-omni on macOS
|
||||
|
||||
set -e
|
||||
|
||||
echo "🐳 vLLM-omni Docker Setup for macOS"
|
||||
echo "=================================="
|
||||
|
||||
# Check if Docker is installed
|
||||
if ! command -v docker &> /dev/null; then
|
||||
echo "❌ Docker is not installed. Please install Docker Desktop for Mac:"
|
||||
echo " https://www.docker.com/products/docker-desktop/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if Docker is running
|
||||
if ! docker info &> /dev/null; then
|
||||
echo "❌ Docker is not running. Please start Docker Desktop."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "✅ Docker is installed and running"
|
||||
|
||||
# Check if NVIDIA Docker runtime is available (optional for macOS)
|
||||
if docker info | grep -q nvidia; then
|
||||
echo "✅ NVIDIA Docker runtime detected"
|
||||
GPU_SUPPORT=true
|
||||
else
|
||||
echo "⚠️ NVIDIA Docker runtime not detected (normal on macOS)"
|
||||
GPU_SUPPORT=false
|
||||
fi
|
||||
|
||||
# Create necessary directories
|
||||
echo "📁 Creating directories..."
|
||||
mkdir -p models logs notebooks
|
||||
|
||||
# Build the Docker image
|
||||
echo "🔨 Building vLLM-omni Docker image..."
|
||||
docker build -f docker/Dockerfile.cpu -t vllm-omni-cpu:latest .
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "✅ Docker image built successfully"
|
||||
else
|
||||
echo "❌ Failed to build Docker image"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Test the image
|
||||
echo "🧪 Testing Docker image..."
|
||||
docker run --rm vllm-omni-cpu:latest python -c "import vllm_omni; print('vLLM-omni imported successfully')"
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "✅ Docker image test passed"
|
||||
else
|
||||
echo "❌ Docker image test failed"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "🎉 Setup complete! You can now use:"
|
||||
echo ""
|
||||
echo " # Start the server:"
|
||||
echo " docker-compose -f docker/docker-compose.cpu.yml up vllm-omni-cpu"
|
||||
echo ""
|
||||
echo " # Start with Jupyter:"
|
||||
echo " docker-compose -f docker/docker-compose.cpu.yml --profile jupyter up"
|
||||
echo ""
|
||||
echo " # Run a single command:"
|
||||
echo " docker run --rm -p 8000:8000 vllm-omni-cpu:latest vllm serve Qwen/Qwen3-0.6B --omni --port 8000"
|
||||
echo ""
|
||||
echo " # Interactive shell:"
|
||||
echo " docker run -it --rm vllm-omni-cpu:latest bash"
|
||||
echo ""
|
||||
|
||||
# Optional: Start the server
|
||||
read -p "🚀 Start the vLLM-omni server now? (y/N): " -n 1 -r
|
||||
echo
|
||||
if [[ $REPLY =~ ^[Yy]$ ]]; then
|
||||
echo "Starting vLLM-omni server..."
|
||||
docker-compose -f docker/docker-compose.cpu.yml up vllm-omni-cpu
|
||||
fi
|
||||
@@ -1,147 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Model download script for vLLM-omni
|
||||
Downloads commonly used AR and DiT models
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
def download_ar_models(models=None):
|
||||
"""Download AR (Autoregressive) models"""
|
||||
if models is None:
|
||||
models = [
|
||||
"Qwen/Qwen3-0.6B",
|
||||
]
|
||||
|
||||
print("🚀 Downloading AR models...")
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
except ImportError:
|
||||
print("❌ transformers not installed. Installing...")
|
||||
os.system("pip install transformers")
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
for model_name in models:
|
||||
print(f"\n📥 Downloading AR model: {model_name}")
|
||||
try:
|
||||
print(f" - Downloading tokenizer...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
print(f" - Downloading model...")
|
||||
model = AutoModel.from_pretrained(model_name)
|
||||
|
||||
print(f"✅ {model_name} downloaded successfully")
|
||||
|
||||
# Print model info
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
size_mb = total_params * 4 / (1024 * 1024)
|
||||
print(f" 📊 Parameters: {total_params:,} ({size_mb:.1f} MB)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to download {model_name}: {e}")
|
||||
|
||||
def download_dit_models(models=None):
|
||||
"""Download DiT (Diffusion Transformer) models"""
|
||||
if models is None:
|
||||
models = [
|
||||
"stabilityai/stable-diffusion-2-1",
|
||||
"runwayml/stable-diffusion-v1-5"
|
||||
]
|
||||
|
||||
print("\n🚀 Downloading DiT models...")
|
||||
|
||||
try:
|
||||
from diffusers import StableDiffusionPipeline
|
||||
except ImportError:
|
||||
print("❌ diffusers not installed. Installing...")
|
||||
os.system("pip install diffusers")
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
for model_name in models:
|
||||
print(f"\n📥 Downloading DiT model: {model_name}")
|
||||
try:
|
||||
print(f" - Downloading pipeline...")
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_name)
|
||||
|
||||
print(f"✅ {model_name} downloaded successfully")
|
||||
|
||||
# Print model info
|
||||
total_params = sum(p.numel() for p in pipe.unet.parameters())
|
||||
size_mb = total_params * 4 / (1024 * 1024)
|
||||
print(f" 📊 UNet parameters: {total_params:,} ({size_mb:.1f} MB)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to download {model_name}: {e}")
|
||||
|
||||
def check_cache_directory():
|
||||
"""Check and display cache directory info"""
|
||||
cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
|
||||
print(f"📁 Cache directory: {cache_dir}")
|
||||
|
||||
if cache_dir.exists():
|
||||
total_size = sum(f.stat().st_size for f in cache_dir.rglob('*') if f.is_file())
|
||||
size_gb = total_size / (1024**3)
|
||||
print(f"📊 Total cache size: {size_gb:.2f} GB")
|
||||
|
||||
# List downloaded models
|
||||
model_dirs = [d for d in cache_dir.iterdir() if d.is_dir() and d.name.startswith("models--")]
|
||||
if model_dirs:
|
||||
print(f"📋 Downloaded models ({len(model_dirs)}):")
|
||||
for model_dir in sorted(model_dirs):
|
||||
model_name = model_dir.name.replace("models--", "").replace("--", "/")
|
||||
print(f" - {model_name}")
|
||||
else:
|
||||
print("📁 Cache directory does not exist yet")
|
||||
|
||||
def clear_cache():
|
||||
"""Clear model cache"""
|
||||
cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
|
||||
|
||||
if cache_dir.exists():
|
||||
print(f"🗑️ Clearing cache directory: {cache_dir}")
|
||||
import shutil
|
||||
shutil.rmtree(cache_dir)
|
||||
print("✅ Cache cleared successfully")
|
||||
else:
|
||||
print("📁 Cache directory does not exist")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Download models for vLLM-omni")
|
||||
parser.add_argument("--ar-models", nargs="+", help="AR models to download")
|
||||
parser.add_argument("--dit-models", nargs="+", help="DiT models to download")
|
||||
parser.add_argument("--all", action="store_true", help="Download all default models")
|
||||
parser.add_argument("--check-cache", action="store_true", help="Check cache directory")
|
||||
parser.add_argument("--clear-cache", action="store_true", help="Clear cache directory")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print("🎯 vLLM-omni Model Downloader")
|
||||
print("=" * 40)
|
||||
|
||||
if args.clear_cache:
|
||||
clear_cache()
|
||||
return
|
||||
|
||||
if args.check_cache:
|
||||
check_cache_directory()
|
||||
return
|
||||
|
||||
if args.all or args.ar_models:
|
||||
download_ar_models(args.ar_models)
|
||||
|
||||
if args.all or args.dit_models:
|
||||
download_dit_models(args.dit_models)
|
||||
|
||||
if not any([args.all, args.ar_models, args.dit_models, args.check_cache, args.clear_cache]):
|
||||
print("❓ No action specified. Use --help for options.")
|
||||
print("\nQuick start:")
|
||||
print(" python scripts/download_models.py --all")
|
||||
print(" python scripts/download_models.py --ar-models Qwen/Qwen3-0.6B")
|
||||
print(" python scripts/download_models.py --check-cache")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,82 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Quick vLLM-omni Test Script
|
||||
# Fast test for basic functionality after changes
|
||||
# Usage: ./scripts/quick_test.sh [port]
|
||||
|
||||
set -e
|
||||
|
||||
PORT=${1:-8000}
|
||||
HOST="localhost"
|
||||
|
||||
# Colors
|
||||
GREEN='\033[0;32m'
|
||||
RED='\033[0;31m'
|
||||
BLUE='\033[0;34m'
|
||||
NC='\033[0m'
|
||||
|
||||
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
|
||||
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
|
||||
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
|
||||
|
||||
# Cleanup
|
||||
cleanup() {
|
||||
log_info "Cleaning up..."
|
||||
pkill -f "vllm serve" 2>/dev/null || true
|
||||
}
|
||||
|
||||
trap cleanup EXIT
|
||||
|
||||
# Quick test
|
||||
main() {
|
||||
log_info "Quick vLLM-omni test on port $PORT"
|
||||
|
||||
# Test imports
|
||||
log_info "Testing imports..."
|
||||
python -c "import vllm_omni; from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM; print('✅ Imports OK')" || {
|
||||
log_error "Import failed"
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Start server
|
||||
log_info "Starting server..."
|
||||
vllm serve ./models/Qwen3-0.6B --omni --port $PORT > /dev/null 2>&1 &
|
||||
SERVER_PID=$!
|
||||
|
||||
# Wait for server
|
||||
sleep 15
|
||||
|
||||
# Test health with retry
|
||||
log_info "Testing health endpoint..."
|
||||
for i in {1..5}; do
|
||||
if curl -s http://$HOST:$PORT/health | grep -q "healthy"; then
|
||||
log_success "Health check passed"
|
||||
break
|
||||
else
|
||||
if [ $i -eq 5 ]; then
|
||||
log_error "Health check failed after 5 attempts"
|
||||
exit 1
|
||||
fi
|
||||
log_info "Health check attempt $i failed, retrying in 3 seconds..."
|
||||
sleep 3
|
||||
fi
|
||||
done
|
||||
|
||||
# Test generation
|
||||
log_info "Testing text generation..."
|
||||
response=$(curl -s -X POST http://$HOST:$PORT/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompts": ["Quick test"], "max_tokens": 10, "temperature": 0.7}')
|
||||
|
||||
if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data else 1)" 2>/dev/null; then
|
||||
log_success "Generation test passed"
|
||||
else
|
||||
log_error "Generation test failed"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
log_success "Quick test completed successfully!"
|
||||
}
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
main "$@"
|
||||
+74
-12
@@ -122,16 +122,22 @@ start_server() {
|
||||
# Test health endpoint
|
||||
test_health() {
|
||||
log_info "Testing health endpoint..."
|
||||
local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/health 2>/dev/null || echo "000")
|
||||
local http_code="${response: -3}"
|
||||
local body="${response%???}"
|
||||
|
||||
if [ "$http_code" = "200" ]; then
|
||||
log_success "Health check passed: $body"
|
||||
else
|
||||
log_error "Health check failed (HTTP $http_code): $body"
|
||||
return 1
|
||||
fi
|
||||
for attempt in {1..10}; do
|
||||
local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/health 2>/dev/null || echo "000")
|
||||
local http_code="${response: -3}"
|
||||
local body="${response%???}"
|
||||
|
||||
if [ "$http_code" = "200" ]; then
|
||||
log_success "Health check passed: $body"
|
||||
return 0
|
||||
fi
|
||||
|
||||
log_warning "Health check attempt $attempt failed (HTTP $http_code). Retrying in 3s..."
|
||||
sleep 3
|
||||
done
|
||||
|
||||
log_error "Health check failed after multiple attempts"
|
||||
return 1
|
||||
}
|
||||
|
||||
# Test info endpoint
|
||||
@@ -186,6 +192,60 @@ test_api_client() {
|
||||
fi
|
||||
}
|
||||
|
||||
# Test AR → DiT pipeline example
|
||||
test_ar_dit_pipeline() {
|
||||
log_info "Testing AR → DiT diffusers pipeline example..."
|
||||
|
||||
if ! python - 2>/dev/null <<'PY'
|
||||
import importlib
|
||||
for pkg in ("torch", "diffusers"):
|
||||
if importlib.util.find_spec(pkg) is None:
|
||||
raise ImportError(pkg)
|
||||
PY
|
||||
then
|
||||
log_warning "Required packages for diffusers pipeline not found, skipping AR → DiT test"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if [ ! -f "examples/omni/ar_dit_diffusers.py" ]; then
|
||||
log_warning "AR → DiT example not found, skipping"
|
||||
return 0
|
||||
fi
|
||||
|
||||
local output_path="logs/ar_dit_pipeline.png"
|
||||
local log_path="logs/ar_dit_pipeline.log"
|
||||
rm -f "$output_path"
|
||||
|
||||
local cmd=(python examples/omni/ar_dit_diffusers.py
|
||||
--prompt "Smoke test prompt of a lighthouse"
|
||||
--seed 0
|
||||
--output "$output_path")
|
||||
|
||||
local env_prefix=()
|
||||
if [ -n "$AR_MODEL" ]; then
|
||||
env_prefix+=("AR_MODEL=$AR_MODEL")
|
||||
fi
|
||||
if [ -n "$DIT_MODEL" ]; then
|
||||
env_prefix+=("DIT_MODEL=$DIT_MODEL")
|
||||
fi
|
||||
|
||||
if ! env "${env_prefix[@]}" "${cmd[@]}" >"$log_path" 2>&1; then
|
||||
log_error "AR → DiT pipeline test failed"
|
||||
log_info "AR → DiT logs:"
|
||||
cat "$log_path"
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [ ! -f "$output_path" ]; then
|
||||
log_error "AR → DiT pipeline did not produce an image"
|
||||
log_info "AR → DiT logs:"
|
||||
cat "$log_path"
|
||||
return 1
|
||||
fi
|
||||
|
||||
log_success "AR → DiT pipeline test passed (output: $output_path)"
|
||||
}
|
||||
|
||||
# Test simple usage example
|
||||
test_simple_usage() {
|
||||
log_info "Testing simple usage example..."
|
||||
@@ -250,10 +310,11 @@ run_tests() {
|
||||
test_generation
|
||||
test_performance
|
||||
test_api_client
|
||||
|
||||
test_ar_dit_pipeline
|
||||
|
||||
# Note: Simple usage test is commented out as it starts its own server
|
||||
# test_simple_usage
|
||||
|
||||
|
||||
log_success "All tests completed successfully!"
|
||||
echo "=========================================="
|
||||
log_info "Test Summary:"
|
||||
@@ -262,6 +323,7 @@ run_tests() {
|
||||
log_info "✅ Text generation functionality"
|
||||
log_info "✅ Performance metrics"
|
||||
log_info "✅ API client integration"
|
||||
log_info "✅ AR → DiT diffusers pipeline"
|
||||
log_info "✅ All imports working correctly"
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user