add accumulator

Author: DNXie

src/forge/observability/__init__.py

@@ -27,6 +27,7 @@
     Reduce,
     reduce_metrics_states,
     RewardThresholdFilter,
+    SampleAccumulator,
     StdAccumulator,
     SumAccumulator,
     TopBottomKFilter,
@@ -62,6 +63,8 @@
     "MaxAccumulator",
     "MinAccumulator",
     "StdAccumulator",
+    "SampleAccumulator",
+    # Filter classes
     "TopBottomKFilter",
     "RewardThresholdFilter",
     "RandomRatioFilter",

src/forge/observability/metrics.py

@@ -24,6 +24,7 @@ class Reduce(Enum):
     MAX = "max"
     MIN = "min"
     STD = "std"
+    SAMPLE = "sample"
 
     @property
     def accumulator_class(self):
@@ -33,6 +34,7 @@ def accumulator_class(self):
             Reduce.MAX: MaxAccumulator,
             Reduce.MIN: MinAccumulator,
             Reduce.STD: StdAccumulator,
+            Reduce.SAMPLE: SampleAccumulator,
         }
         return mapping[self]
 
@@ -188,6 +190,10 @@ def filter_flush(self, samples: List[Dict]) -> List[Dict]:
         """
         return samples
 
+    def reset(self) -> None:
+        """Clears for next accumulation cycle."""
+        pass
+
 
 class RandomRatioFilter:
     """Randomly keep a fraction of samples."""
@@ -260,6 +266,10 @@ def filter_flush(self, samples: List[Dict]) -> List[Dict]:
         bottoms = [s for _, s in self._bottom_heap]
         return bottoms + tops
 
+    def reset(self):
+        self._top_heap = []
+        self._bottom_heap = []
+
 
 ################
 # Accumulators #
@@ -449,6 +459,56 @@ def reset(self) -> None:
         self.count = 0
 
 
+class SampleAccumulator(MetricAccumulator):
+    """Accumulator for sample-level metrics (e.g., prompt/response/reward dicts).
+
+    Optionally uses a SampleFilter to decide what to keep at append/flush time.
+    """
+
+    def __init__(self, reduction: Reduce, filter: SampleFilter | None = None):
+        super().__init__(reduction)
+        self.samples: List[Dict[str, Any]] = []
+        self.filter = filter
+
+    def append(self, value: dict) -> None:
+        assert isinstance(value, dict)
+
+        # If filter is provided, only keep the sample if filter_append returns True
+        if self.filter:
+            if self.filter.filter_append(value):
+                self.samples.append(value)
+        else:
+            self.samples.append(value)
+
+    def get_value(self) -> list[dict]:
+        """Return locally collected (and optionally filtered) samples."""
+        if self.filter:
+            # Apply flush-time filter (e.g. heap selection, threshold trimming)
+            return self.filter.filter_flush(self.samples)
+        return self.samples
+
+    def get_state(self) -> Dict[str, Any]:
+        """Serialize accumulator state for cross-rank reduction."""
+        return {
+            "reduction_type": self.reduction_type.value,
+            "samples": self.get_value(),
+        }
+
+    @classmethod
+    def get_reduced_value_from_states(cls, states: List[Dict[str, Any]]) -> list[dict]:
+        """Merge sample states across ranks."""
+        merged = []
+        for s in states:
+            merged.extend(s.get("samples", []))
+        return merged
+
+    def reset(self) -> None:
+        """Clear local samples and reset filter state."""
+        self.samples.clear()
+        if self.filter:
+            self.filter.reset()
+
+
 #############
 # Collector #
 #############

tests/unit_tests/data/test_metrics_sampler.py

@@ -17,7 +17,9 @@
 
 from forge.observability.metrics import (
     RandomRatioFilter,
+    Reduce,
     RewardThresholdFilter,
+    SampleAccumulator,
     TopBottomKFilter,
 )
 
@@ -53,14 +55,30 @@ def test_reward_threshold_filter_gt():
     assert f.filter_append(make_sample(0.2)) is False
 
 
-def test_top_bottom_k_filter_simple():
-    f = TopBottomKFilter(top_k=2, bottom_k=2, key="reward")
+def test_sample_accumulator_with_topbottom_filter():
+    """Ensure SampleAccumulator integrates with TopBottomKFilter correctly."""
+    f = TopBottomKFilter(top_k=2, bottom_k=1, key="reward")
+    acc = SampleAccumulator(Reduce.SAMPLE, filter=f)
+
     rewards = [0.1, 0.9, 0.5, 0.7, 0.3]
     for r in rewards:
-        f.filter_append(make_sample(r))
+        acc.append(make_sample(r))
+
+    result = acc.get_value()
+    result_rewards = sorted(s["reward"] for s in result)
+
+    # Expect bottom-1 (0.1) and top-2 (0.7, 0.9)
+    assert result_rewards == [0.1, 0.7, 0.9]
+
+
+def test_sample_accumulator_no_filter_returns_all():
+    """Ensure SampleAccumulator without a filter returns all samples."""
+    acc = SampleAccumulator(Reduce.SAMPLE, filter=None)
 
-    samples = f.filter_flush([])
-    values = sorted(s["reward"] for s in samples)
+    samples = [make_sample(r) for r in [0.2, 0.4, 0.6]]
+    for s in samples:
+        acc.append(s)
 
-    # Expect bottom-2 (0.1, 0.3) and top-2 (0.7, 0.9)
-    assert values == [0.1, 0.3, 0.7, 0.9]
+    result = acc.get_value()
+    assert len(result) == len(samples)
+    assert [s["reward"] for s in result] == [0.2, 0.4, 0.6]