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

SIENTIAPDE-1171

Browse Source 
Update GITHUB_BRANCH in values.yaml and refactor MLFlowRepository for model retraining

- Changed GITHUB_BRANCH in values.yaml to reflect the new pipeline for retraining.
- Refactored MLFlowRepository by introducing create_model_experiment and perform_model_retrain methods to streamline model retraining processes.
- Updated test cases to cover new functionalities and ensure proper error handling in model management workflows.
This commit is contained in:
vitor-aignosi
2025-07-23 16:51:33 -03:00
parent 89b9892a5b
commit 4283730e7a
4 changed files with 275 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -116,32 +116,7 @@ class MLFlowRepository():
next_run_number = len(runs) + 1
return f"{model_name}-{next_run_number}"
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
"""
Retrain a model with new data.
Parameters:
data (pandas.DataFrame): The new data to use for retraining.
model_name (str): The name of the model to retrain.
metrics_list (list): The metrics to be used to compare the models.
compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one.
If False, the retrain will always be considered.
split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets.
If False, the data will be used as a unique block for retraining.
update_report (bool): If True, a report will be created with the data of the retrained model.
update_transformation (bool): If True, the model will be updated in the MLflow tracking server.
update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server.
shuffle_data (bool): If True, the data will be shuffled before splitting.
model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'.
Returns:
mlflow.sklearn.Model: The retrained prediction model.
mlflow.sklearn.Model: The retrained data model.
mse (float): The mean squared error of the retrained model.
r2 (float): The R-squared score of the retrained model.
"""
def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
# load predictor model
predictor_uri = f"models:/{model_name}/production"
# load transform model
@@ -156,18 +131,27 @@ class MLFlowRepository():
prediction_model = mlflow.sklearn.load_model(predictor_uri)
data_model = data_model.fit(data)
treated_data = data_model.predict(data)
# align target column with treated_data
target_name = data_model.target_variable
y = data[target_name]
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
# Example usage
experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment)
return prediction_model, data_model, experiment
def perform_model_retrain(self,
prediction_model,
data_model,
experiment: str,
model_name: str,
data: pd.DataFrame):
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
mlflow.set_experiment(experiment)
experiment_description = "Retrain model {model_name} with new data"
experiment_description = f"Retrain model {model_name} with new data"
current_run_name = self.get_next_run_name(experiment)
with mlflow.start_run(
run_name=current_run_name, description=experiment_description
@@ -203,6 +187,14 @@ class MLFlowRepository():
return "Model retrained successfully", experiment
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
prediction_model, data_model, experiment = self.create_model_experiment(
model_name, data)
retrain_result = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data)
return retrain_result
def get_experiment(self, experiment_name: str) -> int:
experiment = mlflow.get_experiment_by_name(experiment_name)

