OmniDreams WebRTC: Add support for GPU-accelerated H.264 (NVENC) encoding#359
OmniDreams WebRTC: Add support for GPU-accelerated H.264 (NVENC) encoding#359ksheth-dev wants to merge 3 commits into
Conversation
Greptile SummaryThis PR moves the OmniDreams WebRTC video encoder from CPU (aiortc software path) to GPU via NVIDIA NVENC (H.264, PyNvVideoCodec 2.1), behind a
Confidence Score: 4/5Safe to merge with one fix: the old NVENCVideoTrack must be closed before being replaced in the H.264 fallback path. The encoder abstraction, codec negotiation, and fallback logic are well-structured and thoroughly tested. The concrete defect is in _enforce_h264_or_fallback: the original NVENCVideoTrack is discarded by overwriting managed_session.video_track without calling close() first, so stop() is never invoked and its asyncio Queue is never drained on every NVENC session where H.264 fails to negotiate. flashdreams/flashdreams/serving/webrtc/manager.py — specifically _enforce_h264_or_fallback where the old video track is replaced without being closed. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant SM as SessionManager
participant SE as select_encoder
participant HW as PyNvHardwareEncoder
participant SW as DefaultRTCEncoder
participant NT as NVENCVideoTrack
participant RTC as RTCPeerConnection
SM->>SE: select_encoder(backend, width, height, fps)
SE->>HW: is_supported() [GetEncoderCaps probe]
alt NVENC supported
HW-->>SE: (True, "")
SE->>HW: CreateEncoder(...)
SM->>HW: create_track(maxsize)
HW-->>SM: NVENCVideoTrack
SM->>RTC: addTransceiver + setCodecPreferences([H264])
else NVENC not supported
HW-->>SE: (False, reason)
SE->>SW: DefaultRTCEncoder(fps)
SM->>SW: create_track(maxsize)
SM->>RTC: addTransceiver(BufferedVideoTrack)
end
Note over SM,RTC: SDP negotiation
SM->>SM: _enforce_h264_or_fallback()
alt H264 not negotiated
SM->>HW: close()
SM->>RTC: replaceTrack(BufferedVideoTrack)
end
loop Generation worker
SM->>SM: runtime.generate_chunk()
SM->>HW: deliver_chunk(tensor, NVENCVideoTrack)
HW->>NT: enqueue_encoded_packet_nowait(Packet)
NT-->>RTC: "recv() -> H264Encoder.pack()"
end
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant SM as SessionManager
participant SE as select_encoder
participant HW as PyNvHardwareEncoder
participant SW as DefaultRTCEncoder
participant NT as NVENCVideoTrack
participant RTC as RTCPeerConnection
SM->>SE: select_encoder(backend, width, height, fps)
SE->>HW: is_supported() [GetEncoderCaps probe]
alt NVENC supported
HW-->>SE: (True, "")
SE->>HW: CreateEncoder(...)
SM->>HW: create_track(maxsize)
HW-->>SM: NVENCVideoTrack
SM->>RTC: addTransceiver + setCodecPreferences([H264])
else NVENC not supported
HW-->>SE: (False, reason)
SE->>SW: DefaultRTCEncoder(fps)
SM->>SW: create_track(maxsize)
SM->>RTC: addTransceiver(BufferedVideoTrack)
end
Note over SM,RTC: SDP negotiation
SM->>SM: _enforce_h264_or_fallback()
alt H264 not negotiated
SM->>HW: close()
SM->>RTC: replaceTrack(BufferedVideoTrack)
end
loop Generation worker
SM->>SM: runtime.generate_chunk()
SM->>HW: deliver_chunk(tensor, NVENCVideoTrack)
HW->>NT: enqueue_encoded_packet_nowait(Packet)
NT-->>RTC: "recv() -> H264Encoder.pack()"
end
Reviews (2): Last reviewed commit: "webrtc: satisfy pre-commit ty checks and..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
Thanks for putting this together. I think this is the right direction, but I would not merge it yet. A few things need to be addressed first:
After those are addressed and the branch has current CI results, I think this is worth a full review. |
|
Agreeing with @jarcherNV request for rebase. I asked an agent for a rebase and I'm seeing issues (the stream freezes after a couple of frames, reporting increasing chunk id but remaining visually frozen). I would delay in-depth review until we have a concrete rebased version. |
|
There were some changes merged that unifies Lingbot and omnidreams WebRTC serving, which caused the conflict. |
|
It seems there are some ordering issues in the chunk generation and encoder
The generation order is sequential, but the packet enqueue order is whichever encode worker finishes/schedules first. |
- VideoEncoder Protocol + ChunkDeliveryResult in
flashdreams/serving/webrtc/encoders.py, with two implementations:
DefaultRTCEncoder (adapter over aiortc's software encoder) and
PyNvHardwareEncoder (NVENC H.264 via PyNvVideoCodec, output as
av.Packet).
- NVENCVideoTrack (media.py) delivers pre-encoded packets on aiortc's
public av.Packet -> H264Encoder.pack() path -- no reach into
aiortc private attributes. Drop-oldest overflow with a bounded
Queue[av.Packet]; recv() paces at 1/fps like BufferedVideoTrack.
- select_encoder() factory with a two-stage capability probe:
Stage 1 GetEncoderCaps (silent fallback to DefaultRTCEncoder on
environmental "not supported"; loud EncoderInitError under
backend='nvenc'); Stage 2 CreateEncoder (hard error, never a
silent fallback -- masking a driver / session-pool / hardware
failure would hide real problems).
- NVENC configured for interactive low-latency H.264: fmt=ABGR
(NV_ENC_BUFFER_FORMAT_ABGR is word-ordered -> little-endian byte
layout [R,G,B,A]; see _chunk_to_abgr_cuda_frames docstring),
preset=P4, ULL tuning, CBR rate control, bf=0, lookahead=0,
repeatspspps=1 (aiortc's H264Encoder.pack does not synthesize SPS
and PPS). Packets carry pts on the RTP 90 kHz clock so
H264Encoder.pack rescales cleanly via convert_timebase.
- PyNvVideoCodec>=2.1,<3 added as a hard dep on
integrations/omnidreams -- omnidreams already requires CUDA at
runtime, so the install matrix is unchanged.
- Startup emits a single INFO log line per session,
'Video encoder ready: backend={pynvvideocodec|aiortc} ...',
giving one grep anchor regardless of which backend was selected.
- ci_cpu tests cover: Protocol conformance, factory branch coverage
(Stage-1 library / caps / bounds failures under both auto and nvenc,
Stage-2 hard-error regression guard under both), pre-encoded packet
plumbing on NVENCVideoTrack (pts / time_base / drop-oldest overflow),
cross-chunk packet-ordering invariant that guards the manager's
sequential-await pattern from a future create_task rewrite, and
compat guards for aiortc + PyNvVideoCodec public surfaces.
- ci_gpu smoke test encodes a 4-frame chunk on real hardware and
asserts Annex-B start code + IDR + SPS + PPS in the first packet,
plus pts=0 and time_base=1/90000.
No runtime path currently consumes the new abstraction; behavior
change lands in the follow-up commit.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
… fallback
- OmnidreamsRuntimeConfig gains encoder_backend
(auto | nvenc | default), encoder_bitrate_bps, encoder_gop.
select_encoder() runs on the runtime executor thread inside
_initialize_video_encoder_sync; the encoder is closed cleanly in
_close_sync so its NVENC session slot is released promptly.
- _generate_one_chunk_sync drops the .cpu() D2H copy on the tensor
and replaces it with torch.cuda.current_stream().synchronize().
Same wall-clock cost as before (both wait for the compute stream to
drain), but the hardware encoder can now read the CUDA tensor
directly via DLPack; the software path picks up the D2H inside
BufferedVideoTrack's worker.
- omnidreams.webrtc.server exposes one new CLI flag,
--prefer_sw_encoder. Defaults false -> encoder_backend='auto'
(probe NVENC and fall back silently to aiortc's software encoder
on Stage-1 unsupported). Setting it maps to encoder_backend='default'
and skips the NVENC probe entirely -- useful for A/B profiling and
known-flaky NVENC hosts. The tri-state config field stays on
OmnidreamsRuntimeConfig for programmatic use and test coverage of
the 'nvenc' loud-on-failure branch.
- Manager wiring: addTransceiver + setCodecPreferences constrain the
SDP to H.264 when the encoder emits pre-encoded packets
(prefers_codec == 'h264'). Post-answer, transceiver._codecs is
inspected; if H.264 did not land (e.g. a browser that will not
offer it), _enforce_h264_or_fallback closes the NVENC session and
installs DefaultRTCEncoder + BufferedVideoTrack via
sender.replaceTrack -- a pre-stream swap, no renegotiation.
- Chunk delivery in the generation worker goes through
encoder.deliver_chunk(...). The await is load-bearing for
cross-chunk packet ordering (guarded by the regression test added
in the previous commit): a create_task rewrite here would allow
chunks to complete out of order and packets to land on the track
with non-monotonic pts.
- Test suite extends _FakeVideoEncoder to satisfy the VideoEncoder
Protocol, threads video_encoder through every ManagedWebRTCSession
construction, and adds ci_cpu tests for _enforce_h264_or_fallback
(swap on VP8-negotiated, keep on H.264-negotiated, swap on empty
codec list) plus a parametrized test that --prefer_sw_encoder
correctly maps to encoder_backend ('default' when set, 'auto' when
unset).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- encoders.py: guard the ``PyNvVideoCodec`` import under ``TYPE_CHECKING`` so ty sees a single ``Any``-typed ``nvc`` name. Without the guard, ty preserved a ``<module PyNvVideoCodec> | None | Any`` union and rejected the ``None`` fallback and every attribute access on ``GetEncoderCaps`` / ``FORCEIDR`` / ``CreateEncoder``. Runtime path (the ``else:`` branch) keeps the same guarded try / except behaviour. - encoders.py: narrow the ``VideoEncoder.create_track`` return type from ``MediaStreamTrack`` to ``BufferedVideoTrack | NVENCVideoTrack`` so the manager can assign the result to the ``ManagedWebRTCSession .video_track`` field without a widening error. - manager.py: change ``ManagedWebRTCSession.video_track`` annotation from ``MediaStreamTrack`` to ``BufferedVideoTrack | NVENCVideoTrack``. The aiortc base type does not advertise ``.close`` / ``.fps`` / ``.qsize``, which the generation worker and liveness watchdog need. Drop the now-unused ``MediaStreamTrack`` import. - test_encoders.py: swap mypy-style ``# type: ignore[misc|arg-type]`` to ty-style ``# ty:ignore[invalid-assignment|invalid-argument-type]`` on the frozen-dataclass and SimpleNamespace test sites. Add ``assert packet.pts is not None`` before ``int(packet.pts)`` in the ordering tests -- ``av.Packet.pts`` is nullable at the type level even though the fake encoder always sets it. - test_webrtc_manager.py: add a local ``_FakeVideoEncoder`` stub and thread ``video_encoder=`` through the ``ManagedWebRTCSession`` construction the shared base test uses. - integrations/lingbot/tests/test_webrtc_runtime.py: same treatment for the three ``_ManagedLingbotSession`` constructions -- the lingbot tests broke when we promoted ``video_encoder`` to a required field on the shared ``ManagedWebRTCSession`` base. - integrations/omnidreams/tests/test_webrtc_runtime.py: drop 8 unused ``# ty:ignore[invalid-argument-type]`` comments ty flagged (5 on ``video_encoder=…``, 3 on ``transceiver=transceiver``). - integrations/omnidreams/ludus-renderer/examples/render_mirror_augmented.py: drop the ``# ty:ignore[unresolved-import]`` on the ``PyNvVideoCodec`` import. The comment was pre-existing but became unused when this feature promoted PyNvVideoCodec from an optional extra to a hard dependency of ``integrations/omnidreams``. - Ruff format touched a few files in passing (test_nvenc_track.py, test_nvenc_smoke.py, test_encoders.py); those changes are formatting-only and included here. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
fa14978 to
f4edab9
Compare
| 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 |
There was a problem hiding this comment.
NVENCVideoTrack orphaned after encoder fallback
When _enforce_h264_or_fallback swaps the encoder, it overwrites managed_session.video_track with the new BufferedVideoTrack but never calls close() on the old NVENCVideoTrack. ManagedWebRTCSession.close() calls await self.video_track.close() — by that point it reaches the fallback track, not the original NVENC track. The old track's stop() (which sets aiortc's readyState to "ended") is never called, leaving the track live, its asyncio Queue unconsumed, and any internal aiortc observer references pinned until the process exits.
|
Thanks for updating this. It looks like it is making good progress. I still see a couple of issues that I think should be addressed before merge:
|
|
The feature changes have been refactored in the latest version of dev/ksheth/pynvvideocodec_v2 branch which is rebased on latest from main branch. With the latest refactored changes, the following above comments are addressed:
|
|
Tested the nvenc code path on RTX6000 pro, and it did improve the omnidreams interactive demo demo in webRTC - from 12 FPS to 14 FPS testing it with and without I think we can get this merged once @jarcherNV 's comments get addressed. The PR is general enough and I think other demos that use webRTC can easily use this nvenc implementation. |
|
Another issue when testing in multi-gpu env - The hardware encoder is initialized on every distributed rank even though only rank 0 serves WebRTC media. In |
4 participants
Summary
Move the OmniDreams interactive-demo WebRTC video encoder off the CPU and onto the
GPU via NVIDIA NVENC (H.264, PyNvVideoCodec 2.1), behind a small
VideoEncoderabstraction that the session picks between at startup:
Summary
Move the OmniDreams interactive-demo WebRTC onto the
GPU via NVIDIA NVENC (H.264, PyNvVideoCodec 2.1), behind a small
VideoEncoderabstraction that the session picks between at startup:
PyNvHardwareEncoder: hardware H.264 viaPyNvVideoCodec. Selected bydefault. Produces pre-encoded
av.Packets that aiortc forwards straight tothe RTP packetizer via a new `EncodedPacke
DefaultRTCVideoEncoder: thin wrapper over aiortc's built-in softwareencoder (raw frames through a
BufferedVideoTrack, aiortc negotiates thecodec via SDP).
Runtime capability probe via
PyNvVideoCodec.GetEncoderCapsdetects environmentswithout NVENC hardware and gracefully falls back to the software backend with
an INFO log -- the server still comes up on non-NVIDIA machines.
New CLI flag
--prefer_sw_encoderlets operators opt out of NVENC explicitlyfor A/B profiling comparisons or on environments where the SW path is preferred.
Measured impact with NVENC
At 1280x704, 30 fps, 10 Mb/s target on a single-GPU host:
enqueue_ms: ~4x reduction (28 ms -> 7 ms)Test plan
uv run --package flashdreams-omnidreams pytest integrations/omnidreams/tests/test_pynvvideocodec_encode.py -v-- encoder unit tests (requires NVENC GPU)uv run --package flashdreams-omnidreams pytest integrations/omnidreams/tests/test_webrtc_codec_prefs.py -v-- SDP H.264 codec preferenceuv run --package flashdreams-omnidreams pytest integrations/omnidreams/tests/test_webrtc_runtime.py -v-- runtime + fake-encoder mocksPyNvVideoCodec H.264 encoder created ...andVideo codec negotiated: video/H264--prefer_sw_encoder; verifyFFmpeg software encoder (aiortc) selectedandVideo codec negotiated: video/VP8