From c3813684e2ad5d1b6908ffe714f14d8b25f2748d Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 1 Apr 2026 13:35:38 +0000 Subject: [PATCH] feat: log metrics to human-readable JSONL files W&B stores metrics in a binary protobuf format that isn't readable locally. This adds a JSONL sidecar file so metrics can be inspected by humans and agents without the W&B API. Each process writes to its own file under {output_dir}/metrics/ to avoid race conditions: in shared W&B mode the filename uses the process label (e.g. trainer.jsonl, orchestrator.jsonl), otherwise it defaults to metrics.jsonl. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/prime_rl/utils/monitor/wandb.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index 96465688f9..cc05fd6c39 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -31,6 +31,7 @@ def __init__( self.logger = get_logger() self.history: list[dict[str, Any]] = [] self.output_dir = output_dir + self._jsonl_file = None rank = int(os.environ.get("RANK", os.environ.get("DP_RANK", "0"))) self.enabled = self.config is not None @@ -46,6 +47,7 @@ def __init__( self._maybe_overwrite_wandb_command() shared_mode = os.environ.get("WANDB_SHARED_MODE") == "1" + self._setup_jsonl(shared_mode) if shared_mode: run_id = os.environ.get("WANDB_SHARED_RUN_ID") label = os.environ.get("WANDB_SHARED_LABEL") @@ -108,6 +110,16 @@ def init_wandb(max_retries: int): log_mode="INCREMENTAL", ) + def _setup_jsonl(self, shared_mode: bool) -> None: + """Open a JSONL file for human-readable metric logging.""" + if self.output_dir is None: + return + label = os.environ.get("WANDB_SHARED_LABEL", "metrics") if shared_mode else "metrics" + path = self.output_dir / "metrics" / f"{label}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + self._jsonl_file = open(path, "a") + self.logger.info(f"Logging metrics to {path}") + def _maybe_overwrite_wandb_command(self) -> None: """Overwrites sys.argv with the start command if it is set in the environment variables.""" wandb_args = os.environ.get("WANDB_ARGS", None) @@ -122,6 +134,9 @@ def log(self, metrics: dict[str, Any], step: int) -> None: if not self.enabled: return wandb.log({**metrics, "step": step}) + if self._jsonl_file is not None: + self._jsonl_file.write(json.dumps({**metrics, "step": step}, default=str) + "\n") + self._jsonl_file.flush() def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: """Logs rollouts to W&B table.""" @@ -228,6 +243,11 @@ def log_distributions(self, distributions: dict[str, list[float]], step: int) -> """Log distributions (no-op for W&B).""" pass + def close(self) -> None: + if self._jsonl_file is not None: + self._jsonl_file.close() + self._jsonl_file = None + def save_final_summary(self, filename: str = "final_summary.json") -> None: """Save final summary to W&B table.""" if not self.is_master or not self.enabled: