SIENTIAPDE-1253: Enforce TrainModelParams object in Training activity and tests, removing dict conversion.

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': {
+                'train_params': TrainModelParams(...)  # Already converted object
-                    'experiment_run_id': 456,
-                    'target_variable': 'price',
-                    'variable_columns': ['feature1', 'feature2'],
-                    'train_size': 80,
-                    'shuffle': True,
-                    'use_scaler': True,
-                    # ... other TrainModelParams fields
-                }
             })
             # 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)

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 = {
+    train_params = TrainModelParams(
-        'experiment_run_id': 123,
+        experiment_run_id=123,
-        'target_variable': 'price',
+        target_variable='price',
-        'variable_columns': ['feature1', 'feature2'],
+        variable_columns=['feature1', 'feature2'],
-        'train_size': 80,
+        train_size=80,
-        'shuffle': True,
+        shuffle=True,
-        'use_scaler': True,
+        use_scaler=True,
-        'include_ar': False,
+        include_ar=False,
-        'bucket_name': 'test-bucket',
+        bucket_name='test-bucket',
-        'file_name': 'test.csv',
+        file_name='test.csv',
-        'line_separator': '\n',
+        line_separator='\n',
-        'decimal_separator': '.',
+        decimal_separator='.',
-        'lag_train': 1,
+        lag_train=1,
-        'lag_val': 1,
+        lag_val=1,
-        'rem_static_win': False,
+        rem_static_win=False,
-        'low_lim': {'feature1': 0.0, 'feature2': 0.0},
+        low_lim={'feature1': 0.0, 'feature2': 0.0},
-        'upp_lim': {'feature1': 100.0, 'feature2': 100.0},
+        upp_lim={'feature1': 100.0, 'feature2': 100.0},
-        'window': 10,
+        window=10,
-        'experiment_name': 'test_experiment',
+        experiment_name='test_experiment',
-        'removed_intervals': [],
+        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': {
+        'train_params': 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': [],
-        },
     }
 
     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': {
+        'train_params': 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': [],
-        },
     }
 
     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': {
+        'train_params': 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': [],
-        },
     }
 
     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': {
+        'train_params': 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': [],
-        },
     }
 
     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()