path: root/packages/meshbay-node/src/meshbay_node/hwaccel.py

"""
Hardware H.264 encoding for the streaming re-encode path.

Streaming copies the video whenever a browser can decode it (`-c:v copy`), and
that path costs nothing. The re-encode path is the expensive one — HEVC, and
the codecs with no MSE string at all (Xvid, MPEG-2, VC-1) — and it is
`libx264 -preset veryfast`, one process per viewer, for the length of a film.
On a laptop that is fine. On an Atom or Celeron mini-PC it does not reach real
time at 1080p, which is why `transcode_incompatible_video` exists as an
operator opt-out: a node that cannot transcode says so instead of serving a
stream that stutters.

**This module is the other answer to the same problem.** The GPU in every
machine of the last decade encodes H.264 in fixed-function silicon and does not
care how weak the CPU beside it is. Where that hardware is there and works, the
node re-encodes on it and the opt-out is not needed.

Nothing here is configured by hand, and nothing is inferred from a CPU model or
from a driver's name. **The capability is established by encoding**: a listed
`h264_qsv` proves only that ffmpeg was built with it, and a render node proves
only that a GPU exists — neither says the driver on this machine can do the
work. Measured on the development VM, where `ffmpeg -encoders` lists every
VAAPI encoder and the virtio-gpu driver then fails to initialise:
`libva: virtio_gpu_drv_video.so init failed`.

**The probe encodes 1080p and then reads the result back with ffprobe**, and
accepts nothing that is not High profile at level 4.1. That is not
belt-and-braces: `stream_init` announces `avc1.640029` and the client checks it
before it trusts a byte, so an encoder that quietly wrote a different level
would make the node a liar. It also means an argument spelled the way one
encoder wants and not another — `-level 4.1` against `-level 41` — is caught
here rather than by a viewer, which is why each candidate below may offer
several variants and the machine picks.

Candidates, in the order they are tried:

  vaapi   Linux, any GPU with a VA-API driver — Intel iGPU, AMD through mesa.
  qsv     Intel Quick Sync, which is how the same iGPU is reached on Windows.
  nvenc   NVIDIA, either platform.

Two gaps, named rather than left to be discovered: **AMD on Windows** (AMF) and
**macOS** (VideoToolbox). Neither set of arguments could be tried anywhere in
this project, and MeshBay ships no macOS package at all; a node on either
re-encodes in software, exactly as every node did before this module existed.
Adding one is a `Candidate` and nothing else — the probe is what decides
whether it works, so a wrong guess costs a rejected candidate, not a broken
stream.

Then three modes per stream, remembered per source codec:

  hw     Hardware decode and encode. The whole pipeline on the GPU, which is
         what makes 1080p HEVC playable on a machine that cannot decode it at
         all in software.
  hwenc  Software decode, hardware encode. The fallback for a source the GPU
         has no decoder for — iHD has none for MPEG-4 Part 2, so an Xvid .avi
         lands here, and an SD Xvid decodes in software for nearly nothing.
  sw     libx264, the path that was always here.
"""

from __future__ import annotations

import asyncio
import logging
import os
import shutil
import sys
import tempfile
from dataclasses import dataclass

from meshbay_node import platform

log = logging.getLogger(__name__)

HW = "hw"
HWENC = "hwenc"
SW = "sw"

# How long one test encode may take before it is considered broken. A driver
# that deadlocks must not hold the first viewer's stream, and 20 s is far
# beyond the fraction of a second the probe needs when it works.
PROBE_TIMEOUT_SECS = 20

# What the probe encodes. 1080p because the level a film gets is the level that
# has to be checked: an encoder asked for 4.1 on a 320x240 clip may well write
# a lower one, and rejecting a good candidate for that would be the probe
# failing rather than the hardware.
PROBE_SIZE = "1920x1080"

# The announced codec string is `avc1.640029` — High (profile_idc 100, no
# constraint flags) at level 4.1 (0x29) — and every encoder here is held to it.
_PROFILE_ARGS = ["-profile:v", "high", "-level", "4.1"]
_WANT_PROFILE = "High"
_WANT_LEVEL = 41

# -pix_fmt yuv420p: a 10-bit or 4:4:4 HEVC source (common for HDR WEB-DLs)
# fails "-profile:v high" outright otherwise — libx264's High profile is 8-bit
# 4:2:0 only. Downsampling loses nothing a browser could show anyway (MSE has
# no HDR path), which is also what each hardware filter below does on the GPU.
_SOFTWARE_ARGS = ["-c:v", "libx264", "-pix_fmt", "yuv420p", *_PROFILE_ARGS,
                  "-preset", "veryfast", "-crf", "21"]


