Skip to content

Commit 85cfb13

Browse files
balthazurbalthazur
committed
Tune aiortc WebRTC bitrate ramp
1 parent 8805940 commit 85cfb13

4 files changed

Lines changed: 256 additions & 0 deletions

File tree

‎inference/core/interfaces/stream_manager/manager_app/webrtc.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,16 @@
3737
WebRTCOffer,
3838
WebRTCTURNConfig,
3939
)
40+
from inference.core.interfaces.webrtc_worker.aiortc_bitrate import (
41+
apply_aiortc_bitrate_tuning,
42+
prefer_h264_for_peer_connection,
43+
)
4044
from inference.core.interfaces.webrtc_worker.h264_nvenc import prefer_h264_nvenc_encoder
4145
from inference.core.utils.async_utils import Queue as SyncAsyncQueue
4246
from inference.core.utils.function import experimental
4347
from inference.core.workflows.execution_engine.entities.base import WorkflowImageData
4448

49+
apply_aiortc_bitrate_tuning()
4550
prefer_h264_nvenc_encoder()
4651
logging.getLogger("aiortc").setLevel(logging.WARNING)
4752

@@ -476,6 +481,7 @@ async def on_connectionstatechange():
476481
await peer_connection.setRemoteDescription(
477482
RTCSessionDescription(sdp=webrtc_offer.sdp, type=webrtc_offer.type)
478483
)
484+
prefer_h264_for_peer_connection(peer_connection)
479485
answer = await peer_connection.createAnswer()
480486
await peer_connection.setLocalDescription(answer)
481487
logger.debug(f"WebRTC connection status: {peer_connection.connectionState}")

