8ea98360c3
This commit refactors the training workflow and associated activities to raise exceptions on failure instead of returning success/failure dictionaries. This allows the Temporal workflow to handle errors more effectively and ensures that the workflow stops when a critical error occurs. Key changes: - The train_model workflow is introduced to orchestrate the entire training process, including parameter validation, data download, model training, and model saving. - The validate_train_params activity is added to validate and convert training parameters. - The train_model and save_model activities are updated to raise exceptions on failure. - The ExperimentStatus enum is updated to include a new status for orchestrator validation errors. - The tests are updated to reflect the new exception-based error handling. - The activities now return the TrainModelResult directly instead of a dictionary.
500 lines
16 KiB
Python
500 lines
16 KiB
Python
"""Unit tests for Training activity."""
|
|
|
|
from io import BytesIO
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
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
|
|
|
|
|
|
@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 = 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,
|
|
}
|
|
|
|
# Execute
|
|
result = await training.train_model(input_data)
|
|
|
|
# Assertions - now returns TrainModelResult directly
|
|
assert result == mock_final_result
|
|
assert result.mse_val == 0.5
|
|
assert result.mae_val == 0.3
|
|
assert result.r2_val == 0.95
|
|
|
|
# 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)
|
|
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,
|
|
}
|
|
|
|
# Should raise ValueError
|
|
with pytest.raises(ValueError, match='uploaded_file must be BytesIO'):
|
|
await training.train_model(input_data)
|
|
|
|
# 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')
|
|
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,
|
|
}
|
|
|
|
# Should raise ValueError
|
|
with pytest.raises(ValueError, match='Training data is empty'):
|
|
await training.train_model(input_data)
|
|
|
|
# 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')
|
|
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,
|
|
}
|
|
|
|
# Should raise Exception
|
|
with pytest.raises(Exception, match='Database connection failed'):
|
|
await training.train_model(input_data)
|
|
|
|
# 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_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')
|
|
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,
|
|
}
|
|
|
|
# Should raise Exception
|
|
with pytest.raises(Exception, match='Metric calculation failed'):
|
|
await training.train_model(input_data)
|
|
|
|
# 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
|
|
}
|
|
|
|
# Should raise ValueError
|
|
with pytest.raises(ValueError, match='train_params must be TrainModelParams.*dict'):
|
|
await training.train_model(input_data)
|
|
|
|
|
|
# ============================================================================
|
|
# Tests for validate_train_params
|
|
# ============================================================================
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_validate_train_params_success():
|
|
"""Test successful validation of training parameters."""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
training.info = MagicMock()
|
|
|
|
input_data = {
|
|
'metadata': {'workflow_id': 'test-123'},
|
|
'experiment_run_id': 456,
|
|
'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': [],
|
|
}
|
|
|
|
result = await training.validate_train_params(input_data)
|
|
|
|
assert isinstance(result, TrainModelParams)
|
|
assert result.experiment_run_id == 456
|
|
assert result.target_variable == 'price'
|
|
assert result.variable_columns == ['feature1', 'feature2']
|
|
assert result.train_size == 80
|
|
assert result.experiment_name == 'test_experiment'
|
|
assert training.info.call_count == 2
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_validate_train_params_missing_required_field():
|
|
"""Test validation fails when required field is missing."""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
training.info = MagicMock()
|
|
training.error = MagicMock()
|
|
|
|
input_data = {
|
|
'metadata': {'workflow_id': 'test-123'},
|
|
'experiment_run_id': 456,
|
|
'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',
|
|
'removed_intervals': [],
|
|
}
|
|
|
|
with pytest.raises(ValueError, match='target_variable'):
|
|
await training.validate_train_params(input_data)
|
|
|
|
notification_handler.send_notification.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_validate_train_params_invalid_type():
|
|
"""Test validation fails when field has invalid type."""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
training.info = MagicMock()
|
|
training.error = MagicMock()
|
|
|
|
input_data = {
|
|
'metadata': {'workflow_id': 'test-invalid'},
|
|
'experiment_run_id': 456,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1'],
|
|
'train_size': 'invalid',
|
|
'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',
|
|
'removed_intervals': [],
|
|
}
|
|
|
|
with pytest.raises((ValueError, TypeError)):
|
|
await training.validate_train_params(input_data)
|
|
|
|
notification_handler.send_notification.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_validate_train_params_empty_input():
|
|
"""Test validation fails with empty input."""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
training.info = MagicMock()
|
|
training.error = MagicMock()
|
|
|
|
input_data = {'metadata': {}}
|
|
|
|
with pytest.raises(ValueError):
|
|
await training.validate_train_params(input_data)
|
|
|
|
notification_handler.send_notification.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_validate_train_params_without_metadata():
|
|
"""Test validation works even without metadata key."""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
training = Training(logger=logger, notification_handler=notification_handler)
|
|
|
|
training.info = MagicMock()
|
|
|
|
input_data = {
|
|
'experiment_run_id': 789,
|
|
'target_variable': 'temperature',
|
|
'variable_columns': ['sensor1'],
|
|
'train_size': 75,
|
|
'shuffle': False,
|
|
'use_scaler': True,
|
|
'include_ar': True,
|
|
'bucket_name': 'sensors',
|
|
'file_name': 'data.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'lag_train': 2,
|
|
'lag_val': 2,
|
|
'rem_static_win': True,
|
|
'low_lim': {'sensor1': -50.0},
|
|
'upp_lim': {'sensor1': 150.0},
|
|
'window': 20,
|
|
'experiment_name': 'sensor_experiment',
|
|
'removed_intervals': [],
|
|
}
|
|
|
|
result = await training.validate_train_params(input_data)
|
|
|
|
assert isinstance(result, TrainModelParams)
|
|
assert result.experiment_run_id == 789
|
|
assert result.target_variable == 'temperature'
|
|
assert result.experiment_name == 'sensor_experiment'
|