如何使用vLLM + Docker Compose 搭建大模型推理服务（基于云端GPU VPS搭建）

import requests
import time
import concurrent.futures
import statistics
import sys
import json
from datetime import datetime

def test_request_simple(id, output_tokens=50):
    """简单测试请求"""
    start = time.time()
    try:
        response = requests.post("http://localhost/v1/chat/completions",
            json={
                "model": "Qwen/Qwen2-1.5B-Instruct",
                "messages": [{"role": "user", "content": f"Request {id}: Explain AI briefly"}],
                "max_tokens": output_tokens
            }, timeout=30)
        success = response.status_code == 200
        latency = time.time() - start
        return {"id": id, "success": success, "latency": latency, "status": response.status_code}
    except Exception as e:
        return {"id": id, "success": False, "latency": time.time() - start, "error": str(e)}

def test_request_varied(id):
    """不同长度的请求测试"""
    # 轮流使用不同的输出长度
    lengths = [10, 50, 100, 200]
    output_tokens = lengths[id % len(lengths)]
    
    start = time.time()
    try:
        response = requests.post("http://localhost/v1/chat/completions",
            json={
                "model": "Qwen/Qwen2-1.5B-Instruct",
                "messages": [{"role": "user", "content": f"Request {id}: Tell me about machine learning"}],
                "max_tokens": output_tokens,
                "temperature": 0.7
            }, timeout=30)
        success = response.status_code == 200
        latency = time.time() - start
        
        # 获取实际生成的token数（如果有）
        actual_tokens = 0
        if success:
            result = response.json()
            actual_tokens = result.get('usage', {}).get('completion_tokens', 0)
        
        return {
            "id": id, 
            "success": success, 
            "latency": latency, 
            "output_tokens": output_tokens,
            "actual_tokens": actual_tokens,
            "status": response.status_code
        }
    except Exception as e:
        return {"id": id, "success": False, "latency": time.time() - start, "error": str(e)}

def run_concurrent_test(name, num_requests, concurrency, request_type="simple"):
    """运行并发测试"""
    print(f"\n{'='*60}")
    print(f"测试: {name}")
    print(f"请求数: {num_requests} | 并发数: {concurrency}")
    print(f"请求类型: {request_type}")
    print(f"{'='*60}")
    
    test_func = test_request_simple if request_type == "simple" else test_request_varied
    
    start_time = time.time()
    results = []
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
        futures = [executor.submit(test_func, i) for i in range(num_requests)]
        
        completed = 0
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            results.append(result)
            completed += 1
            
            # 每10个请求打印一次进度
            if completed % 10 == 0 or completed == num_requests:
                print(f"进度: {completed}/{num_requests} ({completed/num_requests*100:.0f}%)")
    
    total_time = time.time() - start_time
    
    # 统计分析
    successful = [r for r in results if r["success"]]
    failed = [r for r in results if not r["success"]]
    
    if successful:
        latencies = [r["latency"] for r in successful]
        
        print(f"\n{'='*60}")
        print(f"测试结果")
        print(f"{'='*60}")
        print(f"总耗时: {total_time:.2f}s")
        print(f"吞吐量: {len(successful)/total_time:.2f} req/s")
        print(f"成功率: {len(successful)}/{num_requests} ({len(successful)/num_requests*100:.1f}%)")
        
        if failed:
            print(f"失败数: {len(failed)}")
            for f in failed[:5]:  # 只显示前5个错误
                print(f"  - 请求 {f['id']}: {f.get('error', 'Unknown error')}")
        
        print(f"\n延迟统计:")
        print(f"  平均: {statistics.mean(latencies):.3f}s")
        print(f"  中位数: {statistics.median(latencies):.3f}s")
        print(f"  最小: {min(latencies):.3f}s")
        print(f"  最大: {max(latencies):.3f}s")
        
        if len(latencies) >= 10:
            # 计算百分位数
            latencies_sorted = sorted(latencies)
            p95_index = int(len(latencies_sorted) * 0.95)
            p99_index = int(len(latencies_sorted) * 0.99)
            print(f"  P95: {latencies_sorted[p95_index]:.3f}s")
            if p99_index < len(latencies_sorted):
                print(f"  P99: {latencies_sorted[p99_index]:.3f}s")
        
        # 如果是变化长度的测试，显示token统计
        if request_type == "varied" and "actual_tokens" in successful[0]:
            actual_tokens = [r["actual_tokens"] for r in successful]
            print(f"\nToken生成统计:")
            print(f"  平均生成tokens: {statistics.mean(actual_tokens):.1f}")
            print(f"  总生成tokens: {sum(actual_tokens)}")
            print(f"  tokens/秒: {sum(actual_tokens)/total_time:.1f}")
    else:
        print(f"\n❌ 所有请求都失败了！")
    
    return results

