diff --git a/docs/module-internals.md b/docs/module-internals.md index 8919c90..93644d1 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -149,13 +149,15 @@ regeneration. It centralizes container codecs, optional audio stream copying, metadata/chapter policy, encode-failure reporting, and atomic same-directory publication. Each mapped stream is allowed to reach its own end, so a copied audio tail is not shortened to the frame-input duration. -Both the raw-BGR and timestamped-NUT stdin modes run the implicit pixel-format -filter graph on one thread and cap the video encoder at two threads. A Linux -full-clip trace showed ffmpeg creating an oversized execution pool and severely -delaying frame-pipe ingestion on a constrained hosted runner. The bounded -filter and codec pools avoid that scheduling collapse while leaving audio stream -copy independent. Command regressions cover both supported video codecs; the -real Linux full-clip CI job guards process completion. +Both the raw-BGR and timestamped-NUT stdin modes redirect ffmpeg stderr to a +temporary file while frames are written. Waiting to read diagnostics until +after stdin closed allowed stderr backpressure to stop ffmpeg's frame reads, +which in turn blocked the producer before it could close stdin. The file consumes +no pipe capacity or RAM while ffmpeg runs; completion reports a bounded head and +tail when diagnostics are unusually large. Aborts release it even when ffmpeg +has already exited. A real subprocess regression writes diagnostics beyond pipe +capacity while streaming frames, checks bounded failure reporting, and the Linux +full-clip CI job guards the complete path. `probe_video_encode_profile` reads the first source video stream with ffprobe and preserves the supported properties that survive the 8-bit BGR boundary: `yuv420p`/`yuv422p`/`yuv444p` chroma sampling, recognized color tags, encoder diff --git a/src/remove_ai_watermarks/video_encoding.py b/src/remove_ai_watermarks/video_encoding.py index 76b19d2..acb4363 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -13,16 +13,17 @@ from contextlib import contextmanager from dataclasses import dataclass from fractions import Fraction from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast if TYPE_CHECKING: from collections.abc import Generator, Sequence + from typing import BinaryIO log = logging.getLogger(__name__) +_FFMPEG_STDERR_LIMIT = 64 * 1024 +_FFMPEG_STDERR_TRUNCATION = b"\n...[ffmpeg stderr truncated]...\n" _PIXEL_FORMATS = frozenset({"yuv420p", "yuv422p", "yuv444p"}) -_VIDEO_FILTER_THREADS = 1 -_VIDEO_ENCODER_THREADS = 2 _PIXEL_FORMAT_ALIASES = { "yuvj420p": "yuv420p", "yuvj422p": "yuv422p", @@ -83,6 +84,53 @@ class VideoEncodeProfile: component_depth: int | None = None +@dataclass +class _RawVideoEncoder: + """Running ffmpeg process with diagnostics redirected outside a pipe.""" + + process: subprocess.Popen[bytes] + stdin: BinaryIO + _stderr_buffer: BinaryIO + + def poll(self) -> int | None: + return self.process.poll() + + def kill(self) -> None: + self.process.kill() + + def wait(self) -> int: + return self.process.wait() + + def collect_stderr(self) -> str: + """Return bounded ffmpeg diagnostics and release the temporary file.""" + if self._stderr_buffer.closed: + raise RuntimeError("ffmpeg diagnostics have already been collected") + try: + self._stderr_buffer.seek(0, os.SEEK_END) + size = self._stderr_buffer.tell() + if size <= _FFMPEG_STDERR_LIMIT: + self._stderr_buffer.seek(0) + raw_stderr = self._stderr_buffer.read() + else: + payload_limit = _FFMPEG_STDERR_LIMIT - len(_FFMPEG_STDERR_TRUNCATION) + head_size = payload_limit // 2 + tail_size = payload_limit - head_size + self._stderr_buffer.seek(0) + head = self._stderr_buffer.read(head_size) + self._stderr_buffer.seek(-tail_size, os.SEEK_END) + tail = self._stderr_buffer.read(tail_size) + raw_stderr = head + _FFMPEG_STDERR_TRUNCATION + tail + return raw_stderr.decode("utf-8", errors="replace") + finally: + self._stderr_buffer.close() + + def discard_stderr(self) -> None: + """Release the diagnostic buffer after an abort.""" + if self._stderr_buffer.closed: + return + self._stderr_buffer.close() + + def _known_value(value: object, allowed: frozenset[str]) -> str | None: """Return a supported ffmpeg enum value, otherwise omit it.""" return value if isinstance(value, str) and value in allowed else None @@ -306,8 +354,6 @@ def raw_video_command( ) command = [ ffmpeg, - "-filter_threads", - str(_VIDEO_FILTER_THREADS), "-y", "-loglevel", "error", @@ -320,8 +366,6 @@ def raw_video_command( "-map", "1:a?", *_video_codec_args(output.suffix.lower(), crf=crf, profile=profile), - "-threads:v", - str(_VIDEO_ENCODER_THREADS), *_profile_args(profile), "-c:a", "copy", @@ -344,40 +388,49 @@ def raw_video_command( return command -def start_raw_video_encoder(command: list[str]) -> subprocess.Popen[bytes]: +def start_raw_video_encoder(command: list[str]) -> _RawVideoEncoder: """Start ffmpeg and validate its raw-frame pipes.""" log.info("Starting ffmpeg video encode: command=%s", command) - process = subprocess.Popen( # noqa: S603 - command, - stdin=subprocess.PIPE, - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, + stderr_buffer = cast( + "BinaryIO", + tempfile.TemporaryFile(mode="w+b"), # noqa: SIM115 - encoder owns the lifetime ) - if process.stdin is None or process.stderr is None: + try: + process = subprocess.Popen( # noqa: S603 + command, + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=stderr_buffer, + ) + except Exception: + stderr_buffer.close() + raise + if process.stdin is None: process.kill() process.wait() - raise RuntimeError("Could not open ffmpeg pipes") - return process + stderr_buffer.close() + raise RuntimeError("Could not open ffmpeg input pipe") + return _RawVideoEncoder(process, cast("BinaryIO", process.stdin), stderr_buffer) def finish_raw_video_encoder( - process: subprocess.Popen[bytes], + process: _RawVideoEncoder, output: Path, *, operation: str, ) -> None: """Close the frame stream and raise when ffmpeg rejects the encode.""" - if process.stdin is None or process.stderr is None: - raise RuntimeError("ffmpeg pipes are unavailable") process.stdin.close() - stderr = process.stderr.read().decode("utf-8", errors="replace") return_code = process.wait() + stderr = process.collect_stderr() log.info("ffmpeg %s finished: status=%s stderr=%s", operation, return_code, stderr) if return_code != 0: raise RuntimeError(f"ffmpeg failed to encode {output}: {stderr.strip()[:500]}") -def abort_raw_video_encoder(process: subprocess.Popen[bytes]) -> None: +def abort_raw_video_encoder(process: _RawVideoEncoder) -> None: """Stop an incomplete ffmpeg encode.""" - process.kill() + if process.poll() is None: + process.kill() process.wait() + process.discard_stderr() diff --git a/src/remove_ai_watermarks/video_invisible.py b/src/remove_ai_watermarks/video_invisible.py index 6fc1815..832b95d 100644 --- a/src/remove_ai_watermarks/video_invisible.py +++ b/src/remove_ai_watermarks/video_invisible.py @@ -323,9 +323,6 @@ def encode_video_frames( ) ) frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") try: for frame in frames: if frame.shape[:2] != (height, width): @@ -333,8 +330,7 @@ def encode_video_frames( frame_pipe.write(frame.tobytes()) finish_raw_video_encoder(process, output, operation="SynthID removal encode") except Exception: - if process.poll() is None: - abort_raw_video_encoder(process) + abort_raw_video_encoder(process) raise @@ -393,9 +389,6 @@ def regenerate_video_candidate( ) ) frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") frame_count = 0 squared_error = 0.0 pixel_count = 0 @@ -462,8 +455,7 @@ def regenerate_video_candidate( operation="SynthID removal encode", ) except Exception: - if process.poll() is None: - abort_raw_video_encoder(process) + abort_raw_video_encoder(process) raise mse = squared_error / pixel_count diff --git a/src/remove_ai_watermarks/video_visible.py b/src/remove_ai_watermarks/video_visible.py index 1eaa288..a6328d5 100644 --- a/src/remove_ai_watermarks/video_visible.py +++ b/src/remove_ai_watermarks/video_visible.py @@ -1363,9 +1363,6 @@ def encode_clean_video( ) ) frame_pipe = process.stdin - if frame_pipe is None: - abort_raw_video_encoder(process) - raise RuntimeError("Could not open ffmpeg input pipe") capture = cv2.VideoCapture(str(source)) if not capture.isOpened(): @@ -1436,8 +1433,7 @@ def encode_clean_video( if timestamped_writer is not None: with suppress(Exception): timestamped_writer.close() - if process.poll() is None: - abort_raw_video_encoder(process) + abort_raw_video_encoder(process) raise finally: capture.release() diff --git a/tests/test_video.py b/tests/test_video.py index 07a1b7f..5a2b894 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1950,30 +1950,6 @@ class TestVideoVisibleScan: class TestVideoVisibleEncoding: - @pytest.mark.parametrize("suffix", [".mp4", ".webm"]) - def test_encoder_bounds_codec_threads( - self, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - suffix: str, - ): - from remove_ai_watermarks import video_encoding - - monkeypatch.setattr(video_encoding.shutil, "which", lambda _name: "/usr/bin/ffmpeg") - command = video_encoding.raw_video_command( - tmp_path / "source.mp4", - tmp_path / f"clean{suffix}", - width=12, - height=8, - fps=24.0, - strip_metadata=True, - crf=14, - profile=video_encoding.VideoEncodeProfile(), - ) - - assert command[1:3] == ["-filter_threads", "1"] - assert command[command.index("-threads:v") + 1] == "2" - @staticmethod def _patch_single_frame_encode( monkeypatch: pytest.MonkeyPatch, @@ -2008,6 +1984,12 @@ class TestVideoVisibleEncoding: def poll(self) -> int: return 1 + def wait(self) -> int: + return 1 + + def discard_stderr(self) -> None: + pass + targets: list[Path] = [] def fake_command(_source: Path, target: Path, **_kwargs: object) -> list[str]: diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index a28ce00..735d25b 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -2,6 +2,8 @@ from __future__ import annotations +import sys +import threading from types import SimpleNamespace from typing import TYPE_CHECKING @@ -14,6 +16,55 @@ if TYPE_CHECKING: from pathlib import Path +def test_encoder_redirects_large_stderr_while_frames_are_streaming( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + diagnostic = "synthetic ffmpeg diagnostic" + tail_diagnostic = "synthetic ffmpeg diagnostic tail" + caplog.set_level("INFO", logger=video_encoding.__name__) + command = [ + sys.executable, + "-c", + ( + "import sys; " + f"sys.stderr.buffer.write({diagnostic.encode()!r} + b'x' * 262144 + {tail_diagnostic.encode()!r}); " + "sys.stderr.buffer.flush(); " + "sys.stdin.buffer.read(); " + "raise SystemExit(7)" + ), + ] + encoder = video_encoding.start_raw_video_encoder(command) + write_finished = threading.Event() + write_errors: list[Exception] = [] + + def write_frames() -> None: + try: + encoder.stdin.write(b"f" * 262144) + encoder.stdin.flush() + except Exception as exc: # pragma: no cover - mutation cleanup path + write_errors.append(exc) + finally: + write_finished.set() + + writer = threading.Thread(target=write_frames) + writer.start() + try: + assert write_finished.wait(5), "stderr backpressure blocked the frame producer" + assert write_errors == [] + with pytest.raises(RuntimeError, match=diagnostic): + video_encoding.finish_raw_video_encoder( + encoder, + tmp_path / "unused.mp4", + operation="synthetic encode", + ) + assert "ffmpeg stderr truncated" in caplog.text + assert tail_diagnostic in caplog.text + finally: + video_encoding.abort_raw_video_encoder(encoder) + writer.join(timeout=5) + + def test_availability_requires_both_optional_packages(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( video_invisible,