From 79d999f412e144449ac71fd18ce4a05f7607a718 Mon Sep 17 00:00:00 2001 From: shobhitagnihotri69 Date: Mon, 21 Sep 2026 20:13:40 +0530 Subject: [PATCH] fix(video): stream canonical remux with bounded memory (#580) --- src/hflow/episode.py | 31 ++++-- src/hflow/video.py | 195 ++++++++++++++++++++++++---------- tests/test_video.py | 91 ++++++++++++++-- tests/test_video_streaming.py | 183 +++++++++++++++++++++++++++++++ 4 files changed, 427 insertions(+), 73 deletions(-) create mode 100644 tests/test_video_streaming.py diff --git a/src/hflow/episode.py b/src/hflow/episode.py index 5ed798e5..a78ecb7b 100644 --- a/src/hflow/episode.py +++ b/src/hflow/episode.py @@ -556,24 +556,33 @@ def video(self, camera: str | None = None) -> Path: "Transform it first (hflow.write_canonical_episode) or read the " "raw messages yourself via ep.channel()/the mcap package." ) - channel = self.channel(topic) - fps = video_module.estimate_fps_from_log_times(channel.timestamps.tolist(), topic=topic) - self._video_fps[topic] = fps output = self.workdir / f"{_sanitize_topic(topic)}.mp4" + if topic in self._video_fps and output.exists(): + return output + fps = video_module.estimate_fps_from_streaming_log_times( + ( + int(timestamp) + for batch in self._reader.iter_batches( + topics=[topic], channel_ids=[info.channel_id] + ) + for timestamp in batch.log_times + ), + topic=topic, + ) + self._video_fps[topic] = fps if output.exists(): # Sound because write_access_units_to_mp4 replaces atomically: a # file at the final path is always a completed remux. return output def validated_access_units() -> "Iterator[bytes]": - # Stream-decode instead of channel.messages: caching a decoded - # copy of every video payload would double the episode's memory. - for message in channel.iter_decoded(): - if message.format != "h264": - raise ValueError( - f"camera {topic!r} carries {message.format!r}, expected 'h264'" - ) - yield message.data + for batch in self.iter_decoded_batches(topics=[topic], channel_ids=[info.channel_id]): + for message in batch.messages: + if message.format != "h264": + raise ValueError( + f"camera {topic!r} carries {message.format!r}, expected 'h264'" + ) + yield message.data return video_module.write_access_units_to_mp4( validated_access_units(), diff --git a/src/hflow/video.py b/src/hflow/video.py index b944ba95..bf433bfc 100644 --- a/src/hflow/video.py +++ b/src/hflow/video.py @@ -17,10 +17,12 @@ """ import itertools +import sqlite3 import statistics import subprocess import tempfile from collections.abc import Iterable, Iterator, Sequence +from contextlib import closing, suppress from dataclasses import dataclass from pathlib import Path from typing import NamedTuple @@ -92,6 +94,46 @@ def estimate_fps_from_log_times(log_times_ns: Sequence[int], *, topic: str) -> f return 1e9 / median_delta_ns +def estimate_fps_from_streaming_log_times(log_times_ns: Iterable[int], *, topic: str) -> float: + """Infer the same exact median rate without retaining all timestamps in RAM. + + Keep intervals in a temporary disk index with a bounded SQLite page cache. + Selecting its middle one or two entries needs no in-memory sort. Only + timestamps are consumed; callers can release each raw video batch in turn. + """ + with ( + tempfile.TemporaryDirectory(prefix="hflow-video-fps-") as directory, + closing(sqlite3.connect(Path(directory) / "intervals.sqlite")) as connection, + ): + connection.execute("PRAGMA cache_size = -1024") + connection.execute("CREATE TABLE intervals (delta INTEGER NOT NULL)") + connection.execute("CREATE INDEX ordered_intervals ON intervals (delta)") + timestamp_count = 0 + + def intervals() -> Iterator[tuple[int]]: + nonlocal timestamp_count + previous: int | None = None + for timestamp in log_times_ns: + timestamp_count += 1 + if previous is not None: + yield (timestamp - previous,) + previous = timestamp + + connection.executemany("INSERT INTO intervals VALUES (?)", intervals()) + if timestamp_count < 2: + # Reuse the existing error contract for empty/single-frame streams. + return estimate_fps_from_log_times([0] * timestamp_count, topic=topic) + interval_count = timestamp_count - 1 + middle = connection.execute( + "SELECT delta FROM intervals ORDER BY delta LIMIT ? OFFSET ?", + (2 - interval_count % 2, (interval_count - 1) // 2), + ).fetchall() + median_delta_ns = statistics.median(row[0] for row in middle) + if median_delta_ns <= 0: + return estimate_fps_from_log_times([0, 0], topic=topic) + return 1e9 / median_delta_ns + + def source_log_times_for_sampled_frames( source_timestamps_ns: Sequence[int], *, @@ -466,39 +508,52 @@ def scan_picture_coding_types(stream: bytes) -> PictureCodingScan: fail-closed as an error: an uncountable slice means the stream's B-frame freedom cannot be proven. """ - picture_is_b: list[bool] = [] - for slice_header in _iter_slice_headers(stream): - if slice_header.first_mb_in_slice is None or slice_header.slice_type is None: - raise ValueError( - "a slice header is incomplete or truncated; picture coding types " - "cannot be classified" - ) - slice_is_b = slice_header.slice_type % 5 == 1 - if slice_header.first_mb_in_slice == 0: - picture_is_b.append(slice_is_b) - elif picture_is_b: - picture_is_b[-1] = picture_is_b[-1] or slice_is_b - else: - raise ValueError( - "a slice header claims first_mb_in_slice > 0 before any picture starts; " - "picture coding types cannot be classified" - ) - reorder_depth = 0 - current_run = 0 - for picture_is_b_flag in picture_is_b: - current_run = current_run + 1 if picture_is_b_flag else 0 - reorder_depth = max(reorder_depth, current_run) - trailing_b_pictures = 0 - for picture_is_b_flag in reversed(picture_is_b): - if not picture_is_b_flag: - break - trailing_b_pictures += 1 - return PictureCodingScan( - picture_count=len(picture_is_b), - b_picture_count=sum(picture_is_b), - reorder_depth=reorder_depth, - trailing_b_pictures=trailing_b_pictures, - ) + scanner = _PictureCodingScanner() + scanner.update(stream) + return scanner.result() + + +@dataclass +class _PictureCodingScanner: + """Accumulate the existing slice-header evidence across access units.""" + + picture_count: int = 0 + b_picture_count: int = 0 + reorder_depth: int = 0 + trailing_b_pictures: int = 0 + _picture_is_b: bool | None = None + _preceding_b_run: int = 0 + + def update(self, access_unit: bytes) -> None: + for slice_header in _iter_slice_headers(access_unit): + if slice_header.first_mb_in_slice is None or slice_header.slice_type is None: + raise ValueError( + "a slice header is incomplete or truncated; picture coding types " + "cannot be classified" + ) + if slice_header.first_mb_in_slice == 0: + self.picture_count += 1 + self._preceding_b_run = self.trailing_b_pictures + self.trailing_b_pictures = 0 + self._picture_is_b = False + elif self._picture_is_b is None: + raise ValueError( + "a slice header claims first_mb_in_slice > 0 before any picture starts; " + "picture coding types cannot be classified" + ) + if slice_header.slice_type % 5 == 1 and not self._picture_is_b: + self._picture_is_b = True + self.b_picture_count += 1 + self.trailing_b_pictures = self._preceding_b_run + 1 + self.reorder_depth = max(self.reorder_depth, self.trailing_b_pictures) + + def result(self) -> PictureCodingScan: + return PictureCodingScan( + picture_count=self.picture_count, + b_picture_count=self.b_picture_count, + reorder_depth=self.reorder_depth, + trailing_b_pictures=self.trailing_b_pictures, + ) def ensure_access_unit_delimiter(access_unit: bytes) -> bytes: @@ -599,7 +654,7 @@ def write_access_units_to_mp4( ) -> Path: """Losslessly remux H.264 access units into an MP4 file (no re-encode). - Concatenates the access units to a raw Annex B stream and remuxes with + Streams complete access units to ffmpeg stdin and remuxes with ``ffmpeg -r {fps} -f h264 -i - -c:v copy -movflags +faststart``. The resulting file plays in anything; frame timing is constant-rate ``fps`` (callers needing exact per-frame log times use the message timestamps). @@ -610,18 +665,7 @@ def write_access_units_to_mp4( decoded). Canonical video requires ``bframes=0`` (docs/FORMAT.md, "The H.264 bitstream constraints"). """ - annex_b_stream = b"".join(units) - coding_types = scan_picture_coding_types(annex_b_stream) - if coding_types.b_picture_count: - raise ValueError( - f"cannot remux to MP4 without dropping frames: the stream carries " - f"{coding_types.b_picture_count} B picture(s) across " - f"{coding_types.picture_count} (reorder depth " - f"{coding_types.reorder_depth}); a -c:v copy remux drops the reorder " - f"tail, putting the last {coding_types.trailing_b_pictures} frame(s) at " - "risk (measured 301 of 303 in #250). Canonical video requires " - "bframes=0; re-encode upstream -- see docs/FORMAT.md item 4" - ) + executable = ffmpeg_path() # Write to a unique sibling temp path and replace atomically: callers cache # on bare file existence, so the final path must never hold a partial MP4. # A fixed ``.tmp`` name lets concurrent remuxes truncate/unlink the @@ -631,7 +675,7 @@ def write_access_units_to_mp4( ) as temp_file: temporary_output = Path(temp_file.name) command: list[str] = [ - str(ffmpeg_path()), + str(executable), "-hide_banner", "-loglevel", "error", @@ -654,18 +698,59 @@ def write_access_units_to_mp4( str(temporary_output), ] try: - completed = subprocess.run(command, input=annex_b_stream, capture_output=True) - if completed.returncode != 0: - raise VideoEncodeError( - f"ffmpeg remux failed (exit {completed.returncode}): " - f"{_stderr_tail(completed.stderr)}" + # File-backed stderr cannot fill a pipe while stdin is being written, + # and diagnostics never accumulate in Python memory. + with tempfile.TemporaryFile() as stderr: + process = subprocess.Popen( + command, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=stderr ) + assert process.stdin is not None + try: + scanner = _PictureCodingScanner() + pipe_broken = False + for unit in units: + scanner.update(unit) + if not pipe_broken: + try: + process.stdin.write(unit) + process.stdin.flush() + except BrokenPipeError: + # Complete validation even if ffmpeg exits early: + # malformed input/B-frame errors retain precedence. + pipe_broken = True + try: + process.stdin.close() + except BrokenPipeError: + pipe_broken = True + coding_types = scanner.result() + if coding_types.b_picture_count: + raise ValueError( + f"cannot remux to MP4 without dropping frames: the stream carries " + f"{coding_types.b_picture_count} B picture(s) across " + f"{coding_types.picture_count} (reorder depth " + f"{coding_types.reorder_depth}); a -c:v copy remux drops the reorder " + f"tail, putting the last {coding_types.trailing_b_pictures} frame(s) at " + "risk (measured 301 of 303 in #250). Canonical video requires " + "bframes=0; re-encode upstream -- see docs/FORMAT.md item 4" + ) + returncode = process.wait() + if returncode != 0 or pipe_broken: + stderr.seek(0, 2) + stderr.seek(max(0, stderr.tell() - 4 * _STDERR_TAIL_CHARACTER_LIMIT)) + raise VideoEncodeError( + f"ffmpeg remux failed (exit {returncode}): {_stderr_tail(stderr.read())}" + ) + finally: + if process.poll() is None: + process.kill() + process.wait() + with suppress(BrokenPipeError): + process.stdin.close() if not temporary_output.is_file() or temporary_output.stat().st_size == 0: raise VideoEncodeError(f"ffmpeg remux exited 0 but produced no output at {output}") - except BaseException: + temporary_output.replace(output) + finally: temporary_output.unlink(missing_ok=True) - raise - temporary_output.replace(output) return output diff --git a/tests/test_video.py b/tests/test_video.py index 10825199..3b847ac7 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -2,10 +2,15 @@ import subprocess import time +from collections.abc import Iterator from pathlib import Path +from typing import Any import pytest +from foxglove_schemas_protobuf.CompressedVideo_pb2 import CompressedVideo +from mcap_protobuf.writer import Writer +from hflow.episode import Episode from hflow.ffmpeg import ffmpeg_path, ffprobe_path from hflow.video import ( AccessUnit, @@ -132,6 +137,70 @@ def test_remux_to_mp4_preserves_every_frame( assert remux_elapsed_seconds < 5.0 +@pytest.mark.parametrize("invalid_format", [False, True]) +def test_episode_video_remuxes_batches_without_materializing_a_channel( + encoded_units: list[AccessUnit], + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + invalid_format: bool, +) -> None: + source = tmp_path / "camera.mcap" + with Writer(str(source)) as writer: + for index, unit in enumerate(encoded_units): + writer.write_message( + "/camera", + CompressedVideo( + format="h265" if invalid_format and index == FRAME_COUNT - 1 else "h264", + data=unit.data, + ), + log_time=round(index * 1e9 / FPS), + ) + + def refuse_materialization(*args: Any, **kwargs: Any) -> None: + pytest.fail("video must not materialize the camera channel") + + monkeypatch.setattr(Episode, "channel", refuse_materialization) + workdir = tmp_path / "video-cache" + with Episode(source, workdir=workdir) as episode: + original_batches = episode._reader.iter_batches + + def small_batches(*args: Any, **kwargs: Any) -> Iterator[Any]: + kwargs["batch_max_messages"] = 3 + yield from original_batches(*args, **kwargs) + + monkeypatch.setattr(episode._reader, "iter_batches", small_batches) + if invalid_format: + with pytest.raises(ValueError, match="carries 'h265', expected 'h264'"): + episode.video() + assert not list(episode.workdir.glob("*.mp4")) + assert not list(episode.workdir.glob(".*.tmp")) + assert not episode._channel_data_by_id + return + output = episode.video() + assert int(_ffprobe_video_stream_fields(output)["nb_read_frames"]) == FRAME_COUNT + assert episode._video_fps["/camera"] == pytest.approx(FPS) + assert not episode._channel_data_by_id + + def refuse_cached_read(*args: Any, **kwargs: Any) -> None: + pytest.fail("a completed video with cached FPS must not reread the channel") + + with monkeypatch.context() as cached_patch: + cached_patch.setattr(episode._reader, "iter_batches", refuse_cached_read) + assert episode.video() == output + + # FPS alone is insufficient: a removed MP4 must be recreated. + output.unlink() + assert episode.video() == output + assert int(_ffprobe_video_stream_fields(output)["nb_read_frames"]) == FRAME_COUNT + + with Episode(source, workdir=workdir) as reopened: + # A fresh handle must recover FPS, but can reuse the completed MP4. + monkeypatch.setattr(reopened, "iter_decoded_batches", refuse_cached_read) + assert not reopened._video_fps + assert reopened.video() == output + assert reopened._video_fps["/camera"] == pytest.approx(FPS) + + @pytest.fixture(scope="module") def b_frame_stream(jpeg_frames: list[bytes]) -> bytes: """The same frames re-encoded with libx264 defaults (B-frames enabled).""" @@ -155,7 +224,7 @@ def b_frame_stream(jpeg_frames: list[bytes]) -> bytes: "-pix_fmt", "yuv420p", "-x264-params", - "bframes=3:b_adapt=0", + "bframes=3:b_adapt=0:aud=1", "-f", "h264", "-", @@ -219,21 +288,29 @@ def test_scan_refuses_an_unparseable_slice_header() -> None: scan_picture_coding_types(unparseable_slice) +@pytest.mark.parametrize("split_units", [False, True]) def test_remux_refuses_a_b_frame_stream_naming_the_tail( - b_frame_stream: bytes, tmp_path: Path + b_frame_stream: bytes, tmp_path: Path, split_units: bool ) -> None: output_path = tmp_path / "bframe.mp4" + units = ( + (unit.data for unit in split_annex_b_stream(b_frame_stream)) + if split_units + else iter((b_frame_stream,)) + ) with pytest.raises(ValueError, match="reorder depth") as error: - write_access_units_to_mp4((b_frame_stream,), fps=FPS, output=output_path) + write_access_units_to_mp4(units, fps=FPS, output=output_path) message = str(error.value) - assert "B picture" in message - assert "at risk" in message + scan = scan_picture_coding_types(b_frame_stream) + assert f"{scan.b_picture_count} B picture(s) across {scan.picture_count}" in message + assert f"reorder depth {scan.reorder_depth}" in message + assert f"last {scan.trailing_b_pictures} frame(s) at risk" in message assert "docs/FORMAT.md" in message - # The refusal fires before ffmpeg runs, so no partial MP4 is left behind. + # Even a refusal after streaming must remove the partial MP4. assert not output_path.exists() - assert not output_path.with_name(output_path.name + ".tmp").exists() + assert not list(tmp_path.glob(f".{output_path.name}.*.tmp")) def test_encode_guarantees_accept_a_conforming_stream( diff --git a/tests/test_video_streaming.py b/tests/test_video_streaming.py new file mode 100644 index 00000000..5a27e8ac --- /dev/null +++ b/tests/test_video_streaming.py @@ -0,0 +1,183 @@ +"""Bounded remux consumption and cleanup at the subprocess boundary.""" + +import subprocess +import sys +import time +from collections.abc import Iterator +from pathlib import Path +from typing import IO + +import pytest + +from hflow import video + +# A parseable P slice larger than the stdin buffer. The process double consumes +# bytes, while test_video exercises the same remux with real H.264 and ffmpeg. +UNIT = b"\x00\x00\x00\x01\x41\xf0" + b"x" * (64 * 1024) + + +@pytest.fixture +def remux_process( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> tuple[Path, list[subprocess.Popen[bytes]]]: + progress = tmp_path / "progress" + executable = tmp_path / "ffmpeg" + executable.write_text( + f"#!{sys.executable}\n" + "import pathlib, sys\n" + f"progress = pathlib.Path({str(progress)!r})\n" + "with open(sys.argv[-1], 'wb') as output:\n" + " count = 0\n" + " while True:\n" + f" unit = sys.stdin.buffer.read({len(UNIT)})\n" + " if not unit:\n" + " break\n" + " output.write(unit)\n" + " output.flush()\n" + " count += len(unit)\n" + " progress.write_text(str(count))\n" + ) + executable.chmod(0o755) + monkeypatch.setattr(video, "ffmpeg_path", lambda: executable) + processes: list[subprocess.Popen[bytes]] = [] + original_popen = subprocess.Popen + + def start_process( + command: list[str], *, stdin: int, stdout: int, stderr: IO[bytes] + ) -> subprocess.Popen[bytes]: + process = original_popen(command, stdin=stdin, stdout=stdout, stderr=stderr) + processes.append(process) + return process + + monkeypatch.setattr(subprocess, "Popen", start_process) + return progress, processes + + +def _wait_for_consumption(progress: Path, byte_count: int) -> None: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + if progress.exists() and progress.read_text() == str(byte_count): + return + time.sleep(0.01) + pytest.fail("the subprocess did not consume the unit before the next was requested") + + +def test_remux_consumes_each_unit_before_requesting_the_next( + remux_process: tuple[Path, list[subprocess.Popen[bytes]]], tmp_path: Path +) -> None: + progress, processes = remux_process + output = tmp_path / "video.mp4" + + def units() -> Iterator[bytes]: + for index in range(4): + yield UNIT + _wait_for_consumption(progress, (index + 1) * len(UNIT)) + assert not output.exists() + + assert video.write_access_units_to_mp4(units(), fps=30, output=output) == output + assert output.read_bytes() == UNIT * 4 + assert processes[0].returncode == 0 + assert not list(tmp_path.glob(".video.mp4.*.tmp")) + + +@pytest.mark.parametrize("existing_output", [False, True]) +@pytest.mark.parametrize("failure", ["iterator", "malformed", "orphan_slice", "b_frame"]) +def test_streaming_failure_reaps_process_and_preserves_atomic_output( + remux_process: tuple[Path, list[subprocess.Popen[bytes]]], + tmp_path: Path, + existing_output: bool, + failure: str, +) -> None: + progress, processes = remux_process + output = tmp_path / "video.mp4" + if existing_output: + output.write_bytes(b"previous completed MP4") + + def units() -> Iterator[bytes]: + # The orphan slice must be the first picture; start streaming a non-VCL + # NAL first so the validation failure still happens after a write. + yield UNIT if failure != "orphan_slice" else b"\x00\x00\x00\x01\x09" + UNIT[5:] + _wait_for_consumption(progress, len(UNIT)) + if failure == "iterator": + raise RuntimeError("source decode failed") + if failure == "malformed": + yield b"\x00\x00\x00\x01\x41\x00" + elif failure == "orphan_slice": + yield b"\x00\x00\x00\x01\x41\x58" # first_mb=1, slice_type=P + else: + yield b"\x00\x00\x00\x01\x41\xa0" # first_mb=0, slice_type=B + + expected_error = RuntimeError if failure == "iterator" else ValueError + message = { + "iterator": "source decode failed", + "malformed": "incomplete or truncated", + "orphan_slice": "before any picture starts", + "b_frame": r"1 B picture\(s\) across 2 \(reorder depth 1\)", + }[failure] + with pytest.raises(expected_error, match=message): + video.write_access_units_to_mp4(units(), fps=30, output=output) + assert len(processes) == 1 + assert processes[0].poll() is not None + assert processes[0].stdin is not None and processes[0].stdin.closed + assert not list(tmp_path.glob(".video.mp4.*.tmp")) + if existing_output: + assert output.read_bytes() == b"previous completed MP4" + else: + assert not output.exists() + + +@pytest.mark.parametrize("exit_code", [0, 7]) +@pytest.mark.parametrize("early_exit", [False, True]) +def test_ffmpeg_failure_removes_partial_output_and_reports_stderr( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, exit_code: int, early_exit: bool +) -> None: + executable = tmp_path / "ffmpeg" + # More than a pipe buffer of stderr before reading stdin catches deadlocks. + executable.write_text( + f"#!{sys.executable}\n" + "import sys\n" + "sys.stderr.write('x' * 200_000 + 'remux diagnostic')\n" + "sys.stderr.flush()\n" + + ("" if early_exit else "sys.stdin.buffer.read()\n") + + f"sys.exit({exit_code})\n" + ) + executable.chmod(0o755) + monkeypatch.setattr(video, "ffmpeg_path", lambda: executable) + output = tmp_path / "video.mp4" + message = ( + "exited 0 but produced no output" + if exit_code == 0 and not early_exit + else f"exit {exit_code}.*remux diagnostic" + ) + with pytest.raises(video.VideoEncodeError, match=message): + video.write_access_units_to_mp4(iter((UNIT,) * 4), fps=30, output=output) + assert not output.exists() + assert not list(tmp_path.glob(".video.mp4.*.tmp")) + + +def test_process_start_failure_removes_temporary_output( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(video, "ffmpeg_path", lambda: tmp_path / "missing-ffmpeg") + output = tmp_path / "video.mp4" + with pytest.raises(FileNotFoundError): + video.write_access_units_to_mp4(iter((UNIT,)), fps=30, output=output) + assert not output.exists() + assert not list(tmp_path.glob(".video.mp4.*.tmp")) + + +@pytest.mark.parametrize( + "timestamps", [[], [1], [1, 1], [3, 2, 1], [0, 10], [0, 10, 30], [0, 10, 10, 90, 95]] +) +def test_streaming_fps_preserves_exact_median_and_errors(timestamps: list[int]) -> None: + try: + expected = video.estimate_fps_from_log_times(timestamps, topic="/camera") + except ValueError as error: + with pytest.raises(ValueError) as streaming_error: + video.estimate_fps_from_streaming_log_times(iter(timestamps), topic="/camera") + assert str(streaming_error.value) == str(error) + else: + assert ( + video.estimate_fps_from_streaming_log_times(iter(timestamps), topic="/camera") + == expected + )