‎inference/core/interfaces/webrtc_worker/aiortc_bitrate.py‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
import time
2+
from typing import Any, Dict, Mapping, Optional
3+
4+
from inference.core import logger
5+
6+
H264_BITRATE_LIMITS_BPS = {
7+
"DEFAULT_BITRATE": 6_000_000,
8+
"MIN_BITRATE": 2_000_000,
9+
"MAX_BITRATE": 20_000_000,
10+
}
11+
12+
VP8_BITRATE_LIMITS_BPS = {
13+
"DEFAULT_BITRATE": 4_000_000,
14+
"MIN_BITRATE": 1_500_000,
15+
"MAX_BITRATE": 8_000_000,
16+
}
17+
18+
REMB_START_BITRATE_BPS = 6_000_000
19+
REMB_MIN_BITRATE_BPS = 2_000_000
20+
REMB_MAX_BITRATE_BPS = 20_000_000
21+
22+
_BITRATE_TUNING_APPLIED = False
23+
24+
25+
def apply_aiortc_bitrate_tuning() -> None:
26+
"""Apply WebRTC bitrate defaults before aiortc peer connections are created."""
27+
global _BITRATE_TUNING_APPLIED
28+
29+
if _BITRATE_TUNING_APPLIED:
30+
return
31+
32+
import aiortc.codecs.h264 as h264
33+
import aiortc.codecs.vpx as vpx
34+
35+
h264_limits = _apply_codec_bitrate_limits(h264, H264_BITRATE_LIMITS_BPS)
36+
vp8_limits = _apply_codec_bitrate_limits(vpx, VP8_BITRATE_LIMITS_BPS)
37+
_patch_receiver_remb_estimator()
38+
_BITRATE_TUNING_APPLIED = True
39+
40+
logger.info(
41+
"[WEBRTC_BITRATE] applied aiortc bitrate tuning h264=%s vp8=%s "
42+
"remb_start_bps=%s remb_min_bps=%s remb_max_bps=%s",
43+
h264_limits,
44+
vp8_limits,
45+
REMB_START_BITRATE_BPS,
46+
REMB_MIN_BITRATE_BPS,
47+
REMB_MAX_BITRATE_BPS,
48+
)
49+
50+
51+
def prefer_h264_for_peer_connection(peer_connection: Any) -> bool:
52+
"""Prefer H.264 when the browser offers it, keeping other codecs as fallback."""
53+
from aiortc import RTCRtpSender
54+
55+
codecs = RTCRtpSender.getCapabilities("video").codecs
56+
h264_codecs = [codec for codec in codecs if codec.mimeType.lower() == "video/h264"]
57+
if not h264_codecs:
58+
logger.info("[WEBRTC_BITRATE] H.264 codec preference unavailable")
59+
return False
60+
61+
preferred_codecs = h264_codecs + [
62+
codec for codec in codecs if codec not in h264_codecs
63+
]
64+
applied = False
65+
66+
for transceiver in peer_connection.getTransceivers():
67+
if transceiver.kind != "video":
68+
continue
69+
transceiver.setCodecPreferences(preferred_codecs)
70+
applied = True
71+
72+
if applied:
73+
logger.info("[WEBRTC_BITRATE] preferred H.264 for video transceivers")
74+
return applied
75+
76+
77+
def _apply_codec_bitrate_limits(
78+
module: Any,
79+
bitrate_limits: Mapping[str, int],
80+
) -> Dict[str, int]:
81+
applied = {}
82+
for constant_name, value in bitrate_limits.items():
83+
setattr(module, constant_name, value)
84+
applied[constant_name] = value
85+
return applied
86+
87+
88+
def _patch_receiver_remb_estimator(rate_module: Optional[Any] = None) -> None:
89+
if rate_module is None:
90+
import aiortc.rate as rate_module
91+
92+
estimator_cls = rate_module.RemoteBitrateEstimator
93+
if getattr(estimator_cls, "_roboflow_remb_tuning_patched", False):
94+
return
95+
96+
original_init = estimator_cls.__init__
97+
original_add = estimator_cls.add
98+
99+
def __init__(self: Any) -> None:
100+
original_init(self)
101+
now_ms = int(time.time() * 1000)
102+
self.rate_control.set_estimate(REMB_START_BITRATE_BPS, now_ms)
103+
self.rate_control.latest_estimated_throughput = REMB_START_BITRATE_BPS
104+
105+
def add(
106+
self: Any,
107+
arrival_time_ms: int,
108+
abs_send_time: int,
109+
payload_size: int,
110+
ssrc: int,
111+
) -> Any:
112+
remb = original_add(
113+
self,
114+
arrival_time_ms,
115+
abs_send_time,
116+
payload_size,
117+
ssrc,
118+
)
119+
if remb is None:
120+
return None
121+
122+
requested_bitrate, ssrcs = remb
123+
applied_bitrate = _clamp_bitrate(
124+
requested_bitrate,
125+
min_bitrate=REMB_MIN_BITRATE_BPS,
126+
max_bitrate=REMB_MAX_BITRATE_BPS,
127+
)
128+
if applied_bitrate != requested_bitrate:
129+
self.rate_control.current_bitrate = applied_bitrate
130+
return applied_bitrate, ssrcs
131+
132+
estimator_cls.__init__ = __init__
133+
estimator_cls.add = add
134+
estimator_cls._roboflow_remb_tuning_patched = True
135+
136+
137+
def _clamp_bitrate(
138+
bitrate: int,
139+
min_bitrate: int,
140+
max_bitrate: int,
141+
) -> int:
142+
return min(max(bitrate, min_bitrate), max_bitrate)

‎inference/core/interfaces/webrtc_worker/webrtc.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,10 @@
4949
WebRTCData,
5050
WorkflowConfiguration,
5151
)
52+
from inference.core.interfaces.webrtc_worker.aiortc_bitrate import (
53+
apply_aiortc_bitrate_tuning,
54+
prefer_h264_for_peer_connection,
55+
)
5256
from inference.core.interfaces.webrtc_worker.entities import (
5357
DataOutputMode,
5458
StreamOutputMode,
@@ -78,6 +82,7 @@
7882
from inference.core.workflows.execution_engine.entities.base import WorkflowImageData
7983
from inference.usage_tracking.collector import usage_collector
8084

85+
apply_aiortc_bitrate_tuning()
8186
prefer_h264_nvenc_encoder()
8287
logging.getLogger("aiortc").setLevel(logging.WARNING)
8388

@@ -1301,6 +1306,7 @@ def on_message(message):
13011306
sdp=webrtc_request.webrtc_offer.sdp, type=webrtc_request.webrtc_offer.type
13021307
)
13031308
)
1309+
prefer_h264_for_peer_connection(peer_connection)
13041310
answer = await peer_connection.createAnswer()
13051311
await peer_connection.setLocalDescription(answer)
13061312

