test_vllm_connection.py

#!/usr/bin/env python3
"""
vLLM服务器连接测试脚本
测试GPU 3（端口8895）和GPU 4（端口8889）上的vLLM服务器
"""

import sys
import os
sys.path.append('/home/czy/CISC_PRO')

import requests
import json
from src.vllm_models import VLLMModelWrapper, VLLMConfig

def test_vllm_server(name: str, base_url: str, model_name: str):
    """测试单个vLLM服务器"""
    print(f"\n🔍 测试 {name} 服务器...")
    print(f"   URL: {base_url}")
    print(f"   模型: {model_name}")
    
    try:
        # 1. 测试基本连接
        response = requests.get(f"{base_url}/v1/models", timeout=10)
        if response.status_code == 200:
            models_data = response.json()
            available_models = [m.get('id', 'unknown') for m in models_data.get('data', [])]
            print(f"✅ 服务器连接成功")
            print(f"📋 可用模型: {available_models}")
        else:
            print(f"❌ 服务器响应异常: {response.status_code}")
            return False
        
        # 2. 测试模型包装器
        config = VLLMConfig(
            base_url=base_url,
            model_name=model_name,
            timeout=30
        )
        
        wrapper = VLLMModelWrapper(config)
        print(f"✅ 模型包装器初始化成功")
        
        # 3. 测试简单生成
        test_prompt = "What is 2+2? Answer:"
        result = wrapper._single_generate(
            test_prompt,
            max_tokens=10,
            temperature=0.0
        )
        
        generated_text = result["generated_text"].strip()
        print(f"✅ 文本生成测试成功")
        print(f"📝 测试输入: {test_prompt}")
        print(f"📤 生成结果: {generated_text}")
        print(f"📊 Token统计: {result.get('completion_tokens', 0)} 完成tokens")
        
        return True
        
    except requests.ConnectionError:
        print(f"❌ 无法连接到服务器 {base_url}")
        print(f"   请确保vLLM服务器正在运行")
        return False
    except requests.Timeout:
        print(f"❌ 连接超时")
        return False
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False

def test_gsm8k_sample():
    """测试GSM8K样本处理"""
    print(f"\n🧮 测试GSM8K数据集处理...")
    
    try:
        from src.local_datasets import LocalGSM8KProcessor
        
        processor = LocalGSM8KProcessor()
        print(f"✅ GSM8K数据集加载成功: {len(processor.dataset)} 个样本")
        
        # 获取一个样本
        sample = processor.dataset[0]
        print(f"📝 测试样本: {sample['question'][:100]}...")
        
        # 测试置信度prompt生成
        confidence_prompts = processor.get_confidence_prompts(
            sample, 
            "The answer is 42.", 
            num_prompts=2
        )
        print(f"✅ 置信度prompt生成成功: {len(confidence_prompts)} 个")
        print(f"📋 示例prompt: {confidence_prompts[0][:200]}...")
        
        return True
        
    except Exception as e:
        print(f"❌ GSM8K处理测试失败: {e}")
        return False

def main():
    print("🧪 vLLM服务器连接测试")
    print("=" * 60)
    
    # 测试配置
    servers = [
        {
            "name": "Llama (GPU 3)",
            "base_url": "http://localhost:8895",
            "model_name": "Meta-Llama-3.1-8B-Instruct"
        },
        {
            "name": "Mistral (GPU 4)", 
            "base_url": "http://localhost:8889",
            "model_name": "Mistral-7B-Instruct-v0.3"
        }
    ]
    
    results = []
    
    # 测试每个服务器
    for server in servers:
        success = test_vllm_server(
            server["name"],
            server["base_url"], 
            server["model_name"]
        )
        results.append((server["name"], success))
    
    # 测试GSM8K处理
    gsm8k_success = test_gsm8k_sample()
    
    # 总结
    print("\n" + "=" * 60)
    print("🏁 测试总结")
    print("=" * 60)
    
    all_success = True
    for name, success in results:
        status = "✅ 正常" if success else "❌ 失败"
        print(f"{name:20s}: {status}")
        if not success:
            all_success = False
    
    gsm8k_status = "✅ 正常" if gsm8k_success else "❌ 失败"
    print(f"{'GSM8K数据处理':20s}: {gsm8k_status}")
    if not gsm8k_success:
        all_success = False
    
    if all_success:
        print(f"\n🎉 所有测试通过!")
        print(f"✅ 系统已准备好运行置信度prompt测试")
        print(f"\n🚀 运行测试命令:")
        print(f"   ./run_vllm_confidence_test.sh llama 10    # 使用Llama模型")
        print(f"   ./run_vllm_confidence_test.sh mistral 10  # 使用Mistral模型")
    else:
        print(f"\n❌ 部分测试失败，请检查配置")
        print(f"\n🔧 故障排除:")
        print(f"   1. 检查vLLM服务器: ps -ef | grep vllm")
        print(f"   2. 启动服务器: ./start_vllm_servers.sh")
        print(f"   3. 查看日志: tail -f logs/*.log")
        sys.exit(1)

if __name__ == "__main__":
    main()