GPU/CPU 디바이스 자동 감지 유틸리티

"""
GPU/CPU 디바이스 자동 감지 유틸리티

@source voice-to-text-v2
@extracted 2026-02-15
@version 1.0.0

의존성:
    - (선택) torch - CUDA GPU 감지
    - (선택) ctranslate2 - CTranslate2 CUDA 감지

사용법:
    from utils.device_detector import detect_device, get_device_info

    # 최적 디바이스 감지
    device, compute_type = detect_device()
    # → ("cuda", "float16") 또는 ("cpu", "int8")

    # 상세 정보 조회
    info = get_device_info()
    print(info)
    # → {"device": "cuda", "compute_type": "float16", "gpu_name": "RTX 4090", "vram_gb": 24.0}
"""

import logging

logger = logging.getLogger(__name__)


def detect_device(
    prefer_gpu: bool = True,
    cpu_compute_type: str = "int8",
    gpu_compute_type: str = "float16",
) -> tuple[str, str]:
    """
    최적 디바이스 및 컴퓨트 타입 자동 감지.

    감지 순서:
    1. torch.cuda 확인
    2. ctranslate2 CUDA 지원 확인
    3. CPU 폴백

    Args:
        prefer_gpu: GPU 우선 사용 여부 (False면 항상 CPU)
        cpu_compute_type: CPU 사용 시 컴퓨트 타입 (기본: "int8")
        gpu_compute_type: GPU 사용 시 컴퓨트 타입 (기본: "float16")

    Returns:
        (device, compute_type) 튜플
        - GPU 사용 가능: ("cuda", gpu_compute_type)
        - CPU: ("cpu", cpu_compute_type)
    """
    if not prefer_gpu:
        logger.info("GPU 비활성화, CPU 사용")
        return ("cpu", cpu_compute_type)

    # 1. PyTorch CUDA 확인
    try:
        import torch
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            logger.info(f"CUDA GPU 감지: {gpu_name} ({vram_gb:.1f}GB VRAM)")
            return ("cuda", gpu_compute_type)
    except ImportError:
        pass

    # 2. CTranslate2 CUDA 확인
    try:
        import ctranslate2
        cuda_types = ctranslate2.get_supported_compute_types("cuda")
        if cuda_types:
            logger.info(f"CTranslate2 CUDA 지원: {cuda_types}")
            compute = gpu_compute_type if gpu_compute_type in cuda_types else (cuda_types[0] if cuda_types else cpu_compute_type)
            return ("cuda", compute)
    except Exception:
        pass

    logger.info("GPU 미감지, CPU 사용")
    return ("cpu", cpu_compute_type)


def get_device_info() -> dict:
    """
    현재 시스템의 디바이스 상세 정보 조회.

    Returns:
        dict: device, compute_type, gpu_name, vram_gb, cuda_available, torch_version
    """
    info = {
        "device": "cpu",
        "compute_type": "int8",
        "gpu_name": None,
        "vram_gb": None,
        "cuda_available": False,
        "torch_version": None,
        "ctranslate2_available": False,
    }

    # PyTorch 확인
    try:
        import torch
        info["torch_version"] = torch.__version__
        if torch.cuda.is_available():
            info["cuda_available"] = True
            info["device"] = "cuda"
            info["compute_type"] = "float16"
            info["gpu_name"] = torch.cuda.get_device_name(0)
            info["vram_gb"] = round(
                torch.cuda.get_device_properties(0).total_memory / (1024**3), 1
            )
    except ImportError:
        pass

    # CTranslate2 확인
    try:
        import ctranslate2
        info["ctranslate2_available"] = True
        if not info["cuda_available"]:
            try:
                cuda_types = ctranslate2.get_supported_compute_types("cuda")
                if cuda_types:
                    info["cuda_available"] = True
                    info["device"] = "cuda"
                    info["compute_type"] = "float16" if "float16" in cuda_types else cuda_types[0]
            except Exception:
                pass
    except ImportError:
        pass

    return info