Source code for modeva.utils.mlflow

import shutil
from os import environ, makedirs
from os.path import expanduser, join

import mlflow


def get_mlflow_config():
    mlflow_home = get_mlflow_home()
    local_mlflow_uri = "sqlite:///" + mlflow_home + "/modeva.db"
    mlflow_folder = mlflow_home + "/mlruns"
    return mlflow_home, local_mlflow_uri, mlflow_folder




def set_mlflow_home(mlflow_home: str = None):
    """Set the environment variable MLFLOW_HOME, which is the path of the mlflow home directory.

    Parameters
    ----------
    mlflow_home : str, default=None
        The path to mlflow directory. If `None`, the default path
        is `~/modeva_mlflow`.
    """
    if mlflow_home is None:
        mlflow_home = join("~", "modeva_mlflow")

    mlflow_home = expanduser(mlflow_home)
    environ["MLFLOW_HOME"] = mlflow_home
    makedirs(environ.get("MLFLOW_HOME"), exist_ok=True)





def get_mlflow_home():
    """Return the path of the mlflow home directory.

    The directory is set to a folder named 'modeva_mlflow' in the
    user home folder by default.

    Alternatively, it can be set by the 'MLFLOW_HOME' environment
    variable or programmatically by giving an explicit folder path. The '~'
    symbol is expanded to the user home folder.

    If the folder does not already exist, it is automatically created.

    Returns
    -------
    mlflow_home: str
        The path to mlflow directory.
    """
    mlflow_home = environ.get("MLFLOW_HOME", join("~", "modeva_mlflow"))
    mlflow_home = expanduser(mlflow_home)
    makedirs(mlflow_home, exist_ok=True)

    mlflow_folder = mlflow_home + "/mlruns"
    if mlflow_folder[0] == "/":
        # for Linux systems
        mlflow_artifact_path = "file://" + mlflow_folder
    else:
        # for Windows systems
        mlflow_artifact_path = "file:///" + mlflow_folder
    mlflow.store.tracking.DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = mlflow_artifact_path
    mlflow.store.model_registry.DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = mlflow_artifact_path
    mlflow.store.tracking.file_store.DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = mlflow_artifact_path
    mlflow.store.model_registry.file_store.DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = mlflow_artifact_path
    mlflow.tracking._tracking_service.utils.DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = mlflow_artifact_path
    return mlflow_home





def clear_mlflow_home():
    """Delete all the content of the data home cache.
    """
    mlflow_home = get_mlflow_home()
    shutil.rmtree(mlflow_home)