diff --git a/flashdreams/flashdreams/serving/webrtc/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py new file mode 100644 index 00000000..12ef8cc8 --- /dev/null +++ b/flashdreams/flashdreams/serving/webrtc/encoders.py @@ -0,0 +1,569 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Video encoder backends for the WebRTC serving path. + +Two implementations share a single ``VideoEncoder`` protocol: + +- :class:`DefaultRTCEncoder` — thin adapter over aiortc's built-in software + encoder path. Frames leave the model as raw RGB and aiortc's RTP sender + loop drives encoding. +- :class:`PyNvHardwareEncoder` — GPU-resident NVENC H.264 via + ``PyNvVideoCodec`` 2.1. Emits pre-encoded Annex-B packets that aiortc's + ``H264Encoder.pack()`` fragments into RTP payloads (no re-encoding). + +Selection is performed by :func:`select_encoder`, which layers a +capability query (``GetEncoderCaps``) with two failure semantics: an +environmental "no" from Stage 1 falls back silently to the software path, +while a Stage-2 ``CreateEncoder`` failure is logged with traceback and +re-raised so it cannot be masked by a silent fallback. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +from collections.abc import Callable +from dataclasses import dataclass +from fractions import Fraction +from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable + +import torch +from aiortc import MediaStreamTrack +from av.packet import Packet +from loguru import logger + +if TYPE_CHECKING: + from flashdreams.serving.webrtc.media import ( + BufferedVideoTrack, + NVENCVideoTrack, + ) + +# PyNvVideoCodec is a runtime dep. Members like ``nvc.GetEncoderCaps`` / +# ``nvc.FORCEIDR`` are not declared in any type stub, and the module may be +# absent at import time on non-NVENC hosts (or intentionally removed for +# testing the auto-fallback path). Hiding the import behind ``TYPE_CHECKING`` +# gives ty a single ``Any``-typed name to reason about while the runtime +# path keeps the guarded try / except. +if TYPE_CHECKING: + nvc: Any = None + _PYNVVIDEOCODEC_AVAILABLE: bool = False +else: + try: + import PyNvVideoCodec as nvc + + _PYNVVIDEOCODEC_AVAILABLE = True + except ImportError: + nvc = None + _PYNVVIDEOCODEC_AVAILABLE = False + + +EncoderBackend = Literal["auto", "nvenc", "default"] + + +# H.264 NAL type identifiers used when inspecting Annex-B bitstreams. +_H264_NAL_TYPE_IDR = 5 +_H264_NAL_TYPE_SPS = 7 +_H264_NAL_TYPE_PPS = 8 + +# RTP video clock, per RFC 6184. aiortc's ``H264Encoder.pack()`` rescales +# from whatever ``time_base`` the packet carries into this base, so setting +# ``time_base = 1 / _RTP_VIDEO_CLOCK`` on emitted packets is the +# lowest-conversion choice. +_RTP_VIDEO_CLOCK = 90_000 + + +class EncoderInitError(RuntimeError): + """Raised when a forced encoder backend cannot be initialized.""" + + +@dataclass(slots=True, frozen=True) +class ChunkDeliveryResult: + """Uniform result shape for :meth:`VideoEncoder.deliver_chunk`. + + Callers do not need to know which encoder produced the chunk; the + fields here cover both paths. + """ + + backend: str + num_frames: int + num_keyframes: int + encode_ms: float + + +@runtime_checkable +class VideoEncoder(Protocol): + """Encoder backend paired with a compatible :class:`MediaStreamTrack`. + + Each backend owns two responsibilities: creating a fresh media track + sized for one session (:meth:`create_track`), and encoding + enqueueing + one chunk of frames onto that track (:meth:`deliver_chunk`). Callers + pick one backend at startup and branch nowhere else. + """ + + fps: int + backend: str + prefers_codec: str | None + """SDP codec preference hint. ``"h264"`` for the NVENC path (must be + honoured or the pre-encoded bitstream will not decode); ``None`` to + let aiortc pick freely.""" + + def create_track(self, *, maxsize: int) -> BufferedVideoTrack | NVENCVideoTrack: + """Create a fresh session track compatible with this encoder.""" + ... + + async def deliver_chunk( + self, + chunk: torch.Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + """Encode ``chunk`` and enqueue the resulting frames/packets.""" + ... + + def close(self) -> None: + """Release any encoder-owned resources (hardware handles, etc.).""" + ... + + +# --------------------------------------------------------------------------- +# Software path (aiortc's built-in encoder) +# --------------------------------------------------------------------------- + + +class DefaultRTCEncoder: + """Software path: hand raw RGB frames to aiortc and let it encode. + + This class does not encode itself; it wraps a + :class:`~flashdreams.serving.webrtc.media.BufferedVideoTrack` that + carries raw RGB frames, and aiortc's RTP sender loop drives encoding + frame-by-frame. The concrete codec (VP8, VP9, H.264) is chosen by SDP + negotiation. + """ + + prefers_codec: str | None = None + backend = "aiortc" + + def __init__(self, *, fps: int) -> None: + if fps <= 0: + raise ValueError(f"fps must be > 0, got {fps}") + self.fps = fps + logger.info( + "Video encoder ready: backend={} fps={}", + self.backend, + fps, + ) + + def create_track(self, *, maxsize: int) -> BufferedVideoTrack: + from flashdreams.serving.webrtc.media import BufferedVideoTrack + + return BufferedVideoTrack(fps=self.fps, maxsize=maxsize) + + async def deliver_chunk( + self, + chunk: torch.Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + # aiortc's software encoder chooses keyframe cadence internally + # and responds to receiver PLI/FIR feedback, so a caller-supplied + # force_keyframe is neither honoured nor useful on this path. + del force_keyframe + from flashdreams.serving.webrtc.media import BufferedVideoTrack + + if not isinstance(track, BufferedVideoTrack): + raise TypeError( + "DefaultRTCEncoder requires a BufferedVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + enqueued = await track.enqueue_chunk(chunk) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + # No encoder-owned resources; the track is owned by aiortc's PC. + return + + +# --------------------------------------------------------------------------- +# Hardware path (NVENC via PyNvVideoCodec 2.1) +# --------------------------------------------------------------------------- + + +def _payload_contains_nal_type(payload: bytes, nal_type: int) -> bool: + """Scan an Annex-B H.264 payload for the presence of a specific NAL type.""" + i = 0 + while True: + idx = payload.find(b"\x00\x00\x01", i) + if idx < 0: + return False + nal_start = idx + 3 + if nal_start >= len(payload): + return False + if (payload[nal_start] & 0x1F) == nal_type: + return True + i = nal_start + 1 + + +def _chunk_to_abgr_cuda_frames(chunk: torch.Tensor) -> torch.Tensor: + """Convert a model-output chunk to NVENC-``ABGR``-formatted CUDA frames. + + Accepts ``[T, 3, H, W]`` or ``[1, 1, T, 3, H, W]`` (the shape produced + by the omnidreams runtime) in either ``uint8`` or float dtype + (float assumed to be in ``[-1, 1]``). Returns a contiguous + ``[T, H, W, 4]`` ``uint8`` CUDA tensor with alpha=255. + + **NVENC ``NV_ENC_BUFFER_FORMAT_ABGR`` is a word-ordered token, not + memory-ordered.** From ``nvEncodeAPI.h``: "a pixel is represented by + a 32-bit word with R in the lowest 8 bits, G in the next 8 bits, B + in the 8 bits after that and A in the highest 8 bits" (word + ``0xAABBGGRR``). In little-endian memory that is the byte sequence + ``[R, G, B, A]`` — so the channel-last tensor we hand to the encoder + must have channel 0 = R, 1 = G, 2 = B, 3 = A. Writing ``[A, B, G, R]`` + (the naive memory-order reading of the name) makes NVENC interpret + the alpha byte as R, producing a visible RGB↔BGR swap on the wire. + + ABGR (rather than NV12) is chosen so NVENC's driver-side RGB→YUV + conversion handles the colour transform, sparing us a bespoke NV12 + kernel. + """ + if not chunk.is_cuda: + raise ValueError("expected CUDA tensor for hardware encode path") + if chunk.ndim == 6: + if chunk.shape[0] != 1 or chunk.shape[1] != 1: + raise ValueError( + "expected single-batch, single-view chunk [1, 1, T, 3, H, W]; " + f"got {tuple(chunk.shape)}" + ) + chunk = chunk[0, 0] + if chunk.ndim != 4 or chunk.shape[1] != 3: + raise ValueError( + "expected chunk shape [T, 3, H, W] or [1, 1, T, 3, H, W]; " + f"got {tuple(chunk.shape)}" + ) + if chunk.dtype == torch.uint8: + rgb = chunk.permute(0, 2, 3, 1) + else: + rgb = ((chunk.float() + 1.0) / 2.0 * 255.0).clamp(0, 255).byte() + rgb = rgb.permute(0, 2, 3, 1) + t, h, w, _ = rgb.shape + a = torch.full((t, h, w, 1), 255, dtype=torch.uint8, device=rgb.device) + # Channel-last [R, G, B, A] → little-endian bytes [R, G, B, A] → + # NVENC word 0xAABBGGRR = NV_ENC_BUFFER_FORMAT_ABGR. + rgba = torch.cat([rgb, a], dim=-1) + return rgba.contiguous() + + +class PyNvHardwareEncoder: + """NVENC H.264 encoder backed by ``PyNvVideoCodec`` 2.1. + + Accepts CUDA tensors and emits Annex-B H.264 packets streamed onto an + :class:`~flashdreams.serving.webrtc.media.NVENCVideoTrack` as they are + encoded. Packets carry ``pts`` on the RTP 90 kHz video clock so + aiortc's ``H264Encoder.pack()`` can rescale without loss. + """ + + prefers_codec: str | None = "h264" + backend = "pynvvideocodec" + + @classmethod + def is_supported( + cls, + *, + gpu_id: int = 0, + width: int = 0, + height: int = 0, + ) -> tuple[bool, str]: + """Query the driver for NVENC H.264 support at the target resolution. + + Uses ``PyNvVideoCodec.GetEncoderCaps`` (public API) so that no + NVENC session is allocated for the probe itself. Returns + ``(True, "")`` when the environment supports NVENC H.264 at the + requested resolution, else ``(False, human_reason)`` where + ``human_reason`` is a diagnostic suitable for logging. + """ + if not _PYNVVIDEOCODEC_AVAILABLE: + return False, "PyNvVideoCodec library is not installed" + try: + # PyNvVideoCodec 2.1 signature: GetEncoderCaps(gpuid=, codec=). + # Keyword is ``gpuid`` (no underscore); positional order is + # (gpuid, codec). + caps = nvc.GetEncoderCaps(gpuid=gpu_id, codec="h264") + except Exception as exc: + return False, ( + f"GetEncoderCaps gpuid={gpu_id} raised {type(exc).__name__}: {exc}" + ) + if not caps: + return False, (f"GetEncoderCaps gpuid={gpu_id} returned no capabilities") + max_w = int(caps.get("width_max", 0) or 0) + max_h = int(caps.get("height_max", 0) or 0) + min_w = int(caps.get("width_min", 0) or 0) + min_h = int(caps.get("height_min", 0) or 0) + if width > 0 and height > 0: + if max_w > 0 and max_h > 0 and (width > max_w or height > max_h): + return False, ( + f"requested resolution {width}x{height} exceeds driver " + f"maximum {max_w}x{max_h} (gpuid={gpu_id})" + ) + if (min_w > 0 or min_h > 0) and (width < min_w or height < min_h): + return False, ( + f"requested resolution {width}x{height} below driver " + f"minimum {min_w}x{min_h} (gpuid={gpu_id})" + ) + return True, "" + + def __init__( + self, + *, + width: int, + height: int, + fps: int, + bitrate: int, + gpu_id: int, + gop: int = 30, + ) -> None: + if not _PYNVVIDEOCODEC_AVAILABLE: + raise RuntimeError( + "PyNvVideoCodec is not installed; PyNvHardwareEncoder cannot " + "be used. Install with `pip install PyNvVideoCodec` or " + "select DefaultRTCEncoder." + ) + if width <= 0 or height <= 0: + raise ValueError(f"width and height must be > 0, got {width}x{height}") + if fps <= 0: + raise ValueError(f"fps must be > 0, got {fps}") + if bitrate <= 0: + raise ValueError(f"bitrate must be > 0, got {bitrate}") + if gop <= 0: + raise ValueError(f"gop must be > 0, got {gop}") + + # Fail fast with a driver-specific diagnostic before allocating a + # session slot. When called via ``select_encoder`` this is + # redundant with the factory's own probe, but a direct instantiation + # (e.g. from a test) still gets a clear reason. + supported, reason = self.is_supported( + gpu_id=gpu_id, + width=width, + height=height, + ) + if not supported: + raise RuntimeError(f"NVENC H.264 not supported: {reason}") + + self.fps = fps + self._width = width + self._height = height + self._bitrate = bitrate + self._gpu_id = gpu_id + self._gop = gop + self._pts_counter = 0 + self._time_base = Fraction(1, _RTP_VIDEO_CLOCK) + # ``FORCEIDR`` is exposed by PyNvVideoCodec as an integer flag + # bitmask. Cache it so we don't hit ``getattr`` on every frame. + self._force_idr_flag = int(nvc.FORCEIDR) # type: ignore[attr-defined] + + # ``repeatspspps=1`` prepends SPS+PPS to every IDR: aiortc's + # ``H264Encoder.pack()`` does not synthesize parameter sets, so the + # RTP stream must carry them in-band or the receiver cannot lock on. + # ``bf=0`` and ``lookahead=0`` keep the output strictly 1:1 with + # input frames, which is what interactive streaming needs. + self._encoder = nvc.CreateEncoder( # type: ignore[attr-defined] + width=width, + height=height, + fmt="ABGR", + usecpuinputbuffer=False, + codec="h264", + preset="P4", + tuning_info="ultra_low_latency", + rc="cbr", + fps=fps, + bitrate=bitrate, + bf=0, + lookahead=0, + repeatspspps=1, + idrperiod=gop, + ) + logger.info( + "Video encoder ready: backend={} codec=h264 {}x{}@{}fps " + "bitrate={}bps gop={} gpu={}", + self.backend, + width, + height, + fps, + bitrate, + gop, + gpu_id, + ) + + def create_track(self, *, maxsize: int) -> NVENCVideoTrack: + # Lazy import to avoid an ``encoders`` ↔ ``media`` cycle. + from flashdreams.serving.webrtc.media import NVENCVideoTrack + + return NVENCVideoTrack(fps=self.fps, maxsize=maxsize) + + async def deliver_chunk( + self, + chunk: torch.Tensor, + track: MediaStreamTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + from flashdreams.serving.webrtc.media import NVENCVideoTrack + + if not isinstance(track, NVENCVideoTrack): + raise TypeError( + "PyNvHardwareEncoder requires an NVENCVideoTrack; got " + f"{type(track).__name__}. Create it via encoder.create_track()." + ) + loop = asyncio.get_running_loop() + + def _stream(pkt: Packet) -> None: + # Marshal each packet back onto the asyncio loop as soon as + # NVENC produces it, rather than waiting for the whole chunk + # to finish encoding — otherwise the first packet's latency is + # bounded below by ``T / fps``. + try: + loop.call_soon_threadsafe( + track.enqueue_encoded_packet_nowait, + pkt, + ) + except RuntimeError: + # Loop has been closed; drop the packet silently rather + # than raising from the encode worker thread. + return + + num_frames, num_keyframes, encode_ms = await asyncio.to_thread( + self.encode_chunk_sync, + chunk, + force_keyframe=force_keyframe, + on_packet=_stream, + ) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=num_frames, + num_keyframes=num_keyframes, + encode_ms=encode_ms, + ) + + def encode_chunk_sync( + self, + chunk: torch.Tensor, + *, + force_keyframe: bool = False, + on_packet: Callable[[Packet], None] | None = None, + ) -> tuple[int, int, float]: + """Encode a chunk synchronously; returns ``(num_frames, num_keyframes, encode_ms)``. + + Kept public because callers (e.g. tests) that already run on a + worker thread should not have to route through :meth:`deliver_chunk` + just to get access to the emitted packets. + """ + frames = _chunk_to_abgr_cuda_frames(chunk) + num_frames = frames.shape[0] + num_keyframes = 0 + start_s = time.perf_counter() + for i in range(num_frames): + frame = frames[i].contiguous() + if force_keyframe and i == 0: + bs = self._encoder.Encode(frame, self._force_idr_flag) + else: + bs = self._encoder.Encode(frame) + if not bs: + continue + payload = bytes(bs) + packet = Packet(payload) + packet.pts = (self._pts_counter * _RTP_VIDEO_CLOCK) // self.fps + packet.time_base = self._time_base + self._pts_counter += 1 + if _payload_contains_nal_type(payload, _H264_NAL_TYPE_IDR): + num_keyframes += 1 + if on_packet is not None: + on_packet(packet) + encode_ms = (time.perf_counter() - start_s) * 1000.0 + return num_frames, num_keyframes, encode_ms + + def close(self) -> None: + with contextlib.suppress(Exception): + self._encoder.EndEncode() + + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + + +def select_encoder( + *, + backend: EncoderBackend, + width: int, + height: int, + fps: int, + bitrate: int, + gpu_id: int, + gop: int = 30, +) -> VideoEncoder: + """Choose an encoder implementation. + + ``backend == "default"`` returns a :class:`DefaultRTCEncoder` and does + not touch the NVENC probe path at all. + + ``backend == "nvenc"`` requires the hardware path; any probe failure + raises :class:`EncoderInitError` so misconfigured deployments surface + at startup rather than silently degrading. + + ``backend == "auto"`` prefers hardware when the environment supports + it. Selection uses two stages with two different failure semantics: + + * **Stage 1** (``GetEncoderCaps`` + resolution bounds): a "no" here + means the environment cannot do this, which is expected on hosts + without NVENC. Fall back silently to + :class:`DefaultRTCEncoder`, logging the reason at INFO. + * **Stage 2** (``CreateEncoder``): a raise here means Stage 1 said + the environment supports NVENC but session allocation failed anyway + — driver bug, session-pool exhaustion, hardware fault, or + misconfiguration. Log with traceback and re-raise; do not silently + degrade, because doing so would hide a real problem. + """ + if backend == "default": + return DefaultRTCEncoder(fps=fps) + + supported, reason = PyNvHardwareEncoder.is_supported( + gpu_id=gpu_id, + width=width, + height=height, + ) + if not supported: + if backend == "nvenc": + raise EncoderInitError( + f"encoder_backend='nvenc' requested but NVENC H.264 is not " + f"supported on this system: {reason}" + ) + logger.info( + "NVENC H.264 not supported on this system ({}); using aiortc " + "software encoder.", + reason, + ) + return DefaultRTCEncoder(fps=fps) + + try: + return PyNvHardwareEncoder( + width=width, + height=height, + fps=fps, + bitrate=bitrate, + gpu_id=gpu_id, + gop=gop, + ) + except Exception: + logger.exception( + "GetEncoderCaps reported NVENC H.264 supported, but CreateEncoder " + "failed. Not silently falling back to the software encoder — the " + "underlying failure needs to be diagnosed.", + ) + raise diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 5e97813f..5418c90a 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -16,11 +16,20 @@ from typing import Any import torch -from aiortc import RTCConfiguration, RTCPeerConnection, RTCSessionDescription +from aiortc import ( + RTCConfiguration, + RTCPeerConnection, + RTCRtpSender, + RTCSessionDescription, +) from loguru import logger from flashdreams.serving.webrtc.controls import KeyboardResampler -from flashdreams.serving.webrtc.media import BufferedVideoTrack +from flashdreams.serving.webrtc.encoders import ( + DefaultRTCEncoder, + VideoEncoder, +) +from flashdreams.serving.webrtc.media import BufferedVideoTrack, NVENCVideoTrack from flashdreams.serving.webrtc.server import SessionBusyError from flashdreams.serving.webrtc.warmup import ( run_loopback_warmup_session, @@ -61,7 +70,8 @@ class ManagedWebRTCSession: """Per-session state for the single active WebRTC peer connection.""" runtime: Any - video_track: BufferedVideoTrack + video_track: BufferedVideoTrack | NVENCVideoTrack + video_encoder: VideoEncoder peer_connection: Any resampler: KeyboardResampler control_channel: Any | None = None @@ -158,6 +168,69 @@ def _make_resampler(self, *, start_v: float) -> KeyboardResampler: def _register_extra_peer_handlers(self, peer_connection: Any) -> None: """Register optional extra peer-connection event handlers.""" + def _prefer_h264_video_codec(self, *, transceiver: Any) -> None: + """Constrain the transceiver's codec preferences to H.264 variants. + + Required when the selected encoder emits pre-encoded H.264 packets + (``av.Packet`` route through ``H264Encoder.pack()``): if the SDP + negotiates VP8/VP9 instead, aiortc will pack the H.264 bitstream + under the wrong codec header and the receiver will fail to decode. + + If the local aiortc build does not advertise H.264, no preference + is set; the SDP-time fallback in ``_enforce_h264_or_fallback`` + will then swap the encoder to :class:`DefaultRTCEncoder`. + """ + caps = RTCRtpSender.getCapabilities("video") + h264_codecs = [c for c in caps.codecs if c.mimeType.lower() == "video/h264"] + if not h264_codecs: + return + transceiver.setCodecPreferences(h264_codecs) + + def _enforce_h264_or_fallback( + self, + *, + transceiver: Any, + managed_session: ManagedWebRTCSession, + num_frames: int, + ) -> None: + """Verify H.264 was negotiated; swap to the software encoder if not. + + aiortc exposes the negotiated codec set on + ``RTCRtpTransceiver._codecs`` after ``setLocalDescription``. We + read it via that attribute (aiortc-internal, but stable in the + pinned version) and, if H.264 did not land, close the hardware + encoder and install a :class:`DefaultRTCEncoder` with a + :class:`BufferedVideoTrack` on the same sender before the first + RTP packet flies. ``replaceTrack`` does not renegotiate; aiortc's + RTP loop will encode raw ``av.VideoFrame`` output with whatever + codec (VP8/VP9/H.264) actually landed in the SDP. + """ + negotiated = getattr(transceiver, "_codecs", None) or [] + if negotiated and negotiated[0].mimeType.lower() == "video/h264": + logger.info( + "Video codec negotiated: {} (hardware encoder path active).", + negotiated[0].mimeType, + ) + return + + chosen = negotiated[0].mimeType if negotiated else "" + logger.warning( + "H.264 preferred by hardware encoder but SDP negotiation " + "landed on {!r}; swapping to the software encoder before " + "streaming begins.", + chosen, + ) + # Close the hardware encoder so its NVENC session is released + # promptly; the software adapter has no hardware resources to + # release itself. + managed_session.video_encoder.close() + + fallback_encoder = DefaultRTCEncoder(fps=self.fps) + fallback_track = fallback_encoder.create_track(maxsize=num_frames) + transceiver.sender.replaceTrack(fallback_track) + managed_session.video_encoder = fallback_encoder + managed_session.video_track = fallback_track + def _on_offer_received(self, offer_sdp: str) -> None: """Hook invoked with the remote offer SDP before negotiation.""" @@ -283,8 +356,17 @@ async def _create_answer_with_runtime_ready_locked( # frames than steady state; sizing to it would force a per-chunk # stall, so we size to the steady-state count. num_frames = self._runtime.peek_steady_chunk_num_frames() - video_track = BufferedVideoTrack(fps=self.fps, maxsize=num_frames) - peer_connection.addTrack(video_track) + video_encoder: VideoEncoder = self._runtime.video_encoder + video_track = video_encoder.create_track(maxsize=num_frames) + # Use ``addTransceiver`` (not ``addTrack``) so we can constrain the + # SDP m-line's codec list via ``setCodecPreferences`` when the + # encoder emits pre-encoded H.264 packets. + video_transceiver = peer_connection.addTransceiver( + video_track, + direction="sendonly", + ) + if video_encoder.prefers_codec == "h264": + self._prefer_h264_video_codec(transceiver=video_transceiver) # Start the resampler's virtual clock at 0; the real anchor is set # in the ``on_datachannel`` handler so chunk 0's window starts when # input can actually arrive. @@ -293,6 +375,7 @@ async def _create_answer_with_runtime_ready_locked( managed_session = ManagedWebRTCSession( runtime=self._runtime, video_track=video_track, + video_encoder=video_encoder, peer_connection=peer_connection, resampler=resampler, last_client_message_at=loop.time(), @@ -350,6 +433,12 @@ async def on_connectionstatechange() -> None: answer = await peer_connection.createAnswer() await peer_connection.setLocalDescription(answer) await wait_for_ice_gathering_complete(peer_connection) + if video_encoder.prefers_codec == "h264": + self._enforce_h264_or_fallback( + transceiver=video_transceiver, + managed_session=managed_session, + num_frames=num_frames, + ) local_description = peer_connection.localDescription if local_description is None: raise RuntimeError("Peer connection did not produce local description.") @@ -548,6 +637,7 @@ async def _generation_worker( runtime = managed_session.runtime resampler = managed_session.resampler video_track = managed_session.video_track + video_encoder = managed_session.video_encoder # Stay idle until the user interacts. Generating eagerly would burn # GPU cycles on a still scene the viewer never sees. Once an event @@ -617,7 +707,12 @@ async def _generation_worker( return continue t_after_gen = loop.time() - enqueued = await video_track.enqueue_chunk(result.video_chunk) + delivery = await video_encoder.deliver_chunk( + result.video_chunk, + video_track, + force_keyframe=False, + ) + enqueued = delivery.num_frames t_after_enqueue = loop.time() gen_ms = (t_after_gen - t_before_gen) * 1e3 diff --git a/flashdreams/flashdreams/serving/webrtc/media.py b/flashdreams/flashdreams/serving/webrtc/media.py index df5925e6..6d6eba8e 100644 --- a/flashdreams/flashdreams/serving/webrtc/media.py +++ b/flashdreams/flashdreams/serving/webrtc/media.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +import contextlib from collections.abc import Callable from fractions import Fraction @@ -12,6 +13,7 @@ from aiortc import MediaStreamTrack from aiortc.mediastreams import MediaStreamError from av import VideoFrame +from av.packet import Packet from loguru import logger _STALL_THRESHOLD_MS = 1.0 @@ -149,3 +151,117 @@ async def close(self) -> None: break self._frames.put_nowait(None) self.stop() + + +class NVENCVideoTrack(MediaStreamTrack): + """WebRTC video track that delivers pre-encoded H.264 packets. + + Paired with ``PyNvHardwareEncoder``: :meth:`recv` returns + :class:`av.Packet` (not :class:`av.VideoFrame`), which aiortc's + ``RTCRtpSender`` routes through ``H264Encoder.pack()`` for RTP + fragmentation only. The encoder sets ``pts`` and ``time_base`` on + each packet before enqueueing; this track only paces delivery to + ``fps`` and applies a drop-oldest overflow policy so a slow + consumer cannot stall the encode worker. + """ + + kind = "video" + + def __init__(self, *, fps: int, maxsize: int) -> None: + super().__init__() + if fps <= 0: + raise ValueError("fps must be > 0") + if maxsize <= 0: + raise ValueError("maxsize must be > 0") + self._fps = fps + self._frame_interval_s = 1.0 / fps + self._next_deadline_s: float | None = None + self._maxsize = maxsize + self._packets: asyncio.Queue[Packet | None] = asyncio.Queue( + maxsize=maxsize, + ) + self._closed = False + self._dropped_packets = 0 + + @property + def fps(self) -> int: + return self._fps + + @property + def maxsize(self) -> int: + return self._maxsize + + @property + def dropped_packets(self) -> int: + return self._dropped_packets + + def qsize(self) -> int: + return self._packets.qsize() + + def enqueue_encoded_packet_nowait(self, packet: Packet) -> bool: + """Synchronously enqueue one packet on the loop thread. + + Called from the encode worker via ``loop.call_soon_threadsafe`` so + packets become visible to :meth:`recv` as soon as they are + produced, without waiting for the whole chunk to finish encoding. + Drops the oldest queued packet on overflow so real-time streaming + does not stall behind a slow consumer. + """ + if self._closed: + return False + if self._maxsize > 0 and self._packets.full(): + with contextlib.suppress(asyncio.QueueEmpty): + self._packets.get_nowait() + self._dropped_packets += 1 + logger.debug( + "NVENCVideoTrack overflow: dropped oldest packet " + "(total dropped={})", + self._dropped_packets, + ) + self._packets.put_nowait(packet) + return True + + async def recv(self) -> Packet: + if self._closed: + raise MediaStreamError + + loop = asyncio.get_running_loop() + t_get_start = loop.time() + packet = await self._packets.get() + if packet is None: + raise MediaStreamError + get_wait_ms = (loop.time() - t_get_start) * 1000.0 + first_packet = self._next_deadline_s is None + just_stalled = (not first_packet) and get_wait_ms > _STALL_THRESHOLD_MS + + now_s = loop.time() + if first_packet or just_stalled: + self._next_deadline_s = now_s + else: + proposed = self._next_deadline_s + self._frame_interval_s + wait_s = proposed - now_s + if wait_s > 0: + await asyncio.sleep(wait_s) + self._next_deadline_s = proposed + else: + if -wait_s * 1000.0 > _PACING_LAG_LOG_MS: + logger.debug( + "NVENCVideoTrack pacing lag: deadline {:.1f}ms " + "behind walltime; re-anchoring (queue depth {}).", + -wait_s * 1000.0, + self._packets.qsize(), + ) + self._next_deadline_s = now_s + return packet + + async def close(self) -> None: + if self._closed: + return + self._closed = True + while True: + try: + self._packets.get_nowait() + except asyncio.QueueEmpty: + break + self._packets.put_nowait(None) + self.stop() diff --git a/flashdreams/tests/test_encoders.py b/flashdreams/tests/test_encoders.py new file mode 100644 index 00000000..a6680d56 --- /dev/null +++ b/flashdreams/tests/test_encoders.py @@ -0,0 +1,509 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the ``VideoEncoder`` abstraction and ``select_encoder`` factory. + +Covers: + +- Protocol conformance for the software-path adapter. +- ``select_encoder`` branch coverage under ``backend={"default","auto","nvenc"}``. +- The two failure semantics of the capability probe: a Stage-1 no is a + silent fallback; a Stage-2 ``CreateEncoder`` raise is a hard error + (never a silent fallback). The latter is the regression guard that + keeps a future refactor from masking driver problems. +- The cross-chunk packet-ordering invariant: sequential ``await + deliver_chunk(...)`` calls must produce monotonically-ordered packets + on the paired track. A ``create_task``-style refactor would silently + break this — see :class:`TestDeliverChunkOrdering`. +- Compatibility guards for the two upstream libraries we couple to + (``aiortc`` and ``PyNvVideoCodec``). +""" + +from __future__ import annotations + +import asyncio +import threading +from fractions import Fraction +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from av.packet import Packet + +pytestmark = pytest.mark.ci_cpu + +from flashdreams.serving.webrtc import encoders as enc_mod +from flashdreams.serving.webrtc.encoders import ( + ChunkDeliveryResult, + DefaultRTCEncoder, + EncoderInitError, + VideoEncoder, + select_encoder, +) +from flashdreams.serving.webrtc.media import NVENCVideoTrack + +_SELECT_KW = dict( + width=1280, + height=704, + fps=30, + bitrate=6_000_000, + gpu_id=0, + gop=30, +) + + +# --------------------------------------------------------------------------- +# Protocol conformance +# --------------------------------------------------------------------------- + + +class TestVideoEncoderProtocol: + def test_default_encoder_satisfies_protocol(self) -> None: + assert isinstance(DefaultRTCEncoder(fps=30), VideoEncoder) + + def test_default_encoder_advertises_no_codec_preference(self) -> None: + assert DefaultRTCEncoder(fps=30).prefers_codec is None + + def test_default_encoder_rejects_bad_fps(self) -> None: + with pytest.raises(ValueError, match="fps"): + DefaultRTCEncoder(fps=0) + + +# --------------------------------------------------------------------------- +# select_encoder: "default" backend never touches NVENC +# --------------------------------------------------------------------------- + + +class TestSelectDefaultBackend: + def test_returns_default_encoder(self) -> None: + enc = select_encoder(backend="default", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + assert enc.fps == 30 + + def test_does_not_call_getencodercaps(self) -> None: + # Even if the library exists, "default" must not touch it — this + # keeps the software path deterministic and safe on hosts where + # GetEncoderCaps might have side effects. + fake_nvc = MagicMock() + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + select_encoder(backend="default", **_SELECT_KW) + fake_nvc.GetEncoderCaps.assert_not_called() + + def test_works_even_when_library_missing(self) -> None: + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", False), + patch.object(enc_mod, "nvc", None), + ): + enc = select_encoder(backend="default", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + + +# --------------------------------------------------------------------------- +# select_encoder: Stage 1 (GetEncoderCaps) failure paths +# --------------------------------------------------------------------------- + + +class TestSelectStage1Failure: + """Stage 1 says "environment can't do this" — auto silently falls back; + nvenc raises loudly.""" + + def test_library_missing_auto_falls_back(self, caplog) -> None: + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", False), + patch.object(enc_mod, "nvc", None), + ): + enc = select_encoder(backend="auto", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + + def test_library_missing_nvenc_raises(self) -> None: + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", False), + patch.object(enc_mod, "nvc", None), + ): + with pytest.raises(EncoderInitError, match="not installed"): + select_encoder(backend="nvenc", **_SELECT_KW) + + @pytest.mark.parametrize( + "caps_effect, expected_reason_frag", + [ + (RuntimeError("driver comms error"), "driver comms error"), + ({}, "no capabilities"), + ( + { + "width_max": 640, + "height_max": 480, + "width_min": 32, + "height_min": 32, + }, + "exceeds driver maximum", + ), + ( + { + "width_max": 8192, + "height_max": 8192, + "width_min": 1920, + "height_min": 1088, + }, + "below driver minimum", + ), + ], + ids=["caps_raise", "caps_empty", "caps_max_too_small", "caps_min_too_big"], + ) + def test_caps_failure_auto_falls_back( + self, + caps_effect, + expected_reason_frag, + ) -> None: + fake_nvc = MagicMock() + if isinstance(caps_effect, BaseException): + fake_nvc.GetEncoderCaps.side_effect = caps_effect + else: + fake_nvc.GetEncoderCaps.return_value = caps_effect + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + enc = select_encoder(backend="auto", **_SELECT_KW) + assert isinstance(enc, DefaultRTCEncoder) + # ``CreateEncoder`` must not be reached when Stage 1 fails. + fake_nvc.CreateEncoder.assert_not_called() + + def test_caps_failure_nvenc_raises_with_reason(self) -> None: + fake_nvc = MagicMock() + fake_nvc.GetEncoderCaps.side_effect = RuntimeError("driver comms error") + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + with pytest.raises(EncoderInitError, match="driver comms error"): + select_encoder(backend="nvenc", **_SELECT_KW) + + +# --------------------------------------------------------------------------- +# select_encoder: Stage 2 (CreateEncoder) failure — HARD ERROR +# --------------------------------------------------------------------------- + + +class TestSelectStage2HardError: + """The key regression guard: if Stage 1 says supported but Stage 2's + ``CreateEncoder`` raises, ``select_encoder`` must re-raise. Silently + falling back to the software encoder here would mask a real bug + (driver, session pool exhaustion, hardware fault). Any future refactor + that makes this test fail is almost certainly regressing that semantic. + """ + + def _fake_nvc_caps_ok(self) -> MagicMock: + fake = MagicMock() + fake.GetEncoderCaps.return_value = { + "width_max": 8192, + "height_max": 8192, + "width_min": 32, + "height_min": 32, + } + fake.FORCEIDR = 0x1 + return fake + + def test_construct_failure_reraises_under_auto(self) -> None: + fake_nvc = self._fake_nvc_caps_ok() + fake_nvc.CreateEncoder.side_effect = RuntimeError( + "NVENC session pool exhausted" + ) + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + with pytest.raises(RuntimeError, match="session pool exhausted"): + select_encoder(backend="auto", **_SELECT_KW) + + def test_construct_failure_reraises_under_nvenc(self) -> None: + fake_nvc = self._fake_nvc_caps_ok() + fake_nvc.CreateEncoder.side_effect = RuntimeError("hardware fault") + with ( + patch.object(enc_mod, "_PYNVVIDEOCODEC_AVAILABLE", True), + patch.object(enc_mod, "nvc", fake_nvc), + ): + with pytest.raises(RuntimeError, match="hardware fault"): + select_encoder(backend="nvenc", **_SELECT_KW) + + +# --------------------------------------------------------------------------- +# ChunkDeliveryResult +# --------------------------------------------------------------------------- + + +class TestChunkDeliveryResult: + def test_is_frozen_dataclass(self) -> None: + result = ChunkDeliveryResult( + backend="fake", + num_frames=4, + num_keyframes=1, + encode_ms=1.5, + ) + with pytest.raises((AttributeError, Exception)): + result.backend = "other" # ty:ignore[invalid-assignment] + + +# --------------------------------------------------------------------------- +# DefaultRTCEncoder.deliver_chunk delegates to track.enqueue_chunk +# --------------------------------------------------------------------------- + + +class _FakeBufferedVideoTrack: + """Minimal stand-in that ``isinstance(track, BufferedVideoTrack)`` + treats as a real track (subclassing lets us bypass MediaStreamTrack's + aiortc runtime dependencies without breaking the isinstance check).""" + + def __init__(self) -> None: + self.enqueued_chunks: list = [] + + async def enqueue_chunk(self, chunk) -> int: + self.enqueued_chunks.append(chunk) + return 4 + + +class TestDefaultRTCEncoderDeliver: + @pytest.mark.asyncio + async def test_deliver_chunk_returns_frames_from_track(self) -> None: + from flashdreams.serving.webrtc import media as media_mod + + fake_track = _FakeBufferedVideoTrack() + # Patch the isinstance check inside deliver_chunk to accept our fake. + with patch.object(media_mod, "BufferedVideoTrack", _FakeBufferedVideoTrack): + enc = DefaultRTCEncoder(fps=30) + result = await enc.deliver_chunk( + SimpleNamespace(shape=(4, 3, 8, 8)), # ty:ignore[invalid-argument-type] + fake_track, # ty:ignore[invalid-argument-type] + ) + assert result.backend == "aiortc" + assert result.num_frames == 4 + assert result.num_keyframes == 0 + assert len(fake_track.enqueued_chunks) == 1 + + @pytest.mark.asyncio + async def test_deliver_chunk_rejects_wrong_track_type(self) -> None: + enc = DefaultRTCEncoder(fps=30) + with pytest.raises(TypeError, match="BufferedVideoTrack"): + await enc.deliver_chunk( + SimpleNamespace(), # ty:ignore[invalid-argument-type] + SimpleNamespace(), # ty:ignore[invalid-argument-type] + ) + + +# --------------------------------------------------------------------------- +# Compatibility guards for upstream libraries +# --------------------------------------------------------------------------- + + +class TestCompatGuards: + """These tests exist to fail loudly when an upstream dependency + changes its public surface underneath us. They are cheap and + catch dependency drift far earlier than a runtime failure would.""" + + def test_aiortc_h264_encoder_has_pack_and_encode(self) -> None: + from aiortc.codecs.h264 import H264Encoder + + assert callable(getattr(H264Encoder, "encode", None)), ( + "aiortc H264Encoder.encode disappeared — the software encoder " + "contract this design assumes has changed." + ) + assert callable(getattr(H264Encoder, "pack", None)), ( + "aiortc H264Encoder.pack disappeared — the pre-encoded packet " + "path this design relies on has changed." + ) + + def test_aiortc_sender_module_importable(self) -> None: + # Any structural change to rtcrtpsender that breaks import will + # break the runtime; catch it here before the first RTP packet. + import aiortc.rtcrtpsender # noqa: F401 + + def test_getencodercaps_callable_when_library_available(self) -> None: + if not enc_mod._PYNVVIDEOCODEC_AVAILABLE: + pytest.skip("PyNvVideoCodec is not installed on this host") + assert callable(getattr(enc_mod.nvc, "GetEncoderCaps", None)), ( + "PyNvVideoCodec.GetEncoderCaps renamed or removed — Stage 1 " + "capability probe cannot function." + ) + + +# --------------------------------------------------------------------------- +# Cross-chunk packet-ordering invariant +# --------------------------------------------------------------------------- + + +class _OrderingFakeEncoder: + """Minimal encoder that mimics :class:`PyNvHardwareEncoder`'s async + marshaling contract without needing CUDA or PyNvVideoCodec. + + Encoding runs on an ``asyncio.to_thread`` worker and each emitted + packet is delivered to the track via ``loop.call_soon_threadsafe``, + exactly as the real hardware encoder does. This lets us assert the + end-to-end ordering property without any GPU dependency. + + The pts counter is guarded by a lock because the fire-and-forget + test dispatches multiple ``deliver_chunk`` calls concurrently, and + ``self._pts_counter += 1`` is not atomic across Python bytecodes. + """ + + backend = "fake" + prefers_codec: str | None = "h264" + + def __init__(self, *, fps: int, frames_per_chunk: int) -> None: + self.fps = fps + self._frames_per_chunk = frames_per_chunk + self._pts_counter = 0 + self._pts_counter_lock = threading.Lock() + self._time_base = Fraction(1, 90_000) + + def create_track(self, *, maxsize: int) -> NVENCVideoTrack: + return NVENCVideoTrack(fps=self.fps, maxsize=maxsize) + + async def deliver_chunk( + self, + chunk: object, + track: NVENCVideoTrack, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del chunk, force_keyframe + loop = asyncio.get_running_loop() + frames = self._frames_per_chunk + + def _encode_worker() -> None: + # One packet per frame, monotonically-increasing pts. The + # per-iteration call_soon_threadsafe hand-off mirrors what + # PyNvHardwareEncoder does with its on_packet callback. + for _ in range(frames): + packet = Packet(b"\x00\x00\x00\x01\x67") + with self._pts_counter_lock: + pts_frame_index = self._pts_counter + self._pts_counter += 1 + packet.pts = (pts_frame_index * 90_000) // self.fps + packet.time_base = self._time_base + loop.call_soon_threadsafe( + track.enqueue_encoded_packet_nowait, + packet, + ) + + await asyncio.to_thread(_encode_worker) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=frames, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + return + + +# 30 chunks × 8 frames = 240 packets — realistic scale for an interactive +# session (about 8 seconds of 30 fps video, or a few minutes of chunked +# generation at typical omnidreams cadence). At the encoder's real fps=30 +# the track's recv() pacing would make draining take ~8 s per test, so we +# use a much higher fps here to keep the pacing throttle negligible while +# preserving distinct, monotonic pts values (pts = i at fps=90_000 since +# ``(i * 90_000) // 90_000 == i``). +_ORDERING_NUM_CHUNKS = 30 +_ORDERING_FRAMES_PER_CHUNK = 8 +_ORDERING_TOTAL_FRAMES = _ORDERING_NUM_CHUNKS * _ORDERING_FRAMES_PER_CHUNK +_ORDERING_FPS = 90_000 + + +class TestDeliverChunkOrdering: + """Regression guard for the manager's sequential await pattern. + + The manager's ``_generation_worker`` does ``await deliver_chunk(...)`` + in a single loop, which forces chunk N's packets to land on the track + before chunk N+1 starts encoding. If a future refactor swaps that for + ``asyncio.create_task(deliver_chunk(...))`` — or introduces a + producer-consumer queue between generation and encoding without a + reorder buffer — chunks could complete out of order and packets + would land on the track with non-monotonic ``pts``. + + These tests replay the manager's sequential-await pattern against a + fake encoder that mirrors :class:`PyNvHardwareEncoder`'s async + marshaling. They pass under the current ``await``-based design and + would fail under a fire-and-forget rewrite. + """ + + @pytest.mark.asyncio + async def test_sequential_await_produces_monotonic_pts(self) -> None: + encoder = _OrderingFakeEncoder( + fps=_ORDERING_FPS, + frames_per_chunk=_ORDERING_FRAMES_PER_CHUNK, + ) + track = encoder.create_track(maxsize=_ORDERING_TOTAL_FRAMES) + + for _ in range(_ORDERING_NUM_CHUNKS): + await encoder.deliver_chunk(object(), track) + + seen_pts: list[int] = [] + for _ in range(_ORDERING_TOTAL_FRAMES): + packet = await asyncio.wait_for(track.recv(), timeout=1.0) + # ``_OrderingFakeEncoder`` always sets pts before enqueueing; the + # ``av.Packet.pts`` field is nullable at the type level, so narrow. + assert packet.pts is not None + seen_pts.append(int(packet.pts)) + + # Strictly monotonic and matches the exact expected pts sequence. + expected = [ + (i * 90_000) // _ORDERING_FPS for i in range(_ORDERING_TOTAL_FRAMES) + ] + assert seen_pts == expected, ( + "packets emitted out of order across chunks: " + f"first 16 got={seen_pts[:16]} expected={expected[:16]}" + ) + assert seen_pts == sorted(seen_pts) + + @pytest.mark.asyncio + async def test_fire_and_forget_pattern_would_break_ordering(self) -> None: + """Counter-check: if the manager ever spawns deliver_chunk via + ``asyncio.create_task`` (fire-and-forget), packets *can* interleave + across chunks. This test documents that failure mode so a future + refactor that reintroduces the anti-pattern can be caught by + comparing behaviour against the sequential-await test above. + + The test is not a bug — it is a demonstration that the sequential + await in the manager is *load-bearing* for ordering. + """ + encoder = _OrderingFakeEncoder( + fps=_ORDERING_FPS, + frames_per_chunk=_ORDERING_FRAMES_PER_CHUNK, + ) + track = encoder.create_track(maxsize=_ORDERING_TOTAL_FRAMES) + + # Bias interleaving: create_task, then gather. Individual chunks + # each finish quickly; scheduler order within the loop does not + # guarantee packets arrive in the same order the tasks were spawned. + tasks = [ + asyncio.create_task(encoder.deliver_chunk(object(), track)) + for _ in range(_ORDERING_NUM_CHUNKS) + ] + await asyncio.gather(*tasks) + + seen_pts: list[int] = [] + for _ in range(_ORDERING_TOTAL_FRAMES): + packet = await asyncio.wait_for(track.recv(), timeout=1.0) + # ``_OrderingFakeEncoder`` always sets pts before enqueueing; the + # ``av.Packet.pts`` field is nullable at the type level, so narrow. + assert packet.pts is not None + seen_pts.append(int(packet.pts)) + + # Every pts value must be present exactly once; that part is a + # correctness invariant of the fake encoder (guarded by the lock + # around ``_pts_counter``). Duplicate pts here would be a bug in + # the fake, not in the code under test. + expected_set = { + (i * 90_000) // _ORDERING_FPS for i in range(_ORDERING_TOTAL_FRAMES) + } + assert set(seen_pts) == expected_set + assert len(seen_pts) == _ORDERING_TOTAL_FRAMES + + # The full sequence, however, is not necessarily monotonic — the + # fire-and-forget pattern permits interleave. We do NOT assert + # monotonicity here; on the contrary, the moment this test starts + # asserting monotonic pts, the manager's sequential-await guard + # has become unnecessary (and the test above becomes redundant). diff --git a/flashdreams/tests/test_nvenc_track.py b/flashdreams/tests/test_nvenc_track.py new file mode 100644 index 00000000..81b88c95 --- /dev/null +++ b/flashdreams/tests/test_nvenc_track.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for :class:`NVENCVideoTrack` — pre-encoded H.264 packet plumbing. + +The track is a thin queue in front of aiortc's ``RTCRtpSender``. What +matters is that: + +- ``recv()`` returns exactly the enqueued :class:`av.Packet` (aiortc's + ``H264Encoder.pack()`` reads ``payload``, ``pts`` and ``time_base`` + directly), and +- overflow drops the *oldest* packet — a stale I-frame is more useful than + falling behind wall-clock on an interactive stream. +""" + +from __future__ import annotations + +import asyncio +from fractions import Fraction + +import pytest +from aiortc.mediastreams import MediaStreamError +from av.packet import Packet + +pytestmark = pytest.mark.ci_cpu + +from flashdreams.serving.webrtc.media import NVENCVideoTrack + + +def _mk_packet(pts: int, payload: bytes = b"\x00\x00\x00\x01\x67") -> Packet: + pkt = Packet(payload) + pkt.pts = pts + pkt.time_base = Fraction(1, 90_000) + return pkt + + +class TestConstructorValidation: + def test_rejects_bad_fps(self) -> None: + with pytest.raises(ValueError, match="fps"): + NVENCVideoTrack(fps=0, maxsize=8) + + def test_rejects_bad_maxsize(self) -> None: + with pytest.raises(ValueError, match="maxsize"): + NVENCVideoTrack(fps=30, maxsize=0) + + def test_exposes_fps_and_maxsize(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + assert track.fps == 30 + assert track.maxsize == 8 + assert track.qsize() == 0 + assert track.dropped_packets == 0 + + +class TestEnqueue: + def test_enqueue_returns_true_on_open_track(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + assert track.enqueue_encoded_packet_nowait(_mk_packet(0)) is True + assert track.qsize() == 1 + + @pytest.mark.asyncio + async def test_enqueue_returns_false_on_closed_track(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + await track.close() + assert track.enqueue_encoded_packet_nowait(_mk_packet(0)) is False + + +class TestRecv: + @pytest.mark.asyncio + async def test_recv_returns_enqueued_packet_unchanged(self) -> None: + # aiortc's H264Encoder.pack() consumes bytes(packet), packet.pts + # and packet.time_base directly, so recv() must not mutate them. + track = NVENCVideoTrack(fps=30, maxsize=8) + original = _mk_packet(pts=1234, payload=b"\x00\x00\x00\x01\x25\xaa") + track.enqueue_encoded_packet_nowait(original) + got = await asyncio.wait_for(track.recv(), timeout=1.0) + assert bytes(got) == b"\x00\x00\x00\x01\x25\xaa" + assert got.pts == 1234 + assert got.time_base == Fraction(1, 90_000) + + @pytest.mark.asyncio + async def test_recv_on_closed_track_raises_mediastreamerror(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=8) + await track.close() + with pytest.raises(MediaStreamError): + await asyncio.wait_for(track.recv(), timeout=1.0) + + +class TestOverflow: + """When the queue fills up, drop the oldest packet, not the newest. + + Real-time streams tolerate a stale I-frame worse than they tolerate + stale motion, and aiortc will re-request a keyframe via PLI/FIR if + the oldest dropped packet was a keyframe. + """ + + def test_full_queue_drops_oldest(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=2) + p0 = _mk_packet(pts=0) + p1 = _mk_packet(pts=1) + p2 = _mk_packet(pts=2) + assert track.enqueue_encoded_packet_nowait(p0) is True + assert track.enqueue_encoded_packet_nowait(p1) is True + # This must drop p0, not refuse p2. + assert track.enqueue_encoded_packet_nowait(p2) is True + assert track.qsize() == 2 + assert track.dropped_packets == 1 + + @pytest.mark.asyncio + async def test_overflow_preserves_newest_packet(self) -> None: + track = NVENCVideoTrack(fps=30, maxsize=2) + track.enqueue_encoded_packet_nowait(_mk_packet(pts=0)) + track.enqueue_encoded_packet_nowait(_mk_packet(pts=1)) + track.enqueue_encoded_packet_nowait(_mk_packet(pts=2)) + # Queue should now hold pts=1 and pts=2 (p0 dropped). + first = await asyncio.wait_for(track.recv(), timeout=1.0) + second = await asyncio.wait_for(track.recv(), timeout=1.0) + assert first.pts == 1 + assert second.pts == 2 diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 69bb05ce..15b957a1 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -48,6 +48,19 @@ async def close(self) -> None: self.closed = True +class _FakeVideoEncoder: + """Minimal ``VideoEncoder``-shaped stub for ``ManagedWebRTCSession`` + construction. Enough to satisfy the dataclass field; the manager tests + here do not exercise ``create_track`` / ``deliver_chunk`` on it.""" + + fps = 30 + backend = "fake" + prefers_codec: str | None = None + + def close(self) -> None: + return + + class _FakePeerConnection: def __init__(self) -> None: self.closed = False @@ -101,6 +114,7 @@ def _managed_session( managed = ManagedWebRTCSession( runtime=runtime, video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=peer, resampler=_FakeResampler(), # ty:ignore[invalid-argument-type] control_channel=channel, diff --git a/integrations/lingbot/tests/test_webrtc_runtime.py b/integrations/lingbot/tests/test_webrtc_runtime.py index 5fc4f32b..971e4612 100644 --- a/integrations/lingbot/tests/test_webrtc_runtime.py +++ b/integrations/lingbot/tests/test_webrtc_runtime.py @@ -52,6 +52,19 @@ def send(self, payload: str) -> None: self.messages.append(decoded) +class _FakeVideoEncoder: + """Minimal ``VideoEncoder``-shaped stub for ``_ManagedLingbotSession`` + construction. Enough to satisfy the dataclass field; the tests here do + not exercise ``create_track`` / ``deliver_chunk`` on it.""" + + fps = 30 + backend = "fake" + prefers_codec: str | None = None + + def close(self) -> None: + return + + def _fake_runtime_factory(config: LingbotRuntimeConfig) -> object: del config return object() @@ -932,6 +945,7 @@ async def test_heartbeat_message_refreshes_client_liveness( managed_session = session._ManagedLingbotSession( runtime=object(), video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=_FakeCloseable(), resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), @@ -962,6 +976,7 @@ async def test_client_liveness_timeout_closes_active_session( managed_session = session._ManagedLingbotSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] last_client_message_at=asyncio.get_running_loop().time() - 1.0, @@ -993,6 +1008,7 @@ async def test_disconnect_message_closes_active_session( managed_session = session._ManagedLingbotSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), # ty:ignore[invalid-argument-type] peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), diff --git a/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py b/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py index 68253f1d..b4e9e8c3 100644 --- a/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py +++ b/integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py @@ -123,7 +123,7 @@ def main(): import subprocess import tempfile - import PyNvVideoCodec as nvc # ty:ignore[unresolved-import] + import PyNvVideoCodec as nvc output_path = os.path.join( os.path.dirname(__file__), f"../_images/mirror_augmented_bev.mp4" diff --git a/integrations/omnidreams/omnidreams/webrtc/server.py b/integrations/omnidreams/omnidreams/webrtc/server.py index a36fda7a..df0a6b00 100644 --- a/integrations/omnidreams/omnidreams/webrtc/server.py +++ b/integrations/omnidreams/omnidreams/webrtc/server.py @@ -106,6 +106,19 @@ def parse_args() -> argparse.Namespace: type=str, default="camera_front_wide_120fov", ) + parser.add_argument( + "--prefer_sw_encoder", + action="store_true", + help=( + "Prefer the FFmpeg software encoder (aiortc) over the " + "hardware encoder (PyNvVideoCodec/NVENC H.264). Useful on " + "hosts where NVENC is unavailable or misbehaving, and for " + "A/B profiling against the hardware path. Without this flag " + "the encoder is auto-selected at startup: NVENC when the " + "driver reports support at the target resolution, aiortc's " + "software encoder otherwise." + ), + ) return parser.parse_args() @@ -157,6 +170,7 @@ def build_runtime_config( warmup_chunks=args.warmup_chunks, warmup_timeout_s=args.warmup_timeout_s, debug_serve_hdmaps=args.debug_serve_hdmaps, + encoder_backend="default" if args.prefer_sw_encoder else "auto", ) diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index b307c7d2..5852cb3f 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -54,6 +54,11 @@ CameraPoseIntegrator, PoseSegment, ) +from flashdreams.serving.webrtc.encoders import ( + EncoderBackend, + VideoEncoder, + select_encoder, +) from flashdreams.serving.webrtc.manager import ( DEFAULT_CLIENT_LIVENESS_TIMEOUT_S, BaseWebRTCSessionManager, @@ -422,6 +427,14 @@ class OmnidreamsRuntimeConfig: warmup_chunks: int = 10 warmup_timeout_s: float = 600.0 debug_serve_hdmaps: bool = False + # Video encoder selection. ``"auto"`` prefers NVENC when the driver + # reports support at the target resolution (Stage-1 probe via + # ``PyNvVideoCodec.GetEncoderCaps``) and falls back to aiortc's + # software encoder otherwise. ``"nvenc"`` fails startup if NVENC + # cannot be initialized. ``"default"`` skips the probe entirely. + encoder_backend: EncoderBackend = "auto" + encoder_bitrate_bps: int = 6_000_000 + encoder_gop: int = 30 class OmnidreamsInferenceRuntime: @@ -453,6 +466,11 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: self._next_timestamp_us: int = 0 self._closed = False self._clipgt_temp_dir: tempfile.TemporaryDirectory[str] | None = None + # Selected once at initialization; the concrete backend is chosen + # by ``select_encoder`` based on ``config.encoder_backend`` and + # the driver's ``GetEncoderCaps`` response at + # ``config.video_width`` / ``config.video_height``. + self._video_encoder: VideoEncoder | None = None # Pin every blocking runtime call to one OS thread: Omnidreams' CUDA # graph capture/replay state is thread-local, so spreading calls across # workers (e.g. asyncio.to_thread) crashes capture after a few chunks. @@ -474,6 +492,15 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: def is_master(self) -> bool: return self.rank == self.MASTER_RANK + @property + def video_encoder(self) -> VideoEncoder: + """Return the encoder selected at :meth:`initialize` time.""" + if self._video_encoder is None: + raise OmnidreamsRuntimeError( + "Video encoder is not initialized; call runtime.initialize() first." + ) + return self._video_encoder + def wait_for_termination(self) -> None: self.rank_coordinator.worker_loop(exit_signal=WebRTCControlSignal.EXIT) @@ -720,11 +747,39 @@ def _initialize_sync(self) -> None: self._initial_ego_pose = scene_data.ego_poses[0].transformation_matrix self._next_timestamp_us = int(scene_data.ego_poses[0].timestamp) self._reset_rollout_sync() + self._initialize_video_encoder_sync() logger.info( "Omnidreams runtime initialization complete in {:.1f}s.", time.perf_counter() - init_t0, ) + def _initialize_video_encoder_sync(self) -> None: + """Select the video encoder for this runtime. + + Runs on the runtime executor thread so any GPU-side probe + (``CreateEncoder``) sees the same CUDA context the model uses. + """ + if self._video_encoder is not None: + self._video_encoder.close() + self._video_encoder = None + device = ( + self._device + if self._device is not None + else _resolve_cuda_device( + self.config.device, + ) + ) + gpu_id = device.index if device.index is not None else 0 + self._video_encoder = select_encoder( + backend=self.config.encoder_backend, + width=self.config.video_width, + height=self.config.video_height, + fps=self.config.fps, + bitrate=self.config.encoder_bitrate_bps, + gpu_id=gpu_id, + gop=self.config.encoder_gop, + ) + def _prepare_clipgt_dir(self, clipgt_dir: Path) -> Path: def _has_prefixed_parquets(path: Path) -> bool: return any(path.glob("*.calibration_estimate.parquet")) @@ -793,6 +848,10 @@ def _close_sync(self) -> None: self._camera_to_rig = None self._initial_ego_pose = None + if self._video_encoder is not None: + self._video_encoder.close() + self._video_encoder = None + if state is not None and wrapper is not None: wrapper.cleanup(state) if wrapper is not None: @@ -879,10 +938,19 @@ def _generate_one_chunk_sync( else: video_chunk = output.rgb_frames + # Preserve the compute-stream sync barrier that the previous + # ``.cpu()`` provided implicitly. The tensor stays on-device — the + # hardware encoder reads it via DLPack (zero-copy) while the + # software encoder path performs its own D2H copy inside + # ``BufferedVideoTrack``'s worker thread. Either way, the model's + # writes must be visible before those readers run. + if self._device is not None and self._device.type == "cuda": + torch.cuda.current_stream(self._device).synchronize() + result = WebRTCStepResult( chunk_index=self.autoregressive_index, num_frames=int(video_chunk.shape[2]), - video_chunk=video_chunk.detach().cpu(), + video_chunk=video_chunk.detach(), stats=None, ) self.autoregressive_index += 1 diff --git a/integrations/omnidreams/pyproject.toml b/integrations/omnidreams/pyproject.toml index 5f796eb7..5b16daa9 100644 --- a/integrations/omnidreams/pyproject.toml +++ b/integrations/omnidreams/pyproject.toml @@ -52,6 +52,15 @@ dependencies = [ "tqdm>=4.67", "transformers>=5.0,<6", "huggingface_hub>=0.24", + # NVENC hardware H.264 encoding for the WebRTC path via + # ``flashdreams.serving.webrtc.encoders.PyNvHardwareEncoder``. Not + # optional: omnidreams already requires CUDA at runtime, so the + # PyNvVideoCodec install adds no extra platform surface. Hosts + # without an NVENC ASIC fall back to aiortc's software encoder + # automatically at session initialization (Stage-1 ``GetEncoderCaps`` + # probe returns unsupported, ``select_encoder`` logs the reason and + # constructs ``DefaultRTCEncoder``). + "PyNvVideoCodec>=2.1,<3", ] [tool.uv.sources] diff --git a/integrations/omnidreams/tests/test_nvenc_smoke.py b/integrations/omnidreams/tests/test_nvenc_smoke.py new file mode 100644 index 00000000..b72742f9 --- /dev/null +++ b/integrations/omnidreams/tests/test_nvenc_smoke.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU smoke test for :class:`PyNvHardwareEncoder`. + +Requires a host with NVENC-capable hardware and ``PyNvVideoCodec`` +installed. Skips cleanly on any other host so the ``ci_gpu`` job runs +elsewhere aren't confused by a hard failure. + +What we assert (kept small on purpose): + +- The capability probe reports supported. +- Encoding a modest CUDA chunk produces at least one packet. +- The first emitted packet is an Annex-B keyframe carrying SPS + PPS + + IDR NAL units. SPS/PPS presence verifies ``repeatspspps=1`` reached + the hardware — without it, receivers cannot lock on. +- Packet ``pts`` / ``time_base`` are set to the RTP 90 kHz clock, which + is what aiortc's ``H264Encoder.pack()`` expects. +""" + +from __future__ import annotations + +from fractions import Fraction + +import pytest +import torch + +pytestmark = pytest.mark.ci_gpu + +from flashdreams.serving.webrtc.encoders import ( # noqa: E402 + _PYNVVIDEOCODEC_AVAILABLE, + PyNvHardwareEncoder, + _payload_contains_nal_type, +) + +_H264_NAL_TYPE_IDR = 5 +_H264_NAL_TYPE_SPS = 7 +_H264_NAL_TYPE_PPS = 8 + + +def _has_annex_b_start_code(payload: bytes) -> bool: + return payload.startswith(b"\x00\x00\x00\x01") or payload.startswith( + b"\x00\x00\x01" + ) + + +@pytest.fixture(scope="module") +def nvenc_available() -> None: + if not _PYNVVIDEOCODEC_AVAILABLE: + pytest.skip("PyNvVideoCodec is not installed on this host") + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available on this host") + supported, reason = PyNvHardwareEncoder.is_supported( + gpu_id=0, + width=512, + height=288, + ) + if not supported: + pytest.skip(f"NVENC H.264 not supported on this host: {reason}") + + +def test_probe_reports_supported(nvenc_available) -> None: + supported, reason = PyNvHardwareEncoder.is_supported( + gpu_id=0, + width=512, + height=288, + ) + assert supported, reason + + +def test_encode_chunk_produces_annex_b_keyframe_with_sps_pps( + nvenc_available, +) -> None: + encoder = PyNvHardwareEncoder( + width=512, + height=288, + fps=30, + bitrate=2_000_000, + gpu_id=0, + gop=15, + ) + try: + # 4-frame CUDA chunk in the omnidreams runtime's output layout. + chunk = torch.randint( + 0, + 255, + (4, 3, 288, 512), + dtype=torch.uint8, + device="cuda", + ) + packets: list = [] + num_frames, num_keyframes, encode_ms = encoder.encode_chunk_sync( + chunk, + force_keyframe=True, + on_packet=packets.append, + ) + assert num_frames == 4 + assert num_keyframes >= 1 + assert encode_ms > 0.0 + assert len(packets) >= 1 + + first = bytes(packets[0]) + assert _has_annex_b_start_code(first), ( + "first packet does not start with an Annex-B start code" + ) + assert _payload_contains_nal_type(first, _H264_NAL_TYPE_IDR), ( + "first packet does not carry an IDR NAL (nal_type=5) despite " + "force_keyframe=True" + ) + assert _payload_contains_nal_type(first, _H264_NAL_TYPE_SPS), ( + "SPS NAL missing from keyframe packet — repeatspspps=1 not " + "reaching the hardware; receivers will fail to lock on" + ) + assert _payload_contains_nal_type(first, _H264_NAL_TYPE_PPS), ( + "PPS NAL missing from keyframe packet — repeatspspps=1 not " + "reaching the hardware; receivers will fail to lock on" + ) + + # PTS uses the RTP 90 kHz clock so aiortc's H264Encoder.pack() + # rescales cleanly. First packet's pts must be 0. + assert packets[0].pts == 0 + assert packets[0].time_base == Fraction(1, 90_000) + # Second packet, if produced, must be exactly 90_000/fps ticks later. + if len(packets) >= 2: + assert packets[1].pts == 90_000 // 30 + finally: + encoder.close() diff --git a/integrations/omnidreams/tests/test_webrtc_runtime.py b/integrations/omnidreams/tests/test_webrtc_runtime.py index 39589ff4..2e2b2294 100644 --- a/integrations/omnidreams/tests/test_webrtc_runtime.py +++ b/integrations/omnidreams/tests/test_webrtc_runtime.py @@ -27,7 +27,12 @@ WSAD_SUPPORTED_KEYS, CameraPoseIntegrator, ) +from flashdreams.serving.webrtc.encoders import ( + ChunkDeliveryResult, + DefaultRTCEncoder, +) from flashdreams.serving.webrtc.manager import WebRTCStepResult +from flashdreams.serving.webrtc.media import BufferedVideoTrack pytestmark = pytest.mark.ci_cpu @@ -40,6 +45,51 @@ async def close(self) -> None: self.closed = True +class _FakeVideoEncoder: + """Minimal :class:`VideoEncoder`-shaped stub for the manager tests. + + Wraps a real :class:`BufferedVideoTrack` because the manager attaches + the track to a real :class:`RTCPeerConnection` in the warmup path; + aiortc rejects anything that is not a genuine ``MediaStreamTrack``. + """ + + backend = "fake" + prefers_codec: str | None = None + + def __init__(self, *, fps: int = 30) -> None: + self.fps = fps + self.delivered_chunks: list[Any] = [] + self.closed = False + + def create_track(self, *, maxsize: int) -> BufferedVideoTrack: + return BufferedVideoTrack(fps=self.fps, maxsize=max(1, maxsize)) + + async def deliver_chunk( + self, + chunk: Any, + track: Any, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del force_keyframe + self.delivered_chunks.append(chunk) + # If a real BufferedVideoTrack was provided, thread the chunk + # through its enqueue path so downstream consumers see frames. + if isinstance(track, BufferedVideoTrack): + enqueued = await track.enqueue_chunk(chunk) + else: + enqueued = int(chunk.shape[2]) if chunk.ndim == 6 else int(chunk.shape[0]) + return ChunkDeliveryResult( + backend=self.backend, + num_frames=enqueued, + num_keyframes=0, + encode_ms=0.1, + ) + + def close(self) -> None: + self.closed = True + + def _fake_runtime_factory(config: OmnidreamsRuntimeConfig) -> object: del config return object() @@ -397,6 +447,7 @@ def test_build_runtime_config_threads_hf_scene_args(tmp_path: Path) -> None: warmup_chunks=0, warmup_timeout_s=30.0, debug_serve_hdmaps=True, + prefer_sw_encoder=False, ) cfg = webrtc_server.build_runtime_config(args, device_override="cuda:7") @@ -408,6 +459,43 @@ def test_build_runtime_config_threads_hf_scene_args(tmp_path: Path) -> None: assert cfg.video_height == 360 assert cfg.video_width == 640 assert cfg.debug_serve_hdmaps is True + # ``--prefer_sw_encoder`` unset maps to the ``auto`` backend, which + # still probes NVENC and only falls back to software when the driver + # reports it unsupported. + assert cfg.encoder_backend == "auto" + + +@pytest.mark.parametrize( + "prefer_sw_encoder, expected_backend", + [(False, "auto"), (True, "default")], +) +def test_build_runtime_config_maps_prefer_sw_encoder_to_backend( + tmp_path: Path, + prefer_sw_encoder: bool, + expected_backend: str, +) -> None: + """--prefer_sw_encoder is the single CLI switch that toggles between + the auto-probe path and the forced-software path. Any regression in + this mapping would silently disable the hardware encoder (or worse, + fail to disable it when explicitly asked).""" + args = argparse.Namespace( + pipeline_config_name="omnidreams-sv-2steps-chunk2-loc6-lightvae-lighttae-perf", + scene_dir=tmp_path / "local-scene", + scene_uuid=None, + scene_variant="default", + seed=1, + device="cuda:0", + video_height=360, + video_width=640, + fps=24, + camera_name="camera_front_wide_120fov", + warmup_chunks=0, + warmup_timeout_s=30.0, + debug_serve_hdmaps=False, + prefer_sw_encoder=prefer_sw_encoder, + ) + cfg = webrtc_server.build_runtime_config(args) + assert cfg.encoder_backend == expected_backend def test_parse_args_omits_scene_dir_by_default( @@ -499,6 +587,7 @@ def test_build_runtime_config_clears_scene_uuid_for_local_scene(tmp_path: Path) warmup_chunks=0, warmup_timeout_s=30.0, debug_serve_hdmaps=True, + prefer_sw_encoder=False, ) cfg = webrtc_server.build_runtime_config(args) @@ -570,6 +659,9 @@ def __init__(self, config: OmnidreamsRuntimeConfig) -> None: self.generated_segments: list[ list[tuple[float, float, frozenset[str]]] ] = [] + # The manager reads ``runtime.video_encoder`` when it wires the + # peer connection during the warmup loopback session. + self.video_encoder = _FakeVideoEncoder(fps=config.fps) async def initialize(self) -> None: self.initialize_calls += 1 @@ -638,6 +730,7 @@ async def test_heartbeat_message_refreshes_client_liveness( managed_session = session._ManagedOmnidreamsSession( runtime=object(), video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=_FakeCloseable(), resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), @@ -668,6 +761,7 @@ async def test_client_liveness_timeout_closes_active_session( managed_session = session._ManagedOmnidreamsSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] last_client_message_at=asyncio.get_running_loop().time() - 1.0, @@ -699,6 +793,7 @@ async def test_disconnect_message_closes_active_session( managed_session = session._ManagedOmnidreamsSession( runtime=object(), video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=peer_connection, resampler=object(), # ty:ignore[invalid-argument-type] control_channel=object(), @@ -783,6 +878,7 @@ def send(self, message: str) -> None: managed_session = session._ManagedOmnidreamsSession( runtime=runtime, video_track=video_track, # ty:ignore[invalid-argument-type] + video_encoder=_FakeVideoEncoder(), peer_connection=peer_connection, resampler=_FakeResampler(), # ty:ignore[invalid-argument-type] control_channel=control_channel, @@ -803,3 +899,132 @@ def send(self, message: str) -> None: assert video_track.closed assert peer_connection.closed assert len(control_channel.messages) == 1 + + +class _HardwareEncoderStub: + """A stand-in that ``_enforce_h264_or_fallback`` should recognize as a + hardware encoder (``prefers_codec == "h264"``) and, when H.264 fails to + negotiate, close and replace with :class:`DefaultRTCEncoder`.""" + + backend = "pynvvideocodec" + prefers_codec: str | None = "h264" + + def __init__(self, *, fps: int = 30) -> None: + self.fps = fps + self.closed = False + + def create_track(self, *, maxsize: int) -> Any: + del maxsize + return _FakeCloseable() + + async def deliver_chunk( + self, + chunk: Any, + track: Any, + *, + force_keyframe: bool = False, + ) -> ChunkDeliveryResult: + del chunk, track, force_keyframe + return ChunkDeliveryResult( + backend=self.backend, + num_frames=0, + num_keyframes=0, + encode_ms=0.0, + ) + + def close(self) -> None: + self.closed = True + + +@dataclass +class _FakeSdpCodec: + mimeType: str + + +class _FakeSender: + def __init__(self) -> None: + self.replaced_with: Any = None + + def replaceTrack(self, track: Any) -> None: + self.replaced_with = track + + +class _FakeTransceiver: + def __init__(self, negotiated: list[_FakeSdpCodec]) -> None: + self._codecs = negotiated + self.sender = _FakeSender() + + +def _sdp_fallback_managed_session( + hw_encoder: _HardwareEncoderStub, +) -> session._ManagedOmnidreamsSession: + return session._ManagedOmnidreamsSession( + runtime=object(), + video_track=_FakeCloseable(), # ty:ignore[invalid-argument-type] + video_encoder=hw_encoder, + peer_connection=_FakeCloseable(), + resampler=object(), # ty:ignore[invalid-argument-type] + ) + + +def test_enforce_h264_or_fallback_swaps_when_negotiation_lands_on_non_h264() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + original_track = _FakeCloseable() + managed_session = _sdp_fallback_managed_session(hw_encoder) + managed_session.video_track = original_track # ty:ignore[invalid-assignment] + transceiver = _FakeTransceiver([_FakeSdpCodec(mimeType="video/VP8")]) + + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert hw_encoder.closed, "hardware encoder should be closed on fallback" + assert isinstance(managed_session.video_encoder, DefaultRTCEncoder) + assert isinstance(managed_session.video_track, BufferedVideoTrack) + assert transceiver.sender.replaced_with is managed_session.video_track + + +def test_enforce_h264_or_fallback_keeps_hardware_when_h264_negotiated() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + original_track = _FakeCloseable() + managed_session = _sdp_fallback_managed_session(hw_encoder) + managed_session.video_track = original_track # ty:ignore[invalid-assignment] + transceiver = _FakeTransceiver([_FakeSdpCodec(mimeType="video/H264")]) + + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert not hw_encoder.closed + assert managed_session.video_encoder is hw_encoder + assert managed_session.video_track is original_track + assert transceiver.sender.replaced_with is None + + +def test_enforce_h264_or_fallback_swaps_when_no_codecs_negotiated() -> None: + manager = OmnidreamsWebRTCSessionManager( + runtime_config=OmnidreamsRuntimeConfig(device="cpu", warmup_chunks=0), + ) + hw_encoder = _HardwareEncoderStub(fps=30) + managed_session = _sdp_fallback_managed_session(hw_encoder) + transceiver = _FakeTransceiver([]) + + manager._enforce_h264_or_fallback( + transceiver=transceiver, + managed_session=managed_session, + num_frames=4, + ) + + assert hw_encoder.closed + assert isinstance(managed_session.video_encoder, DefaultRTCEncoder) + assert isinstance(managed_session.video_track, BufferedVideoTrack) diff --git a/uv.lock b/uv.lock index 2b0d0321..e34d4569 100644 --- a/uv.lock +++ b/uv.lock @@ -1553,6 +1553,7 @@ requires-dist = [ { name = "opencv-python-headless", specifier = ">=4.5" }, { name = "pillow", specifier = ">=10.0" }, { name = "pyarrow", specifier = ">=16.0" }, + { name = "pynvvideocodec", specifier = ">=2.1,<3" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, { name = "pyvirtualdisplay", marker = "extra == 'dev'", specifier = ">=3.0" },