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..46886dc443 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: @@ -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