hamel commited on
Commit
4f4d638
·
unverified ·
1 Parent(s): ba043a3

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

Browse files
Files changed (1) hide show
  1. src/axolotl/utils/callbacks.py +11 -4
src/axolotl/utils/callbacks.py CHANGED
@@ -4,6 +4,8 @@ from __future__ import annotations
4
 
5
  import logging
6
  import os
 
 
7
  from typing import TYPE_CHECKING, Dict, List
8
 
9
  import evaluate
@@ -561,10 +563,15 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
561
  ):
562
  if is_main_process():
563
  try:
564
- artifact = wandb.Artifact(name="axolotl-config", type="config")
565
- artifact.add_file(local_path=self.axolotl_config_path)
566
- wandb.run.log_artifact(artifact)
567
- LOG.info("Axolotl config has been saved to WandB as an artifact.")
 
 
 
 
 
568
  except (FileNotFoundError, ConnectionError) as err:
569
  LOG.warning(f"Error while saving Axolotl config to WandB: {err}")
570
  return control
 
4
 
5
  import logging
6
  import os
7
+ from shutil import copyfile
8
+ from tempfile import NamedTemporaryFile
9
  from typing import TYPE_CHECKING, Dict, List
10
 
11
  import evaluate
 
563
  ):
564
  if is_main_process():
565
  try:
566
+ # 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.
567
+ with NamedTemporaryFile(
568
+ mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
569
+ ) as temp_file:
570
+ copyfile(self.axolotl_config_path, temp_file.name)
571
+ wandb.save(temp_file.name)
572
+ LOG.info(
573
+ "The Axolotl config has been saved to the WandB run under files."
574
+ )
575
  except (FileNotFoundError, ConnectionError) as err:
576
  LOG.warning(f"Error while saving Axolotl config to WandB: {err}")
577
  return control