File size: 543 Bytes
090c24d
 
 
 
 
 
 
 
 
d756534
090c24d
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
"""Module for mlflow utilities"""

import os

from axolotl.utils.dict import DictDefault


def setup_mlflow_env_vars(cfg: DictDefault):
    for key in cfg.keys():
        if key.startswith("mlflow_") or key.startswith("hf_mlflow_"):
            value = cfg.get(key, "")

            if value and isinstance(value, str) and len(value) > 0:
                os.environ[key.upper()] = value

    # Enable mlflow if experiment name is present
    if cfg.mlflow_experiment_name and len(cfg.mlflow_experiment_name) > 0:
        cfg.use_mlflow = True