import pickle
import tempfile
import time
import uuid
from pathlib import Path
from typing import Union, Dict, Any, Optional
import mlflow
import pandas as pd
import tqdm
from mlflow.exceptions import MlflowException
from mlflow.tracking import MlflowClient
from sklearn.base import BaseEstimator
from .base import ModelBase
from ..auth import auth
from ..exceptions import ModelNotSupported
from ..testsuite.utils.metric_utils import get_scorer_funcs
from ..utils.constants import (MODEVA_MODEL_TAG,
REGRESSION_METRICS,
CLASSIFICATION_METRICS,
REGRESSION,
CLASSIFICATION)
from ..utils.mlflow import get_mlflow_config
from ..utils.registry import list_registered_models_, delete_registered_model_
[docs]
class LocalModelZoo:
"""A class for managing multiple models, their training, and evaluation.
LocalModelZoo provides functionality to:
- Add and manage multiple models
- Train models on a given dataset
- Track model performance metrics
- Register models with MLflow for experiment tracking
- Compare model performance through a leaderboard
Parameters
----------
dataset : Dataset
The dataset object to be used for model training and evaluation.
models : Dict, default=None
Dictionary of models to initialize with. If None, starts with empty dict.
name : str, default=None
Name for the model zoo. If None, a UUID-based name will be generated.
random_state : int, default=0
Random seed for reproducibility.
experiment_name : str, default=None
Name for the MLflow experiment. If None, will use name parameter or generate UUID-based name.
Examples
--------
.. minigallery:: ../galleries/*/*modelzoo*.py
"""
def __init__(
self,
dataset,
models: Dict = None,
name: str = None,
random_state: int = 0,
experiment_name: str = None):
auth.run()
self.__dataset = dataset
self.__models = models if models is not None else dict()
self.random_state = random_state
self.experiment_name = experiment_name
if self.experiment_name is None:
self.experiment_name = f"modelzoo-{str(uuid.uuid4())[:8]}" if name is None else name
self._init_mlflow()
self.run_logs = {}
def __validate_model(self, model: Union[str, ModelBase, BaseEstimator]):
if not isinstance(model, ModelBase):
if self.dataset.task_type == REGRESSION:
module = "modeva_sklearn_regressor or modeva_arbitrary_regressor"
elif self.dataset.task_type == CLASSIFICATION:
module = "modeva_sklearn_classifier or modeva_arbitrary_classifier"
raise ModelNotSupported(f"Model not supported. Make sure to wrap your model"
f" using modeva.models.wrappers.api.%s" % module)
@property
def dataset(self):
"""Return the dataset object."""
return self.__dataset
@property
def models(self):
"""Return the model dictionary."""
return self.__models
[docs]
def add_model(self, model, name: str = None, replace: bool = False):
"""Add a new model together with its name to models dictionary.
Parameters
----------
name : str
The name of model.
model : object
The model object.
replace : bool
Whether to replace old model when new model with same name
"""
if name is None:
name = model.name
if not replace:
if name in self.__models:
raise NameError(f"Model name {name} exists, please provide a new name.")
self.__validate_model(model)
self.__models[name] = model
[docs]
def list_model_names(self):
"""Return the list of model names."""
return list(self.models.keys())
[docs]
def get_model(self, name: str):
"""Return a model object.
Parameters
----------
name : str
The name of model to be extracted.
"""
if name in self.models.keys():
return self.models[name]
else:
raise ValueError(f"Model {name} not found")
[docs]
def train(self, name: str):
"""Train a model object.
Parameters
----------
name : str
The name of model to be trained.
"""
if self.dataset.task_type == REGRESSION:
metric_list = REGRESSION_METRICS
elif self.dataset.task_type == CLASSIFICATION:
metric_list = CLASSIFICATION_METRICS
run_logs = {}
model = self.get_model(name=name)
train_x, train_y, train_sample_weight = self.__dataset.get_X_y_data(dataset="train")
test_x, test_y, test_sample_weight = self.__dataset.get_X_y_data(dataset="test")
start_time = time.time()
model.fit(train_x, train_y, sample_weight=train_sample_weight)
end_time = time.time()
run_logs["start_time"] = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(int(start_time)))
run_logs["end_time"] = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(int(end_time)))
run_logs["Duration"] = end_time - start_time
estimator = self.get_model(name)
for score_name in metric_list:
score_sign, scoring_func, predict_func = get_scorer_funcs(estimator, score_name)
run_logs["train %s" % score_name] = scoring_func(y=train_y,
prediction=predict_func(train_x),
sample_weight=train_sample_weight)
run_logs["test %s" % score_name] = scoring_func(y=test_y,
prediction=predict_func(test_x),
sample_weight=test_sample_weight)
self.run_logs[name] = run_logs
[docs]
def train_all(self, silent: bool = False):
"""Train all models.
Parameters
----------
silent : bool, default=False
If False, will show the progress bar during model training.
"""
if silent:
for name, model in self.models.items():
self.train(name)
else:
for name, model in tqdm.tqdm(self.models.items(), leave=True):
self.train(name)
[docs]
def leaderboard(self, order_by: str = None, ascending=False):
"""Show the leaderboard of all models.
Parameters
----------
order_by : str, default=None
The leaderboard will be ordered by this metric.
If None, will show the results by the order of model training.
ascending : bool, default=False
The ordering parameter used when order_by is not None.
"""
df = pd.DataFrame(self.run_logs).T
if order_by is not None:
df = df.sort_values(order_by, ascending=ascending)
return df
# ------------------
# MLFLOW
# ------------------
def _init_mlflow(self):
mlflow_home, local_mlflow_uri, mlflow_folder = get_mlflow_config()
mlflow.set_tracking_uri(local_mlflow_uri)
self.mlflow_client = MlflowClient()
self.mlflow_tag = MODEVA_MODEL_TAG
try:
self.experiment_id = mlflow.create_experiment(self.experiment_name, tags={
"register_purpose": True,
})
except MlflowException:
mlflow.set_experiment(self.experiment_name)
self.experiment_id = dict(mlflow.get_experiment_by_name(self.experiment_name))["experiment_id"]
def _register_model(
self,
name,
model,
register_name: str = None,
description: str = None,
tags: Optional[Dict[str, Any]] = None,
experiment_id: str = None,
run_id: str = None
):
# use default model name if no name provided
if register_name is None:
register_name = name
if tags is None:
tags = {'dataset': self.dataset.name}
else:
tags['dataset'] = self.dataset.name
# Create a new run
if run_id is None:
run = self.mlflow_client.create_run(self.experiment_id)
run_id = run.info.run_id
# Save dataset and metainfo in temp directory and log them in current new run
with tempfile.TemporaryDirectory() as tmpdirname:
model.save(f'{tmpdirname}/{register_name}.pkl')
self.mlflow_client.log_artifact(run_id, f'{tmpdirname}/{register_name}.pkl')
try:
self.mlflow_client.create_registered_model(register_name,
tags={"artifact_type": self.mlflow_tag,
"experiment_id": experiment_id})
except MlflowException:
# Register dataset existed, create a new version down below
pass
if experiment_id is None:
model_uri = f"runs:/{run_id}"
else:
model_uri = f"{experiment_id}/{run_id}"
mv = self.mlflow_client.create_model_version(register_name, model_uri,
run_id, description=description, tags=tags)
[docs]
def register(
self,
name: str,
register_name: str = None,
description: str = None,
tags: Optional[Dict[str, Any]] = {},
run_id: str = None,
):
"""Register a model into MLFlow.
Parameters
----------
name : str
The current name of the model to be registered.
register_name : str, default=None
The register name of the model in MLFlow.
If None, will be the same as name.
description : str, default=None
The description of this model.
tags : dict, default=None
The tags.
run_id : str, default=None
The run id in MLFLow.
"""
model = self.get_model(name)
# Run id is needed to keep model registered under experiment
if isinstance(tags, list):
# Convert tag from list to dict, required by mlflow
ts = {}
for tag in tags:
ts[tag] = tag
tags = ts
self._register_model(
name=name, # name
model=model, # model object
register_name=register_name,
description=description,
tags=tags,
experiment_id=self.experiment_id,
run_id=run_id)
[docs]
def load_registered_model(self, name: str, version: int = None):
"""Return the list all registered models.
Parameters
----------
name : str
The name of model used for filtering.
version : int, default=None
Model version.
"""
mlflow_home, local_mlflow_uri, mlflow_folder = get_mlflow_config()
try:
latest_version = self.mlflow_client.get_latest_versions(name=name)
except MlflowException:
raise ValueError(f"Cannot find registered model {name}.")
if version is None:
# No version provided, get the latest version
uri = self.mlflow_client.get_model_version_download_uri(name, latest_version[0].version)
version = latest_version[0].version
else:
try:
uri = self.mlflow_client.get_model_version_download_uri(name, version)
except MlflowException:
raise ValueError(f"Cannot find registered dataset version {version}.")
prefix = uri.split("/")[0]
run_id = uri.split("/")[1]
if prefix == "runs:":
experiment_id = self.experiment_id
else:
experiment_id = prefix
file_path = Path(mlflow_folder, f"{experiment_id}", run_id, "artifacts", f"{name}.pkl")
model = pickle.load(open(file_path, "rb"))
model.name = f"{name}-{version}" # Update model name to registered name
return model
[docs]
def list_registered_models(self, name: str = None, format: str = 'frame', flat: bool = False):
"""Return the list all registered models.
Parameters
----------
name : str
The name of model used for filtering.
format : str, default="frame"
The format of displayed model list.
flat : bool, default=False
Whether to flatting the results.
"""
return list_registered_models_(mlflow_client=self.mlflow_client,
name=name,
experiment_id=self.experiment_id,
format=format,
flat=flat)
[docs]
def delete_registered_model(self, name: str, dataset: str):
"""Delete registered model.
Parameters
----------
name : str
The model name to delete
dataset : str
The dataset name to trained model
"""
delete_registered_model_(mlflow_client=self.mlflow_client,
name=name,
dataset=dataset,
experiment_id=self.experiment_id)