feat(wandb): log full RLConfig to W&B instead of per-component subconfig

src/prime_rl/entrypoints/rl.py

@@ -38,6 +38,7 @@
 
 RL_TOML = "rl.toml"
 RL_SBATCH = "rl.sbatch"
+FULL_CONFIG_JSON = "full_config.json"
 
 TRAINER_TOML = "trainer.toml"
 ORCHESTRATOR_TOML = "orchestrator.toml"
@@ -77,6 +78,13 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None:
             tomli_w.dump(to_toml_dict(config.inference, exclude=exclude_inference), f)
 
 
+def write_full_config_json(config: RLConfig, output_dir: Path) -> None:
+    """Write the full resolved RLConfig as JSON for wandb logging."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    with open(output_dir / FULL_CONFIG_JSON, "w") as f:
+        json.dump(config.model_dump(mode="json"), f, indent=2)
+
+
 def rl_local(config: RLConfig):
     assert config.deployment.type == "single_node"
 
@@ -558,6 +566,13 @@ def rl(config: RLConfig):
 
         pre_download_model(config.trainer.model.name)
 
+    # Write the full resolved config as JSON so the WandbMonitor can log it to W&B.
+    # Set the path in os.environ so all locally-launched subprocesses inherit it;
+    # the multi-node SLURM template exports it separately (srun doesn't inherit os.environ).
+    config_dir = config.output_dir / "configs"
+    write_full_config_json(config, config_dir)
+    os.environ["WANDB_FULL_CONFIG_PATH"] = str(config_dir / FULL_CONFIG_JSON)
+
     if config.slurm is not None:
         rl_slurm(config)
     else:

src/prime_rl/templates/multi_node_rl.sbatch.j2

@@ -66,6 +66,7 @@ ln -sfn inference/node_0.log $OUTPUT_DIR/logs/inference.log
 
 export WANDB_SHARED_MODE=1
 export WANDB_SHARED_RUN_ID=${WANDB_SHARED_RUN_ID:-$(python3 -c "import uuid; print(uuid.uuid4().hex)")}
+export WANDB_FULL_CONFIG_PATH={{ config_dir }}/full_config.json
 
 # Networking
 HOSTNAMES=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )

src/prime_rl/utils/monitor/wandb.py

@@ -99,6 +99,18 @@ def __init__(
 
         retryable_errors = (CommError, ServerResponseError) if shared_mode else (CommError,)
 
+        # Prefer the full RLConfig (written by the launcher) over the per-component
+        # subconfig, so the W&B run captures the entire resolved configuration.
+        wandb_config = run_config.model_dump() if run_config else None
+        full_config_path = os.environ.get("WANDB_FULL_CONFIG_PATH")
+        if full_config_path and Path(full_config_path).exists():
+            try:
+                with open(full_config_path) as f:
+                    wandb_config = json.load(f)
+                self.logger.info(f"Using full config from {full_config_path} for W&B")
+            except Exception as e:
+                self.logger.warning(f"Failed to read full config from {full_config_path}: {e}")
+
         def init_wandb(max_retries: int):
             for attempt in range(max_retries):
                 try:
@@ -111,7 +123,7 @@ def init_wandb(max_retries: int):
                         group=config.group,
                         tags=config.tags,
                         dir=output_dir,
-                        config=run_config.model_dump() if run_config else None,
+                        config=wandb_config,
                         settings=settings,
                     )
                 except retryable_errors as e:
@@ -137,6 +149,13 @@ def init_wandb(max_retries: int):
         self.wandb = init_wandb(max_retries)
         self.run_id = self.wandb.id
 
+        # Upload config files to the W&B run so they're downloadable from the UI.
+        config_dir = Path(full_config_path).parent if full_config_path else None
+        if config_dir and config_dir.is_dir():
+            for config_file in config_dir.iterdir():
+                if config_file.is_file():
+                    wandb.save(str(config_file), policy="now")
+
         wandb.define_metric("*", step_metric="step")
 
         # Provision the curated "overview" saved view once per project (the run's primary process