diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index bc0113e..63206a7 100644 --- a/laborious/activities/mlflow.py +++ b/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 diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index 5158f05..bcbb49b 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -47,12 +47,13 @@ def build_opc_config(): def build_mongodb_config(): - username = getenv('MONGODB_USERNAME', 'sientia') - password = getenv('MONGODB_PASSWORD', 'sientia') - uri = getenv('MONGODB_URL', 'localhost:27017') + username = getenv('MONGODB_USERNAME', 'root') + password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c') + 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 } diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index b3473c6..2f10a9d 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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 diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 005bb18..375c84f 100644 --- a/laborious/worker/worker.py +++ b/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', diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py new file mode 100644 index 0000000..3d6f8d6 --- /dev/null +++ b/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) + ) diff --git a/requirements.txt b/requirements.txt index 8b1535e..fed750d 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index a92f3f3..3d83736 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/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" diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 6d10c1c..62f4fc3 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,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', + } 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/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py new file mode 100644 index 0000000..bcb2288 --- /dev/null +++ b/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 + ) + ]) diff --git a/values.yaml b/values.yaml index 3a54067..03f681b 100644 --- a/values.yaml +++ b/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