Perf/parallel partial movie encoding

docs/source/guides/configuration.rst

@@ -354,8 +354,9 @@ A list of all config options
    'ffmpeg_loglevel', 'flush_cache', 'frame_height', 'frame_rate',
    'frame_size', 'frame_width', 'frame_x_radius', 'frame_y_radius',
    'from_animation_number', `fullscreen`, 'images_dir', 'input_file', 'left_side',
-   'log_dir', 'log_to_file', 'max_files_cached', 'media_dir', 'media_width',
-   'movie_file_extension', 'notify_outdated_version', 'output_file', 'partial_movie_dir',
+   'log_dir', 'log_to_file', 'max_files_cached', 'max_inflight_encoders',
+   'media_dir', 'media_width', 'movie_file_extension', 'notify_outdated_version',
+   'output_file', 'partial_movie_dir',
    'pixel_height', 'pixel_width', 'plugins', 'preview',
    'progress_bar', 'quality', 'right_side', 'save_as_gif', 'save_last_frame',
    'save_pngs', 'scene_names', 'show_in_file_browser', 'sound', 'tex_dir',

manim/_config/default.cfg

@@ -129,6 +129,13 @@ use_projection_stroke_shaders = False
 
 movie_file_extension = .mp4
 
+# Maximum number of partial movie files being encoded concurrently while the
+# scene continues rendering. 1 encodes each animation's file before the next
+# animation starts; values > 1 overlap encoding with rendering (bounds pipeline
+# depth, not CPU use; each encoder already slice-threads internally).
+# 4 is a good value on typical hardware.
+max_inflight_encoders = 1
+
 # These now override the --quality option.
 frame_rate = 60
 pixel_height = 1080

manim/_config/utils.py

@@ -283,6 +283,7 @@ class MyScene(Scene): ...
         "log_dir",
         "log_to_file",
         "max_files_cached",
+        "max_inflight_encoders",
         "media_dir",
         "movie_file_extension",
         "notify_outdated_version",
@@ -605,6 +606,7 @@ def digest_parser(self, parser: configparser.ConfigParser) -> Self:
             "from_animation_number",
             "upto_animation_number",
             "max_files_cached",
+            "max_inflight_encoders",
             # the next two must be set BEFORE digesting frame_width and frame_height
             "pixel_height",
             "pixel_width",
@@ -1227,6 +1229,22 @@ def max_files_cached(self) -> int:
     def max_files_cached(self, value: int) -> None:
         self._set_pos_number("max_files_cached", value, True)
 
+    @property
+    def max_inflight_encoders(self) -> int:
+        """Maximum number of partial movie files encoded concurrently while the
+        scene continues rendering. 1 encodes each animation's file before the
+        next animation starts; values > 1 overlap encoding with rendering
+        (4 is a good value on typical hardware). No flag.
+        """
+        return self._d["max_inflight_encoders"]
+
+    @max_inflight_encoders.setter
+    def max_inflight_encoders(self, value: int) -> None:
+        if isinstance(value, int) and value >= 1:
+            self._d.__setitem__("max_inflight_encoders", value)
+        else:
+            raise ValueError("max_inflight_encoders must be a positive integer")
+
     @property
     def window_monitor(self) -> int:
         """The monitor on which the scene will be rendered."""

manim/scene/scene.py

@@ -264,6 +264,9 @@ def render(self, preview: bool = False) -> bool:
             # TODO: The CairoRenderer does not have the method clear_screen()
             self.renderer.clear_screen()  # type: ignore[union-attr]
             self.renderer.num_plays = 0
+            # The rerun replaces the file writer; drain its encode jobs so no
+            # worker is still writing a partial file the new writer may reuse.
+            self.renderer.file_writer.join_all_encode_jobs()
             return True
         self.tear_down()
         # We have to reset these settings in case of multiple renders.

manim/scene/scene_file_writer.py

@@ -90,6 +90,93 @@
             output_audio.mux(packet)
 
 
