SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment.
This commit is contained in:
Bruno Domingues
481
model-manager/utils/repository/model_repository.py
Normal file
481
model-manager/utils/repository/model_repository.py
Normal file
@@ -0,0 +1,481 @@
|
||||
"""
|
||||
Model Monitoring Repository
|
||||
|
||||
This module contains the ModelMonitoringRepository class,
|
||||
which is responsible for handling the communication with the Model Monitoring API.
|
||||
|
||||
It includes the methods that are used to answer ModelMonitoringService
|
||||
requests using the Model Monitoring API functions.
|
||||
|
||||
By Monitoring we mean the evaluation of the performance of models, the generation of reports.
|
||||
|
||||
"""
|
||||
from datetime import datetime
|
||||
import traceback
|
||||
import pandas as pd
|
||||
import mlflow
|
||||
from os import makedirs, path, remove
|
||||
from sientia.ModelServing import ModelServing
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
|
||||
class MLFlowRepository():
|
||||
def __init__(self, host, username, password, logger: Logger):
|
||||
|
||||
self.model_serving = ModelServing(tracking_uri=host,
|
||||
username=username, password=password,
|
||||
logger=logger)
|
||||
self.logger = logger
|
||||
|
||||
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
|
||||
"""
|
||||
Detect and parse datetime index from data. index must be a timestamp like column.
|
||||
This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
|
||||
If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
|
||||
If another type or format, must raise an error.
|
||||
"""
|
||||
index = data.index
|
||||
|
||||
# Get type of first element of index
|
||||
index_type = type(index[0])
|
||||
|
||||
self.logger.custom_info(f"Index type: {index_type}", metadata)
|
||||
|
||||
message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}"
|
||||
|
||||
# Check if all in index are of the same type
|
||||
if not all(isinstance(i, index_type) for i in index):
|
||||
raise ValueError(
|
||||
f"{message}")
|
||||
|
||||
# Check type and converts to DATETIME_FORMAT_WITH_TZ
|
||||
if index_type == str:
|
||||
# Validate format of string and return error if not valid
|
||||
try:
|
||||
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
|
||||
except ValueError:
|
||||
raise ValueError(
|
||||
f"{message}")
|
||||
|
||||
elif index_type == datetime or index_type == pd.Timestamp:
|
||||
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{message}")
|
||||
|
||||
return data
|
||||
|
||||
def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
|
||||
"""
|
||||
Transform data using a model.
|
||||
|
||||
Parameters:
|
||||
- model_name (str): The name of the model to use for transformation.
|
||||
- data (pandas.DataFrame): The data to transform.
|
||||
- model_retention (int): The number of minutes to keep the model.
|
||||
|
||||
Returns:
|
||||
- dict: A dictionary containing the transformed data.
|
||||
"""
|
||||
|
||||
try:
|
||||
self.logger.custom_debug(
|
||||
f"Data received for model transformation: {data.to_csv()}", metadata)
|
||||
|
||||
model_retention = model_config.get('retention_minutes', 0)
|
||||
flavor = model_config.get('transform_flavor', 'sklearn')
|
||||
compressed = model_config.get('is_compressed', False)
|
||||
retention_target = model_config.get('retention_target', 'model')
|
||||
transform_keyword = model_config.get(
|
||||
'transform_function_keyword', 'predict')
|
||||
|
||||
transformed_data = self.model_serving.get_cached_transform(
|
||||
model_name, data, model_retention, flavor,
|
||||
compressed, retention_target, transform_keyword
|
||||
)
|
||||
|
||||
self.logger.custom_debug(
|
||||
f"Data received from model transformation: {transformed_data.to_csv()}", metadata)
|
||||
|
||||
transformed_data = self.detect_and_parse_datetime_index(
|
||||
transformed_data, metadata)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'content': transformed_data.to_dict()
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'content': {
|
||||
'message': str(e),
|
||||
'traceback': traceback.format_exc()
|
||||
}
|
||||
}
|
||||
|
||||
def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
|
||||
"""
|
||||
Predict data using a model.
|
||||
|
||||
Parameters:
|
||||
- model_name (str): The name of the model to use for prediction.
|
||||
- data (pandas.DataFrame): The data to predict.
|
||||
- model_retention (int): The number of minutes to keep the model.
|
||||
|
||||
Returns:
|
||||
- dict: A dictionary containing the predicted data.
|
||||
"""
|
||||
try:
|
||||
model_retention = model_config.get('retention_minutes', 0)
|
||||
flavor = model_config.get('predict_flavor', 'pyfunc')
|
||||
compressed = model_config.get('is_compressed', False)
|
||||
retention_target = model_config.get('retention_target', 'model')
|
||||
|
||||
input_index = data.index
|
||||
start_time = datetime.now()
|
||||
|
||||
self.logger.custom_debug(
|
||||
f"Data received for model prediction: {data.to_csv()}", metadata)
|
||||
data = self.model_serving.get_cached_predict(
|
||||
model_name, data, model_retention, flavor,
|
||||
compressed, retention_target
|
||||
)
|
||||
|
||||
end_time = datetime.now()
|
||||
data = pd.DataFrame(data, columns=['prediction'])
|
||||
self.logger.custom_debug(
|
||||
f"Data received from model prediction: {data.to_csv()}", metadata)
|
||||
data.index = input_index
|
||||
data['response_time'] = (end_time - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'content': data.to_dict()
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'content': {
|
||||
'message': str(e),
|
||||
'traceback': traceback.format_exc()
|
||||
}
|
||||
}
|
||||
|
||||
def get_experiment_by_run_id(self, run_id: str) -> dict:
|
||||
# Get the run information using the run_id
|
||||
run = mlflow.get_run(run_id)
|
||||
|
||||
# Extract the experiment ID from the run
|
||||
experiment_id = run.info.experiment_id
|
||||
|
||||
# Get the experiment details using the experiment ID
|
||||
experiment = mlflow.get_experiment(experiment_id)
|
||||
experiment_name = experiment.name
|
||||
return experiment_name
|
||||
|
||||
def get_next_run_name(self, model_name: str) -> str:
|
||||
"""
|
||||
Generate the next run name for a specific MLFlow model.
|
||||
|
||||
This method calculates the next sequential run number for a model
|
||||
by searching existing runs and incrementing the count. It ensures
|
||||
unique run names for model training and retraining operations.
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the MLFlow model
|
||||
|
||||
Returns:
|
||||
str: The next run name in format 'model_name-run_number'
|
||||
"""
|
||||
runs = mlflow.search_runs(
|
||||
experiment_names=[model_name], order_by=["start_time desc"])
|
||||
next_run_number = len(runs) + 1
|
||||
return f"{model_name}-{next_run_number}"
|
||||
|
||||
def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
|
||||
"""
|
||||
Create a new MLFlow experiment for model retraining.
|
||||
|
||||
This method sets up the complete environment for model retraining by:
|
||||
1. Loading the current production prediction model
|
||||
2. Loading the current production transformation model
|
||||
3. Fitting the transformation model with new data
|
||||
4. Preparing data for prediction model retraining
|
||||
5. Setting up the MLFlow experiment context
|
||||
|
||||
Args:
|
||||
model_name (str): Name of the MLFlow model to retrain
|
||||
data (pd.DataFrame): Training data for model retraining
|
||||
|
||||
Returns:
|
||||
tuple: (prediction_model, data_model, experiment)
|
||||
- prediction_model: Loaded prediction model for retraining
|
||||
- data_model: Fitted transformation model
|
||||
- experiment: MLFlow experiment name
|
||||
"""
|
||||
# load predictor model
|
||||
predictor_uri = f"models:/{model_name}/production"
|
||||
# load transform model
|
||||
latest_production_id = self.model_serving.get_model_run_id(
|
||||
model_name, stage="Production"
|
||||
)
|
||||
transform_uri = self.model_serving.get_model_uri(
|
||||
latest_production_id, prediction=False
|
||||
)
|
||||
# load
|
||||
data_model = mlflow.sklearn.load_model(transform_uri)
|
||||
prediction_model = mlflow.sklearn.load_model(predictor_uri)
|
||||
data_model = data_model.fit(data)
|
||||
treated_data = data_model.predict(data)
|
||||
|
||||
target_name = data_model.target_variable
|
||||
y = data[target_name]
|
||||
treated_data = pd.merge(
|
||||
treated_data, y, left_index=True, right_index=True)
|
||||
prediction_model = prediction_model.fit(treated_data)
|
||||
experiment = self.get_experiment_by_run_id(latest_production_id)
|
||||
mlflow.set_experiment(experiment)
|
||||
|
||||
return prediction_model, data_model, experiment
|
||||
|
||||
def perform_model_retrain(self,
|
||||
prediction_model,
|
||||
data_model,
|
||||
experiment: str,
|
||||
model_name: str,
|
||||
data: pd.DataFrame):
|
||||
"""
|
||||
Execute the complete model retraining process in MLFlow.
|
||||
|
||||
This method performs the actual model retraining by:
|
||||
1. Starting a new MLFlow run with descriptive metadata
|
||||
2. Logging model parameters and hyperparameters
|
||||
3. Retraining both prediction and transformation models
|
||||
4. Logging training data as artifacts
|
||||
5. Saving retrained models to MLFlow registry
|
||||
|
||||
Args:
|
||||
prediction_model: MLFlow prediction model to retrain
|
||||
data_model: MLFlow transformation model to retrain
|
||||
experiment (str): MLFlow experiment name for the retraining
|
||||
model_name (str): Name of the model being retrained
|
||||
data (pd.DataFrame): Training data used for retraining
|
||||
|
||||
Returns:
|
||||
tuple: (status_message, experiment_name)
|
||||
- status_message (str): Success confirmation message
|
||||
- experiment_name (str): Name of the experiment
|
||||
"""
|
||||
pred_model_atributes = vars(prediction_model) # load class attributes
|
||||
data_model_atributes = vars(data_model) # load class attributes
|
||||
experiment_description = f"Retrain model {model_name} with new data"
|
||||
current_run_name = self.get_next_run_name(experiment)
|
||||
with mlflow.start_run(
|
||||
run_name=current_run_name, description=experiment_description
|
||||
) as _run:
|
||||
# update transfomation model
|
||||
# fixed parameters
|
||||
for name_atribute, val_atribute in pred_model_atributes.items():
|
||||
if name_atribute != "model":
|
||||
mlflow.log_param(name_atribute, val_atribute)
|
||||
# update prediction model
|
||||
for name_atribute, val_atribute in data_model_atributes.items():
|
||||
if name_atribute != "model":
|
||||
mlflow.log_param(name_atribute, val_atribute)
|
||||
# dynamic parameters, including model itself
|
||||
mlflow.sklearn.log_model(data_model, "data_model")
|
||||
|
||||
makedirs("temp", exist_ok=True)
|
||||
|
||||
file_path = f"temp/raw_data_{model_name}.csv"
|
||||
data.to_csv(file_path, index=True)
|
||||
|
||||
# log the data raw
|
||||
mlflow.log_artifact(file_path)
|
||||
|
||||
# dynamic parameters, including model itself
|
||||
mlflow.sklearn.log_model(prediction_model, "prediction_model")
|
||||
mlflow.log_param("retrain", True)
|
||||
|
||||
# clear temp file
|
||||
if path.exists(file_path):
|
||||
remove(file_path)
|
||||
|
||||
return "Model retrained successfully", experiment
|
||||
|
||||
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
|
||||
"""
|
||||
Orchestrate the complete model retraining workflow.
|
||||
|
||||
This method coordinates the entire model retraining process by:
|
||||
1. Creating the MLFlow experiment environment
|
||||
2. Loading existing production models
|
||||
3. Executing the retraining process
|
||||
4. Returning comprehensive retraining results
|
||||
|
||||
Args:
|
||||
data (pd.DataFrame): Training data for model retraining
|
||||
model_name (str): Name of the MLFlow model to retrain
|
||||
|
||||
Returns:
|
||||
tuple: (status_message, experiment_name)
|
||||
- status_message (str): Retraining operation status
|
||||
- experiment_name (str): MLFlow experiment identifier
|
||||
"""
|
||||
prediction_model, data_model, experiment = self.create_model_experiment(
|
||||
model_name, data)
|
||||
retrain_result = self.perform_model_retrain(
|
||||
prediction_model, data_model, experiment, model_name, data)
|
||||
return retrain_result
|
||||
|
||||
def get_experiment(self, experiment_name: str) -> int:
|
||||
"""
|
||||
Retrieve MLFlow experiment ID by experiment name.
|
||||
|
||||
This method searches for an MLFlow experiment by name and
|
||||
returns its unique identifier. It provides error handling
|
||||
for non-existent experiments.
|
||||
|
||||
Args:
|
||||
experiment_name (str): Name of the MLFlow experiment
|
||||
|
||||
Returns:
|
||||
int: MLFlow experiment ID
|
||||
|
||||
Raises:
|
||||
ValueError: If the experiment name is not found
|
||||
"""
|
||||
experiment = mlflow.get_experiment_by_name(experiment_name)
|
||||
|
||||
if experiment is None:
|
||||
raise ValueError(f'Experiment {experiment_name} not found')
|
||||
|
||||
return int(experiment.experiment_id)
|
||||
|
||||
def get_experiment_last_run(self, experiment_id: int) -> str:
|
||||
"""
|
||||
Retrieve the most recent retraining run ID for an experiment.
|
||||
|
||||
This method searches for the latest run in an MLFlow experiment
|
||||
that has been marked as a retraining run. It filters runs by
|
||||
the 'retrain' parameter and orders them by completion time.
|
||||
|
||||
Args:
|
||||
experiment_id (int): MLFlow experiment ID
|
||||
|
||||
Returns:
|
||||
str: MLFlow run ID of the most recent retraining run
|
||||
|
||||
Raises:
|
||||
ValueError: If runs data is not in expected DataFrame format
|
||||
"""
|
||||
runs = mlflow.search_runs(
|
||||
experiment_ids=[experiment_id],
|
||||
filter_string="", # Sem filtro no MLflow ainda
|
||||
output_format="pandas"
|
||||
)
|
||||
|
||||
if not isinstance(runs, pd.DataFrame):
|
||||
raise ValueError('Runs is not a pandas DataFrame')
|
||||
|
||||
# Filtrar apenas as runs onde params.retrain == True
|
||||
filtered_runs = runs[runs["params.retrain"] == 'True']
|
||||
|
||||
# Converter a coluna 'end_time' para datetime
|
||||
filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
|
||||
|
||||
# Ordenar o DataFrame de forma descendente pela coluna 'end_time'
|
||||
filtered_runs = filtered_runs.sort_values(
|
||||
by='end_time', ascending=False)
|
||||
|
||||
# Pegar a última run_id do DataFrame filtrado e ordenado
|
||||
latest_run_id = filtered_runs.iloc[0]['run_id']
|
||||
|
||||
return latest_run_id
|
||||
|
||||
def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
|
||||
"""
|
||||
Update production model with a specific MLFlow run.
|
||||
|
||||
This method promotes a model from a specific MLFlow run to
|
||||
production stage. It handles model registration, versioning,
|
||||
and stage transitions with proper error handling.
|
||||
|
||||
Args:
|
||||
run_id (str): MLFlow run ID containing the model to promote
|
||||
model_name (str): Name of the MLFlow model
|
||||
|
||||
Returns:
|
||||
dict: Model update metadata containing:
|
||||
- model_name (str): Name of the updated model
|
||||
- version (str): New model version number
|
||||
- mlflow_run_id (str): Source run ID
|
||||
|
||||
Update Process:
|
||||
1. Registers the model from the specified run
|
||||
2. Retrieves the latest model version
|
||||
3. Transitions the model to 'Production' stage
|
||||
4. Archives existing production versions
|
||||
"""
|
||||
# Registrar o modelo
|
||||
# Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
|
||||
# Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
|
||||
mlflow.register_model(
|
||||
f"runs:/{run_id}/prediction_model", model_name)
|
||||
|
||||
# Colocar a versão do modelo em produção
|
||||
# Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
|
||||
client = mlflow.tracking.MlflowClient()
|
||||
|
||||
# Obter a versão mais recente registrada do modelo
|
||||
model_versions = client.get_registered_model(
|
||||
model_name).latest_versions
|
||||
|
||||
if not isinstance(model_versions, list):
|
||||
raise ValueError('Model versions is not a list')
|
||||
|
||||
max_version = max(model_versions, key=lambda x: int(x.version)).version
|
||||
|
||||
# Mover a versão mais recente do modelo para o estágio de 'Production'
|
||||
client.transition_model_version_stage(
|
||||
name=model_name,
|
||||
version=max_version,
|
||||
stage="Production",
|
||||
archive_existing_versions=True
|
||||
)
|
||||
|
||||
return {
|
||||
'model_name': model_name,
|
||||
'version': max_version,
|
||||
'mlflow_run_id': run_id
|
||||
}
|
||||
|
||||
def update_production_model(self, experiment: str, model_name: str) -> dict:
|
||||
"""
|
||||
Update production model using the latest retraining run.
|
||||
|
||||
This method orchestrates the complete production model update
|
||||
process by identifying the most recent retraining run and
|
||||
promoting it to production stage.
|
||||
|
||||
Args:
|
||||
experiment (str): MLFlow experiment name
|
||||
model_name (str): Name of the MLFlow model
|
||||
|
||||
Returns:
|
||||
dict: Complete model update metadata containing:
|
||||
- model_name (str): Name of the updated model
|
||||
- version (str): New model version number
|
||||
- mlflow_run_id (str): Source run ID
|
||||
- mlflow_experiment_id (int): Experiment ID
|
||||
"""
|
||||
experiment_id = self.get_experiment(experiment)
|
||||
run_id = self.get_experiment_last_run(experiment_id)
|
||||
metadata = self.update_production_model_by_run_id(run_id, model_name)
|
||||
|
||||
metadata['mlflow_experiment_id'] = experiment_id
|
||||
|
||||
return metadata
|
||||
Reference in New Issue
Block a user