from __future__ import annotations import subprocessimport sysimport textwrapimport threadingimport timefrom pathlib import Pathfrom unittest.mock import ANY, Mock import avimport numpy as npimport pytestfrom click.testing import CliRunner from manim import FadeIn, Scene, Square, capture, tempconfigfrom manim._config import configfrom manim._config.output import OutputFormat, OutputSpecfrom manim._config.output_plan import (    resolve_media_layout,    resolve_module_name,    resolve_output_plan,    resolve_requested_output_name,)from manim._config.video_encoder import resolve_video_encoderfrom manim.cli.render.commands import renderfrom manim.scene.scene_file_writer import SceneFileWriter, _SceneFileWriterSettingsfrom manim.utils.exceptions import RerunSceneException _ENCODER_THREAD_PREFIX = "partial-movie-encoder-"_VIDEO_OUTPUT = OutputSpec(    OutputFormat.MP4,    transparent=False,    save_sections=False,    fallback_to_still=False,)_NO_OUTPUT = OutputSpec(    OutputFormat.NONE,    transparent=False,    save_sections=False,    fallback_to_still=False,)_UNIQUE_PLAYS = 6_TOTAL_PLAYS = _UNIQUE_PLAYS + 2  def _make_writer(    scene_name: str,    output: OutputSpec = _VIDEO_OUTPUT,) -> SceneFileWriter:    module_name = resolve_module_name(config)    layout = resolve_media_layout(        config,        output,        module_name=module_name,        scene_name=scene_name,        working_directory=Path.cwd(),    )    plan = resolve_output_plan(        layout,        output,        scene_name=scene_name,        requested_output_name=resolve_requested_output_name(config),    )    settings = _SceneFileWriterSettings(        plan=plan,        video_encoder=resolve_video_encoder(            output,            width=config.pixel_width,            height=config.pixel_height,            frame_rate=config.frame_rate,        ),        max_inflight_encoders=config.max_inflight_encoders,        encoder_queue_size=config.encoder_queue_size,        max_files_cached=config.max_files_cached,        assets_dir=Path.cwd(),    )    return SceneFileWriter(settings)  _SCENE_NAME = "ParallelEncodingCacheScene"_SCENE_SOURCE = textwrap.dedent(    f"""\    from manim import FadeIn, Scene, Square      class {_SCENE_NAME}(Scene):        def construct(self):            for index in range({_UNIQUE_PLAYS}):                square = Square(side_length=0.2 + index / 100)                self.play(FadeIn(square), run_time=0.1)                self.clear()             self.play(FadeIn(Square(side_length=0.75)), run_time=0.1)            self.clear()            self.play(FadeIn(Square(side_length=0.75)), run_time=0.1)    """)  def _alive_encoder_threads():    return [        thread        for thread in threading.enumerate()        if thread.name.startswith(_ENCODER_THREAD_PREFIX) and thread.is_alive()    ]  def _render_scene(media_dir, scene_file):    command = [        sys.executable,        "-m",        "manim",        "-ql",        "--max-inflight-encoders",        "3",        "--media_dir",        str(media_dir),        str(scene_file),        _SCENE_NAME,    ]    _, err, exit_code = capture(command)    assert exit_code == 0, err    quality_directory = media_dir / "videos" / scene_file.stem / "480p15"    return quality_directory, quality_directory / "partial_movie_files" / _SCENE_NAME  @pytest.mark.slowdef test_parallel_encoding_cache_behavior(tmp_path):    scene_file = tmp_path / "parallel_encoding_scenes.py"    scene_file.write_text(_SCENE_SOURCE)     # The duplicate tail play must cache-hit within the run (one file fewer    # than the number of plays); that hit exercises the writer's in-flight    # lookup while the first tail file may still be encoding.    media_dir = tmp_path / "media"    quality_directory, partial_directory = _render_scene(media_dir, scene_file)    partial_movies = sorted(partial_directory.glob("*.mp4"))    assert len(partial_movies) == _TOTAL_PLAYS - 1    assert (quality_directory / f"{_SCENE_NAME}.mp4").exists()    for partial_movie in partial_movies:        with av.open(partial_movie) as container:            next(container.decode(video=0))     partial_snapshot = {        partial_movie.name: partial_movie.stat().st_mtime_ns        for partial_movie in partial_movies    }     _render_scene(media_dir, scene_file)     second_snapshot = {        partial_movie.name: partial_movie.stat().st_mtime_ns        for partial_movie in sorted(partial_directory.glob("*.mp4"))    }    assert second_snapshot == partial_snapshot  @pytest.mark.slowdef test_disable_caching_is_evaluated_per_play(config, tmp_path):    class DynamicCachePolicyScene(Scene):        def construct(self):            self.play(FadeIn(Square()), run_time=0.1)            self.clear()            with tempconfig({"disable_caching": True}):                self.play(FadeIn(Square()), run_time=0.1)            self.clear()            self.play(FadeIn(Square()), run_time=0.1)            self.clear()            self.play(FadeIn(Square()), run_time=0.1)     with tempconfig(        {            "format": "mp4",            "media_dir": tmp_path,            "quality": "low_quality",            "disable_caching": False,        },    ):        scene = DynamicCachePolicyScene()        scene.render()     initial_hash, uncached_hash, restored_hash, cached_hash = (        scene.renderer.animations_hashes    )    assert initial_hash is not None    assert uncached_hash == "uncached_00001"    assert restored_hash is not None    assert cached_hash == restored_hash  @pytest.mark.slowdef test_no_encoder_threads_survive_render(config, tmp_path):    class ThreadSweepScene(Scene):        def construct(self):            for index in range(3):                square = Square(side_length=0.3 + index / 10)                self.play(FadeIn(square), run_time=0.1)                self.clear()     with tempconfig(        {            "media_dir": tmp_path,            "quality": "low_quality",            "max_inflight_encoders": 3,            "encoder_queue_size": 2,        },    ):        scene = ThreadSweepScene()        scene.render()     assert not _alive_encoder_threads()  def _frame():    return np.zeros((4, 4, 4), dtype=np.uint8)  def _fake_segment_encoder(tmp_path, name):    target = tmp_path / f"{name}.mp4"    encoder = Mock(target=target)    encoder.abort.side_effect = lambda: target.unlink(missing_ok=True)    return encoder  def _new_encode_job(tmp_path, name, encoder=None, frame_queue_size=8):    from manim.scene.scene_file_writer import _PartialMovieEncodeJob     if encoder is None:        encoder = _fake_segment_encoder(tmp_path, name)    return _PartialMovieEncodeJob(        animation_index=0,        encoder=encoder,        frame_queue_size=frame_queue_size,    )  def _assert_failed_join(job, expected_exception):    job.thread.join(timeout=5)    assert not job.thread.is_alive(), "Partial movie encoder did not finish"     with pytest.raises(type(expected_exception)) as exc_info:        job.join()     assert exc_info.value is expected_exception    assert not _alive_encoder_threads()  def test_encode_failure_propagates_and_drains_bounded_queue(    tmp_path,    manim_caplog,):    expected_exception = RuntimeError("encode failed")    encode_failed = threading.Event()    encoder = _fake_segment_encoder(tmp_path, "encode_failure")     def fail_write(*args, **kwargs):        encode_failed.set()        raise expected_exception     encoder.write_frame.side_effect = fail_write    job = _new_encode_job(tmp_path, "encode_failure", encoder)    job.put(1, _frame())    assert encode_failed.wait(timeout=2), "Encode failure was not triggered"     def fill_queue_and_seal():        for _ in range(job.queue.maxsize + 1):            job.put(1, _frame())        job.seal()     producer = threading.Thread(target=fill_queue_and_seal, daemon=True)    producer.start()    producer.join(timeout=5)    assert not producer.is_alive(), "Producer deadlocked on the bounded queue"     _assert_failed_join(job, expected_exception)    encoder.abort.assert_called_once_with()    assert "Partial movie file written" not in manim_caplog.text  def test_finish_failure_aborts_segment(tmp_path, manim_caplog):    expected_exception = RuntimeError("finish failed")    encoder = _fake_segment_encoder(tmp_path, "finish_failure")    encoder.finish.side_effect = expected_exception    job = _new_encode_job(tmp_path, "finish_failure", encoder)    job.put(1, _frame())    job.seal()     _assert_failed_join(job, expected_exception)    encoder.abort.assert_called_once_with()    assert "Partial movie file written" not in manim_caplog.text  def test_encode_failure_precedes_abort_failure(tmp_path, manim_caplog):    expected_exception = RuntimeError("encode failed")    abort_exception = RuntimeError("abort failed")    encoder = _fake_segment_encoder(tmp_path, "encode_and_abort_failure")    encoder.write_frame.side_effect = expected_exception    encoder.abort.side_effect = abort_exception    job = _new_encode_job(tmp_path, "encode_and_abort_failure", encoder)    job.put(1, _frame())    job.seal()     _assert_failed_join(job, expected_exception)    assert "Failed to clean up incomplete segment" in manim_caplog.text    assert "abort failed" in manim_caplog.text    assert "Partial movie file written" not in manim_caplog.text  def test_successful_encode_job_logs_partial_movie_written(tmp_path, manim_caplog):    encoder = _fake_segment_encoder(tmp_path, "encode_success")    job = _new_encode_job(tmp_path, "encode_success", encoder)    job.put(1, _frame())    job.seal()     job.join()     encoder.write_frame.assert_called_once_with(ANY, repeat=1)    encoder.finish.assert_called_once_with()    encoder.abort.assert_not_called()    assert "Partial movie file written" in manim_caplog.text    assert not _alive_encoder_threads()  def test_write_frame_fails_fast_after_encoder_failure(config, tmp_path):    expected_exception = RuntimeError("encode failed")    encoder = _fake_segment_encoder(tmp_path, "fail_fast")    encoder.write_frame.side_effect = expected_exception    config.media_dir = str(tmp_path)    writer = _make_writer("FailFastScene")    job = _new_encode_job(tmp_path, "fail_fast", encoder)    job.path.write_bytes(b"stale")    writer._current_encode_job = job     try:        writer.write_frame(_frame())        for _ in range(500):            if job.failed:                break            time.sleep(0.01)        assert job.failed, "Encoder failure was not captured in time"         with pytest.raises(RuntimeError) as exc_info:            writer.write_frame(_frame())         assert exc_info.value is expected_exception        assert writer._current_encode_job is None        assert not job.path.exists()        assert not _alive_encoder_threads()    finally:        if writer._current_encode_job is not None:            job.abort()            writer._current_encode_job = None        job.thread.join(timeout=5)  @pytest.mark.parametrize(    ("max_inflight_encoders", "encoder_queue_size", "expected_queue_size"),    [(1, 8, 0), (1, 3, 0), (2, 8, 8), (2, 3, 3)],)def test_frame_queue_configuration(    config,    tmp_path,    max_inflight_encoders,    encoder_queue_size,    expected_queue_size,):    config.max_inflight_encoders = max_inflight_encoders    config.encoder_queue_size = encoder_queue_size    writer = _make_writer("FrameQueueSizeScene")     writer.open_partial_movie_stream(        animation_index=0,        file_path=tmp_path / "partial.mp4",    )    job = writer._current_encode_job    assert job is not None    assert job.queue.maxsize == expected_queue_size     writer.close_partial_movie_stream()    writer.join_all_encode_jobs()    assert not _alive_encoder_threads()  def test_pool_settings_are_not_read_from_mutated_config(config, tmp_path):    config.max_inflight_encoders = 2    config.encoder_queue_size = 3    writer = _make_writer("CapturedPoolSettingsScene")    config.max_inflight_encoders = 1    config.encoder_queue_size = 9     writer.open_partial_movie_stream(        animation_index=0,        file_path=tmp_path / "partial.mp4",    )    job = writer._current_encode_job    assert job is not None    assert job.queue.maxsize == 3     writer.close_partial_movie_stream()    assert writer._inflight_encode_jobs == [job]    writer.join_all_encode_jobs()    assert not _alive_encoder_threads()  def test_cache_limit_is_not_read_from_mutated_config(config, monkeypatch):    config.max_files_cached = -1    writer = _make_writer("CapturedCacheLimitScene")    config.max_files_cached = 0    writer.combine_to_movie = Mock()    prune = Mock()    monkeypatch.setattr(        "manim.scene.scene_file_writer.prune_segment_cache",        prune,    )     writer.finish()     prune.assert_called_once_with(writer.partial_movie_directory, -1)  @pytest.mark.parametrize("max_inflight_encoders", [1, 2, 3])def test_close_partial_movie_stream_respects_cap_and_joins_fifo(    config,    tmp_path,    max_inflight_encoders,):    config.max_inflight_encoders = max_inflight_encoders    writer = _make_writer("EncoderCapScene")    jobs = [Mock(path=tmp_path / f"partial_{index}.mp4") for index in range(3)]     for index, job in enumerate(jobs):        writer._current_encode_job = job        writer.close_partial_movie_stream()         closed_jobs = jobs[: index + 1]        expected_joined = max(            0,            len(closed_jobs) - max_inflight_encoders + 1,        )        expected_inflight = closed_jobs[expected_joined:]        assert writer._inflight_encode_jobs == expected_inflight        assert writer._inflight_by_path == {            str(inflight_job.path): inflight_job for inflight_job in expected_inflight        }        assert len(writer._inflight_encode_jobs) < max_inflight_encoders         for joined_job in closed_jobs[:expected_joined]:            joined_job.join.assert_called_once_with()            assert str(joined_job.path) not in writer._inflight_by_path        for inflight_job in expected_inflight:            inflight_job.join.assert_not_called()            assert writer._inflight_by_path[str(inflight_job.path)] is inflight_job     for job in jobs:        job.seal.assert_called_once_with()  def test_cap_join_failure_drains_all_inflight_jobs(config, tmp_path):    primary_exception = RuntimeError("first join failed")    secondary_exception = RuntimeError("second join failed")    config.max_inflight_encoders = 3    writer = _make_writer("EncoderCapFailureScene")    jobs = [Mock(path=tmp_path / f"partial_{index}.mp4") for index in range(3)]    jobs[0].join.side_effect = primary_exception    jobs[1].join.side_effect = secondary_exception     for job in jobs[:2]:        writer._current_encode_job = job        writer.close_partial_movie_stream()     writer._current_encode_job = jobs[2]    with pytest.raises(RuntimeError) as exc_info:        writer.close_partial_movie_stream()     assert exc_info.value is primary_exception    for job in jobs:        job.seal.assert_called_once_with()        job.join.assert_called_once_with()    assert writer._current_encode_job is None    assert writer._inflight_encode_jobs == []    assert writer._inflight_by_path == {}  def test_is_already_cached_joins_same_path_inflight_job(config, tmp_path):    writer = _make_writer("CachedInflightScene")    hash_invocation = "same_path_hash"    path = (        writer.partial_movie_directory        / f"{hash_invocation}{writer.output_spec.segment_extension}"    )    job = Mock(path=path)    writer._inflight_encode_jobs.append(job)    writer._inflight_by_path[str(path)] = job     writer.is_already_cached(hash_invocation)     job.join.assert_called_once_with()    assert writer._inflight_encode_jobs == []    assert writer._inflight_by_path == {}

Lines 1–500 of 1102. Load more