Skip to content

Commit

Permalink
[WandB] Push axolotl config to top level wandb files (#1014)
Browse files Browse the repository at this point in the history
  • Loading branch information
hamelsmu authored Dec 29, 2023
1 parent ba043a3 commit 4f4d638
Showing 1 changed file with 11 additions and 4 deletions.
15 changes: 11 additions & 4 deletions src/axolotl/utils/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

import logging
import os
from shutil import copyfile
from tempfile import NamedTemporaryFile
from typing import TYPE_CHECKING, Dict, List

import evaluate
Expand Down Expand Up @@ -561,10 +563,15 @@ def on_train_begin(
):
if is_main_process():
try:
artifact = wandb.Artifact(name="axolotl-config", type="config")
artifact.add_file(local_path=self.axolotl_config_path)
wandb.run.log_artifact(artifact)
LOG.info("Axolotl config has been saved to WandB as an artifact.")
# sync config to top level in run, cannot delete file right away because wandb schedules it to be synced even w/policy = 'now', so let OS delete it later.
with NamedTemporaryFile(
mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
) as temp_file:
copyfile(self.axolotl_config_path, temp_file.name)
wandb.save(temp_file.name)
LOG.info(
"The Axolotl config has been saved to the WandB run under files."
)
except (FileNotFoundError, ConnectionError) as err:
LOG.warning(f"Error while saving Axolotl config to WandB: {err}")
return control

0 comments on commit 4f4d638

Please sign in to comment.