fix wandb hang

Author: Felipe Mello

apps/grpo/main.py

@@ -322,6 +322,7 @@ async def main(cfg: DictConfig):
         )
     metric_logging_cfg = cfg.get("metric_logging", {"console": {"log_per_rank": False}})
     mlogger = await get_or_create_metric_logger(process_name="Controller")
+    await mlogger.init_backends.call_one(metric_logging_cfg)
     await ts.initialize(strategy=ts.ControllerStorageVolumes())
 
     # ---- Setup services ---- #
@@ -350,11 +351,6 @@ async def main(cfg: DictConfig):
         ),
     )
 
-    # Call after services are initialized
-    # TODO (felipemello): if called before, and per_rank_share_run=True, it hangs
-    # probably wandb requires primary runs to finish before shared runs can be initialized
-    await mlogger.init_backends.call_one(metric_logging_cfg)
-
     print("All services initialized successfully!")
 
     # ---- Core RL loops ---- #

apps/toy_rl/toy_metrics/main.py

@@ -7,7 +7,7 @@
 import asyncio
 
 import logging
-import time
+from datetime import datetime
 
 from forge.controller.actor import ForgeActor
 from forge.controller.provisioner import shutdown
@@ -17,8 +17,13 @@
 
 from monarch.actor import current_rank, endpoint
 
-logging.basicConfig(level=logging.DEBUG)
-logging.getLogger("forge.observability.metrics").setLevel(logging.DEBUG)
+logging.basicConfig(
+    level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logging.getLogger("forge.observability.metrics").setLevel(logging.INFO)
+logging.getLogger("forge.observability.metric_actors").setLevel(logging.INFO)
+# Reduce wandb logging noise
+logging.getLogger("wandb").setLevel(logging.WARNING)
 
 
 class TrainActor(ForgeActor):
@@ -79,8 +84,7 @@ async def generate_step(self, step: int, substep: int):
 
 # Main
 async def main():
-    """Example demonstrating distributed metric logging with different backends."""
-    group = f"grpo_exp_{int(time.time())}"
+    group = "time" + str(int(datetime.now().timestamp()))
 
     # Config format: {backend_name: backend_config_dict}
     config = {
@@ -89,22 +93,18 @@ async def main():
             "project": "toy_metrics",
             "group": group,
             "logging_mode": "per_rank_no_reduce",
-            "per_rank_share_run": False,
+            "per_rank_share_run": True,
         },
     }
 
     service_config = {"procs": 2, "num_replicas": 2, "with_gpus": False}
     mlogger = await get_or_create_metric_logger(process_name="Controller")
+    await mlogger.init_backends.call_one(config)
 
-    # Spawn services first (triggers registrations via provisioner hook)
+    # Spawn services (will register fetchers)
     trainer = await TrainActor.options(**service_config).as_service()
     generator = await GeneratorActor.options(**service_config).as_service()
 
-    # Call after services are initialized
-    # TODO (felipemello): if called before, and per_rank_share_run=True, it hangs
-    # probably wandb requires primary runs to finish before shared runs can be initialized
-    await mlogger.init_backends.call_one(config)
-
     for i in range(3):
         print(f"\n=== Global Step {i} ===")
         record_metric("main/global_step", 1, Reduce.MEAN)

apps/vllm/main.py

@@ -27,7 +27,8 @@
 
 async def run(cfg: DictConfig):
     metric_logging_cfg = cfg.get("metric_logging", {"console": {"log_per_rank": False}})
-    mlogger = await get_or_create_metric_logger(actor_name="Controller")
+    mlogger = await get_or_create_metric_logger(process_name="Controller")
+    await mlogger.init_backends.call_one(metric_logging_cfg)
 
     if (prompt := cfg.get("prompt")) is None:
         gd = cfg.policy.get("sampling_config", {}).get("guided_decoding", False)
@@ -36,11 +37,6 @@ async def run(cfg: DictConfig):
     print("Spawning service...")
     policy = await Policy.options(**cfg.services.policy).as_service(**cfg.policy)
 
-    # Call after services are initialized
-    # TODO (felipemello): if called before, and per_rank_share_run=True, it hangs
-    # probably wandb requires primary runs to finish before shared runs can be initialized
-    await mlogger.init_backends.call_one(metric_logging_cfg)
-
     import time
 
     print("Requesting generation...")

src/forge/observability/metrics.py

@@ -568,9 +568,6 @@ async def flush(
             states[key] = acc.get_state()
             acc.reset()
 
-        # Update step (used by NO_REDUCE backends in push)
-        self.step = step
-
         # Log to PER_RANK_REDUCE backends only (NO_REDUCE already logged in push)
         if self.per_rank_reduce_backends:
             metrics_for_backends = reduce_metrics_states([states])
@@ -579,6 +576,9 @@ async def flush(
             for backend in self.per_rank_reduce_backends:
                 await backend.log_batch(metrics_for_backends, step)
 
+        # Update step (used by NO_REDUCE backends in push)
+        self.step = step + 1
+
         return states if return_state else {}
 
     async def shutdown(self):
@@ -768,22 +768,32 @@ async def _init_shared_global(self):
         settings = wandb.Settings(
             mode="shared", x_primary=True, x_label="controller_primary"
         )
-        self.run = wandb.init(project=self.project, group=self.group, settings=settings)
+
+        self.run = wandb.init(
+            project=self.project,
+            group=self.group,
+            settings=settings,
+        )
 
     async def _init_shared_local(self, primary_metadata: Dict[str, Any]):
         import wandb
+        from wandb.sdk.lib.service import service_token
 
         shared_id = primary_metadata.get("shared_run_id")
         if shared_id is None:
             raise ValueError(
                 f"Shared ID required but not provided for {self.name} backend init"
             )
+
+        # Clear any stale service tokens that might be pointing to dead processes
+        # In multiprocessing environments, WandB service tokens can become stale and point
+        # to dead service processes. This causes wandb.init() to hang indefinitely trying
+        # to connect to non-existent services. Clearing forces fresh service connection.
+        service_token.clear_service_in_env()
+
         settings = wandb.Settings(mode="shared", x_primary=False, x_label=self.name)
         self.run = wandb.init(
-            id=shared_id,
-            project=self.project,
-            group=self.group,
-            settings=settings,
+            id=shared_id, project=self.project, group=self.group, settings=settings
         )
 
     async def log_batch(