S2.32: tracker memory and time budgets; determinism test

VideoTracker::set_budget (default 2 GiB of state, optional seconds per
propagate call) stops with Error::Budget; state_bytes reports what it holds.
Two runs of the same prompts give bit-identical masks.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/sam2/src/session.rs

@@ -133,11 +133,34 @@ struct FrameFeatures {
     high: (Tensor, Tensor),
 }
 
+/// What a tracker may hold and take. A call that would exceed it stops with [`Error::Budget`].
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct Budget {
+    /// Memory banks, results and cached frame features (bytes).
+    pub max_state_bytes: usize,
+    /// Wall-clock seconds one `propagate` call may take (`None`: no limit).
+    pub max_seconds: Option<f64>,
+}
+
+impl Default for Budget {
+    fn default() -> Budget {
+        Budget { max_state_bytes: 2 << 30, max_seconds: None }
+    }
+}
+
+/// Bytes of one stored frame: memory features and position encoding `[4096, 64]`, the pointer.
+const FRAME_BYTES: usize = 2 * GRID * GRID * 64 * 4 + 256 * 4;
+/// Bytes of a conditioning frame's kept result (`256 x 256` logits).
+const RESULT_BYTES: usize = 4 * GRID * 4 * GRID * 4;
+/// Bytes of one frame's cached neck features (lowest level and position, two skip levels).
+const FEATURE_BYTES: usize = (2 * 256 * GRID * GRID + 32 * 16 * GRID * GRID + 64 * 4 * GRID * GRID) * 4;
+
 pub struct VideoTracker<'a> {
     model: Arc<VideoModel>,
     frames: Box<dyn FrameSource + 'a>,
     objects: BTreeMap<ObjectId, ObjectState>,
     cache: VecDeque<(usize, Arc<FrameFeatures>)>,
+    budget: Budget,
 }
 
 fn sine_embedding_1d(positions: &[f32], dim: usize) -> Result<Tensor> {
@@ -161,13 +184,32 @@ impl<'a> VideoTracker<'a> {
         if n == 0 || n > MAX_FRAMES {
             return Err(Error::Input(format!("a clip needs between 1 and {MAX_FRAMES} frames, got {n}")));
         }
-        Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new() })
+        Ok(VideoTracker { model, frames, objects: BTreeMap::new(), cache: VecDeque::new(), budget: Budget::default() })
     }
 
     pub fn num_frames(&self) -> usize {
         self.frames.num_frames()
     }
 
+    /// Replace the limits (see [`Budget`]).
+    pub fn set_budget(&mut self, budget: Budget) {
+        self.budget = budget;
+    }
+
+    /// Bytes held now: stored frames, kept results and cached features.
+    pub fn state_bytes(&self) -> usize {
+        let frames: usize = self.objects.values().map(|o| (o.cond.len() + o.non_cond.len()) * FRAME_BYTES + o.cond.len() * RESULT_BYTES).sum();
+        frames + self.cache.len() * FEATURE_BYTES
+    }
+
+    fn check_budget(&self) -> Result<()> {
+        let used = self.state_bytes();
+        if used > self.budget.max_state_bytes {
+            return Err(Error::Budget(format!("the tracker holds {} MB, over its {} MB limit", used >> 20, self.budget.max_state_bytes >> 20)));
+        }
+        Ok(())
+    }
+
     /// Forget every object and cached frame (the clip stays).
     pub fn reset(&mut self) {
         self.objects.clear();
@@ -298,6 +340,7 @@ impl<'a> VideoTracker<'a> {
             state.non_cond.insert(frame, entry);
             Self::prune(state, frame);
         }
+        self.check_budget()?;
         Ok(result)
     }
 
@@ -383,6 +426,7 @@ impl<'a> VideoTracker<'a> {
             state.non_cond.insert(frame, entry);
             Self::prune(state, frame);
         }
+        self.check_budget()?;
         Ok(result)
     }
 
