From 193eeae0e2453522125e31ba419811ec5d01cae0 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Wed, 22 Jul 2026 16:48:03 +0000 Subject: [PATCH 1/2] feat(wandb): log full RLConfig to W&B instead of per-component subconfig The rl entrypoint splits the config into subconfigs (trainer.toml, orchestrator.toml, inference.toml) before launching subprocesses. Each process only loads its subconfig, so wandb.init only received the orchestrator config (the primary in shared W&B mode). The trainer, inference, slurm, deployment, and shared sections were missing entirely. Fix: the rl() entrypoint writes the full resolved RLConfig as JSON (output_dir/configs/full_config.json) and sets PRIME_FULL_CONFIG_PATH in os.environ so all locally-launched subprocesses inherit it. The multi-node SLURM template exports the same env var separately (since srun tasks don't inherit the launcher's os.environ). The WandbMonitor reads this file and uses it as the wandb run config when available, falling back to the per-component subconfig otherwise. --- src/prime_rl/entrypoints/rl.py | 15 +++++++++++++++ src/prime_rl/templates/multi_node_rl.sbatch.j2 | 1 + src/prime_rl/utils/monitor/wandb.py | 14 +++++++++++++- 3 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 08e9f481ed..1ce1b78e6a 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/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: diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index f401f46a62..23143816c9 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/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 ) ) diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index ef8de745e7..c2cf436540 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/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: From 542734160b06a95f0f0614b886d27b452193c67a Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Wed, 22 Jul 2026 20:59:19 +0000 Subject: [PATCH 2/2] feat(wandb): upload config files to W&B run After wandb.init, save all files in the config_dir (full_config.json + subconfig TOMLs) to the W&B run so they are downloadable from the UI Files tab. Uses WANDB_FULL_CONFIG_PATH parent dir to locate the config directory. --- src/prime_rl/utils/monitor/wandb.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index c2cf436540..46886dc443 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -149,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