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 <[email protected]>
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."
|