From 4283730e7ac99e824798ededa0892c73e7ee72b1 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 23 Jul 2025 16:51:33 -0300 Subject: [PATCH] SIENTIAPDE-1171 Update GITHUB_BRANCH in values.yaml and refactor MLFlowRepository for model retraining - Changed GITHUB_BRANCH in values.yaml to reflect the new pipeline for retraining. - Refactored MLFlowRepository by introducing create_model_experiment and perform_model_retrain methods to streamline model retraining processes. - Updated test cases to cover new functionalities and ensure proper error handling in model management workflows. --- .../utils/repository/model_repository.py | 52 ++-- .../utils/repository/test_model_repository.py | 246 +++++++++++++++++- .../laborious/utils/test_connectors_config.py | 11 +- values.yaml | 2 +- 4 files changed, 275 insertions(+), 36 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 86f1499..187e324 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -116,32 +116,7 @@ class MLFlowRepository(): 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. - """ - + def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple: # load predictor model predictor_uri = f"models:/{model_name}/production" # load transform model @@ -156,18 +131,27 @@ class MLFlowRepository(): 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) + 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): + 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" + 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 @@ -203,6 +187,14 @@ class MLFlowRepository(): return "Model retrained successfully", experiment + def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: + + 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: experiment = mlflow.get_experiment_by_name(experiment_name) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 6d10c1c..f3978cb 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,6 +1,8 @@ -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import ANY, MagicMock, call, patch import numpy as np +from pandas import DataFrame import pytest +from laborious.utils.repository import model_repository from laborious.utils.repository.model_repository import MLFlowRepository @@ -92,3 +94,245 @@ def test_predict_error(mlflow_repository): 'traceback': ANY } } + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_by_run_id(mlflow, mlflow_repository): + mlflow.get_run.return_value = MagicMock( + info=MagicMock( + experiment_id='0', + ) + ) + mlflow.get_experiment.return_value = MagicMock() + mlflow.get_experiment.return_value.name = 'test' + + output = mlflow_repository.get_experiment_by_run_id('0') + assert output == 'test' + mlflow.get_run.assert_called_once_with('0') + mlflow.get_experiment.assert_called_once_with('0') + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_next_run_name(mlflow, mlflow_repository): + mlflow.search_runs.return_value = [1, 2, 3] + output = mlflow_repository.get_next_run_name('run') + assert output == 'run-4' + mlflow.search_runs.assert_called_once_with( + experiment_names=['run'], + order_by=['start_time desc'], + ) + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_success(mlflow, mlflow_repository): + mlflow.get_experiment_by_name.return_value = MagicMock( + experiment_id='0') + + output = mlflow_repository.get_experiment('test') + + assert output == 0 + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_error(mlflow, mlflow_repository): + mlflow.get_experiment_by_name.return_value = None + + try: + mlflow_repository.get_experiment('test') + except ValueError as e: + assert str(e) == 'Experiment test not found' + else: + assert False + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_last_run(mlflow, mlflow_repository): + mlflow.search_runs.return_value = DataFrame({ + 'params.retrain': ['True', 'False', 'True', 'False'], + 'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'], + 'run_id': ['0', '1', '2', '3'], + }) + + output = mlflow_repository.get_experiment_last_run(0) + + mlflow.search_runs.assert_called_once_with( + experiment_ids=[0], + filter_string="", + output_format="pandas", + ) + + assert output == '2' + + +@patch('laborious.utils.repository.model_repository.mlflow.sklearn') +@patch('laborious.utils.repository.model_repository.mlflow.set_experiment') +def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): + + mlflow_repository.model_serving.get_model_run_id = MagicMock( + return_value='0') + mlflow_repository.model_serving.get_model_uri = MagicMock( + return_value='test') + mlflow_repository.get_experiment_by_run_id = MagicMock() + + data_model_mock = MagicMock() + prediction_model_mock = MagicMock() + + sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock] + + data_model_mock.fit.return_value = data_model_mock + data_model_mock.predict.return_value = DataFrame({ + 'x': [10, 20, 30], + }) + data_model_mock.target_variable = 'y' + + prediction_model_mock.fit.return_value = prediction_model_mock + + data = DataFrame({ + 'x': [1, 2, 3], + 'y': [4, 5, 6] + }) + + output = mlflow_repository.create_model_experiment('test', data) + + mlflow_repository.model_serving.get_model_run_id.assert_called_once_with( + 'test', stage='Production') + mlflow_repository.model_serving.get_model_uri.assert_called_once_with( + '0', prediction=False) + + sklearn.load_model.assert_has_calls([ + call(mlflow_repository.model_serving.get_model_uri.return_value), + call("models:/test/production"), + ]) + assert sklearn.load_model.call_count == 2 + + data_model_mock.fit.assert_called_once_with(data) + data_model_mock.predict.assert_called_once_with(data) + + fit_args = prediction_model_mock.fit.call_args[0][0] + assert fit_args.equals( + DataFrame({ + 'x': [10, 20, 30], + 'y': [4, 5, 6], + }) + ) + + mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0') + + set_experiment.assert_called_once_with( + mlflow_repository.get_experiment_by_run_id.return_value + ) + + assert output == (prediction_model_mock, + data_model_mock, + mlflow_repository.get_experiment_by_run_id.return_value) + + +@patch('laborious.utils.repository.model_repository.mlflow.start_run') +@patch('laborious.utils.repository.model_repository.mlflow.log_param') +@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model') +@patch('laborious.utils.repository.model_repository.mlflow.log_artifact') +def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository): + + _os_remove = patch( + 'laborious.utils.repository.model_repository.remove') + _os_makedirs = patch( + 'laborious.utils.repository.model_repository.makedirs') + _os_path_exists = patch( + 'laborious.utils.repository.model_repository.path.exists', + MagicMock(return_value=False)) + + prediction_model_mock = MagicMock() + data_model_mock = MagicMock() + experiment = 'test' + model_name = 'test' + data = MagicMock() + + mlflow_repository.get_next_run_name = MagicMock( + return_value='test-1') + run = MagicMock() + start_run.__enter__.return_value = run + + output = mlflow_repository.perform_model_retrain( + prediction_model_mock, data_model_mock, experiment, model_name, data) + + mlflow_repository.get_next_run_name.assert_called_once_with(experiment) + start_run.assert_called_once_with( + run_name='test-1', description='Retrain model test with new data') + + log_model.assert_has_calls([ + call(data_model_mock, "data_model"), + call(prediction_model_mock, "prediction_model"), + ]) + + data.to_csv.assert_called_once_with( + "temp/raw_data_test.csv", index=True) + + log_artifact.assert_called_once_with( + "temp/raw_data_test.csv") + + log_param.assert_has_calls([ + call("retrain", True), + ]) + + assert output == ("Model retrained successfully", experiment) + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_update_production_model_by_run_id(mlflow, mlflow_repository): + client_mock = MagicMock() + mlflow.tracking.MlflowClient.return_value = client_mock + + client_mock.get_registered_model.return_value = MagicMock( + latest_versions=[ + MagicMock(version='1'), + MagicMock(version='2'), + MagicMock(version='3'), + ] + ) + output = mlflow_repository.update_production_model_by_run_id('0', 'test') + + mlflow.register_model.assert_called_once_with( + "runs:/0/prediction_model", + 'test', + ) + + mlflow.tracking.MlflowClient.assert_called_once() + client_mock.get_registered_model.assert_called_once_with('test') + client_mock.transition_model_version_stage.assert_called_once_with( + name='test', + version='3', + stage='Production', + archive_existing_versions=True, + ) + + assert output == { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + } + + +def test_update_production_model(mlflow_repository): + connector = mlflow_repository + + with patch.object(connector, 'get_experiment', + return_value='0') as get_experiment: + with patch.object(connector, 'get_experiment_last_run', + return_value='2') as get_experiment_last_run: + with patch.object(connector, 'update_production_model_by_run_id', + return_value={'model_name': 'test', 'version': '3', + 'mlflow_run_id': '0'}) as update_production_model_by_run_id: + + output = connector.update_production_model('0', 'test') + + get_experiment.assert_called_once_with('0') + get_experiment_last_run.assert_called_once_with('0') + update_production_model_by_run_id.assert_called_once_with( + '2', 'test') + + assert output == { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + 'mlflow_experiment_id': '0', + } diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index acb5c21..910439c 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -139,10 +139,12 @@ def test_build_mongo_db_config_with_env_vars(): environ['MONGODB_PASSWORD'] = 'sientia1' environ['MONGODB_URL'] = 'localhost:27018' environ['MONGODB_DATABASE_NAME'] = 'test_db' + environ['MONGODB_TTL_INDEX_HOURS'] = '1' assert build_mongodb_config() == { 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', - 'database_name': 'test_db' + 'database_name': 'test_db', + 'ttl_index_seconds': 3600 } @@ -151,8 +153,9 @@ def test_build_mongo_db_config_with_defaults(): environ.pop('MONGODB_PASSWORD', None) environ.pop('MONGODB_DATABASE_NAME', None) environ.pop('MONGODB_URL', None) - + environ.pop('MONGODB_TTL_INDEX_HOURS', None) assert build_mongodb_config() == { - 'connection_string': 'mongodb://sientia:sientia@localhost:27017', - 'database_name': 'sientia' + 'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018', + 'database_name': 'sientia', + 'ttl_index_seconds': 3600 } diff --git a/values.yaml b/values.yaml index bac9395..03f681b 100644 --- a/values.yaml +++ b/values.yaml @@ -123,7 +123,7 @@ env: - name: GITHUB_REPO_URL value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git" - name: GITHUB_BRANCH - value: "SIENTIAPDE-1151-criar-testes-de-stress" + value: "SIENTIAPDE-1171-criar-pipeline-de-retreino-laborious" - name: PYTHON_APP value: "laborious.worker.worker"