MLflow utils Databricks API Specs

Methods

mlflow_load_model (model_uri, type_of_model, model_func_dict)	load an ML model from MLflow run, raises an exception if type_of_model is not in dictionary	mlflow_load_model
mlflow_load_artifact (artifact_uri, artifact_name, type_of_artifact)	load an artifact from MLflow run, accepts joblib, pkl, dict and yaml file types	mlflow_load_artifact
mlflow_get_model_metrics (run_id, key_value_metrics)	gets all model evaluation metrics logged in MLflow run, or a specified key value evaluation metric	mlflow_get_model_metrics
mlflow_get_both_registered_model_info_run_id (name, mlflow_client, run_id, stage)	log model evaluation metrics to mlflow run	mlflow_get_both_registered_model_info_run_id
mlflow_promote_model (name, retrained_run_id, retrained_metric, start_date, eval_date, env, mlflow_client, metrics_name, prev_run_id, prev_metric)	log model evaluation metrics to mlflow run	mlflow_promote_model

mlflow_load_model

mlflow_load_model(model_uri: str, type_of_model: str, model_func_dict: dict)

load an ML model from MLflow run, raises an exception if type_of_model is not in dictionary

Parameters:

• model_uri (str) – model location uri from MLflow run

• type_of_model (str) – name of type of model; sk_model, pytorch, python_model, keras, etc

• model_func_dict (dict) – dictionary of allowed models to be loaded

Returns:

any allowed model callable

Return type:

Any

mlflow_load_artifact

mlflow_load_artifact(artifact_uri: str, artifact_name: str, type_of_artifact: str = 'joblib')

load an artifact from MLflow run, accepts `joblib, pkl, dict and yaml` file types

Parameters:

• artifact_uri (str) – artifact uri location from MLflow run

• artifact_name (str) – name of artifact

• type_of_artifact (str) – filetype, accepts joblib, pkl, dict and yaml file types

Returns:

returns a callable python object; dictionary, pandas dataframe, list

Return type:

Any

mlflow_get_model_metrics

mlflow_get_model_metrics(run_id: str, key_value_metrics: str = None)

gets all model evaluation metrics logged in MLflow run, or a specified key value evaluation metric

Parameters:

• run_id (str) – unique identifier of MLflow run

• key_value_metrics (str) – name of evaluation metric

Returns:

all of evaluation metric (Dict), or singular evaluation metric as float or int

Return type:

Union[float, int, Dict]

mlflow_get_both_registered_model_info_run_id

mlflow_get_both_registered_model_info_run_id(name: str, mlflow_client: mlflow.tracking.client.MlflowClient, run_id: str = None, stage: str = 'Production')

returns the registered model information from the specified MLflow run_id, and the MLflow run_id of the specified staging tag; Staging, Archived or Production

Parameters:

• name (str) – name of model

• mlflow_client (mlflow.tracking.client.MlflowClient) – mlflow client

• run_id (float) – unique identifier of MLflow run

• stage (str) – name of stage

Returns:

run_id, registered model information

Return type:

Tuple[str, Dict]

mlflow_promote_model

mlflow_promote_model(name: str, retrained_run_id: str, retrained_metric: float, start_date: str, eval_date: str, env: str, mlflow_client: mlflow.tracking.client.MlflowClient, metrics_name: str, prev_run_id: str = None, prev_metric: float = 0.0)

function that decides if we need to promote model to the staging tag if there is no model in the specified staging tag, and

Parameters:

• key (str) – name of evaluation metric

• retrained_run_id (str) – unique identifier of retrained MLflow run

• retrained_metric (float) – retrained model primary evaluation metric

• start_date (str) – start date of dataset

• eval_date (str) – end date of dataset

• env (str) – name of environment; staging, dev, prod

• mlflow_client (mlflow.tracking.client.MlflowClient) – initialised MLflow client

• metrics_name (str) – name of evaluation metrics

• prev_run_id (str) – unique identifier of previous MLflow run

• prev_metric (float) – previous MLflow run evaluation metric

Returns:

string response of the promotion of model

Return type:

str