@dataclass(frozen=True)
class Candidate:
    """One way a machine might encode H.264 without its CPU."""

    name: str
    encoder: str
    platforms: tuple[str, ...]
    # `-hwaccel <this>`, and the pixel format decoded frames stay in.
    hwaccel: str
    # Filters for frames already on the GPU, and for frames in system memory.
    hw_filter: str
    sw_filter: str
    # Rate control and anything else, tried in order until one produces High
    # at level 4.1. A later `-level` overrides the shared one, which is how a
    # spelling one encoder rejects is offered without being imposed.
    variants: tuple[tuple[str, tuple[str, ...]], ...]
    # VA-API alone names a device; QSV and NVENC find their own.
    needs_render_node: bool = False


CANDIDATES: tuple[Candidate, ...] = (
    Candidate(
        name="vaapi", encoder="h264_vaapi", platforms=("linux",),
        hwaccel="vaapi",
        hw_filter="scale_vaapi=format=nv12", sw_filter="format=nv12,hwupload",
        needs_render_node=True,
        # CQP is the constant-quality mode every VA-API driver implements (iHD
        # and i965 alike); ICQ and QVBR are not always there. Some Intel
        # generations implement H.264 encoding only on the low-power (VDEnc)
        # path and others only on the full one, which is a property of the chip
        # and the driver — so it is measured rather than looked up.
        variants=(
            ("standard", ("-rc_mode", "CQP", "-qp", "23")),
            ("low-power", ("-rc_mode", "CQP", "-qp", "23", "-low_power", "1")),
        ),
    ),
    Candidate(
        name="qsv", encoder="h264_qsv", platforms=("win32", "linux"),
        hwaccel="qsv",
        hw_filter="vpp_qsv=format=nv12", sw_filter="format=nv12",
        # `-global_quality` is QSV's constant-quality knob. The second variant
        # exists because h264_qsv has no `level` option of its own and takes
        # the generic integer one, which may or may not read "4.1" as 41 —
        # ffmpeg accepts both spellings without complaint, and only the encoded
        # file says which one was understood.
        variants=(
            ("icq", ("-global_quality", "23")),
            ("icq, integer level", ("-global_quality", "23", "-level", "41")),
        ),
    ),
    Candidate(
        name="nvenc", encoder="h264_nvenc", platforms=("win32", "linux"),
        hwaccel="cuda",
        hw_filter="scale_cuda=format=nv12", sw_filter="format=nv12",
        variants=(
            ("vbr", ("-rc", "vbr", "-cq", "23", "-b:v", "0")),
        ),
    ),
)


@dataclass(frozen=True)
class Encoder:
    """A hardware encoder that has been seen to produce what the node claims."""

    candidate: Candidate
    variant: str
    extra: tuple[str, ...]
    device: str | None = None


_enabled = os.environ.get("MESHBAY_HW_VIDEO_ENCODE", "1") not in ("0", "false", "no")
_encoder: Encoder | None = None
_probed = False
_probe_lock: asyncio.Lock | None = None
# (source codec, mode) pairs that have failed once and are not tried again.
_demoted: set[tuple[str, str]] = set()


def set_enabled(enabled: bool) -> None:
    """Operator switch (`node.toml`, `hardware_video_encode`)."""
    global _enabled
    _enabled = enabled


def render_node() -> str | None:
    """The first DRM render node this process can actually open.

    Readable *and* writable, because VA-API maps buffers on it. On a desktop
    session logind grants that through an ACL rather than through membership of
    the `render` group, so asking the kernel answers "can this process use it"
    — which is the question — rather than "is this user in a group", which is
    not.
    """
    try:
        names = sorted(os.listdir("/dev/dri"))
    except OSError:
        return None
    for name in names:
        if not name.startswith("renderD"):
            continue
        dev = f"/dev/dri/{name}"
        if os.access(dev, os.R_OK | os.W_OK):
            return dev
    return None


