ncmr-songs-videos/scripts/hwprobe.py

"""Detect which video encoders this machine can actually use.

Probes ffmpeg for hardware and software encoders and decides which are usable
here, then emits either a human-readable report or a shell ``export`` block.

The reference machine has no NVIDIA GPU, so ``*_nvenc`` encoders exist in the
ffmpeg build but cannot run. Detecting that up front is cheaper than discovering
it in the middle of a render.

Usage::

    python3 scripts/hwprobe.py                 # human-readable report
    python3 scripts/hwprobe.py --json          # machine-readable
    python3 scripts/hwprobe.py --shell         # eval-able export block
"""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
from dataclasses import dataclass, field, asdict

if __package__ in (None, ""):
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from _common import (  # type: ignore[import-not-found]
        Die,
        Report,
        add_common_arguments,
        add_ffmpeg_argument,
        find_ffmpeg,
        main_wrapper,
        report as default_report,
        run,
        text_table,
    )
else:  # pragma: no cover - exercised via `python3 -m scripts.hwprobe`
    from ._common import (
        Die,
        Report,
        add_common_arguments,
        add_ffmpeg_argument,
        find_ffmpeg,
        main_wrapper,
        report as default_report,
        run,
        text_table,
    )


# --------------------------------------------------------------------------- #
# Candidate tables
# --------------------------------------------------------------------------- #

#: Hardware encoder families we know how to enable, keyed by ffmpeg suffix.
HW_FAMILIES = ("nvenc", "vaapi", "qsv", "amf", "v4l2m2m", "vulkan", "videotoolbox")

#: The encoders worth offering as a "best available" pick, in preference order.
#: Hardware first, then software fallbacks that always exist in a full build.
ENCODER_PREFERENCE = (
    "h264_nvenc",
    "hevc_nvenc",
    "av1_nvenc",
    "h264_qsv",
    "hevc_qsv",
    "h264_vaapi",
    "hevc_vaapi",
    "av1_vaapi",
    "h264_amf",
    "hevc_amf",
    "h264_v4l2m2m",
    "hevc_v4l2m2m",
    "hevc_vulkan",
    "libx264",
    "libx265",
    "libsvtav1",
    "libvpx-vp9",
)

#: vaapi/vulkan/qsv devices we look for, in preference order.
DEVICE_CANDIDATES = (
    "/dev/dri/renderD128",
    "/dev/dri/renderD129",
    "/dev/dri/card0",
    "/dev/dri/card1",
)


@dataclass
class Encoder:
    name: str
    description: str

    @property
    def family(self) -> str:
        for family in HW_FAMILIES:
            if self.name.endswith(f"_{family}"):
                return family

        return "software"

    @property
    def is_hardware(self) -> bool:
        return self.family != "software"


@dataclass
class ProbeResult:
    ffmpeg: str
    ffmpeg_version: str
    encoders: list[Encoder] = field(default_factory=list)
    hwaccels: list[str] = field(default_factory=list)
    devices: dict[str, bool] = field(default_factory=dict)
    tested: dict[str, dict] = field(default_factory=dict)
    #: Encoder families that are compiled in but have no usable device.
    blocked: dict[str, str] = field(default_factory=dict)
    recommended: str | None = None
    notes: list[str] = field(default_factory=list)

    def encryptors(self) -> list[Encoder]:
        return [e for e in self.encoders if e.is_hardware]


def parse_encoders(output: str) -> list[Encoder]:
    """Parse ``ffmpeg -encoders`` output into video encoders only."""
    encoders: list[Encoder] = []

    for line in output.splitlines():
        match = re.match(r"^\s*([VAS][\w.]{5})\s+(\S+)\s+(.*\S)\s*$", line)
        if not match:
            continue

        flags, name, description = match.groups()
        if not flags.startswith("V"):
            continue

        encoders.append(Encoder(name=name, description=description))

    return encoders


def parse_hwaccels(output: str) -> list[str]:
    """Parse ``ffmpeg -hwaccels`` output."""
    methods: list[str] = []
    seen_header = False

    for line in output.splitlines():
        line = line.strip()
        if not line:
            continue
        if line.lower().startswith("hardware acceleration"):
            seen_header = True
            continue
        if seen_header:
            methods.append(line)

    return methods


def probe_devices() -> dict[str, bool]:
    """Check which DRM devices exist *and* are readable by this process."""
    found: dict[str, bool] = {}

    for device in DEVICE_CANDIDATES:
        if not os.path.exists(device):
            continue

        found[device] = os.access(device, os.R_OK | os.W_OK)

    return found


