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

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

Browse Source 
Sientiapde 1171 criar pipeline de retreino laborious
This commit is contained in:
vitor-aignosi
2025-07-24 11:46:47 -03:00
committed by GitHub
parent 9410de5f82 a34bdcd839
commit 26887c820c
11 changed files with 936 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
laborious/activities/mlflow.py
View File

@@ -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

9
laborious/utils/connectors_config.py
View File

@@ -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
}

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

@@ -13,6 +13,8 @@ By Monitoring we mean the evaluation of the performance of models, the generatio
from datetime import datetime
import traceback
import pandas as pd
import mlflow
from os import makedirs, path, remove
from sientia.ModelServing import ModelServing
@@ -85,3 +87,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

19
laborious/worker/worker.py
View File

@@ -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',

93
laborious/workflows/minimal_retrain.py Normal file
View File

@@ -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)
)

2
requirements.txt
View File

@@ -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

133
tests/laborious/activities/test_mlflow.py
View File

@@ -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"

289
tests/laborious/utils/repository/test_model_repository.py
View File

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

11
tests/laborious/utils/test_connectors_config.py
View File

@@ -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
}

97
tests/laborious/workflows/test_minimal_retrain.py Normal file
View File

@@ -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
)
])

6
values.yaml
View File

@@ -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