246
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,245 @@ def test_predict_error(mlflow_repository):
'traceback': ANY
}
}
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
mlflow.get_run.return_value = MagicMock(
info=MagicMock(
experiment_id='0',
)
)
mlflow.get_experiment.return_value = MagicMock()
mlflow.get_experiment.return_value.name = 'test'
output = mlflow_repository.get_experiment_by_run_id('0')
assert output == 'test'
mlflow.get_run.assert_called_once_with('0')
mlflow.get_experiment.assert_called_once_with('0')
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_next_run_name(mlflow, mlflow_repository):
mlflow.search_runs.return_value = [1, 2, 3]
output = mlflow_repository.get_next_run_name('run')
assert output == 'run-4'
mlflow.search_runs.assert_called_once_with(
experiment_names=['run'],
order_by=['start_time desc'],
)
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(
experiment_id='0')
output = mlflow_repository.get_experiment('test')
assert output == 0
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_error(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = None
try:
mlflow_repository.get_experiment('test')
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
assert False
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository):
mlflow.search_runs.return_value = DataFrame({
'params.retrain': ['True', 'False', 'True', 'False'],
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
'run_id': ['0', '1', '2', '3'],
})
output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with(
experiment_ids=[0],
filter_string="",
output_format="pandas",
)
assert output == '2'
@patch('laborious.utils.repository.model_repository.mlflow.sklearn')
@patch('laborious.utils.repository.model_repository.mlflow.set_experiment')
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
mlflow_repository.model_serving.get_model_run_id = MagicMock(
return_value='0')
mlflow_repository.model_serving.get_model_uri = MagicMock(
return_value='test')
mlflow_repository.get_experiment_by_run_id = MagicMock()
data_model_mock = MagicMock()
prediction_model_mock = MagicMock()
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
data_model_mock.fit.return_value = data_model_mock
data_model_mock.predict.return_value = DataFrame({
'x': [10, 20, 30],
})
data_model_mock.target_variable = 'y'
prediction_model_mock.fit.return_value = prediction_model_mock
data = DataFrame({
'x': [1, 2, 3],
'y': [4, 5, 6]
})
output = mlflow_repository.create_model_experiment('test', data)
mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
'test', stage='Production')
mlflow_repository.model_serving.get_model_uri.assert_called_once_with(
'0', prediction=False)
sklearn.load_model.assert_has_calls([
call(mlflow_repository.model_serving.get_model_uri.return_value),
call("models:/test/production"),
])
assert sklearn.load_model.call_count == 2
data_model_mock.fit.assert_called_once_with(data)
data_model_mock.predict.assert_called_once_with(data)
fit_args = prediction_model_mock.fit.call_args[0][0]
assert fit_args.equals(
DataFrame({
'x': [10, 20, 30],
'y': [4, 5, 6],
})
)
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
set_experiment.assert_called_once_with(
mlflow_repository.get_experiment_by_run_id.return_value
)
assert output == (prediction_model_mock,
data_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value)
@patch('laborious.utils.repository.model_repository.mlflow.start_run')
@patch('laborious.utils.repository.model_repository.mlflow.log_param')
@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model')
@patch('laborious.utils.repository.model_repository.mlflow.log_artifact')
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
_os_remove = patch(
'laborious.utils.repository.model_repository.remove')
_os_makedirs = patch(
'laborious.utils.repository.model_repository.makedirs')
_os_path_exists = patch(
'laborious.utils.repository.model_repository.path.exists',
MagicMock(return_value=False))
prediction_model_mock = MagicMock()
data_model_mock = MagicMock()
experiment = 'test'
model_name = 'test'
data = MagicMock()
mlflow_repository.get_next_run_name = MagicMock(
return_value='test-1')
run = MagicMock()
start_run.__enter__.return_value = run
output = mlflow_repository.perform_model_retrain(
prediction_model_mock, data_model_mock, experiment, model_name, data)
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
start_run.assert_called_once_with(
run_name='test-1', description='Retrain model test with new data')
log_model.assert_has_calls([
call(data_model_mock, "data_model"),
call(prediction_model_mock, "prediction_model"),
])
data.to_csv.assert_called_once_with(
"temp/raw_data_test.csv", index=True)
log_artifact.assert_called_once_with(
"temp/raw_data_test.csv")
log_param.assert_has_calls([
call("retrain", True),
])
assert output == ("Model retrained successfully", experiment)
@patch('laborious.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
client_mock = MagicMock()
mlflow.tracking.MlflowClient.return_value = client_mock
client_mock.get_registered_model.return_value = MagicMock(
latest_versions=[
MagicMock(version='1'),
MagicMock(version='2'),
MagicMock(version='3'),
]
)
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model",
'test',
)
mlflow.tracking.MlflowClient.assert_called_once()
client_mock.get_registered_model.assert_called_once_with('test')
client_mock.transition_model_version_stage.assert_called_once_with(
name='test',
version='3',
stage='Production',
archive_existing_versions=True,
)
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
}
def test_update_production_model(mlflow_repository):
connector = mlflow_repository
with patch.object(connector, 'get_experiment',
return_value='0') as get_experiment:
with patch.object(connector, 'get_experiment_last_run',
return_value='2') as get_experiment_last_run:
with patch.object(connector, 'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3',
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with(
'2', 'test')
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
'mlflow_experiment_id': '0',
}

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
}

2
values.yaml
View File

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