+class _PartialMovieEncodeJob:
+    """Encode and write the frames for one partial movie file."""
+
+    def __init__(
+        self,
+        path: StrPath,
+        animation_index: int,
+        container: OutputContainer,
+        stream: Stream,
+    ) -> None:
+        self.path = path
+        self.animation_index = animation_index
+        self.container = container
+        self.stream = stream
+        # About 66 MB of 1080p RGBA frames per job. The worker drains through
+        # the sentinel after an exception, so a bounded queue cannot deadlock.
+        self.queue: Queue[tuple[int, PixelArray | None]] = Queue(maxsize=8)
+        self._exception: BaseException | None = None
+        self.thread = Thread(
+            target=self._listen_and_write,
+            name=f"partial-movie-encoder-{animation_index}",
+        )
+        self.thread.start()
+
+    def _capture_exception(self, exception: BaseException) -> None:
+        if self._exception is None:
+            self._exception = exception
+
+    def _listen_and_write(self) -> None:
+        while True:
+            num_frames, frame_data = self.queue.get()
+            if frame_data is None:
+                break
+            if self._exception is not None:
+                continue
+
+            try:
+                self._encode_and_write_frame(frame_data, num_frames)
+            except BaseException as exception:
+                self._capture_exception(exception)
+
+        try:
+            for packet in self.stream.encode():
+                self.container.mux(packet)
+        except BaseException as exception:
+            self._capture_exception(exception)
+        finally:
+            try:
+                self.container.close()
+            except BaseException as exception:
+                self._capture_exception(exception)
+
+    def _encode_and_write_frame(self, frame: PixelArray, num_frames: int) -> None:
+        for _ in range(num_frames):
+            # Notes: precomputing reusing packets does not work!
+            # I.e., you cannot do `packets = encode(...)`
+            # and reuse it, as it seems that `mux(...)`
+            # consumes the packet.
+            # The same issue applies for `av_frame`,
+            # reusing it renders weird-looking frames.
+            av_frame = av.VideoFrame.from_ndarray(frame, format="rgba")
+            for packet in self.stream.encode(av_frame):
+                self.container.mux(packet)
+
+    def put(self, num_frames: int, frame: PixelArray) -> None:
+        """Add a frame to the encoding queue."""
+        self.queue.put((num_frames, frame))
+
+    def seal(self) -> None:
+        """Signal that no more frames will be added."""
+        self.queue.put((-1, None))
+
+    def join(self) -> None:
+        """Wait for encoding to finish and propagate worker failures."""
+        self.thread.join()
+        if self._exception is not None:
+            # A failed encode may leave a structurally valid but truncated
+            # file behind; remove it so a later run cannot cache-hit it.
+            Path(self.path).unlink(missing_ok=True)
+            raise self._exception
+
+        logger.info(
+            f"Animation {self.animation_index} : Partial movie file written in %(path)s",
+            {"path": f"'{self.path}'"},
+        )
+
+
 class SceneFileWriter:
     """SceneFileWriter is the object that actually writes the animations
     played, into video files, using FFMPEG.
@@ -127,6 +214,9 @@
         **kwargs: Any,
     ) -> None:
         self.renderer = renderer
+        self._inflight_encode_jobs: list[_PartialMovieEncodeJob] = []
+        self._inflight_by_path: dict[str, _PartialMovieEncodeJob] = {}
+        self._current_encode_job: _PartialMovieEncodeJob | None = None
         self.init_output_directories(scene_name)
         self.init_audio()
         self.frame_count = 0
@@ -426,30 +516,6 @@
         if write_to_movie() and allow_write:
             self.close_partial_movie_stream()
 
-    def listen_and_write(self) -> None:
-        """For internal use only: blocks until new frame is available on the queue."""
-        while True:
-            num_frames, frame_data = self.queue.get()
-            if frame_data is None:
-                break
-
-            self.encode_and_write_frame(frame_data, num_frames)
-
-    def encode_and_write_frame(self, frame: PixelArray, num_frames: int) -> None:
-        """For internal use only: takes a given frame in ``np.ndarray`` format and
-        writes it to the stream
-        """
-        for _ in range(num_frames):
-            # Notes: precomputing reusing packets does not work!
-            # I.e., you cannot do `packets = encode(...)`
-            # and reuse it, as it seems that `mux(...)`
-            # consumes the packet.
-            # The same issue applies for `av_frame`,
-            # reusing it renders weird-looking frames.
-            av_frame = av.VideoFrame.from_ndarray(frame, format="rgba")
-            for packet in self.video_stream.encode(av_frame):
-                self.video_container.mux(packet)
-
     def write_frame(
         self, frame_or_renderer: PixelArray | OpenGLRenderer, num_frames: int = 1
     ) -> None:
@@ -472,8 +538,12 @@
                     else frame_or_renderer
                 )
 
-            msg = (num_frames, frame)
-            self.queue.put(msg)
+            job = self._current_encode_job
+            if job is None:
+                # Interactive OpenGL rendering emits frames outside an open
+                # partial movie stream; drop them silently.
+                return
+            job.put(num_frames, frame)
 
         if is_png_format() and not config["dry_run"]:
             if isinstance(frame_or_renderer, np.ndarray):
@@ -524,9 +594,11 @@
         If save_last_frame is True, saves the last frame in the default image directory.
         """
         if write_to_movie():
