Limit number of metrics shown in gather_trajectory_groups

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "openpipe-art"
-version = "0.3.9"
+version = "0.3.10"
 description = "The OpenPipe Agent Reinforcement Training (ART) library"
 readme = "README.md"
 requires-python = ">=3.10"

src/art/gather.py

@@ -16,12 +16,14 @@ async def gather_trajectory_groups(
     pbar_desc: str | None = "gather",
     pbar_total_completion_tokens: bool = True,
     max_exceptions: int | float = 0,
+    max_metrics: int | None = None,
 ) -> list[TrajectoryGroup]:
     groups = list(groups)
     context = GatherContext(
         pbar=None,
         pbar_total_completion_tokens=pbar_total_completion_tokens,
         max_exceptions=max_exceptions,
+        max_metrics=max_metrics,
     )
     with set_gather_context(context):
         future = asyncio.gather(*[wrap_group_awaitable(g) for g in groups])
@@ -167,14 +169,18 @@ class GatherContext:
     pbar: tqdm.tqdm | None = None
     metric_sums: Counter[str] = field(default_factory=Counter)
     metric_divisors: Counter[str] = field(default_factory=Counter)
+    max_metrics: int | None = None
     pbar_total_completion_tokens: bool = False
     max_exceptions: int | float = 0
 
     def update_pbar(self, n: int) -> None:
         if self.pbar is not None:
             self.pbar.update(n)
             postfix = {}
-            for metric in self.metric_sums:
+            included_metrics = self.metric_sums.keys()
+            if self.max_metrics is not None:
+                included_metrics = list(self.metric_sums.keys())[: self.max_metrics]
+            for metric in included_metrics:
                 sum = self.metric_sums[metric]
                 divisor = max(1, self.metric_divisors[metric])
                 postfix[metric] = sum / divisor