gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1171

Browse Source 
Update dependencies and enhance ML model retraining functionality

- Updated sientia-dataops-library version in requirements.txt from 1.3.3 to 1.3.4.
- Incremented image tag in values.yaml from 0.2.4 to 0.2.5 and added a new environment variable MONGODB_TTL_INDEX_HOURS.
- Introduced new methods in MLFlowRepository for model retraining and production model updates, including error handling and logging.
- Added retrain_model and update_production_model activities in mlflow.py to support model management workflows.
- Modified MongoDB connection settings in connectors_config.py for improved security and configuration flexibility.
This commit is contained in:
vitor-aignosi
2025-07-23 12:02:16 -03:00
parent 9410de5f82
commit 89b9892a5b
7 changed files with 421 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

194
laborious/utils/repository/model_repository.py
View File

@@ -13,6 +13,8 @@ By Monitoring we mean the evaluation of the performance of models, the generatio
from datetime import datetime
import traceback
import pandas as pd
import mlflow
from os import makedirs, path, remove
from sientia.ModelServing import ModelServing
@@ -85,3 +87,195 @@ class MLFlowRepository():
'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:
"""
Function to get the next run number of a specific model
Parameters:
model_name (str): the name of the model
Returns:
str: the next 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 retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
"""
Retrain a model with new data.
Parameters:
data (pandas.DataFrame): The new data to use for retraining.
model_name (str): The name of the model to retrain.
metrics_list (list): The metrics to be used to compare the models.
compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one.
If False, the retrain will always be considered.
split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets.
If False, the data will be used as a unique block for retraining.
update_report (bool): If True, a report will be created with the data of the retrained model.
update_transformation (bool): If True, the model will be updated in the MLflow tracking server.
update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server.
shuffle_data (bool): If True, the data will be shuffled before splitting.
model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'.
Returns:
mlflow.sklearn.Model: The retrained prediction model.
mlflow.sklearn.Model: The retrained data model.
mse (float): The mean squared error of the retrained model.
r2 (float): The R-squared score of the retrained model.
"""
# 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)
# align target column with treated_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)
# Example usage
experiment = self.get_experiment_by_run_id(latest_production_id)
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
mlflow.set_experiment(experiment)
experiment_description = "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")
if not path.exists("temp"):
makedirs("temp")
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 get_experiment(self, experiment_name: str) -> int:
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:
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:
# 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:
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