From 61267ec49dea0c13c57b7ed1a69e52a45c6be28a Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Tue, 14 Oct 2025 10:03:34 -0300 Subject: [PATCH] SIENTIAPDE-1253: Enforce TrainModelParams object in Training activity and tests, removing dict conversion. --- model_manager/activities/training.py | 38 ++-- tests/activities/test_training.py | 255 ++++++++++++++++----------- 2 files changed, 167 insertions(+), 126 deletions(-) diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index e022b48..8b9dee7 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -71,7 +71,7 @@ class Training(BaseActivity): Required keys: - metadata (dict): Workflow execution metadata - uploaded_file (BytesIO): Training data file (already downloaded from MinIO) - - train_params (dict): Training parameters (converted to TrainModelParams) + - train_params (TrainModelParams): Training parameters object Returns: dict: Training result with the following structure: @@ -86,15 +86,7 @@ class Training(BaseActivity): result = await train_model({ 'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456}, 'uploaded_file': BytesIO(csv_data), - 'train_params': { - 'experiment_run_id': 456, - 'target_variable': 'price', - 'variable_columns': ['feature1', 'feature2'], - 'train_size': 80, - 'shuffle': True, - 'use_scaler': True, - # ... other TrainModelParams fields - } + 'train_params': TrainModelParams(...) # Already converted object }) # Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None} @@ -103,21 +95,22 @@ class Training(BaseActivity): """ metadata = input_data.get('metadata', {}) uploaded_file = input_data['uploaded_file'] - train_params_dict = input_data['train_params'] + train_params = input_data['train_params'] try: - self.info( - f'Starting model training for target: {train_params_dict.get("target_variable")}', - metadata, - ) - - # Convert dict to TrainModelParams - train_params = TrainModelParams.from_dict(train_params_dict) - # Validate uploaded_file is BytesIO if not isinstance(uploaded_file, BytesIO): raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}') + # Validate train_params is TrainModelParams + if not isinstance(train_params, TrainModelParams): + raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}') + + self.info( + f'Starting model training for target: {train_params.target_variable}', + metadata, + ) + # Step 1: Train the model self.info('Training model with TrainingRepository', metadata) train_result = self.training_repository.train(uploaded_file, train_params) @@ -141,7 +134,12 @@ class Training(BaseActivity): } except Exception as e: # noqa: BLE001 - error_msg = f'Error training model - Target: {train_params_dict.get("target_variable", "unknown")}, Error: {str(e)}' + target = ( + train_params.target_variable + if hasattr(train_params, 'target_variable') + else 'unknown' + ) + error_msg = f'Error training model - Target: {target}, Error: {str(e)}' trace = traceback.format_exc() # Send notification (MongoDB) diff --git a/tests/activities/test_training.py b/tests/activities/test_training.py index 9616ee1..df2affc 100644 --- a/tests/activities/test_training.py +++ b/tests/activities/test_training.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch from pytest import mark from model_manager.activities.training import Training +from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.models.train_model_result import TrainModelResult @@ -42,32 +43,32 @@ async def test_train_model_success(mock_training_repository_class): # 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', - 'removed_intervals': [], - } + train_params = TrainModelParams( + 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', + removed_intervals=[], + ) input_data = { 'metadata': {'workflow_id': 'test-123'}, 'uploaded_file': uploaded_file, - 'train_params': train_params_dict, + 'train_params': train_params, } # Execute @@ -95,30 +96,32 @@ async def test_train_model_invalid_file_type(mock_training_repository_class): training = Training(logger=logger, notification_handler=notification_handler) # Invalid file type (string instead of BytesIO) + train_params = TrainModelParams( + 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', + removed_intervals=[], + ) + 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', - 'removed_intervals': [], - }, + 'train_params': train_params, } result = await training.train_model(input_data) @@ -145,30 +148,32 @@ async def test_train_model_training_error(mock_training_repository_class): training = Training(logger=logger, notification_handler=notification_handler) uploaded_file = BytesIO(b'test,data\n') + train_params = TrainModelParams( + 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', + removed_intervals=[], + ) + 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', - 'removed_intervals': [], - }, + 'train_params': train_params, } result = await training.train_model(input_data) @@ -194,30 +199,32 @@ async def test_train_model_sends_notification_on_error(mock_training_repository_ training = Training(logger=logger, notification_handler=notification_handler) uploaded_file = BytesIO(b'test,data\n1,2') + train_params = TrainModelParams( + 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', + removed_intervals=[], + ) + 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', - 'removed_intervals': [], - }, + 'train_params': train_params, } result = await training.train_model(input_data) @@ -248,30 +255,32 @@ async def test_train_model_after_calculation_error(mock_training_repository_clas training = Training(logger=logger, notification_handler=notification_handler) uploaded_file = BytesIO(b'test,data\n1,2\n3,4') + train_params = TrainModelParams( + 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', + removed_intervals=[], + ) + 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', - 'removed_intervals': [], - }, + 'train_params': train_params, } result = await training.train_model(input_data) @@ -281,3 +290,37 @@ async def test_train_model_after_calculation_error(mock_training_repository_clas assert 'Metric calculation failed' 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_invalid_train_params_type(mock_training_repository_class): + """Test training with invalid train_params type (dict instead of TrainModelParams).""" + mock_repository = MagicMock() + mock_training_repository_class.return_value = mock_repository + + logger = MagicMock() + notification_handler = MagicMock() + training = Training(logger=logger, notification_handler=notification_handler) + + uploaded_file = BytesIO(b'test,data\n1,2\n3,4') + + # Invalid train_params type (dict instead of TrainModelParams object) + input_data = { + 'metadata': {'workflow_id': 'test-invalid'}, + 'uploaded_file': uploaded_file, + 'train_params': { + 'experiment_run_id': 123, + 'target_variable': 'price', + 'variable_columns': ['feature1'], + }, # This is a dict, not TrainModelParams + } + + result = await training.train_model(input_data) + + assert result['success'] is False + assert result['result'] is None + assert 'train_params must be TrainModelParams' in result['error_message'] + assert 'dict' in result['error_message'] + # Verify notification was sent (via BaseActivity) + notification_handler.send_notification.assert_called_once()