def run_stress_gradual(base_requests=50):
    """逐步增加负载测试"""
    print(f"\n{'='*60}")
    print("逐步增加负载测试")
    print(f"{'='*60}")
    
    concurrency_levels = [1, 2, 5, 10, 20, 30, 40, 50]
    results_summary = []
    
    for concurrency in concurrency_levels:
        if concurrency > base_requests:
            break
            
        print(f"\n--- 并发数: {concurrency} ---")
        start = time.time()
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
            futures = [executor.submit(test_request_simple, i, 30) for i in range(base_requests)]
            results = [f.result() for f in concurrent.futures.as_completed(futures)]
        
        total_time = time.time() - start
        successful = [r for r in results if r["success"]]
        
        throughput = len(successful) / total_time if total_time > 0 else 0
        avg_latency = statistics.mean([r["latency"] for r in successful]) if successful else 0
        
        results_summary.append({
            "concurrency": concurrency,
            "throughput": throughput,
            "avg_latency": avg_latency,
            "success_rate": len(successful) / base_requests * 100
        })
        
        print(f"吞吐量: {throughput:.2f} req/s | 平均延迟: {avg_latency:.3f}s | 成功率: {results_summary[-1]['success_rate']:.1f}%")
        
        # 在两次测试之间稍作休息，让系统恢复
        time.sleep(2)
    
    # 打印汇总表
    print(f"\n{'='*60}")
    print("逐步负载测试汇总")
    print(f"{'='*60}")
    print(f"{'并发数':<8} {'吞吐量(req/s)':<15} {'平均延迟(s)':<12} {'成功率':<10}")
    print(f"{'-'*60}")
    for r in results_summary:
        print(f"{r['concurrency']:<8} {r['throughput']:<15.2f} {r['avg_latency']:<12.3f} {r['success_rate']:<10.1f}%")
    
    return results_summary

def main():
    print("vLLM 高级压力测试工具")
    print("=" * 60)
    print("1. 简单并发测试 (20, 30, 50 并发)")
    print("2. 变化负载测试 (不同输出长度)")
    print("3. 逐步增加负载测试 (找到系统瓶颈)")
    print("4. 自定义测试")
    print("5. 快速测试 (50并发, 输出50 tokens)")
    print("6. 退出")
    
    choice = input("\n请选择测试类型 (1-6): ").strip()
    
    if choice == "1":
        for concurrency in [20, 30, 50]:
            run_concurrent_test(f"{concurrency}并发测试", concurrency, concurrency, "simple")
            time.sleep(2)
    
    elif choice == "2":
        run_concurrent_test("变化负载测试", 50, 20, "varied")
    
    elif choice == "3":
        run_stress_gradual(50)
    
    elif choice == "4":
        num = int(input("请求数量: "))
        concurrency = int(input("并发数: "))
        test_type = input("测试类型 (simple/varied): ").strip().lower()
        if test_type not in ['simple', 'varied']:
            test_type = "simple"
        run_concurrent_test(f"自定义测试", num, concurrency, test_type)
    
    elif choice == "5":
        print("\n快速测试: 50个请求, 50并发, 输出50 tokens")
        run_concurrent_test("快速测试", 50, 50, "simple")
    
    else:
        print("再见！")

if __name__ == "__main__":
    main()