from openai import OpenAI
import time
import concurrent.futures
import statistics

# 配置客户端
client = OpenAI(
    api_key="EMPTY",  # vLLM不需要API key
    base_url="http://localhost:8000/v1"
)

def test_vllm_request(prompt, model="DeepSeek-R1-Distill-Qwen-7B"):
    """测试vLLM单个请求"""
    start_time = time.time()
    
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=512,
            temperature=0.7
        )
        end_time = time.time()
        
        response_time = end_time - start_time
        token_count = response.usage.completion_tokens
        
        return response_time, token_count, response.choices[0].message.content
    except Exception as e:
        print(f"请求失败: {e}")
        return None, 0, ""

def benchmark_vllm(num_requests=50, max_workers=10):
    """vLLM性能基准测试"""
    prompts = [
        "请解释什么是机器学习",
        "介绍一下深度学习的基本概念",
        "什么是自然语言处理",
        "解释一下Transformer架构",
        "什么是强化学习"
    ] * (num_requests // 5 + 1)
    
    prompts = prompts[:num_requests]
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        start_time = time.time()
        futures = [executor.submit(test_vllm_request, prompt) for prompt in prompts]
        results = [future.result() for future in futures]
        end_time = time.time()
    
    # 分析结果
    successful_results = [r for r in results if r[0] is not None]
    response_times = [r[0] for r in successful_results]
    token_counts = [r[1] for r in successful_results]
    
    print(f"=== vLLM性能测试结果 ===")
    print(f"总请求数: {num_requests}")
    print(f"成功请求数: {len(successful_results)}")
    print(f"成功率: {len(successful_results)/num_requests*100:.1f}%")
    print(f"总耗时: {end_time - start_time:.2f}秒")
    print(f"平均响应时间: {statistics.mean(response_times):.2f}秒")
    print(f"响应时间中位数: {statistics.median(response_times):.2f}秒")
    print(f"95%分位数: {sorted(response_times)[int(len(response_times)*0.95)]:.2f}秒")
    print(f"平均生成token数: {statistics.mean(token_counts):.1f}")
    print(f"吞吐量: {len(successful_results)/(end_time - start_time):.2f} 请求/秒")
    print(f"Token吞吐量: {sum(token_counts)/(end_time - start_time):.1f} tokens/秒")

if __name__ == "__main__":
    benchmark_vllm(num_requests=100, max_workers=20)