tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[WandB] Push axolotl config to top level wandb files (#1014)

Browse Source
This commit is contained in:
Hamel Husain
2023-12-29 10:52:12 -08:00
committed by GitHub
parent ba043a361e
commit 4f4d638b84
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/axolotl/utils/callbacks.py
View File

@@ -4,6 +4,8 @@ from __future__ import annotations
import logging import logging
import os import os
from shutil import copyfile
from tempfile import NamedTemporaryFile
from typing import TYPE_CHECKING, Dict, List from typing import TYPE_CHECKING, Dict, List
import evaluate import evaluate
@@ -561,10 +563,15 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
): ):
if is_main_process(): if is_main_process():
try: try:
artifact = wandb.Artifact(name="axolotl-config", type="config") # 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.
artifact.add_file(local_path=self.axolotl_config_path) with NamedTemporaryFile(
wandb.run.log_artifact(artifact) mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
LOG.info("Axolotl config has been saved to WandB as an artifact.") ) 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: except (FileNotFoundError, ConnectionError) as err:
LOG.warning(f"Error while saving Axolotl config to WandB: {err}") LOG.warning(f"Error while saving Axolotl config to WandB: {err}")
return control return control