FrameTour-RenderWorker/util/system.py

# -*- coding: utf-8 -*-
"""
系统信息工具

提供系统信息采集功能。
"""

import os
import platform
import subprocess
from typing import Optional, Dict, Any

import psutil
from constant import SOFTWARE_VERSION, DEFAULT_CAPABILITIES, HW_ACCEL_NONE, HW_ACCEL_QSV, HW_ACCEL_CUDA


def get_sys_info():
    """
    获取系统信息

    Returns:
        dict: 系统信息字典
    """
    mem = psutil.virtual_memory()

    info = {
        'os': platform.system(),
        'cpu': f"{os.cpu_count()} cores",
        'memory': f"{mem.total // (1024**3)}GB",
        'cpuUsage': f"{psutil.cpu_percent()}%",
        'memoryAvailable': f"{mem.available // (1024**3)}GB",
        'platform': platform.system(),
        'pythonVersion': platform.python_version(),
        'version': SOFTWARE_VERSION,
    }

    # 尝试获取 GPU 信息
    gpu_info = get_gpu_info()
    if gpu_info:
        info['gpu'] = gpu_info

    return info


def get_capabilities():
    """
    获取 Worker 支持的能力列表

    Returns:
        list: 能力列表
    """
    return DEFAULT_CAPABILITIES.copy()


def get_gpu_info() -> Optional[str]:
    """
    尝试获取 GPU 信息

    Returns:
        str: GPU 信息，失败返回 None
    """
    try:
        # 尝试使用 nvidia-smi
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode == 0:
            gpu_name = result.stdout.strip().split('\n')[0]
            return gpu_name
    except Exception:
        pass

    return None


def get_ffmpeg_version() -> str:
    """
    获取 FFmpeg 版本

    Returns:
        str: FFmpeg 版本号
    """
    try:
        result = subprocess.run(
            ['ffmpeg', '-version'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode == 0:
            first_line = result.stdout.split('\n')[0]
            # 解析版本号，例如 "ffmpeg version 6.0 ..."
            parts = first_line.split()
            for i, part in enumerate(parts):
                if part == 'version' and i + 1 < len(parts):
                    return parts[i + 1]
    except Exception:
        pass

    return 'unknown'


def check_ffmpeg_encoder(encoder: str) -> bool:
    """
    检查 FFmpeg 是否支持指定的编码器

    Args:
        encoder: 编码器名称，如 'h264_nvenc', 'h264_qsv'

    Returns:
        bool: 是否支持该编码器
    """
    try:
        result = subprocess.run(
            ['ffmpeg', '-hide_banner', '-encoders'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode == 0:
            return encoder in result.stdout
    except Exception:
        pass

    return False


def check_ffmpeg_decoder(decoder: str) -> bool:
    """
    检查 FFmpeg 是否支持指定的解码器

    Args:
        decoder: 解码器名称，如 'h264_cuvid', 'h264_qsv'

    Returns:
        bool: 是否支持该解码器
    """
    try:
        result = subprocess.run(
            ['ffmpeg', '-hide_banner', '-decoders'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode == 0:
            return decoder in result.stdout
    except Exception:
        pass

    return False


def check_ffmpeg_hwaccel(hwaccel: str) -> bool:
    """
    检查 FFmpeg 是否支持指定的硬件加速方法

    Args:
        hwaccel: 硬件加速方法，如 'cuda', 'qsv', 'dxva2', 'd3d11va'

    Returns:
        bool: 是否支持该硬件加速方法
    """
    try:
        result = subprocess.run(
            ['ffmpeg', '-hide_banner', '-hwaccels'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode == 0:
            return hwaccel in result.stdout
    except Exception:
        pass

    return False


def detect_hw_accel_support() -> Dict[str, Any]:
    """
    检测系统的硬件加速支持情况

    Returns:
        dict: 硬件加速支持信息
        {
            'cuda': {
                'available': bool,
                'gpu': str or None,
                'encoder': bool,  # h264_nvenc
                'decoder': bool,  # h264_cuvid
            },
            'qsv': {
                'available': bool,
                'encoder': bool,  # h264_qsv
                'decoder': bool,  # h264_qsv
            },
            'recommended': str  # 推荐的加速方式: 'cuda', 'qsv', 'none'
        }
    """
    result = {
        'cuda': {
            'available': False,
            'gpu': None,
            'encoder': False,
            'decoder': False,
        },
        'qsv': {
            'available': False,
            'encoder': False,
            'decoder': False,
        },
        'recommended': HW_ACCEL_NONE
    }

    # 检测 CUDA/NVENC 支持
    gpu_info = get_gpu_info()
    if gpu_info:
        result['cuda']['gpu'] = gpu_info
        result['cuda']['available'] = check_ffmpeg_hwaccel('cuda')
        result['cuda']['encoder'] = check_ffmpeg_encoder('h264_nvenc')
        result['cuda']['decoder'] = check_ffmpeg_decoder('h264_cuvid')

    # 检测 QSV 支持
    result['qsv']['available'] = check_ffmpeg_hwaccel('qsv')
    result['qsv']['encoder'] = check_ffmpeg_encoder('h264_qsv')
    result['qsv']['decoder'] = check_ffmpeg_decoder('h264_qsv')

    # 推荐硬件加速方式（优先 CUDA，其次 QSV）
    if result['cuda']['available'] and result['cuda']['encoder']:
        result['recommended'] = HW_ACCEL_CUDA
    elif result['qsv']['available'] and result['qsv']['encoder']:
        result['recommended'] = HW_ACCEL_QSV

    return result


def get_hw_accel_info_str() -> str:
    """
    获取硬件加速支持信息的可读字符串

    Returns:
        str: 硬件加速支持信息描述
    """
    support = detect_hw_accel_support()

    parts = []

    if support['cuda']['available']:
        gpu = support['cuda']['gpu'] or 'Unknown GPU'
        status = 'encoder+decoder' if support['cuda']['encoder'] and support['cuda']['decoder'] else (
            'encoder only' if support['cuda']['encoder'] else 'decoder only' if support['cuda']['decoder'] else 'hwaccel only'
        )
        parts.append(f"CUDA({gpu}, {status})")

    if support['qsv']['available']:
        status = 'encoder+decoder' if support['qsv']['encoder'] and support['qsv']['decoder'] else (
            'encoder only' if support['qsv']['encoder'] else 'decoder only' if support['qsv']['decoder'] else 'hwaccel only'
        )
        parts.append(f"QSV({status})")

    if not parts:
        return "No hardware acceleration available"

    return ', '.join(parts) + f" [recommended: {support['recommended']}]"