metric logger - backend (1/N)

apps/grpo/qwen3_1_7b.yaml

@@ -1,5 +1,5 @@
 # Grouped Relative Policy Optimization (GRPO)
-# >>> python -m apps.grpo.qwen3_1_7b --config apps/grpo/qwen3_1_7b.yaml
+# >>> python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml
 
 # Global configuration
 group_size: 8

apps/toy_metrics/main.py

@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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 logging
+import sys
+import time
+
+from forge.controller.actor import ForgeActor
+from forge.controller.provisioner import shutdown
+from forge.observability.metric_actors import GlobalLoggingActor
+from forge.observability.metrics import record_metric, ReductionType
+
+from monarch.actor import current_rank, endpoint, get_or_spawn_controller
+
+logging.basicConfig(level=logging.INFO)
+
+
+class TrainActor(ForgeActor):
+    @endpoint
+    async def train_step(self, step: int):
+        rank = current_rank().rank
+        value = rank * 1000 + 100 * step
+        print(f"🔧 Train rank {rank}: Step {step}, loss={value}")
+        await record_metric("train/loss", value)
+
+
+class GeneratorActor(ForgeActor):
+    @endpoint
+    async def generate_step(self, step: int, substep: int):
+        rank = current_rank().rank
+        value = rank * 1000 + step * 100 + substep * 10
+        print(f"🎯 Gen rank {rank}: Step {step}.{substep}, tokens={value}")
+        await record_metric("generate/tokens", value, ReductionType.SUM)
+
+
+# Main
+async def main(mode: str = "wandb_all_log_all"):
+    group = f"experiment_group_{int(time.time())}"
+    if mode == "wandb_all_log_all":
+        backends = [
+            {"class": "console", "log_per_rank": True},
+            {
+                "class": "wandb",
+                "project": "my_project",
+                "group": group,
+                "mode": "wandb_all_log_all",
+                "log_per_rank": True,
+            },
+        ]
+    elif mode == "wandb_rank_0_reduce_all":
+        backends = [
+            {"class": "console", "log_per_rank": False},
+            {
+                "class": "wandb",
+                "project": "my_project",
+                "group": group,
+                "mode": "wandb_rank_0_reduce_all",
+                "log_per_rank": False,
+            },
+        ]
+    else:  # wandb_rank_0_log_all
+        backends = [
+            {
+                "class": "wandb",
+                "project": "my_project",
+                "group": group,
+                "mode": "wandb_rank_0_log_all",
+                "log_per_rank": True,
+            },
+        ]
+
+    logging_config = {
+        "backends": backends,
+    }
+    service_config = {"procs": 2, "num_replicas": 2, "with_gpus": False}
+
+    # Spawn services first (triggers registrations via provisioner hook)
+    trainer = await TrainActor.options(**service_config).as_service()
+    generator = await GeneratorActor.options(**service_config).as_service()
+
+    # Now init config on global (inits backends eagerly across fetchers)
+    global_logger = await get_or_spawn_controller("global_logger", GlobalLoggingActor)
+    await global_logger.initialize_backends.call_one(logging_config)
+
+    for i in range(3):
+        print(f"\n=== Global Step {i} ===")
+        await trainer.train_step.call(i)
+        for sub in range(3):
+            await generator.generate_step.call(i, sub)
+        await global_logger.flush.call_one(i)
+
+    # shutdown
+    await asyncio.gather(global_logger.shutdown.call_one())
+
+    await asyncio.gather(
+        trainer.shutdown(),
+        generator.shutdown(),
+    )
+
+    await shutdown()
+
+
+if __name__ == "__main__":
+    mode = sys.argv[1] if len(sys.argv) > 1 else "wandb_all_log_all"
+    valid_modes = [
+        "wandb_all_log_all",
+        "wandb_rank_0_log_all",
+        "wandb_rank_0_reduce_all",
+    ]
+    if mode not in valid_modes:
+        print(f"Invalid mode: {mode}. Use {valid_modes}")
+        sys.exit(1)
+    asyncio.run(main(mode))

src/forge/controller/__init__.py

@@ -3,7 +3,6 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-
 from .actor import ForgeActor
 from .proc_mesh import get_proc_mesh, stop_proc_mesh
 
@@ -24,9 +23,4 @@ async def spawn_actors(
     return actors
 
 
-__all__ = [
-    "spawn_actors",
-    "stop_proc_mesh",
-    "get_proc_mesh",
-    "ForgeActor",
-]
+__all__ = ["spawn_actors", "stop_proc_mesh", "get_proc_mesh", "ForgeActor"]

src/forge/controller/provisioner.py

@@ -16,7 +16,14 @@
 import monarch
 from monarch._src.actor.allocator import RemoteAllocator, TorchXRemoteAllocInitializer
 from monarch._src.actor.shape import NDSlice, Shape
-from monarch.actor import Actor, endpoint, HostMesh, ProcMesh, this_host
+from monarch.actor import (
+    Actor,
+    endpoint,
+    get_or_spawn_controller,
+    HostMesh,
+    ProcMesh,
+    this_host,
+)
 from monarch.tools import commands
 from monarch.tools.components import hyperactor
 from monarch.tools.config import Config
@@ -142,6 +149,28 @@ async def create_host_mesh(self, name: str, num_hosts: int) -> HostMesh:
             server_name,
         )
 
+    async def _setup_logging(self, procs: ProcMesh) -> None:
+        """Spawn and register local fetcher for metric logging on each process.
+        When a service is spawned, we create for each rank a LocalFetcherActor and
+        store it at GlobalLoggingActor. Backends (e.g. wandb) should be eagerly instantiated
+        later in main by calling `global_logger.initialize_backends.call_one(logging_config)`
+        """
+        from forge.observability.metric_actors import (
+            GlobalLoggingActor,
+            LocalFetcherActor,
+        )
+
+        local_fetcher_actor = await procs.spawn(
+            "local_fetcher_actor", LocalFetcherActor
+        )
+        procs._local_fetcher = local_fetcher_actor
+
+        global_logger = await get_or_spawn_controller(
+            "global_logger", GlobalLoggingActor
+        )
+        process_name = f"proc_mesh_{id(procs)}"
+        await global_logger.register_fetcher.call_one(local_fetcher_actor, process_name)
+
     async def get_proc_mesh(
         self, num_procs: int, with_gpus: bool = False, num_hosts: int | None = None
     ):
@@ -215,11 +244,24 @@ def bootstrap(gpu_ids: list[str]):
                 self._server_names.append(server_name)
                 self._proc_server_map[procs] = server_name
 
+        # Spawn local logging actor on each process and register with global logger
+        await self._setup_logging(procs)
+
         return procs
 
     async def stop_proc_mesh(self, proc_mesh: ProcMesh):
         """Stops a proc mesh."""
         async with self._lock:
+            # Deregister local logger from global logger
+            if hasattr(proc_mesh, "_local_fetcher"):
+                from forge.observability.metric_actors import GlobalLoggingActor
+
+                global_logger = await get_or_spawn_controller(
+                    "global_logger", GlobalLoggingActor
+                )
+                process_name = f"proc_mesh_{id(proc_mesh)}"
+                await global_logger.deregister_fetcher.call_one(process_name)
+
             if hasattr(proc_mesh, "_gpu_ids"):
                 gpu_manager = self._host_gpu_map[proc_mesh._host._host_id]
                 gpu_manager.release_gpus(proc_mesh._gpu_ids)

src/forge/observability/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.