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.
2025-10-09 15:07:27 -03:00
parent c970754e1b
commit bee6036205
7 changed files with 1046 additions and 2 deletions
--- a/tests/utils/repository/test_training_repository.py
+++ b/tests/utils/repository/test_training_repository.py
@@ -0,0 +1,327 @@
+"""Unit tests for TrainingRepository."""
+
+from io import BytesIO
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pandas as pd
+from pytest import fixture, raises
+
+from model_manager.utils.models.train_model_params import TrainModelParams
+from model_manager.utils.models.train_model_result import TrainModelResult
+from model_manager.utils.repository.training_repository import TrainingRepository
+
+
+@fixture
+def logger():
+    """Create a mock logger."""
+    return MagicMock()
+
+
+@fixture
+def training_repository(logger):
+    """Create a TrainingRepository instance."""
+    return TrainingRepository(logger)
+
+
+@fixture
+def train_params():
+    """Create sample training parameters."""
+    return TrainModelParams(
+        variable_columns=['feature1', 'feature2'],
+        lag_train=1,
+        lag_val=1,
+        target_variable='target',
+        rem_static_win=False,
+        low_lim={'feature1': 0.0, 'feature2': 0.0},
+        upp_lim={'feature1': 100.0, 'feature2': 100.0},
+        window=10,
+        use_scaler=True,
+        include_ar=False,
+        bucket_name='test-bucket',
+        file_name='test.csv',
+        line_separator='\n',
+        decimal_separator='.',
+        train_size=80,
+        shuffle=True,
+        experiment_run_id=123,
+        experiment_name='test_experiment',
+        experiment_description='Test experiment',
+        removed_intervals=[],
+    )
+
+
+@fixture
+def sample_csv_data():
+    """Create sample CSV data."""
+    csv_content = """feature1,feature2,target
+1.0,2.0,10.0
+2.0,3.0,15.0
+3.0,4.0,20.0
+4.0,5.0,25.0
+5.0,6.0,30.0
+6.0,7.0,35.0
+7.0,8.0,40.0
+8.0,9.0,45.0
+9.0,10.0,50.0
+10.0,11.0,55.0
+"""
+    return BytesIO(csv_content.encode())
+
+
+@patch('model_manager.utils.repository.training_repository.load_data')
+@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
+@patch('model_manager.utils.repository.training_repository.split_train_test')
+@patch('model_manager.utils.repository.training_repository.LinearRegressionModel')
+def test_train_success(
+    mock_linear_model,
+    mock_split,
+    mock_preprocessor_class,
+    mock_load_data,
+    training_repository,
+    train_params,
+    sample_csv_data,
+):
+    """Test successful model training."""
+    # Setup mocks
+    mock_data = pd.DataFrame(
+        {'feature1': [1, 2, 3, 4, 5], 'feature2': [2, 3, 4, 5, 6], 'target': [10, 15, 20, 25, 30]}
+    )
+    mock_load_data.return_value = mock_data
+
+    mock_preprocessor = MagicMock()
+    mock_preprocessor_class.return_value = mock_preprocessor
+    mock_preprocessor.transform.return_value = mock_data
+
+    x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
+    x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
+    y_train = pd.Series([10, 15, 20], name='target')
+    y_test = pd.Series([25, 30], name='target')
+    mock_split.return_value = (x_train, x_test, y_train, y_test)
+
+    mock_model = MagicMock()
+    mock_linear_model.return_value = mock_model
+
+    mock_scaler = MagicMock()
+    mock_preprocessor.get_scaler.return_value = mock_scaler
+
+    # Execute
+    result = training_repository.train(sample_csv_data, train_params)
+
+    # Assertions
+    assert isinstance(result, TrainModelResult)
+    assert result.params == train_params
+    assert result.process_data == mock_preprocessor
+    assert result.regr == mock_model
+    mock_load_data.assert_called_once()
+    mock_preprocessor.fit.assert_called_once()
+    mock_model.fit.assert_called_once()
+
+
+@patch('model_manager.utils.repository.training_repository.load_data')
+def test_train_empty_data_after_transform(
+    mock_load_data, training_repository, train_params, sample_csv_data
+):
+    """Test training with empty data after transformation."""
+    mock_data = pd.DataFrame({'feature1': [], 'feature2': [], 'target': []})
+    mock_load_data.return_value = mock_data
+
+    with patch.object(training_repository, 'init_data_preprocessor') as mock_init:
+        mock_preprocessor = MagicMock()
+        mock_init.return_value = mock_preprocessor
+        mock_preprocessor.transform.return_value = pd.DataFrame()
+
+        with raises(ValueError, match='Data view is empty after transformation'):
+            training_repository.train(sample_csv_data, train_params)
+
+
+def test_init_scaler_dict_with_minmax_scaler(training_repository, train_params):
+    """Test scaler dict initialization with MinMaxScaler."""
+    mock_preprocessor = MagicMock()
+    mock_scaler = MagicMock()
+    mock_scaler.x_min = [0.0, 1.0]
+    mock_scaler.x_max = [10.0, 11.0]
+    mock_scaler.y_min = 5.0
+    mock_scaler.y_max = 50.0
+    mock_preprocessor.get_scaler.return_value = mock_scaler
+
+    # Patch isinstance to return True for MinMaxScaler
+    with patch(
+        'model_manager.utils.repository.training_repository.isinstance',
+        side_effect=lambda obj, cls: cls.__name__ == 'MinMaxScaler',
+    ):
+        result = training_repository.init_scaler_dict(mock_preprocessor, train_params)
+
+    assert result is not None
+    assert 'feature1' in result
+    assert 'feature2' in result
+    assert 'target' in result
+    assert result['feature1'] == {'min': 0.0, 'max': 10.0}
+    assert result['feature2'] == {'min': 1.0, 'max': 11.0}
+    assert result['target'] == {'min': 5.0, 'max': 50.0}
+
+
+def test_init_scaler_dict_with_z_scaler(training_repository, train_params):
+    """Test scaler dict initialization with Z_Scaler."""
+    mock_preprocessor = MagicMock()
+    mock_scaler = MagicMock()
+    mock_scaler.create_dict.return_value = {'mean': 5.0, 'std': 2.0}
+    mock_preprocessor.get_scaler.return_value = mock_scaler
+
+    # Patch isinstance to return True for Z_Scaler
+    with patch(
+        'model_manager.utils.repository.training_repository.isinstance',
+        side_effect=lambda obj, cls: cls.__name__ == 'Z_Scaler',
+    ):
+        result = training_repository.init_scaler_dict(mock_preprocessor, train_params)
+
+    assert result == {'mean': 5.0, 'std': 2.0}
+    mock_scaler.create_dict.assert_called_once()
+
+
+def test_init_scaler_dict_without_scaler(training_repository):
+    """Test scaler dict initialization when use_scaler is False."""
+    train_params_no_scaler = TrainModelParams(
+        variable_columns=['feature1'],
+        lag_train=1,
+        lag_val=1,
+        target_variable='target',
+        rem_static_win=False,
+        low_lim={'feature1': 0.0},
+        upp_lim={'feature1': 100.0},
+        window=10,
+        use_scaler=False,
+        include_ar=False,
+        bucket_name='test',
+        file_name='test.csv',
+        line_separator='\n',
+        decimal_separator='.',
+        train_size=80,
+        shuffle=True,
+        experiment_run_id=123,
+        experiment_name='test',
+        experiment_description='test',
+        removed_intervals=[],
+    )
+
+    mock_preprocessor = MagicMock()
+    result = training_repository.init_scaler_dict(mock_preprocessor, train_params_no_scaler)
+
+    assert result == {}
+
+
+@patch('model_manager.utils.repository.training_repository.mse')
+@patch('model_manager.utils.repository.training_repository.mae')
+@patch('model_manager.utils.repository.training_repository.r2')
+def test_after_train_calculation_with_scaler(
+    mock_r2, mock_mae, mock_mse, training_repository, train_params
+):
+    """Test post-training calculations with scaler."""
+    # Setup mock train result
+    mock_train_result = MagicMock(spec=TrainModelResult)
+    mock_train_result.params = train_params
+    mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
+    mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
+    mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
+    mock_train_result.y_test = pd.Series([25, 30], name='target')
+
+    mock_regr = MagicMock()
+    mock_regr.predict.return_value = np.array([24.5, 29.5])
+    mock_train_result.regr = mock_regr
+
+    mock_scaler = MagicMock()
+    mock_scaler.denormalize_single_input.side_effect = lambda x, col: x
+    mock_scaler.denormalize_predictions.side_effect = lambda x, col: x
+
+    mock_process_data = MagicMock()
+    mock_process_data.get_scaler.return_value = mock_scaler
+    mock_train_result.process_data = mock_process_data
+
+    # Setup metric mocks
+    mock_mse.return_value = 0.5
+    mock_mae.return_value = 0.3
+    mock_r2.return_value = 0.95
+
+    # Execute
+    result = training_repository.after_train_calculation(train_params, mock_train_result)
+
+    # Assertions
+    assert result == mock_train_result
+    assert result.mse_val == 0.5
+    assert result.mae_val == 0.3
+    assert result.r2_val == 0.95
+    assert result.y_pred is not None
+    mock_regr.predict.assert_called_once()
+    mock_mse.assert_called_once()
+    mock_mae.assert_called_once()
+    mock_r2.assert_called_once()
+
+
+@patch('model_manager.utils.repository.training_repository.mse')
+@patch('model_manager.utils.repository.training_repository.mae')
+@patch('model_manager.utils.repository.training_repository.r2')
+def test_after_train_calculation_without_scaler(mock_r2, mock_mae, mock_mse, training_repository):
+    """Test post-training calculations without scaler."""
+    train_params_no_scaler = TrainModelParams(
+        variable_columns=['feature1'],
+        lag_train=1,
+        lag_val=1,
+        target_variable='target',
+        rem_static_win=False,
+        low_lim={'feature1': 0.0},
+        upp_lim={'feature1': 100.0},
+        window=10,
+        use_scaler=False,
+        include_ar=False,
+        bucket_name='test',
+        file_name='test.csv',
+        line_separator='\n',
+        decimal_separator='.',
+        train_size=80,
+        shuffle=True,
+        experiment_run_id=123,
+        experiment_name='test',
+        experiment_description='test',
+        removed_intervals=[],
+    )
+
+    mock_train_result = MagicMock(spec=TrainModelResult)
+    mock_train_result.params = train_params_no_scaler
+    mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3]})
+    mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5]})
+    mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
+    mock_train_result.y_test = pd.Series([25, 30], name='target')
+
+    mock_regr = MagicMock()
+    mock_regr.predict.return_value = np.array([24.5, 29.5])
+    mock_train_result.regr = mock_regr
+
+    # Setup metric mocks
+    mock_mse.return_value = 0.5
+    mock_mae.return_value = 0.3
+    mock_r2.return_value = 0.95
+
+    # Execute
+    result = training_repository.after_train_calculation(train_params_no_scaler, mock_train_result)
+
+    # Assertions
+    assert result.mse_val == 0.5
+    assert result.mae_val == 0.3
+    assert result.r2_val == 0.95
+
+
+@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
+def test_init_data_preprocessor(mock_preprocessor_class, training_repository, train_params):
+    """Test DataPreprocessor initialization."""
+    mock_preprocessor = MagicMock()
+    mock_preprocessor_class.return_value = mock_preprocessor
+
+    result = training_repository.init_data_preprocessor(train_params)
+
+    assert result == mock_preprocessor
+    mock_preprocessor_class.assert_called_once()
+    call_kwargs = mock_preprocessor_class.call_args[1]
+    assert call_kwargs['target_variable'] == 'target'
+    assert call_kwargs['input_columns'] == ['feature1', 'feature2']
+    assert call_kwargs['low_lim'] == {'feature1': 0.0, 'feature2': 0.0}
+    assert call_kwargs['upp_lim'] == {'feature1': 100.0, 'feature2': 100.0}