SIENTIAPDE-1251: Implement ML model training activity and repository
This commit introduces the 'Training' activity and 'TrainingRepository' for handling ML model training operations within the Model Manager system. - Added model_manager/activities/training.py for the Training activity, which extends BaseActivity and integrates with Temporal workflows. - Added model_manager/utils/repository/training_repository.py for the TrainingRepository, which encapsulates the core training logic. - Updated model_manager/activities/activities.py to include the Training activity in the main activities orchestrator. - Updated README.md to document the new 'Training' component. - Added unit tests for the new activity and repository.
This commit is contained in:
Bruno Domingues
@@ -6,14 +6,20 @@ from model_manager.activities.activities import Activities
|
||||
from model_manager.activities.experiment_tracking import ExperimentTracking
|
||||
from model_manager.activities.gates import Gates
|
||||
from model_manager.activities.mlflow import MLFlow
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
|
||||
@patch('model_manager.activities.activities.MLFlow.__init__')
|
||||
@patch('model_manager.activities.activities.MinIO.__init__')
|
||||
@patch('model_manager.activities.activities.Gates.__init__')
|
||||
@patch('model_manager.activities.activities.Training.__init__')
|
||||
def test___init__(
|
||||
mock_gates_init, mock_minio_init, mock_mlflow_init, mock_experiment_tracking_init
|
||||
mock_training_init,
|
||||
mock_gates_init,
|
||||
mock_minio_init,
|
||||
mock_mlflow_init,
|
||||
mock_experiment_tracking_init,
|
||||
):
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
@@ -54,6 +60,7 @@ def test___init__(
|
||||
assert isinstance(activities, ExperimentTracking)
|
||||
assert isinstance(activities, MLFlow)
|
||||
assert isinstance(activities, Gates)
|
||||
assert isinstance(activities, Training)
|
||||
|
||||
mock_experiment_tracking_init.assert_called_once_with(
|
||||
ANY,
|
||||
@@ -97,6 +104,10 @@ def test___init__(
|
||||
ANY, logger=logger, notification_handler=notification_handler
|
||||
)
|
||||
|
||||
mock_training_init.assert_called_once_with(
|
||||
ANY, logger=logger, notification_handler=notification_handler
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.activities.ExperimentTracking', return_value=MagicMock())
|
||||
|
||||
288
tests/activities/test_training.py
Normal file
288
tests/activities/test_training.py
Normal file
@@ -0,0 +1,288 @@
|
||||
"""Unit tests for Training activity."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from pytest import mark
|
||||
|
||||
from model_manager.activities.training import Training
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
async def test_train_model_success(mock_training_repository_class):
|
||||
"""Test successful model training."""
|
||||
# Create mock repository instance
|
||||
mock_repository = MagicMock()
|
||||
mock_training_repository_class.return_value = mock_repository
|
||||
|
||||
# Create mock train result
|
||||
mock_train_result = MagicMock(spec=TrainModelResult)
|
||||
mock_train_result.mse_val = 0.5
|
||||
mock_train_result.mae_val = 0.3
|
||||
mock_train_result.r2_val = 0.95
|
||||
|
||||
mock_final_result = MagicMock(spec=TrainModelResult)
|
||||
mock_final_result.mse_val = 0.5
|
||||
mock_final_result.mae_val = 0.3
|
||||
mock_final_result.r2_val = 0.95
|
||||
|
||||
# Setup repository mocks
|
||||
mock_repository.train.return_value = mock_train_result
|
||||
mock_repository.after_train_calculation.return_value = mock_final_result
|
||||
|
||||
# Create Training instance
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
# Mock inherited methods
|
||||
training.info = MagicMock()
|
||||
|
||||
# Test data
|
||||
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
|
||||
train_params_dict = {
|
||||
'experiment_run_id': 123,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1', 'feature2'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': True,
|
||||
'include_ar': False,
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'lag_train': 1,
|
||||
'lag_val': 1,
|
||||
'rem_static_win': False,
|
||||
'low_lim': {'feature1': 0.0, 'feature2': 0.0},
|
||||
'upp_lim': {'feature1': 100.0, 'feature2': 100.0},
|
||||
'window': 10,
|
||||
'experiment_name': 'test_experiment',
|
||||
'experiment_description': 'Test experiment',
|
||||
'removed_intervals': [],
|
||||
}
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'uploaded_file': uploaded_file,
|
||||
'train_params': train_params_dict,
|
||||
}
|
||||
|
||||
# Execute
|
||||
result = await training.train_model(input_data)
|
||||
|
||||
# Assertions
|
||||
assert result['success'] is True
|
||||
assert result['result'] == mock_final_result
|
||||
assert result['error_message'] is None
|
||||
|
||||
# Verify repository calls
|
||||
mock_repository.train.assert_called_once()
|
||||
mock_repository.after_train_calculation.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
async def test_train_model_invalid_file_type(mock_training_repository_class):
|
||||
"""Test training with invalid file type."""
|
||||
mock_repository = MagicMock()
|
||||
mock_training_repository_class.return_value = mock_repository
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
# Invalid file type (string instead of BytesIO)
|
||||
input_data = {
|
||||
'metadata': {},
|
||||
'uploaded_file': 'not_a_bytesio',
|
||||
'train_params': {
|
||||
'experiment_run_id': 123,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': False,
|
||||
'include_ar': False,
|
||||
'bucket_name': 'test',
|
||||
'file_name': 'test.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'lag_train': 1,
|
||||
'lag_val': 1,
|
||||
'rem_static_win': False,
|
||||
'low_lim': {'feature1': 0.0},
|
||||
'upp_lim': {'feature1': 100.0},
|
||||
'window': 10,
|
||||
'experiment_name': 'test_experiment',
|
||||
'experiment_description': 'Test experiment',
|
||||
'removed_intervals': [],
|
||||
},
|
||||
}
|
||||
|
||||
result = await training.train_model(input_data)
|
||||
|
||||
assert result['success'] is False
|
||||
assert result['result'] is None
|
||||
assert 'uploaded_file must be BytesIO' in result['error_message']
|
||||
# Verify notification was sent (via BaseActivity)
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
async def test_train_model_training_error(mock_training_repository_class):
|
||||
"""Test training failure during model training."""
|
||||
mock_repository = MagicMock()
|
||||
mock_training_repository_class.return_value = mock_repository
|
||||
|
||||
# Setup repository to raise error
|
||||
mock_repository.train.side_effect = ValueError('Training data is empty')
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
uploaded_file = BytesIO(b'test,data\n')
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-456'},
|
||||
'uploaded_file': uploaded_file,
|
||||
'train_params': {
|
||||
'experiment_run_id': 456,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': False,
|
||||
'include_ar': False,
|
||||
'bucket_name': 'test',
|
||||
'file_name': 'test.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'lag_train': 1,
|
||||
'lag_val': 1,
|
||||
'rem_static_win': False,
|
||||
'low_lim': {'feature1': 0.0},
|
||||
'upp_lim': {'feature1': 100.0},
|
||||
'window': 10,
|
||||
'experiment_name': 'test_experiment',
|
||||
'experiment_description': 'Test experiment',
|
||||
'removed_intervals': [],
|
||||
},
|
||||
}
|
||||
|
||||
result = await training.train_model(input_data)
|
||||
|
||||
assert result['success'] is False
|
||||
assert result['result'] is None
|
||||
assert 'Training data is empty' in result['error_message']
|
||||
# Verify notification was sent (via BaseActivity)
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
async def test_train_model_sends_notification_on_error(mock_training_repository_class):
|
||||
"""Test that notification is sent when training fails."""
|
||||
mock_repository = MagicMock()
|
||||
mock_training_repository_class.return_value = mock_repository
|
||||
|
||||
mock_repository.train.side_effect = Exception('Database connection failed')
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
uploaded_file = BytesIO(b'test,data\n1,2')
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-789', 'experiment_run_id': 789},
|
||||
'uploaded_file': uploaded_file,
|
||||
'train_params': {
|
||||
'experiment_run_id': 789,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': False,
|
||||
'include_ar': False,
|
||||
'bucket_name': 'test',
|
||||
'file_name': 'test.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'lag_train': 1,
|
||||
'lag_val': 1,
|
||||
'rem_static_win': False,
|
||||
'low_lim': {'feature1': 0.0},
|
||||
'upp_lim': {'feature1': 100.0},
|
||||
'window': 10,
|
||||
'experiment_name': 'test_experiment',
|
||||
'experiment_description': 'Test experiment',
|
||||
'removed_intervals': [],
|
||||
},
|
||||
}
|
||||
|
||||
result = await training.train_model(input_data)
|
||||
|
||||
# Verify notification was sent (via BaseActivity)
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
|
||||
# Verify result - the important part is that error was caught and returned
|
||||
assert result['success'] is False
|
||||
assert result['result'] is None
|
||||
assert 'Database connection failed' in result['error_message']
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
async def test_train_model_after_calculation_error(mock_training_repository_class):
|
||||
"""Test training failure during post-training calculations."""
|
||||
mock_repository = MagicMock()
|
||||
mock_training_repository_class.return_value = mock_repository
|
||||
|
||||
# Train succeeds but after_calculation fails
|
||||
mock_train_result = MagicMock(spec=TrainModelResult)
|
||||
mock_repository.train.return_value = mock_train_result
|
||||
mock_repository.after_train_calculation.side_effect = Exception('Metric calculation failed')
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
training = Training(logger=logger, notification_handler=notification_handler)
|
||||
|
||||
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
|
||||
input_data = {
|
||||
'metadata': {},
|
||||
'uploaded_file': uploaded_file,
|
||||
'train_params': {
|
||||
'experiment_run_id': 999,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': False,
|
||||
'include_ar': False,
|
||||
'bucket_name': 'test',
|
||||
'file_name': 'test.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'lag_train': 1,
|
||||
'lag_val': 1,
|
||||
'rem_static_win': False,
|
||||
'low_lim': {'feature1': 0.0},
|
||||
'upp_lim': {'feature1': 100.0},
|
||||
'window': 10,
|
||||
'experiment_name': 'test_experiment',
|
||||
'experiment_description': 'Test experiment',
|
||||
'removed_intervals': [],
|
||||
},
|
||||
}
|
||||
|
||||
result = await training.train_model(input_data)
|
||||
|
||||
assert result['success'] is False
|
||||
assert result['result'] is None
|
||||
assert 'Metric calculation failed' in result['error_message']
|
||||
# Verify notification was sent (via BaseActivity)
|
||||
notification_handler.send_notification.assert_called_once()
|
||||
Reference in New Issue
Block a user