Merge pull request #12 from Aignosi/SIENTIAPDE-1171-criar-pipeline-de-retreino-laborious

Sientiapde 1171 criar pipeline de retreino laborious

laborious/activities/mlflow.py

@@ -1,5 +1,3 @@
-import numpy as np
-from pandas import DataFrame
 from temporalio import activity, workflow
 
 
@@ -9,6 +7,9 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.temporal.utils.logger import Logger
     from laborious.utils.repository.model_repository import MLFlowRepository
     from typing import Any
+    import numpy as np
+    from pandas import DataFrame
+    import traceback
 
 
 class MLFlow(BaseActivity):
@@ -96,3 +97,107 @@ class MLFlow(BaseActivity):
         self.debug(response_data, metadata)
 
         return response_data
+
+    @activity.defn(name="retrain_model")
+    async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Retrain the model.
+        Args:
+            - input_data (dict): The input data. Contains:
+                - model_name (str): The name of the model.
+                - data (dict[str, Any]): The data to retrain the model.
+        """
+        metadata = input_data['metadata']
+        data = DataFrame(input_data['data'])
+        model_name = input_data['model_name']
+
+        self.info(f'Retraining model {model_name}...', metadata)
+
+        timestamp = data['timestamp'].max()
+        self.debug(f'Timestamp: {timestamp}', metadata)
+
+        data.drop(columns=['model_id'], inplace=True, errors='ignore')
+        data.drop(columns=['created_at'], inplace=True, errors='ignore')
+
+        data = data.pivot(index='timestamp', columns='variable',
+                          values='value')
+        data.sort_index(inplace=True)
+        data.reset_index(inplace=True)
+
+        data = data.dropna()
+        data.columns.name = None
+
+        try:
+            retrain_output, experiment = self.model_monitoring_repository.retrain_model(
+                data=data,
+                model_name=model_name
+            )
+
+            return {
+                'status': retrain_output,
+                'timestamp': timestamp,
+                'experiment': experiment
+            }
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.send_notification(
+                metadata=metadata,
+                notification_id='RETRAIN_MODEL_ERROR',
+                message=f'Error retraining model {model_name}: {e}',
+                block='retrain_model',
+                attachment_content=trace
+            )
+            self.error(trace, metadata=metadata)
+            raise e
+
+    @activity.defn(name="update_production_model")
+    async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+        """
+        Update the production model.
+        Args:
+            - input_data (dict): The input data. Contains:
+                - model_name (str): The name of the model.
+                - experiment (str): The name of the experiment.
+                - model_id (str): The id of the model.
+                - timestamp (str): The timestamp of the model.
+                - status (str): The status of the model.
+        Returns:
+            dict[Any, Any]: The report of the model.
+        """
+        metadata = input_data['metadata']
+        model_name = input_data['model_name']
+        model_id = input_data['model_id']
+        experiment = input_data['experiment']
+        timestamp = input_data['timestamp']
+        status = input_data['status']
+
+        self.info(
+            f'Updating production model {model_name} from experiment {experiment}...', metadata)
+
+        try:
+            response = self.model_monitoring_repository.update_production_model(
+                experiment=experiment,
+                model_name=model_name
+            )
+
+            report = DataFrame([response])
+            report['model_id'] = model_id
+            report['model_name'] = model_name
+            report['timestamp'] = timestamp
+            report['status'] = status
+
+            self.info(
+                f'Production model {model_name} updated successfully', metadata)
+            return report.to_dict()
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.send_notification(
+                metadata=metadata,
+                notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
+                message=f'Error updating production model {model_name}: {e}',
+                block='update_production_model',
+                attachment_content=trace
+            )
+            self.error(trace, metadata=metadata)
+            raise e

laborious/utils/connectors_config.py

@@ -47,12 +47,13 @@ def build_opc_config():
 
 
 def build_mongodb_config():
-    username = getenv('MONGODB_USERNAME', 'sientia')
+    username = getenv('MONGODB_USERNAME', 'root')
-    password = getenv('MONGODB_PASSWORD', 'sientia')
+    password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c')
-    uri = getenv('MONGODB_URL', 'localhost:27017')
+    uri = getenv('MONGODB_URL', 'localhost:27018')
 
     connection_string = f'mongodb://{username}:{password}@{uri}'
     return {
         'connection_string': connection_string,
-        'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia')
+        'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
+        'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
     }

laborious/utils/repository/model_repository.py

@@ -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,186 @@ 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 create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
+        # 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):
+
+        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:
+
+        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)
+
+        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

laborious/worker/worker.py

@@ -6,6 +6,7 @@ with workflow.unsafe.imports_passed_through():
     import os
     import sys
     import asyncio
+    from laborious.workflows.minimal_retrain import MinimalRetrain
     from laborious.workflows.predictions_batch import PredictionsBatch
     from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
     from laborious.workflows.sub_workflows.format_and_export_prediction import \
@@ -51,12 +52,28 @@ async def main():
 
     temporal_client = await client.Client.connect(
         target_host=host,
-        namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')
+        namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious')
     )
 
     logger.info('Starting Workers...')
 
     workers = [
+        Worker(
+            temporal_client,
+            task_queue='minimal_retrain-queue',
+            workflows=[MinimalRetrain],
+            activities=[
+                activities.load_custom_query,
+                activities.retrain_model,
+                activities.update_production_model,
+                activities.export_data_to_postgres
+            ],
+            max_concurrent_workflow_tasks=100,
+            max_concurrent_activities=100,
+            max_concurrent_local_activities=100,
+            max_concurrent_workflow_task_polls=100,
+            max_cached_workflows=50,
+        ),
         Worker(
             temporal_client,
             task_queue='predictions_batch-queue',

laborious/workflows/minimal_retrain.py

@@ -0,0 +1,93 @@
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from laborious.activities.activities import Activities
+    from typing import Any
+    from sientia_do.temporal.utils.policies import retry_policy
+    from datetime import timedelta
+
+
+@workflow.defn(name="minimal_retrain")
+class MinimalRetrain():
+    @workflow.run
+    async def run(self, input_data: dict[str, Any]):
+        """
+        This workflow runs a minimal retrain of a model.
+
+        The workflow executes in four steps:
+        1. Loads the data from the database
+        2. Formats the data and perform the retrain
+        3. Updates the production model
+        4. Saves a model
+
+        Args:
+            - input_data (dict[str, Any]): The input data for the workflow.
+                - schedule_name (str): The name of the schedule.
+                - model_name (str): The name of the model.
+                - model_id (int): The id of the model.
+                - query (str): The SQL query to be executed to load data.
+                - schema (dict, optional): The schema to store the report.
+                - table_name (str, optional): The name of the table to store report.
+
+        Returns:
+            None
+
+        Raises:
+            Exception: If any of the required parameters are missing or if the workflow fails.
+        """
+
+        metadata = {
+            'metadata': {
+                'schedule_name': input_data['schedule_name'],
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id'],
+                'workflow_name': 'minimal_retrain'
+            }
+        }
+
+        model_name = input_data['model_name']
+
+        data = await workflow.execute_local_activity_method(
+            Activities.load_custom_query,
+            {
+                **metadata,
+                'query': input_data['query'],
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )
+
+        experiment_response = await workflow.execute_activity_method(
+            Activities.retrain_model,
+            {
+                **metadata,
+                'data': data,
+                'model_name': model_name
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )
+
+        report = await workflow.execute_activity_method(
+            Activities.update_production_model,
+            {
+                **metadata,
+                'model_name': model_name,
+                'model_id': input_data['model_id'],
+                **experiment_response
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )
+
+        await workflow.execute_activity_method(
+            Activities.export_data_to_postgres,
+            {
+                **metadata,
+                'data': report,
+                'schema': input_data['schema'],
+                'table_name': input_data['table_name']
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )

requirements.txt

@@ -3,5 +3,5 @@ psycopg2-binary
 sqlalchemy
 asyncua
 redis
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.3.3
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.3.4
 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.1

tests/laborious/activities/test_mlflow.py

@@ -1,6 +1,7 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import ANY, MagicMock, patch
 
 import numpy as np
+from pandas import DataFrame
 from pytest import fixture, mark
 from laborious.activities.mlflow import MLFlow
 
@@ -29,7 +30,7 @@ def test___init__(mock_mlflow_repository):
 @fixture
 @patch("laborious.activities.mlflow.MLFlowRepository")
 def mlflow(mock_mlflow_repository):
-    return MLFlow(
+    mlflow = MLFlow(
         mlflow_host="http://localhost:5000",
         mlflow_port=5000,
         mlflow_username="admin",
@@ -38,6 +39,10 @@ def mlflow(mock_mlflow_repository):
         notification_handler=MagicMock()
     )
 
+    mlflow.send_notification = MagicMock()
+
+    return mlflow
+
 
 metadata = {
     "metadata": {
@@ -141,3 +146,127 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow):
     mlflow.model_monitoring_repository.predict.assert_called_once_with(
         'test_model', mock_dataframe.return_value, 30
     )
+
+
+@mark.asyncio
+async def test_retrain_model(mlflow):
+    data = {
+        "model_id": [4, 5, 6, 7],
+        "created_at": [1, 2, 3, 4],
+        "timestamp": [1, 1, 2, 2],
+        "variable": ["var1", "var2", "var1", "var2"],
+        "value": [1, 2, 3, 4]
+    }
+
+    mlflow.model_monitoring_repository.retrain_model.return_value = (
+        'Model retrained successfully', 'test')
+
+    response = await mlflow.retrain_model({
+        **metadata,
+        'data': data,
+        'model_name': 'test_model'
+    })
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once()
+
+    assert response == {
+        "status": 'Model retrained successfully',
+        "timestamp": 2,
+        "experiment": 'test'
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_error(mlflow):
+    mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
+        'Error retraining model'
+    )
+
+    data = {
+        "model_id": [4, 5, 6, 7],
+        "created_at": [1, 2, 3, 4],
+        "timestamp": [1, 1, 2, 2],
+        "variable": ["var1", "var2", "var1", "var2"],
+        "value": [1, 2, 3, 4]
+    }
+
+    try:
+        await mlflow.retrain_model({
+            **metadata,
+            'data': data,
+            'model_name': 'test_model'
+        })
+    except Exception as e:
+        assert str(e) == 'Error retraining model'
+        mlflow.send_notification.assert_called_once_with(
+            metadata=metadata['metadata'],
+            notification_id='RETRAIN_MODEL_ERROR',
+            message='Error retraining model test_model: Error retraining model',
+            block='retrain_model',
+            attachment_content=ANY
+        )
+    else:
+        assert False, "No exception raised"
+
+
+@mark.asyncio
+async def test_update_production_model(mlflow):
+    mlflow.model_monitoring_repository.update_production_model.return_value = (
+        {
+            "data1": 1,
+            "data2": 2
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 1,
+        'experiment': 'test',
+        'timestamp': 2,
+        'status': 'success'
+    }
+
+    response = await mlflow.update_production_model(input_data)
+
+    mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
+        experiment='test', model_name='test_model')
+
+    assert response == {
+        'data1': {0: 1},
+        'data2': {0: 2},
+        'model_id': {0: 1},
+        'model_name': {0: 'test_model'},
+        'timestamp': {0: 2},
+        'status': {0: 'success'}
+    }
+
+
+@mark.asyncio
+async def test_update_production_model_error(mlflow):
+    mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
+        'Error updating production model'
+    )
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 1,
+        'experiment': 'test',
+        'timestamp': 2,
+        'status': 'success'
+    }
+
+    try:
+        await mlflow.update_production_model(input_data)
+    except Exception as e:
+        assert str(e) == 'Error updating production model'
+        mlflow.send_notification.assert_called_once_with(
+            metadata=metadata['metadata'],
+            notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
+            message='Error updating production model test_model: Error updating production model',
+            block='update_production_model',
+            attachment_content=ANY
+        )
+    else:
+        assert False, "No exception raised"

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,288 @@ 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')
+def test_get_experiment_last_run_error(mlflow, mlflow_repository):
+    mlflow.search_runs.return_value = []
+
+    try:
+        mlflow_repository.get_experiment_last_run(0)
+    except ValueError as e:
+        assert str(e) == 'Runs is not a pandas DataFrame'
+    else:
+        assert False
+
+
+@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):
+
+    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)
+
+
+def test_retrain_model(mlflow_repository):
+    data = MagicMock()
+    model_name = 'test'
+
+    mlflow_repository.create_model_experiment = MagicMock(
+        return_value=('data_model', 'prediction_model', '0'))
+
+    mlflow_repository.perform_model_retrain = MagicMock(
+        return_value='Model retrained successfully')
+
+    output = mlflow_repository.retrain_model(data, model_name)
+
+    mlflow_repository.create_model_experiment.assert_called_once_with(
+        model_name, data)
+
+    mlflow_repository.perform_model_retrain.assert_called_once_with(
+        'data_model', 'prediction_model', '0', model_name, data)
+
+    assert output == 'Model retrained successfully'
+
+
+@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',
+    }
+
+
+@patch('laborious.utils.repository.model_repository.mlflow')
+def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
+    mlflow.tracking.MlflowClient.return_value = MagicMock(
+        get_registered_model=MagicMock(
+            return_value=MagicMock(
+                latest_versions={}
+            )
+        )
+    )
+
+    try:
+        mlflow_repository.update_production_model_by_run_id('0', 'test')
+    except Exception as e:
+        assert str(e) == 'Model versions is not a list'
+    else:
+        assert False
+
+
+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
     }

tests/laborious/workflows/test_minimal_retrain.py

@@ -0,0 +1,97 @@
+from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
+from pytest import fixture, mark
+from laborious.activities.activities import Activities
+from laborious.workflows.minimal_retrain import MinimalRetrain
+
+
+@fixture
+def minimal_retrain() -> MinimalRetrain:
+    return MinimalRetrain()
+
+
+metadata = {
+    "metadata": {
+        "model_id": "test_model_id",
+        "model_name": "test_model",
+        "workflow_name": "minimal_retrain",
+        "schedule_name": "test_schedule",
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        "model_id": "test_model_id",
+        "model_name": "test_model",
+        "workflow_name": "minimal_retrain",
+        "schedule_name": "test_schedule",
+        "query": "test_query",
+        "schema": "test_schema",
+        "table_name": "test_table",
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        return_value={
+            "data1": "1",
+            "data2": "2",
+        }
+    )
+
+    await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                {
+                    **metadata,
+                    "query": input_data["query"],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.retrain_model,
+            {
+                **metadata,
+                'data': workflow_mock.execute_local_activity_method.return_value,
+                'model_name': input_data['model_name'],
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )
+    ])
+
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.update_production_model,
+            {
+                **metadata,
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id'],
+                **workflow_mock.execute_activity_method.return_value,
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )
+    ])
+
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.export_data_to_postgres,
+            {
+                **metadata,
+                'data': workflow_mock.execute_activity_method.return_value,
+                'schema': input_data['schema'],
+                'table_name': input_data['table_name'],
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )
+    ])

values.yaml

@@ -11,7 +11,7 @@ image:
   # This sets the pull policy for images.
   pullPolicy: Always
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.2.4"
+  tag: "0.2.5"
 
 # This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
 imagePullSecrets:
@@ -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"
 
@@ -178,6 +178,8 @@ env:
     value: "my-release-mongodb.mongodb.svc.cluster.local:27017"
   - name: MONGODB_DATABASE
     value: "sientia"
+  - name: MONGODB_TTL_INDEX_HOURS
+    value: "1"
 
 ssh:
   enabled: true