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.

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:
+    def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> 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
@@ -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 = f"Retrain model {model_name} with new data"
-        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
@@ -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)
 

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',
+                }

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',
+        'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
-        'database_name': 'sientia'
+        'database_name': 'sientia',
+        'ttl_index_seconds': 3600
     }

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"