+            self.join_all_encode_jobs()
             self.combine_to_movie()
             if config.save_sections:
                 self.combine_to_section_videos()
+            # Cache cleanup runs after the in-flight encode jobs have been drained.
             if config["flush_cache"]:
                 self.flush_cache_directory()
             else:
@@ -545,6 +617,14 @@
         """
         if file_path is None:
             file_path = self.partial_movie_files[self.renderer.num_plays]
+            if file_path is None:
+                raise RuntimeError(
+                    "open_partial_movie_stream() called for a play that has no "
+                    "partial movie file path.",
+                )
+        path_key = str(file_path)
+        if path_key in self._inflight_by_path:
+            self._join_job(self._inflight_by_path[path_key])
         self.partial_movie_file_path = file_path
 
         fps = to_av_frame_rate(config.frame_rate)
@@ -576,12 +656,34 @@
         stream.width = config.pixel_width
         stream.height = config.pixel_height
 
-        self.video_container: OutputContainer = video_container
-        self.video_stream: Stream = stream
+        self._current_encode_job = _PartialMovieEncodeJob(
+            path=file_path,
+            animation_index=self.renderer.num_plays,
+            container=video_container,
+            stream=stream,
+        )
 
-        self.queue: Queue[tuple[int, PixelArray | None]] = Queue()
-        self.writer_thread = Thread(target=self.listen_and_write, args=())
-        self.writer_thread.start()
+    def _join_job(self, job: _PartialMovieEncodeJob) -> None:
+        """Remove and join an in-flight partial movie encode job."""
+        if job in self._inflight_encode_jobs:
+            self._inflight_encode_jobs.remove(job)
+        self._inflight_by_path.pop(str(job.path), None)
+        job.join()
+
+    def join_all_encode_jobs(self) -> None:
+        """Join every in-flight encode job, re-raising the first failure."""
+        first_exception: BaseException | None = None
+        for job in list(self._inflight_encode_jobs):
+            try:
+                self._join_job(job)
+            except BaseException as exception:
+                if first_exception is None:
+                    first_exception = exception
+
+        self._inflight_encode_jobs.clear()
+        self._inflight_by_path.clear()
+        if first_exception is not None:
+            raise first_exception
 
     def close_partial_movie_stream(self) -> None:
         """Close the currently opened video container.
@@ -590,18 +692,19 @@
         in the video stream holding a partial file, and then close
         the corresponding container.
         """
-        self.queue.put((-1, None))
-        self.writer_thread.join()
-
-        for packet in self.video_stream.encode():
-            self.video_container.mux(packet)
-
-        self.video_container.close()
+        job = self._current_encode_job
+        if job is None:
+            raise RuntimeError(
+                "close_partial_movie_stream() called without an open partial "
+                "movie stream.",
+            )
+        job.seal()
+        self._inflight_encode_jobs.append(job)
+        self._inflight_by_path[str(job.path)] = job
+        self._current_encode_job = None
 
-        logger.info(
-            f"Animation {self.renderer.num_plays} : Partial movie file written in %(path)s",
-            {"path": f"'{self.partial_movie_file_path}'"},
-        )
+        while len(self._inflight_encode_jobs) >= config.max_inflight_encoders:
+            self._join_job(self._inflight_encode_jobs[0])
 
     def is_already_cached(self, hash_invocation: str) -> bool:
         """Will check if a file named with `hash_invocation` exists.
@@ -622,6 +725,9 @@
             self.partial_movie_directory
             / f"{hash_invocation}{config['movie_file_extension']}"
         )
+        path_key = str(path)
+        if path_key in self._inflight_by_path:
+            self._join_job(self._inflight_by_path[path_key])
         return path.exists()
 
     def combine_files(