Metric Logging updates 4/N - better actor name (#351)

Co-authored-by: Felipe Mello <[email protected]>

Author: 2 people

apps/grpo/main.py

@@ -305,7 +305,7 @@ async def main(cfg: DictConfig):
         provisioner = await init_provisioner()
 
     metric_logging_cfg = cfg.get("metric_logging", {"console": {"log_per_rank": False}})
-    mlogger = await get_or_create_metric_logger()
+    mlogger = await get_or_create_metric_logger(process_name="Controller")
     await mlogger.init_backends.call_one(metric_logging_cfg)
 
     # ---- Setup services ---- #

src/forge/controller/provisioner.py

@@ -313,7 +313,7 @@ def bootstrap(env: dict[str, str]):
         if not FORGE_DISABLE_METRICS.get_value():
             from forge.observability.metric_actors import get_or_create_metric_logger
 
-            _ = await get_or_create_metric_logger(procs)
+            _ = await get_or_create_metric_logger(procs, process_name=mesh_name)
         return procs
 
     async def host_mesh_from_proc(self, proc_mesh: ProcMesh):

src/forge/observability/__init__.py

@@ -12,8 +12,6 @@
 from .metrics import (
     BackendRole,
     ConsoleBackend,
-    get_actor_name_with_rank,
-    get_logger_backend_class,
     LoggerBackend,
     MaxAccumulator,
     MeanAccumulator,
@@ -29,12 +27,12 @@
     WandbBackend,
 )
 from .perf_tracker import trace, Tracer
+from .utils import get_proc_name_with_rank
 
 __all__ = [
     # Main API functions
     "record_metric",
     "reduce_metrics_states",
-    "get_actor_name_with_rank",
     "get_logger_backend_class",
     "get_or_create_metric_logger",
     # Performance tracking
@@ -45,6 +43,8 @@
     "BackendRole",
     # Enums
     "Reduce",
+    # Utility functions
+    "get_proc_name_with_rank",
     # Actor classes
     "GlobalLoggingActor",
     "LocalFetcherActor",

src/forge/observability/metric_actors.py

@@ -8,7 +8,14 @@
 import logging
 from typing import Any, Union
 
-from monarch.actor import Actor, endpoint, get_or_spawn_controller, ProcMesh, this_proc
+from monarch.actor import (
+    Actor,
+    context,
+    endpoint,
+    get_or_spawn_controller,
+    ProcMesh,
+    this_proc,
+)
 
 from forge.env import FORGE_DISABLE_METRICS
 from forge.observability.metrics import (
@@ -27,6 +34,7 @@
 
 async def get_or_create_metric_logger(
     proc_mesh: ProcMesh | None = None,
+    process_name: str | None = None,
 ) -> "GlobalLoggingActor":
     """Initializes a LocalFetcherActor in the specified process mesh (or current process if None),
     if not already initialized, registers it with the GlobalLoggingActor and returns the
@@ -40,6 +48,9 @@ async def get_or_create_metric_logger(
     Args:
         proc_mesh: Optional ProcMesh to spawn LocalFetcherActor on. If None,
             uses `monarch.actor.this_proc()`.
+        process_name: Optional process name (e.g., "TrainActor", "GeneratorActor") for logging.
+            If None, will be auto-detected from the mesh_name provided during actor initialization or
+                a generic mesh name if one was not provided.
 
     Returns:
         GlobalLoggingActor: The global logging controller.
@@ -53,7 +64,7 @@ async def get_or_create_metric_logger(
         from forge.observability.metrics import record_metric
 
         # Main process setup
-        mlogger = await get_or_create_metric_logger()
+        mlogger = await get_or_create_metric_logger(process_name="Controller")
 
         # Initialize logging backends
         await mlogger.init_backends({
@@ -66,13 +77,14 @@ async def get_or_create_metric_logger(
 
         # Training loop
         for step in range(max_steps):
-            record_metric("loss", 1.2, step, reduction_type=Reduce.MEAN)
+            record_metric("loss", 1.2, reduction_type=Reduce.MEAN)
             # ... training code with record_metric() calls ...
             await mlogger.flush(step)  # Log metrics for this step
 
         # Shutdown
         await mlogger.shutdown()
     """
+
     # Get or create the singleton global logger
     global _global_logger
     if _global_logger is None:
@@ -84,14 +96,19 @@ async def get_or_create_metric_logger(
     # Determine process context
     proc = proc_mesh if proc_mesh is not None else this_proc()
 
+    # Auto-detect process_name from proc mesh if not provided
+    if process_name is None:
+        ctx = context()
+        process_name = ctx.actor_instance.actor_id.actor_name
+
     # Check current state for consistency
     proc_has_local_fetcher = hasattr(proc, "_local_fetcher")
     global_logger_has_local_fetcher = await global_logger.has_fetcher.call_one(proc)
 
     # Consistency check: both should be in sync
     if proc_has_local_fetcher != global_logger_has_local_fetcher:
         raise ValueError(
-            f"Inconsistent logging state for proc {proc}: "
+            f"Inconsistent logging state for {proc=} with {process_name=}: "
             f"proc has _local_fetcher={proc_has_local_fetcher}, "
             f"but global_logger has registration={global_logger_has_local_fetcher}. "
             f"This indicates a bug in logging setup/teardown. "
@@ -101,7 +118,7 @@ async def get_or_create_metric_logger(
     # Setup local_fetcher_actor if needed (unless disabled by environment flag)
     if not proc_has_local_fetcher and not FORGE_DISABLE_METRICS.get_value():
         local_fetcher_actor = proc.spawn(
-            "local_fetcher_actor", LocalFetcherActor, global_logger
+            "local_fetcher_actor", LocalFetcherActor, global_logger, process_name
         )
         await global_logger.register_fetcher.call_one(local_fetcher_actor, proc)
         proc._local_fetcher = local_fetcher_actor  # pyre-ignore
@@ -117,8 +134,13 @@ class LocalFetcherActor(Actor):
     GlobalLoggingActor -> per-rank LocalFetcherActor -> per-rank MetricCollector
     """
 
-    def __init__(self, global_logger: Union["GlobalLoggingActor", None] = None) -> None:
+    def __init__(
+        self,
+        global_logger: Union["GlobalLoggingActor", None] = None,
+        process_name: str | None = None,
+    ) -> None:
         self.global_logger = global_logger
+        self.process_name = process_name  # Passed to MetricCollector for logging
         _is_initialized = False
 
     @endpoint
@@ -145,10 +167,22 @@ async def init_backends(
         self,
         metadata_per_primary_backend: dict[str, dict[str, Any]],
         config: dict[str, Any],
+        global_step: int = 0,
     ) -> None:
-        """Init local (per-rank) logger backends and MetricCollector."""
+        """Init local (per-rank) logger backends and MetricCollector.
+
+        Args:
+            metadata_per_primary_backend (dict[str, dict[str, Any]]): Metadata from primary backends for shared state.
+            config (dict[str, Any]): Backend configurations with logging modes and settings.
+            global_step (int): Initial step for metrics.
+        """
         collector = MetricCollector()
-        await collector.init_backends(metadata_per_primary_backend, config)
+        await collector.init_backends(
+            metadata_per_primary_backend,
+            config,
+            global_step,
+            process_name=self.process_name,
+        )
 
     @endpoint
     async def shutdown(self) -> None: