gitlab-instance-0a899031_douyinyoutube/scripts/hardware_probe.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import platform
import shutil
import subprocess
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
CONFIG_DIR = ROOT / "config"
JSON_PATH = CONFIG_DIR / "hardware-profile.json"
ENV_PATH = CONFIG_DIR / "hardware.env"


def run_text(cmd: list[str]) -> str:
    try:
        proc = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="ignore",
            timeout=8,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        return ""
    return (proc.stdout or proc.stderr or "").strip()


def detect_ffmpeg_hwaccels() -> list[str]:
    if not shutil.which("ffmpeg"):
        return []
    text = run_text(["ffmpeg", "-hide_banner", "-hwaccels"])
    accels: list[str] = []
    for line in text.splitlines():
        item = line.strip()
        if not item or item.lower().startswith("hardware acceleration methods"):
            continue
        if item not in accels:
            accels.append(item)
    return accels


def detect_gpu() -> dict:
    dri_dir = Path("/dev/dri")
    render_nodes = sorted(str(path) for path in dri_dir.glob("renderD*")) if dri_dir.is_dir() else []
    cards = sorted(str(path) for path in dri_dir.glob("card*")) if dri_dir.is_dir() else []
    driver = ""
    if cards:
        card0 = Path(cards[0])
        driver_link = Path(f"/sys/class/drm/{card0.name}/device/driver")
        if driver_link.exists():
            try:
                driver = driver_link.resolve().name
            except OSError:
                driver = driver_link.name
    return {
        "present": bool(cards or render_nodes),
        "cards": cards,
        "render_nodes": render_nodes,
        "driver": driver,
        "vainfo": bool(shutil.which("vainfo")),
    }


def detect_npu() -> dict:
    hints = []
    for path in (
        Path("/dev/rknpu"),
        Path("/dev/npu"),
        Path("/dev/galcore"),
        Path("/sys/class/misc/rknpu"),
        Path("/sys/class/misc/verisilicon-npu"),
    ):
        if path.exists():
            hints.append(str(path))
    return {
        "present": bool(hints),
        "devices": hints,
    }


def detect_kernel_features() -> dict:
    features = {
        "drm": Path("/dev/dri/card0").exists(),
        "render_node": any(Path("/dev/dri").glob("renderD*")) if Path("/dev/dri").is_dir() else False,
        "v4l2loopback": Path("/sys/module/v4l2loopback").exists(),
        "kvm": Path("/dev/kvm").exists(),
        "nvidia": shutil.which("nvidia-smi") is not None,
    }
    return features


def recommend(profile: dict) -> dict:
    hwaccels = set(profile["ffmpeg_hwaccels"])
    gpu = profile["gpu"]
    browser_hw = "auto" if gpu["present"] else "disabled"
    decoder = "cpu"
    if "vaapi" in hwaccels and gpu["render_nodes"]:
        decoder = "vaapi"
    elif "v4l2m2m" in hwaccels:
        decoder = "v4l2m2m"
    elif "cuda" in hwaccels or "nvdec" in hwaccels:
        decoder = "cuda"

    hdmi_player = "cpu"
    if shutil.which("mpv") and gpu["cards"]:
        hdmi_player = "mpv-drm"
    elif shutil.which("ffplay") and (os.environ.get("DISPLAY") or shutil.which("startx")):
        hdmi_player = "ffplay-x11"

    risk = []
    if not gpu["present"]:
        risk.append("No DRM/VA render node detected, browser and HDMI playback will stay on conservative settings")
    if decoder == "cpu":
        risk.append("No safe hardware decoder detected, ffmpeg workloads will remain on CPU")
    if hdmi_player == "cpu":
        risk.append("No direct HDMI playback stack detected, install mpv or X11+ffplay for full-screen output")

    return {
        "video_decoder": decoder,
        "browser_hwaccel": browser_hw,
        "hdmi_player": hdmi_player,
        "stability_mode": "safe" if risk else "balanced",
        "degradation_reasons": risk,
    }


def build_profile() -> dict:
    profile = {
        "status": "ok",
        "arch": platform.machine(),
        "kernel": platform.release(),
        "platform": platform.platform(),
        "ffmpeg_hwaccels": detect_ffmpeg_hwaccels(),
        "gpu": detect_gpu(),
        "npu": detect_npu(),
        "kernel_features": detect_kernel_features(),
    }
    profile["recommendation"] = recommend(profile)
    return profile


def write_outputs(profile: dict) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    JSON_PATH.write_text(json.dumps(profile, indent=2, ensure_ascii=False), encoding="utf-8")
    env_lines = [
        f"LIVE_HW_DECODER={profile['recommendation']['video_decoder']}",
        f"LIVE_HDMI_PLAYER={profile['recommendation']['hdmi_player']}",
        f"LIVE_BROWSER_HWACCEL={profile['recommendation']['browser_hwaccel']}",
        f"LIVE_STABILITY_MODE={profile['recommendation']['stability_mode']}",
    ]
    ENV_PATH.write_text("\n".join(env_lines) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Probe Linux hardware capabilities for safe runtime defaults")
    parser.add_argument("--write", action="store_true", help="Write JSON and env files into config/")
    args = parser.parse_args()

    profile = build_profile()
    if args.write:
        write_outputs(profile)
    print(json.dumps(profile, ensure_ascii=False))


if __name__ == "__main__":
    main()