tcapelle winglian commited on
Commit
5ed2939
1 Parent(s): da9b1a3

Update SaveAxolotlConfigtoWandBCallback to use artifact instead of save (#1483)

Browse files

* deprecated wandb.save

* also use wandb.save for axolotl yaml

* chore: lint

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>

src/axolotl/utils/callbacks/__init__.py CHANGED
@@ -761,6 +761,11 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
761
  mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
762
  ) as temp_file:
763
  copyfile(self.axolotl_config_path, temp_file.name)
 
 
 
 
 
764
  wandb.save(temp_file.name)
765
  LOG.info(
766
  "The Axolotl config has been saved to the WandB run under files."
 
761
  mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
762
  ) as temp_file:
763
  copyfile(self.axolotl_config_path, temp_file.name)
764
+ artifact = wandb.Artifact(
765
+ f"config-{wandb.run.id}", type="axolotl-config"
766
+ )
767
+ artifact.add_file(temp_file.name)
768
+ wandb.log_artifact(artifact)
769
  wandb.save(temp_file.name)
770
  LOG.info(
771
  "The Axolotl config has been saved to the WandB run under files."