def probe_device_driver(device: str) -> str | None:
    """Read the kernel driver name backing a DRM node, if sysfs is readable."""
    node = os.path.basename(device)
    uevent = f"/sys/class/drm/{node}/device/uevent"

    try:
        with open(uevent, "r", encoding="utf-8", errors="replace") as handle:
            content = handle.read()
    except OSError:
        return None

    match = re.search(r"^DRIVER=(\S+)", content, re.MULTILINE)

    return match.group(1) if match else None


def missing_render_group(devices: dict[str, bool]) -> bool:
    """True when a DRM device exists but this process cannot open it."""
    return any(not ok for ok in devices.values())


def test_encoder(ffmpeg: str, encoder: str, report: Report, device: str | None = None) -> tuple[bool, str]:
    """Encode two throwaway frames with *encoder*. Returns (worked, detail)."""
    cmd = [ffmpeg, "-hide_banner", "-loglevel", "error"]

    if encoder.endswith("_vaapi"):
        if not device:
            return False, "no vaapi device available"
        cmd += [
            "-init_hw_device", f"vaapi=va:{device}",
            "-filter_hw_device", "va",
            "-f", "lavfi",
            "-i", "color=c=black:s=256x256:r=30:d=1",
            "-vf", "format=nv12,hwupload",
        ]
    elif encoder.endswith("_vulkan"):
        cmd += ["-f", "lavfi", "-i", "color=c=black:s=256x256:r=30:d=1"]
    else:
        cmd += ["-f", "lavfi", "-i", "color=c=black:s=256x256:r=30:d=1"]

    cmd += ["-c:v", encoder, "-frames:v", "2", "-f", "null", "-"]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
    except subprocess.TimeoutExpired:
        return False, "timed out after 60s"
    except FileNotFoundError:
        return False, "ffmpeg binary missing"

    if result.returncode == 0:
        return True, "ok"

    stderr = (result.stderr or "").strip().splitlines()

    return False, stderr[-1] if stderr else f"exit status {result.returncode}"


def probe(ffmpeg: str, report: Report, *, test: bool = True) -> ProbeResult:
    """Gather everything we can learn about this machine's encoders."""

    def capture(args: list[str]) -> str:
        try:
            result = subprocess.run(
                [ffmpeg, "-hide_banner", *args],
                capture_output=True,
                text=True,
                timeout=30,
            )
        except (OSError, subprocess.TimeoutExpired) as exc:
            raise Die(f"could not run {ffmpeg}: {exc}") from exc

        return result.stdout

    banner = capture(["-version"])
    version = banner.splitlines()[0] if banner else "unknown"

    result = ProbeResult(
        ffmpeg=ffmpeg,
        ffmpeg_version=version,
        encoders=parse_encoders(capture(["-encoders"])),
        hwaccels=parse_hwaccels(capture(["-hwaccels"])),
        devices=probe_devices(),
    )

    usable_device = next((d for d, ok in result.devices.items() if ok), None)
    blocked_families: dict[str, str] = {}

    if usable_device:
        driver = probe_device_driver(usable_device)
        if driver:
            result.notes.append(f"{usable_device} is driven by {driver}")
    elif result.devices:
        blocked_families["vaapi"] = "DRM device present but not openable"
        blocked_families["vulkan"] = "DRM device present but not openable"
    else:
        blocked_families["vaapi"] = "no DRM render device found"

    if not test:
        return result

    for encoder in result.encoders:
        if not encoder.is_hardware and encoder.name not in ("libx264",):
            continue
        if encoder.family in blocked_families:
            result.tested[encoder.name] = {
                "ok": False,
                "detail": blocked_families[encoder.family],
                "tested": False,
            }
            continue

        if encoder.family in ("qsv", "amf", "videotoolbox", "v4l2m2m") and not result.devices:
            result.tested[encoder.name] = {
                "ok": False,
                "detail": "no matching device",
                "tested": False,
            }
            continue

        ok, detail = test_encoder(ffmpeg, encoder.name, report, usable_device)
        result.tested[encoder.name] = {"ok": ok, "detail": detail, "tested": True}

    for family, reason in blocked_families.items():
        if any(e.family == family for e in result.encoders):
            result.blocked[family] = reason

    for name in ENCODER_PREFERENCE:
        entry = result.tested.get(name)
        if entry and entry["ok"]:
            result.recommended = name
            break

    return result


