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 main():
    print("=" * 60)
    print("FineTuneX 环境检查")
    print("=" * 60)
    print()
    
    required_packages = {
        "torch": "2.0.0",
        "transformers": "4.40.0",
        "datasets": "2.14.0",
        "accelerate": "0.25.0",
        "peft": "0.7.0",
        "bitsandbytes": "0.41.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
    
    print()
    print("检查 CUDA:")
    print("-" * 60)
    if not check_cuda():
        print("⚠ 警告：CPU 训练速度较慢，建议使用 GPU")
    
    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())