Update SaveAxolotlConfigtoWandBCallback to use artifact instead of sa…

…ve (#1483) * deprecated wandb.save * also use wandb.save for axolotl yaml * chore: lint --------- Co-authored-by: Wing Lian <[email protected]>
axolotl-ai-cloud · Apr 9, 2024 · 5ed2939 · 5ed2939
1 parent da9b1a3
commit 5ed2939
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py
@@ -761,6 +761,11 @@ def on_train_begin(
                     mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
                 ) as temp_file:
                     copyfile(self.axolotl_config_path, temp_file.name)
+                    artifact = wandb.Artifact(
+                        f"config-{wandb.run.id}", type="axolotl-config"
+                    )
+                    artifact.add_file(temp_file.name)
+                    wandb.log_artifact(artifact)
                     wandb.save(temp_file.name)
                 LOG.info(
                     "The Axolotl config has been saved to the WandB run under files."