# --------------------------------------------------------------------------- #
# Output
# --------------------------------------------------------------------------- #


def render_report(result: ProbeResult) -> str:
    lines: list[str] = []
    lines.append(f"ffmpeg   {result.ffmpeg_version}")
    lines.append(f"binary   {result.ffmpeg}")
    lines.append("")

    lines.append("DRM devices")
    if result.devices:
        for device in sorted(result.devices):
            ok = result.devices[device]
            driver = probe_device_driver(device) or "?"
            state = "readable" if ok else "NOT READABLE"
            lines.append(f"  {device.ljust(20)} {driver.ljust(8)} {state}")
    else:
        lines.append("  none found — hardware encoding is unavailable")

    lines.append("")
    lines.append(f"hwaccels  {', '.join(result.hwaccels) or '(none)'}")
    lines.append("")

    hw = result.encryptors()
    sw = [e for e in result.encoders if not e.is_hardware]

    lines.append(f"Hardware encoders in this build ({len(hw)})")
    if hw:
        rows = []
        for encoder in hw:
            entry = result.tested.get(encoder.name)
            if entry is None:
                verdict = "untested"
            elif entry["ok"]:
                verdict = "WORKS"
            elif not entry.get("tested", True):
                verdict = "unavailable"
            else:
                verdict = "failed"
            rows.append((encoder.name, encoder.family, verdict))
        lines.append(text_table(("encoder", "family", "status"), rows))
    else:
        lines.append("  none")

    lines.append("")
    lines.append(f"Software encoders ({len(sw)})")
    notable = [e for e in sw if e.name in ("libx264", "libx265", "libsvtav1", "libvpx-vp9", "libxvid")]
    lines.append("  " + (", ".join(e.name for e in notable) or "none"))

    if result.blocked:
        lines.append("")
        lines.append("Blocked")
        for family, reason in sorted(result.blocked.items()):
            lines.append(f"  {family}: {reason}")

    if result.notes:
        lines.append("")
        for note in result.notes:
            lines.append(f"note: {note}")

    lines.append("")
    if result.recommended:
        lines.append(f"recommended  -c:v {result.recommended}")
    else:
        lines.append("recommended  (nothing usable detected)")

    if missing_render_group(result.devices):
        lines.append("")
        lines.append("hint: a DRM device exists but this user cannot open it.")
        lines.append("      add the user to the 'render' group and re-login:")
        lines.append("        sudo usermod -aG render,video $USER")

    return "\n".join(lines)


def render_shell(result: ProbeResult) -> str:
    """Export block for ``eval \"$(python3 scripts/hwprobe.py --shell)\"``."""
    device = next((d for d, ok in result.devices.items() if ok), "")

    lines = [
        "# generated by scripts/hwprobe.py — eval to apply",
        f"export FFMPEG={result.ffmpeg}",
    ]

    for encoder in result.encryptors():
        entry = result.tested.get(encoder.name, {})
        if entry.get("ok"):
            lines.append(f"export ENCODER_{encoder.family.upper()}={encoder.name}")

    lines.append(f"export ENCODER={result.recommended or 'libx264'}")

    if device:
        lines.append(f"export HW_DEVICE={device}")

    return "\n".join(lines)


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="hwprobe",
        description="Report which ffmpeg encoders this machine can actually use.",
        epilog=(
            "Handy:  eval \"$(python3 scripts/hwprobe.py --shell)\"  "
            "to put the detected encoder in $ENCODER."
        ),
    )
    add_common_arguments(parser)
    add_ffmpeg_argument(parser)

    parser.add_argument("--json", action="store_true", help="emit JSON instead of text")
    parser.add_argument("--shell", action="store_true", help="emit an eval-able export block")
    parser.add_argument(
        "--no-test",
        action="store_true",
        help="only read ffmpeg's capability lists; do not try encoding",
    )

    return parser


@main_wrapper
def main() -> int:
    args = build_parser().parse_args()
    report = Report(quiet=args.quiet)

    ffmpeg = find_ffmpeg(args.ffmpeg)
    result = probe(ffmpeg, report, test=not args.no_test)

    if args.json:
        payload = asdict(result)
        payload["blocked"] = result.blocked
        payload["recommended"] = result.recommended
        print(json.dumps(payload, indent=2))
    elif args.shell:
        print(render_shell(result))
    else:
        print(render_report(result))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())