Merge pull request #1763 from Higgs32584/fix/1615-cpu-inflight-cap

Export: cap CPU frames in flight with EFFECTCRAFT_CPU_INFLIGHT (#1615)

crates/export/src/pipeline.rs

@@ -39,6 +39,57 @@ pub(crate) fn gpu_inflight(var: Option<&str>) -> usize {
     }
 }
 
+/// Caps CPU frames in flight (#1615): a whole number of frames, at least 1. Unset, CPU renders
+/// keep one frame per thread in flight as before. A heavy comp can need gigabytes of working
+/// memory per frame, so 16 threads at once can get the process OOM-killed; the headless CLI
+/// reads no `memory.*` preferences, so this is how a render there is held to fewer frames.
+pub(crate) const CPU_INFLIGHT_ENV: &str = "EFFECTCRAFT_CPU_INFLIGHT";
+
+/// CPU frames in flight from the `EFFECTCRAFT_CPU_INFLIGHT` value: `None` (no cap) when unset or
+/// blank, a positive whole number, or `None` with a warning when it is 0, negative, not a
+/// number or too large.
+pub(crate) fn cpu_inflight(var: Option<&str>) -> Option<usize> {
+    let v = var?;
+    if v.trim().is_empty() {
+        return None;
+    }
+    match v.trim().parse::<usize>() {
+        Ok(n) if n >= 1 => Some(n),
+        _ => {
+            log::warn!(target: "effectcraft_export", "{CPU_INFLIGHT_ENV}={v:?} is not a whole number of frames of at least 1: CPU frames are not capped");
+            None
+        }
+    }
+}
+
+/// The frames-in-flight caps of a batch: `gpu` (#879) for the GPU's memory, `cpu` (#1615) for
+/// the process's.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub(crate) struct Inflight {
+    pub gpu: Option<usize>,
+    pub cpu: Option<usize>,
+}
+
+impl Inflight {
+    /// Frames rendered at once: the smaller cap while the GPU renders, and only the CPU cap once
+    /// it has stopped for good (the CPU fallback renders in parallel, #1511). `None`: uncapped.
+    pub fn renders(&self, gpu_live: bool) -> Option<usize> {
+        let gpu = self.gpu.filter(|_| gpu_live);
+        let n = match (gpu, self.cpu) {
+            (Some(g), Some(c)) => Some(g.min(c)),
+            (g, c) => g.or(c),
+        };
+        n.map(|n| n.max(1))
+    }
+
+    /// Frames per window, rendered (`renders` at a time) and then encoded in parallel: with a CPU
+    /// cap at most that many, so finished images held before encoding count against it too;
+    /// with a GPU cap alone, one per thread (the images are in RAM, not VRAM).
+    pub fn window(&self, renders: usize) -> usize {
+        self.cpu.map_or(renders.max(rayon::current_num_threads()), |c| c.max(1))
+    }
+}
+
 /// `render` then `encode` over the frames `ks`, results in `ks` order.
 ///
 /// Without `cap` every frame renders and encodes in one parallel batch (the CPU path, as
@@ -55,38 +106,41 @@ pub(crate) fn map_frames<I: Send, T: Send>(
     render: impl Fn(u64) -> I + Sync + Send,
     encode: impl Fn(u64, I) -> T + Sync + Send,
 ) -> Vec<T> {
-    map_frames_while(ks, cap, || true, render, encode)
+    map_frames_while(ks, Inflight { gpu: cap, cpu: None }, || true, render, encode)
 }
 
