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

.pylintrc

@@ -9,6 +9,6 @@ generated-members=numpy.*, torch.*
 
 
 [pylint.messages_control]
-disable=missing-function-docstring, line-too-long, import-error,
+disable=missing-function-docstring, line-too-long, import-error, too-many-lines,
     too-many-arguments, too-many-locals, too-many-statements, too-many-branches, too-few-public-methods,
     too-many-instance-attributes, fixme, import-outside-toplevel, logging-fstring-interpolation,

src/axolotl/cli/__init__.py

@@ -41,6 +41,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__), ".."))
@@ -384,6 +385,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

@@ -874,6 +874,11 @@ def get_callbacks(self) -> List[TrainerCallback]:
                 SaveAxolotlConfigtoWandBCallback(self.cfg.axolotl_config_path)
             )
 
+        if self.cfg.vessl_credential_path:
+            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,29 @@
+"""Vessl module for trainer callbacks"""
+import logging
+from typing import Dict
+
+import vessl
+from transformers import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+LOG = logging.getLogger("axolotl.callbacks")
+
+
+class VesslLogMetricsCallback(TrainerCallback):
+    """Callback to send training metrics to VESSL AI"""
+
+    def __init__(self, credential_path: str) -> None:
+        vessl.configure(credentials_file=credential_path)
+
+    def on_log(
+        self,
+        args: TrainingArguments,  # pylint: disable=unused-argument
+        state: TrainerState,
+        control: TrainerControl,  # pylint: disable=unused-argument
+        logs: Dict[str, float],
+        **kwargs  # pylint: disable=unused-argument
+    ):
+        # is_world_process_zero: Whether or not this process is the global main process (when training in a
+        # distributed fashion on several machines, this is only going to be `True` for one process).
+        if state.is_world_process_zero:
+            vessl.log(logs, state.global_step)

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -408,6 +408,12 @@ def check_wandb_run(cls, data):
         return data
 
 
+class VesslConfig(BaseModel):
+    """Vessl AI configuration subset"""
+
+    vessl_credential_path: Optional[str] = None
+
+
 # pylint: disable=too-many-public-methods,too-many-ancestors
 class AxolotlInputConfig(
     ModelInputConfig,
@@ -417,6 +423,7 @@ class AxolotlInputConfig(
     HyperparametersConfig,
     WandbConfig,
     MLFlowConfig,
+    VesslConfig,
     LISAConfig,
     RemappedParameters,
     DeprecatedParameters,

src/axolotl/utils/vessl_.py

@@ -0,0 +1,17 @@
+"""Module for vessl utilities"""
+
+import os
+
+from axolotl.utils.dict import DictDefault
+
+
+def setup_vessl_env_vars(cfg: DictDefault):
+    if cfg.vessl_credential_path:
+        return
+
+    # VESSL_RUN_INITIAL_CONFIG is a variable that contain path to default credential inside a VESSL Run.
+    # Currently there is no docs regarding this variable, but it exists inside the container.
+    # Ref: https://screen.yanolja.in/lrTGow4Pr8eXhAai.png
+    credential_path = os.environ.get("VESSL_RUN_INITIAL_CONFIG")
+    if credential_path:
+        cfg.vessl_credential_path = credential_path