[#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

@@ -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

@@ -854,6 +854,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,27 @@
+"""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
+ ):
+ if state.is_world_process_zero:
+ vessl.log(logs, state.global_step)

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

@@ -395,6 +395,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,
@@ -404,6 +410,7 @@ class AxolotlInputConfig(
  HyperparametersConfig,
  WandbConfig,
  MLFlowConfig,
+ VesslConfig,
  RemappedParameters,
  DeprecatedParameters,
  BaseModel,

src/axolotl/utils/vessl_.py

@@ -0,0 +1,13 @@
+"""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.vessl_credential_path = credential_path