-/// [`map_frames`] whose `cap` holds only while `capped()`: once it turns false (the GPU stopped
-/// for good, so every frame falls back to the CPU), the remaining frames render in parallel
-/// like the uncapped CPU path instead of one chunk of `cap` at a time.
+/// [`map_frames`] with the caps of `caps`, the GPU cap holding only while `gpu_live()`: once it
+/// turns false (the GPU stopped for good, so every frame falls back to the CPU), the remaining
+/// frames render under the CPU cap alone, in parallel when there is none, instead of one chunk
+/// of the GPU cap at a time. A CPU cap also bounds each window ([`Inflight::window`]).
 pub(crate) fn map_frames_while<I: Send, T: Send>(
     ks: Vec<u64>,
-    cap: Option<usize>,
-    capped: impl Fn() -> bool,
+    caps: Inflight,
+    gpu_live: impl Fn() -> bool,
     render: impl Fn(u64) -> I + Sync + Send,
     encode: impl Fn(u64, I) -> T + Sync + Send,
 ) -> Vec<T> {
     use rayon::prelude::*;
-    let Some(n) = cap.map(|n| n.max(1)).filter(|_| capped()) else {
+    if caps.renders(gpu_live()).is_none() {
         return ks.into_par_iter().map(|k| encode(k, render(k))).collect();
-    };
-    let window = n.max(rayon::current_num_threads());
+    }
     let mut out = Vec::with_capacity(ks.len());
-    for (w, win) in ks.chunks(window).enumerate() {
-        if !capped() {
-            out.par_extend(ks[w * window..].par_iter().map(|&k| encode(k, render(k))));
+    let mut rest = &ks[..];
+    while !rest.is_empty() {
+        let Some(n) = caps.renders(gpu_live()) else {
+            out.par_extend(rest.par_iter().map(|&k| encode(k, render(k))));
             break;
-        }
+        };
+        let (win, tail) = rest.split_at(caps.window(n).min(rest.len()));
         let mut imgs = Vec::with_capacity(win.len());
         let mut i = 0;
         while i < win.len() {
-            let end = if capped() { (i + n).min(win.len()) } else { win.len() };
+            let end = caps.renders(gpu_live()).map_or(win.len(), |n| (i + n).min(win.len()));
             imgs.par_extend(win[i..end].par_iter().map(|&k| render(k)));
             i = end;
         }
         out.par_extend(win.par_iter().copied().zip(imgs.into_par_iter()).map(|(k, img)| encode(k, img)));
+        rest = tail;
     }
     out
 }
@@ -226,10 +280,10 @@ impl<'a> Cx<'a> {
             }
             return out;
         }
-        // The cap is for the GPU's memory: once it has stopped for good, frames render on the
-        // CPU in parallel.
+        // The GPU cap is for the GPU's memory: once it has stopped for good, frames render on the
+        // CPU, in parallel or under the CPU cap (#1615), which holds throughout.
         let gpu_live = || self.job.accel.is_none_or(|a| a.available());
-        map_frames_while(ks, self.gpu_cap(), gpu_live, |k| self.frame(comp, k), f)
+        map_frames_while(ks, self.inflight(), gpu_live, |k| self.frame(comp, k), f)
     }
 
     /// [`Cx::frames`] for `ks` while `encode` runs beside them on `beside`'s threads, so the
@@ -286,6 +340,11 @@ impl<'a> Cx<'a> {
         (self.job.accel.is_some() && self.project.settings.gpu_acceleration).then(|| gpu_inflight(std::env::var(GPU_INFLIGHT_ENV).ok().as_deref()))
     }
 
