per_rank_no_reduce mode

Author: DNXie

src/forge/observability/metrics.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import asyncio
 import heapq
 import itertools
 import logging
@@ -715,12 +716,13 @@ def push(self, metric: Metric) -> None:
 
         # For PER_RANK_NO_REDUCE backends: stream without reduce
         for backend in self.per_rank_no_reduce_backends:
-            # if metric.reduction == Reduce.SAMPLE:
-            #     # Wrap singleton Metric into expected {key: [list_of_dicts]} format
-            #     sample = {metric.key: [metric.value]}
-            #     asyncio.create_task(backend.log_samples(sample, self.global_step))
-            # else:
-            backend.log_stream(metric=metric, global_step=self.global_step)
+
+            if metric.reduction == Reduce.SAMPLE:
+                # Wrap singleton Metric into expected {key: [list_of_dicts]} format
+                sample = {metric.key: [metric.value]}
+                asyncio.create_task(backend.log_samples(sample, self.global_step))
+            else:
+                backend.log_stream(metric=metric, global_step=self.global_step)
 
         # Always accumulate for reduction and state return
         key = metric.key