and another one

Author: Felipe Mello

src/forge/observability/metric_actors.py

@@ -14,6 +14,7 @@
 from forge.env import FORGE_DISABLE_METRICS
 from forge.observability.metrics import (
     BackendRole,
+    context,
     get_logger_backend_class,
     LoggerBackend,
     MetricCollector,

src/forge/observability/metrics.py

@@ -465,7 +465,7 @@ async def init_backends(
             return
 
         self.global_step = global_step
-        self.proc_name_with_rank = get_proc_name_with_rank(self.process_name)
+        self.proc_name_with_rank = get_proc_name_with_rank(process_name)
 
         # instantiate local backends if any
         for backend_name, backend_config in config.items():
@@ -482,7 +482,7 @@ async def init_backends(
             await logger_backend.init(
                 role=BackendRole.LOCAL,
                 primary_logger_metadata=primary_metadata,
-                name=proc_name_with_rank,
+                name=self.proc_name_with_rank,
             )
             self.logger_backends.append(logger_backend)
 
@@ -632,6 +632,7 @@ async def log(self, metrics: list[Metric], global_step: int) -> None:
         """
         pass
 
+    @abstractmethod
     async def finish(self) -> None:
         pass
 

tests/unit_tests/observability/conftest.py

@@ -9,33 +9,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from forge.observability.metrics import LoggerBackend, MetricCollector
-
-
-class MockBackend(LoggerBackend):
-    """Mock backend for testing metrics logging without external dependencies."""
-
-    def __init__(self, logger_backend_config=None):
-        super().__init__(logger_backend_config or {})
-        self.logged_metrics = []
-        self.init_called = False
-        self.finish_called = False
-        self.metadata = {}
-
-    async def init(self, role="local", primary_logger_metadata=None, process_name=None):
-        self.init_called = True
-        self.role = role
-        self.primary_logger_metadata = primary_logger_metadata or {}
-        self.process_name = process_name
-
-    async def log(self, metrics, global_step):
-        self.logged_metrics.append((metrics, global_step))
-
-    async def finish(self):
-        self.finish_called = True
-
-    def get_metadata_for_secondary_ranks(self):
-        return self.metadata
+from forge.observability.metrics import MetricCollector
 
 
 @pytest.fixture(autouse=True)

tests/unit_tests/observability/test_metric_actors.py

@@ -62,7 +62,7 @@ async def test_global_logger_basic_ops(self, global_logger):
     async def test_backend_init(self, local_fetcher):
         """Test backend initialization and shutdown."""
         metadata = {"wandb": {"shared_run_id": "test123"}}
-        config = {"console": {"logging_mode": "per_rank_reduce"}}
+        config = {"console": {"reduce_across_ranks": False}}
 
         await local_fetcher.init_backends.call_one(metadata, config, global_step=5)
         await local_fetcher.shutdown.call_one()
@@ -108,9 +108,9 @@ async def test_valid_backend_configs(self, global_logger):
         # Empty config
         await global_logger.init_backends.call_one({})
 
-        # Valid configs for all logging modes
-        for mode in ["per_rank_reduce", "per_rank_no_reduce", "global_reduce"]:
-            config = {"console": {"logging_mode": mode}}
+        # Valid configs for different reduce_across_ranks modes
+        for reduce_across_ranks in [True, False]:
+            config = {"console": {"reduce_across_ranks": reduce_across_ranks}}
             await global_logger.init_backends.call_one(config)
 
     @pytest.mark.timeout(3)
@@ -124,7 +124,7 @@ async def test_invalid_backend_configs(self, global_logger):
         config_with_project = {"console": {"project": "test_project"}}
         await global_logger.init_backends.call_one(config_with_project)
 
-        # Config with reduce_across_ranks should work (Diff 3 doesn't validate logging_mode yet)
+        # Config with reduce_across_ranks should work
         config_with_reduce = {"console": {"reduce_across_ranks": True}}
         await global_logger.init_backends.call_one(config_with_reduce)