FineTurnX/scripts/check_env.py

"""
检查环境依赖
"""

import sys
import subprocess
from packaging import version


def check_package(package_name, min_version=None):
    """检查包是否安装及版本"""
    try:
        import importlib
        module = importlib.import_module(package_name)
        
        if hasattr(module, "__version__"):
            installed_version = module.__version__
            if min_version and version.parse(installed_version) < version.parse(min_version):
                print(f"❌ {package_name}: {installed_version} (需要 >= {min_version})")
                return False
            else:
                print(f"✓ {package_name}: {installed_version}")
        else:
            print(f"✓ {package_name}: 已安装")
        return True
    except ImportError:
        print(f"❌ {package_name}: 未安装")
        return False


def check_cuda():
    """检查 CUDA 是否可用"""
    try:
        import torch
        if torch.cuda.is_available():
            print(f"✓ CUDA: 可用 ({torch.cuda.device_count()} 个 GPU)")
            for i in range(torch.cuda.device_count()):
                print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
            return True
        else:
            print("⚠ CUDA: 不可用 (将使用 CPU 训练)")
            return False
    except ImportError:
        print("❌ PyTorch: 未安装")
        return False


def check_npu():
    """检查华为升腾 NPU 是否可用"""
    try:
        import torch
        if hasattr(torch, 'npu') and torch.npu.is_available():
            print(f"✓ NPU (华为升腾): 可用 ({torch.npu.device_count()} 个 NPU)")
            return True
        else:
            print("⚠ NPU (华为升腾): 不可用")
            return False
    except ImportError:
        print("⚠ torch-npu: 未安装 (华为升腾需要)")
        return False


def main():
    print("=" * 60)
    print("FineTuneX 环境检查")
    print("=" * 60)
    print()
    
    required_packages = {
        "torch": "2.1.0",
        "transformers": "4.40.0",
        "datasets": "2.14.0",
        "accelerate": "0.25.0",
        "peft": "0.7.0",
        "trl": "0.7.0",
        "fastapi": "0.104.0",
        "uvicorn": "0.24.0",
        "pydantic": "2.0.0",
    }
    
    all_ok = True
    
    print("检查 Python 包:")
    print("-" * 60)
    for package, min_ver in required_packages.items():
        if not check_package(package, min_ver):
            all_ok = False
    
    # 检查 bitsandbytes（可选，仅 NVIDIA GPU）
    print("\n检查量化工具（可选，仅 NVIDIA GPU）:")
    print("-" * 60)
    check_package("bitsandbytes", "0.41.0")
    print("（华为升腾 NPU 不支持 bitsandbytes 量化，可忽略此项）")
    
    print()
    print("检查计算设备:")
    print("-" * 60)
    has_cuda = check_cuda()
    has_npu = check_npu()
    
    if not has_cuda and not has_npu:
        print("⚠ 警告：未检测到 GPU/NPU，将使用 CPU 训练（速度较慢）")
    
    print()
    print("=" * 60)
    if all_ok:
        print("✓ 所有依赖已安装！")
    else:
        print("❌ 部分依赖缺失，请运行：pip install -r requirements.txt")
    print("=" * 60)
    
    return 0 if all_ok else 1


if __name__ == "__main__":
    sys.exit(main())