[#4] Add VESSL Callback to Post Metrics to VESSL AI

src/axolotl/cli/__init__.py

@@ -40,6 +40,7 @@
 from axolotl.utils.models import load_tokenizer
 from axolotl.utils.tokenization import check_dataset_labels
 from axolotl.utils.trainer import prepare_optim_env
+from axolotl.utils.vessl_ import setup_vessl_env_vars
 from axolotl.utils.wandb_ import setup_wandb_env_vars
 
 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -369,6 +370,8 @@ def load_cfg(config: Union[str, Path] = Path("examples/"), **kwargs):
 
     setup_mlflow_env_vars(cfg)
 
+    setup_vessl_env_vars(cfg)
+
     return cfg
 
 

src/axolotl/core/trainer_builder.py

@@ -836,6 +836,13 @@ def get_callbacks(self) -> List[TrainerCallback]:
                 SaveAxolotlConfigtoWandBCallback(self.cfg.axolotl_config_path)
             )
 
+        if self.cfg.use_vessl:
+            from axolotl.utils.callbacks.vessl_ import VesslLogMetricsCallback
+
+            callbacks.append(
+                VesslLogMetricsCallback(self.cfg.vessl_credential_path)
+            )
+
         return callbacks
 
     @abstractmethod

src/axolotl/utils/callbacks/vessl_.py

@@ -0,0 +1,19 @@
+"""Vessl module for trainer callbacks"""
+import logging
+from typing import Dict, List
+
+import vessl
+from transformers import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+LOG = logging.getLogger("axolotl.callbacks")
+
+class VesslLogMetricsCallback(TrainerCallback):
+
+    def __init__(self, credential_path: str, metrics: List[str]) -> None:
+        vessl.configure(credentials_file=credential_path)
+
+    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs: Dict[str, float] = None, **kwargs):     
+        if state.is_world_process_zero:
+            vessl.log(logs, state.global_step)
+

src/axolotl/utils/vessl_.py

@@ -0,0 +1,14 @@
+"""Module for vessl utilities"""
+
+import os
+
+from axolotl.utils.dict import DictDefault
+
+
+def setup_vessl_env_vars(cfg: DictDefault):
+    # VESSL_RUN_INITIAL_CONFIG is a variable that contain path to
+    # default credential inside a VESSL Run
+    credential_path = os.environ.get("VESSL_RUN_INITIAL_CONFIG")
+    if credential_path:
+        cfg.use_vessl = True
+        cfg.vessl_credential_path = credential_path