async def _run(args: list[str], timeout: int) -> tuple[int | None, str]:
    proc = await asyncio.create_subprocess_exec(
        *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
    try:
        out, err = await asyncio.wait_for(proc.communicate(), timeout)
    except TimeoutError:
        proc.kill()
        await proc.wait()
        return None, "timed out"
    text = (out or b"").decode("utf-8", "replace") or \
        (err or b"").decode("utf-8", "replace")
    return proc.returncode, text.strip()


async def _test_encode(cand: Candidate, device: str | None,
                       extra: tuple[str, ...]) -> str | None:
    """Encode 1080p and read the result back. Returns why it was rejected.

    `None` means it worked — the encoder exists, the driver initialises, and
    the file it wrote really is High at level 4.1.
    """
    ffmpeg = platform.ffmpeg_cmd()
    if not (os.path.isabs(ffmpeg) or shutil.which(ffmpeg)):
        return "ffmpeg not found"
    fd, out = tempfile.mkstemp(suffix=".mp4", prefix="meshbay-hwprobe-")
    os.close(fd)
    try:
        rc, err = await _run([
            ffmpeg, "-hide_banner", "-loglevel", "error", "-nostdin", "-y",
            *device_args(cand, device),
            "-f", "lavfi", "-i", f"testsrc=size={PROBE_SIZE}:rate=25:duration=0.2",
            "-vf", cand.sw_filter,
            "-c:v", cand.encoder, *_PROFILE_ARGS, *extra,
            out,
        ], PROBE_TIMEOUT_SECS)
        if rc != 0:
            # The *first* line: ffmpeg's last word is usually "nothing was
            # written into the output file", which is the consequence. The
            # cause is at the top — "Cannot load libcuda.so.1", "Error
            # creating a MFX session", "libva ... init failed".
            return (err.splitlines() or ["no output"])[0]
        rc, info = await _run([
            platform.ffprobe_cmd(), "-v", "error", "-select_streams", "v:0",
            "-show_entries", "stream=profile,level", "-of", "csv=p=0", out,
        ], PROBE_TIMEOUT_SECS)
        profile, _, level = info.partition(",")
        if rc != 0 or profile.strip() != _WANT_PROFILE or level.strip() != str(_WANT_LEVEL):
            return (f"produced {info or 'nothing'}, and the node announces "
                    f"{_WANT_PROFILE} at level {_WANT_LEVEL}")
        return None
    finally:
        try:
            os.unlink(out)
        except OSError:
            pass


async def encoder() -> Encoder | None:
    """The working hardware encoder, established once per process."""
    global _encoder, _probed, _probe_lock
    if not _enabled:
        return None
    if _probed:
        return _encoder
    if _probe_lock is None:
        _probe_lock = asyncio.Lock()
    async with _probe_lock:
        if _probed:
            return _encoder
        _encoder = await _find_encoder()
        _probed = True
    return _encoder


async def _find_encoder() -> Encoder | None:
    device = render_node()
    for cand in CANDIDATES:
        if sys.platform not in cand.platforms:
            continue
        if cand.needs_render_node and device is None:
            continue
        for variant, extra in cand.variants:
            why = await _test_encode(cand, device, extra)
            if why is None:
                log.info("hwaccel: %s (%s) — video re-encoding runs on the GPU",
                         cand.encoder, variant)
                return Encoder(candidate=cand, variant=variant, extra=extra,
                               device=device if cand.needs_render_node else None)
            log.debug("hwaccel: %s (%s) rejected: %s", cand.encoder, variant, why)
    log.info("hwaccel: no hardware H264 encoder works here — encoding in software")
    return None


async def modes_for(source_codec: str | None) -> list[str]:
    """Which re-encode modes to try for this source, best first.

    Always ends in `SW`: libx264 is the one that needs no hardware, so a plan
    that ran out of hardware modes is still a plan that plays the film.
    """
    enc = await encoder()
    if enc is None:
        return [SW]
    key = source_codec or "?"
    plan = [m for m in (HW, HWENC) if (key, m) not in _demoted]
    plan.append(SW)
    return plan


def demote(source_codec: str | None, mode: str, detail: str) -> None:
    """This mode does not work for this source codec; stop trying it."""
    if mode == SW:
        return
    key = source_codec or "?"
    if (key, mode) not in _demoted:
        _demoted.add((key, mode))
        log.info("hwaccel: %s re-encoding does not work for %s on this machine "
                 "(%s) — not trying it again", mode, key, detail)


def device_args(cand: Candidate, device: str | None) -> list[str]:
    """How this candidate is told which GPU to use, where it needs telling."""
    if cand.needs_render_node and device:
        return ["-vaapi_device", device]
    return []


def input_args(mode: str, enc: Encoder | None) -> list[str]:
    """ffmpeg options that must precede `-i`."""
    if enc is None or mode == SW:
        return []
    cand = enc.candidate
    if mode == HWENC:
        return device_args(cand, enc.device)
    # `-hwaccel_output_format` keeps decoded frames on the GPU rather than
    # reading them back to system memory, which is the whole saving: a readback
    # of every 1080p frame costs more on a weak machine than the encode it
    # feeds.
    args = ["-hwaccel", cand.hwaccel]
    if enc.device:
        args += ["-hwaccel_device", enc.device]
    return args + ["-hwaccel_output_format", cand.hwaccel]


def codec_args(mode: str, enc: Encoder | None) -> list[str]:
    """The `-vf`/`-c:v` half, where libx264's arguments used to be written."""
    if enc is None or mode == SW:
        return list(_SOFTWARE_ARGS)
    cand = enc.candidate
    # The hardware filter converts a 10-bit HDR source to the 8-bit NV12 the
    # encoder takes, on the GPU — the same downsampling the software path does
    # with `-pix_fmt yuv420p`, and for the same reason.
    filters = cand.hw_filter if mode == HW else cand.sw_filter
    return ["-vf", filters, "-c:v", cand.encoder, *_PROFILE_ARGS, *enc.extra]


def _reset_for_tests() -> None:
    global _encoder, _probed, _probe_lock
    _encoder, _probed, _probe_lock = None, False, None
    _demoted.clear()