[WandB] Push axolotl config to top level wandb files (#1014)
Browse files
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 |
-
|
| 565 |
-
|
| 566 |
-
|
| 567 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|