mlflow_provider.hooks.base

Module Contents

Classes

MLflowBaseHook

Base for MLflow hooks that interacts with a Python API.

class mlflow_provider.hooks.base.MLflowBaseHook(mlflow_conn_id=default_conn_name)

Bases: airflow.hooks.base.BaseHook

Base for MLflow hooks that interacts with a Python API. This is not used by the client hook which uses the requests library.

Parameters:

mlflow_conn_id (str) – mlflow http connection

conn_name_attr = 'mlflow_conn_id'

default_conn_name = 'mlflow_default'

conn_type = 'http'

mlflow_env_variables = ['MLFLOW_TRACKING_URI', 'MLFLOW_TRACKING_TOKEN', 'MLFLOW_TRACKING_USERNAME',...

unset_env_variables(other_env=None)

Precautionary function to be used after the hook or operator has finished executing to clear MLflow credentials from ENV variables.

Parameters:

other_env (dict) – Optionally provide additional creds for other systems like Sagemaker