‎tests/inference/unit_tests/core/interfaces/test_aiortc_bitrate.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import sys
2+
from types import SimpleNamespace
3+
4+
from inference.core.interfaces.webrtc_worker.aiortc_bitrate import (
5+
H264_BITRATE_LIMITS_BPS,
6+
REMB_MAX_BITRATE_BPS,
7+
REMB_MIN_BITRATE_BPS,
8+
REMB_START_BITRATE_BPS,
9+
_apply_codec_bitrate_limits,
10+
_clamp_bitrate,
11+
_patch_receiver_remb_estimator,
12+
prefer_h264_for_peer_connection,
13+
)
14+
15+
16+
def test_apply_codec_bitrate_limits_sets_aiortc_constants():
17+
module = SimpleNamespace(
18+
DEFAULT_BITRATE=1,
19+
MIN_BITRATE=1,
20+
MAX_BITRATE=1,
21+
)
22+
23+
applied = _apply_codec_bitrate_limits(module, H264_BITRATE_LIMITS_BPS)
24+
25+
assert applied == H264_BITRATE_LIMITS_BPS
26+
assert module.DEFAULT_BITRATE == H264_BITRATE_LIMITS_BPS["DEFAULT_BITRATE"]
27+
assert module.MIN_BITRATE == H264_BITRATE_LIMITS_BPS["MIN_BITRATE"]
28+
assert module.MAX_BITRATE == H264_BITRATE_LIMITS_BPS["MAX_BITRATE"]
29+
30+
31+
def test_clamp_bitrate_applies_min_and_max():
32+
assert (
33+
_clamp_bitrate(
34+
REMB_MIN_BITRATE_BPS - 1,
35+
min_bitrate=REMB_MIN_BITRATE_BPS,
36+
max_bitrate=REMB_MAX_BITRATE_BPS,
37+
)
38+
== REMB_MIN_BITRATE_BPS
39+
)
40+
assert (
41+
_clamp_bitrate(
42+
REMB_MAX_BITRATE_BPS + 1,
43+
min_bitrate=REMB_MIN_BITRATE_BPS,
44+
max_bitrate=REMB_MAX_BITRATE_BPS,
45+
)
46+
== REMB_MAX_BITRATE_BPS
47+
)
48+
49+
50+
def test_receiver_remb_estimator_is_seeded_and_clamped():
51+
class FakeRateControl:
52+
def __init__(self):
53+
self.current_bitrate = None
54+
self.latest_estimated_throughput = None
55+
56+
def set_estimate(self, bitrate, now_ms):
57+
self.current_bitrate = bitrate
58+
59+
class FakeRemoteBitrateEstimator:
60+
def __init__(self):
61+
self.rate_control = FakeRateControl()
62+
63+
def add(self, arrival_time_ms, abs_send_time, payload_size, ssrc):
64+
return REMB_MIN_BITRATE_BPS - 1, [ssrc]
65+
66+
rate_module = SimpleNamespace(RemoteBitrateEstimator=FakeRemoteBitrateEstimator)
67+
68+
_patch_receiver_remb_estimator(rate_module)
69+
estimator = rate_module.RemoteBitrateEstimator()
70+
71+
assert estimator.rate_control.current_bitrate == REMB_START_BITRATE_BPS
72+
assert estimator.rate_control.latest_estimated_throughput == REMB_START_BITRATE_BPS
73+
assert estimator.add(1, 1, 1, 1234) == (REMB_MIN_BITRATE_BPS, [1234])
74+
assert estimator.rate_control.current_bitrate == REMB_MIN_BITRATE_BPS
75+
76+
77+
def test_prefer_h264_for_peer_connection_keeps_other_codecs_as_fallback(monkeypatch):
78+
h264 = SimpleNamespace(mimeType="video/H264")
79+
vp8 = SimpleNamespace(mimeType="video/VP8")
80+
av1 = SimpleNamespace(mimeType="video/AV1")
81+
82+
class FakeSender:
83+
@staticmethod
84+
def getCapabilities(kind):
85+
assert kind == "video"
86+
return SimpleNamespace(codecs=[vp8, h264, av1])
87+
88+
class FakeTransceiver:
89+
kind = "video"
90+
91+
def __init__(self):
92+
self.preferred_codecs = None
93+
94+
def setCodecPreferences(self, codecs):
95+
self.preferred_codecs = codecs
96+
97+
transceiver = FakeTransceiver()
98+
peer_connection = SimpleNamespace(getTransceivers=lambda: [transceiver])
99+
monkeypatch.setitem(sys.modules, "aiortc", SimpleNamespace(RTCRtpSender=FakeSender))
100+
101+
assert prefer_h264_for_peer_connection(peer_connection) is True
102+
assert transceiver.preferred_codecs == [h264, vp8, av1]

0 commit comments

Comments
 (0)