+    /// The GPU and CPU frames-in-flight caps every frame loop shares (#1615).
+    fn inflight(&self) -> Inflight {
+        Inflight { gpu: self.gpu_cap(), cpu: cpu_inflight(std::env::var(CPU_INFLIGHT_ENV).ok().as_deref()) }
+    }
+
     /// Output frame `i`: sampled (fields, pulldown), at the colour depth, cropped and resized.
     pub fn frame(&self, comp: &Comp, i: u64) -> Image {
         effectcraft_render::passes::block_on(self.frame_async(comp, i))
@@ -363,7 +422,7 @@ impl<'a> Cx<'a> {
         let gpu_live = || self.job.accel.is_none_or(|a| a.available());
         map_frames_while(
             ks,
-            self.gpu_cap(),
+            self.inflight(),
             gpu_live,
             |k| self.yuv420_frame(comp, k, w, h).ok_or_else(|| self.frame(comp, k)),
             |k, r| match r {
@@ -575,6 +634,7 @@ fn premultiplied8(img: &Image) -> Vec<u8> {
 mod tests {
     use super::*;
     use rayon::prelude::*;
+    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 
     /// Frames `0..n` through [`map_frames`] with `cap`, recording the most running at once.
     fn run_capped(n: u64, cap: Option<usize>) -> (Vec<u64>, usize) {
@@ -607,7 +667,7 @@ mod tests {
         let out = pool.install(|| {
             map_frames_while(
                 (0..24).collect(),
-                Some(1),
+                Inflight { gpu: Some(1), cpu: None },
                 || live.load(Ordering::SeqCst),
                 |k| {
                     if k == 1 {
@@ -767,6 +827,103 @@ mod tests {
         assert_eq!(Cx::new(&job).gpu_cap(), None, "no accelerator");
     }
 
+    #[test]
+    fn cpu_inflight_is_off_unless_a_positive_number() {
+        assert_eq!(cpu_inflight(None), None);
+        assert_eq!(cpu_inflight(Some("2")), Some(2));
+        assert_eq!(cpu_inflight(Some(" 4 ")), Some(4));
+        for bad in ["0", "", "  ", "-2", "three", "2.5", "99999999999999999999999"] {
+            assert_eq!(cpu_inflight(Some(bad)), None, "{bad:?}");
+        }
+    }
+
+    #[test]
+    fn cpu_cap_holds_and_combines_with_the_gpu_cap() {
+        // #1615: CPU renders had no cap; a GPU cap alone still lapses when the GPU stops.
+        let caps = |gpu, cpu| Inflight { gpu, cpu };
+        for live in [true, false] {
+            assert_eq!(caps(None, None).renders(live), None);
+            assert_eq!(caps(None, Some(2)).renders(live), Some(2));
+        }
+        assert_eq!(caps(Some(1), None).renders(true), Some(1));
+        assert_eq!(caps(Some(1), None).renders(false), None);
+        assert_eq!(caps(Some(3), Some(2)).renders(true), Some(2));
+        assert_eq!(caps(Some(1), Some(4)).renders(true), Some(1));
+        // Review of #1615: once the GPU stops, the CPU cap applies, not the GPU's 1.
+        assert_eq!(caps(Some(1), Some(4)).renders(false), Some(4));
+        // A CPU cap of 2 on 8 threads renders at most 2 frames at once.
+        let pool = rayon::ThreadPoolBuilder::new().num_threads(8).build().unwrap();
+        let (now, max) = (AtomicUsize::new(0), AtomicUsize::new(0));
+        let out = pool.install(|| map_frames_while((0..16).collect(), caps(None, Some(2)), || true, |k| track(&now, &max, k), |_, k| k));
+        assert_eq!(out, (0..16).collect::<Vec<_>>());
+        assert!(max.load(Ordering::SeqCst) <= 2, "{max:?} frames in flight");
+    }
+
+    /// Counts `k` in flight in `now`, records the most in `max`, and holds it for a moment.
+    fn track(now: &AtomicUsize, max: &AtomicUsize, k: u64) -> u64 {
+        let c = now.fetch_add(1, Ordering::SeqCst) + 1;
+        max.fetch_max(c, Ordering::SeqCst);
+        std::thread::sleep(std::time::Duration::from_millis(5));
+        now.fetch_sub(1, Ordering::SeqCst);
+        k
+    }
+
+    #[test]
+    fn a_cpu_cap_bounds_the_images_waiting_to_encode() {
+        // Review of #1615: a window of one image per thread was rendered before any encoded, so
+        // a cap of 2 on 8 threads still held 8 finished images.
+        let pool = rayon::ThreadPoolBuilder::new().num_threads(8).build().unwrap();
+        for gpu in [None, Some(1)] {
+            let (held, max) = (AtomicUsize::new(0), AtomicUsize::new(0));
+            let out = pool.install(|| {
+                map_frames_while(
+                    (0..20).collect(),
+                    Inflight { gpu, cpu: Some(2) },
+                    || true,
+                    |k| {
+                        let c = held.fetch_add(1, Ordering::SeqCst) + 1;
+                        max.fetch_max(c, Ordering::SeqCst);
+                        k
+                    },
+                    |k, img| {
+                        held.fetch_sub(1, Ordering::SeqCst);
+                        assert_eq!(k, img);
+                        k
+                    },
+                )
+            });
+            assert_eq!(out, (0..20).collect::<Vec<_>>());
+            assert!(max.load(Ordering::SeqCst) <= 2, "gpu {gpu:?}: {max:?} images held");
+        }
+    }
+
+    #[test]
+    fn the_cpu_cap_takes_over_once_the_gpu_stops() {
+        // Review of #1615: with both caps, frames after the GPU stopped rendered one at a time.
+        let pool = rayon::ThreadPoolBuilder::new().num_threads(8).build().unwrap();
+        let (live, now, max) = (AtomicBool::new(true), AtomicUsize::new(0), AtomicUsize::new(0));
+        let out = pool.install(|| {
+            map_frames_while(
+                (0..40).collect(),
+                Inflight { gpu: Some(1), cpu: Some(4) },
+                || live.load(Ordering::SeqCst),
+                |k| {
+                    if k == 1 {
+                        live.store(false, Ordering::SeqCst);
+                    }
+                    if k < 4 {
+                        return k;
+                    }
+                    track(&now, &max, k)
+                },
+                |_, k| k,
+            )
+        });
+        assert_eq!(out, (0..40).collect::<Vec<_>>());
+        let max = max.load(Ordering::SeqCst);
+        assert!((2..=4).contains(&max), "{max} frames in flight after the GPU stopped");
+    }
+
     #[test]
     fn gpu_inflight_override_is_a_positive_number_or_the_default() {
         assert_eq!(gpu_inflight(None), GPU_INFLIGHT_DEFAULT);