From 166dab9d03537e490c9701627d53650d34a28804 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 31 Jul 2026 14:33:50 -0700 Subject: [PATCH] Bound ffmpeg video encoder threads --- .github/workflows/test.yml | 20 ++++---------------- docs/module-internals.md | 11 ++++++----- src/remove_ai_watermarks/video_encoding.py | 12 ++++++------ tests/test_video.py | 15 +++++---------- tests/test_video_invisible.py | 6 +----- 5 files changed, 22 insertions(+), 42 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 09389b4..cf3d1f8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -55,22 +55,10 @@ jobs: with: python-version: "3.12" - name: Install ffmpeg - run: sudo apt-get update && sudo apt-get install --yes ffmpeg strace + run: sudo apt-get update && sudo apt-get install --yes ffmpeg - name: Sync dev environment run: uv sync --frozen --extra dev - name: Run full-clip video test - run: | - set +e - timeout --kill-after=10s 90s \ - strace -ff -tt -s 256 -o /tmp/video-e2e.strace \ - uv run pytest -vv -o faulthandler_timeout=45 \ - tests/test_video.py::TestVideoVisibleFullClip - status=$? - if [ "$status" -ne 0 ]; then - for trace in /tmp/video-e2e.strace*; do - echo "::group::$trace" - tail -n 100 "$trace" - echo "::endgroup::" - done - fi - exit "$status" + run: >- + uv run pytest -vv -o faulthandler_timeout=60 + tests/test_video.py::TestVideoVisibleFullClip diff --git a/docs/module-internals.md b/docs/module-internals.md index 4f8019c..4484965 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -149,11 +149,12 @@ 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 inputs disable ffmpeg probing before -`pipe:0`: their format is already explicit, and ffmpeg 6 on Linux can otherwise -wait for its normal analysis window while the producer blocks on a full pipe -before the source-audio input opens. Command-order regressions cover both stdin -modes; the real Linux full-clip CI job guards process completion. +Both the raw-BGR and timestamped-NUT stdin modes cap the video encoder at two +threads. A Linux full-clip trace showed ffmpeg creating dozens of worker stacks +and severely delaying frame-pipe ingestion on a constrained hosted runner. The +bounded codec pool avoids 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. `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 1f129d3..9f18278 100644 --- a/src/remove_ai_watermarks/video_encoding.py +++ b/src/remove_ai_watermarks/video_encoding.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: log = logging.getLogger(__name__) _PIXEL_FORMATS = frozenset({"yuv420p", "yuv422p", "yuv444p"}) +_VIDEO_ENCODER_THREADS = 2 _PIXEL_FORMAT_ALIASES = { "yuvj420p": "yuv420p", "yuvj422p": "yuv422p", @@ -286,12 +287,8 @@ def raw_video_command( ffmpeg = shutil.which("ffmpeg") if ffmpeg is None: raise RuntimeError("Video processing requires ffmpeg on PATH") - # The pipe format and stream geometry are already explicit. FFmpeg 6 can - # otherwise wait for its normal analysis window while the producer blocks - # on a full pipe, before the second (audio) input has been opened. - pipe_input = ["-analyzeduration", "0", "-probesize", "32", "-i", "pipe:0"] frame_input = ( - ["-f", "nut", *pipe_input] + ["-f", "nut", "-i", "pipe:0"] if timestamped_input else [ "-f", @@ -302,7 +299,8 @@ def raw_video_command( f"{width}x{height}", "-r", f"{fps:.12g}", - *pipe_input, + "-i", + "pipe:0", ] ) command = [ @@ -319,6 +317,8 @@ 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", diff --git a/tests/test_video.py b/tests/test_video.py index 6df876f..3b3e674 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -1950,17 +1950,19 @@ class TestVideoVisibleScan: class TestVideoVisibleEncoding: - def test_raw_pipe_input_disables_redundant_ffmpeg_probing( + @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 / "clean.mp4", + tmp_path / f"clean{suffix}", width=12, height=8, fps=24.0, @@ -1969,14 +1971,7 @@ class TestVideoVisibleEncoding: profile=video_encoding.VideoEncodeProfile(), ) - pipe_position = command.index("pipe:0") - assert command[pipe_position - 5 : pipe_position] == [ - "-analyzeduration", - "0", - "-probesize", - "32", - "-i", - ] + assert command[command.index("-threads:v") + 1] == "2" @staticmethod def _patch_single_frame_encode( diff --git a/tests/test_video_invisible.py b/tests/test_video_invisible.py index 61ddaf3..a28ce00 100644 --- a/tests/test_video_invisible.py +++ b/tests/test_video_invisible.py @@ -102,13 +102,9 @@ def test_timestamped_encoder_reads_nut_and_passes_pts_through( ) assert "-copyts" in command - assert command[command.index("-f") : command.index("-f") + 8] == [ + assert command[command.index("-f") : command.index("-f") + 4] == [ "-f", "nut", - "-analyzeduration", - "0", - "-probesize", - "32", "-i", "pipe:0", ]