@@ -410,6 +454,7 @@ impl<'a> VideoTracker<'a> {
         if ids.is_empty() {
             return Err(Error::Input("no object has been prompted".into()));
         }
+        let started = std::time::Instant::now();
         for k in 0..=span {
             let frame = if reverse { start.checked_sub(k) } else { Some(start + k).filter(|f| *f < n) };
             let Some(frame) = frame else { break };
@@ -417,6 +462,9 @@ impl<'a> VideoTracker<'a> {
                 if cancel.is_cancelled() {
                     return Err(Error::Cancelled);
                 }
+                if let Some(limit) = self.budget.max_seconds.filter(|l| started.elapsed().as_secs_f64() > *l) {
+                    return Err(Error::Budget(format!("tracking took more than {limit} s")));
+                }
                 let cached = self.objects.get(id).is_some_and(|o| o.cond.contains_key(&frame));
                 let stored = self.objects.get(id).and_then(|o| o.cond.get(&frame)).and_then(|o| o.result.clone());
                 let mask = match stored {

crates/sam2/tests/reference.rs

@@ -544,3 +544,51 @@ fn a_mask_prompt_through_the_object_tracker_interface() {
     let bad = TrackRequest::single(0, n, vec![(1, Prompt { points: vec![], bbox: None, mask: Some(vec![1.0; 5]) })]);
     assert!(model.tracker().unwrap().track(&bad, &mut get, &|| false, &mut |_, _, _| {}).is_err());
 }
+
+/// The same prompts on the same frames give the same masks, bit for bit, run after run (CPU).
+#[test]
+fn tracking_is_deterministic() {
+    use effectcraft_sam2::prompt::Prompt;
+    use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker};
+    let Some((w, r)) = load_video() else { return };
+    let model = std::sync::Arc::new(VideoModel::load(&w).unwrap());
+    let clip = r.get("clip.pixel_values").unwrap().narrow(0, 1, 3).unwrap();
+    let run = || {
+        let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap();
+        t.prompt(0, ObjectId(1), &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap();
+        let mut out = vec![];
+        t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |m| out.push(m)).unwrap();
+        out
+    };
+    let (a, b) = (run(), run());
+    assert_eq!(a.len(), 3);
+    for (x, y) in a.iter().zip(&b) {
+        assert_eq!(x.logits, y.logits, "frame {} differs between runs", x.frame);
+        assert_eq!((x.iou.to_bits(), x.object_score.to_bits()), (y.iou.to_bits(), y.object_score.to_bits()));
+    }
+}
+
+/// Hard limits: a tracker over its memory budget or past its time budget stops with Error::Budget.
+#[test]
+fn the_tracker_stops_at_its_budget() {
+    use effectcraft_sam2::Error;
+    use effectcraft_sam2::prompt::Prompt;
+    use effectcraft_sam2::session::{Budget, Cancel, Direction, ObjectId, VideoModel, VideoTracker};
+    let Some((w, r)) = load_video() else { return };
+    let model = std::sync::Arc::new(VideoModel::load(&w).unwrap());
+    let clip = r.get("clip.pixel_values").unwrap().narrow(0, 1, 3).unwrap();
+    let click = Prompt { points: vec![([416.0, 512.0], true)], bbox: None };
+    // memory: one frame's worth is more than a 1 MB limit
+    let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap();
+    t.set_budget(Budget { max_state_bytes: 1 << 20, max_seconds: None });
+    assert!(matches!(t.prompt(0, ObjectId(1), &click), Err(Error::Budget(_))));
+    // default limit: fine, and the bytes held are reported
+    let mut t = VideoTracker::new(model.clone(), Box::new(Clip(clip.clone()))).unwrap();
+    t.prompt(0, ObjectId(1), &click).unwrap();
+    let held = t.state_bytes();
+    assert!(held > 2 << 20 && held < 200 << 20, "{held}");
+    // time: a limit of zero stops the first propagation
+    t.set_budget(Budget { max_state_bytes: 2 << 30, max_seconds: Some(0.0) });
+    let r = t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |_| {});
+    assert!(matches!(r, Err(Error::Budget(_))), "{r:?}");
+}