metric logger - backend (1/N)

apps/toy_metrics/main.py

@@ -0,0 +1,109 @@
+# 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.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_per_replica": 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.init_config.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_global.call_one(i)
+
+    await global_logger.shutdown.call_one()
+
+
+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,16 +3,14 @@
 #
 # 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
 
 
 # TODO - remove this once everything has moved to
 # service
 async def spawn_actors(
     name: str,
-    actor_cls: ForgeActor,
+    actor_cls,
     cfg,
     processes,
     set_address: bool = False,
@@ -28,5 +26,4 @@ async def spawn_actors(
     "spawn_actors",
     "stop_proc_mesh",
     "get_proc_mesh",
-    "ForgeActor",
 ]

src/forge/controller/provisioner.py

@@ -14,7 +14,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
@@ -119,6 +126,24 @@ 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 metrics on each process."""
+        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
     ):
@@ -184,6 +209,9 @@ def bootstrap(gpu_ids: int):
 
             procs._host = host_mesh
 
+            # Spawn local logging actor on each process and register with global logger
+            await self._setup_logging(procs)
+
             # If we created a server, track so we can tear it down later.
             if server_name:
                 self._server_names.append(server_name)
@@ -194,6 +222,16 @@ def bootstrap(gpu_ids: int):
     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.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/controller/service/replica.py

@@ -15,7 +15,7 @@
 
 from monarch.actor import ActorError
 
-from forge.controller import ForgeActor
+from forge.controller.actor import ForgeActor
 from forge.types import ProcessConfig
 
 logger = logging.getLogger(__name__)

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.