#!/usr/bin/env python3
"""
测试Milvus GPU配置的脚本
"""
import sys
import os

def test_gpu_availability():
    """测试GPU可用性"""
    print("=== GPU可用性测试 ===")
    
    # 测试PyTorch CUDA
    try:
        import torch
        print(f"PyTorch版本: {torch.__version__}")
        print(f"CUDA可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"CUDA设备数量: {torch.cuda.device_count()}")
            print(f"设备名称: {torch.cuda.get_device_name(0)}")
            print(f"CUDA版本: {torch.version.cuda}")
            return True
        else:
            print("CUDA不可用")
            return False
    except Exception as e:
        print(f"PyTorch测试失败: {e}")
        return False

def test_milvus_connection():
    """测试Milvus连接"""
    print("\n=== Milvus连接测试 ===")
    
    try:
        from pymilvus import connections, utility
        print("PyMilvus导入成功")
        
        # 尝试连接本地Milvus
        try:
            connections.connect("default", host="localhost", port="19530")
            print("Milvus连接成功")
            
            # 检查Milvus版本
            version = utility.get_server_version()
            print(f"Milvus版本: {version}")
            
            return True
        except Exception as e:
            print(f"Milvus连接失败: {e}")
            print("注意: 这需要先启动Milvus服务")
            return False
            
    except Exception as e:
        print(f"PyMilvus导入失败: {e}")
        return False

def test_gpu_memory():
    """测试GPU内存"""
    print("\n=== GPU内存测试 ===")
    
    try:
        import torch
        if torch.cuda.is_available():
            # 获取GPU内存信息
            device = torch.cuda.current_device()
            total_memory = torch.cuda.get_device_properties(device).total_memory
            allocated_memory = torch.cuda.memory_allocated(device)
            cached_memory = torch.cuda.memory_reserved(device)
            
            print(f"GPU总内存: {total_memory / 1024**3:.2f} GB")
            print(f"已分配内存: {allocated_memory / 1024**3:.2f} GB")
            print(f"缓存内存: {cached_memory / 1024**3:.2f} GB")
            
            # 测试GPU计算
            print("\n=== GPU计算测试 ===")
            x = torch.randn(1000, 1000, device='cuda')
            y = torch.randn(1000, 1000, device='cuda')
            z = torch.mm(x, y)
            print("GPU矩阵乘法测试成功!")
            print(f"结果张量形状: {z.shape}")
            print(f"结果张量设备: {z.device}")
            
            return True
        else:
            print("CUDA不可用，无法测试GPU内存")
            return False
    except Exception as e:
        print(f"GPU内存测试失败: {e}")
        return False

def main():
    print("Milvus GPU配置测试")
    print("=" * 50)
    
    # 测试GPU可用性
    gpu_available = test_gpu_availability()
    
    # 测试Milvus连接
    milvus_available = test_milvus_connection()
    
    # 测试GPU内存和计算
    gpu_memory_ok = False
    if gpu_available:
        gpu_memory_ok = test_gpu_memory()
    
    # 总结
    print("\n" + "=" * 50)
    print("测试总结:")
    print(f"GPU可用: {'✓' if gpu_available else '✗'}")
    print(f"Milvus连接: {'✓' if milvus_available else '✗'}")
    print(f"GPU计算: {'✓' if gpu_memory_ok else '✗'}")
    
    if gpu_available and gpu_memory_ok:
        print("\n🎉 GPU配置正常，可以用于Milvus等GPU加速应用!")
    else:
        print("\n⚠️  GPU配置存在问题，可能影响Milvus等应用的GPU加速功能")

if __name__ == "__main__":
    main()