gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1253: Refactor training workflow and activities to raise exceptions on failure

Browse Source 
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.
This commit is contained in:
Bruno Domingues
2025-10-15 15:19:29 -03:00
parent 61267ec49d
commit 8ea98360c3
8 changed files with 1472 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
tests/activities/test_mlflow.py
View File

@@ -1,6 +1,7 @@
from unittest.mock import ANY, MagicMock, patch
import numpy as np
import pytest
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
@@ -333,10 +334,9 @@ async def test_save_model_success(mlflow):
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
# Verify response
assert response['success'] is True
assert response['result'] == train_result
assert response['error_message'] is None
# Verify response - now returns TrainModelResult directly
assert response == train_result
assert response.run_name == 'test_experiment-1'
assert train_result.run_name == 'test_experiment-1'
@@ -362,13 +362,9 @@ async def test_save_model_get_next_run_name_error(mlflow):
'train_result': train_result,
}
# Call the method
response = await mlflow.save_model(input_data)
# Verify error handling
assert response['success'] is False
assert response['result'] is None
assert 'MLflow connection error' in response['error_message']
# Call the method - should raise exception
with pytest.raises(Exception, match='MLflow connection error'):
await mlflow.save_model(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
@@ -404,13 +400,9 @@ async def test_save_model_generate_artifacts_error(mlflow):
'train_result': train_result,
}
# Call the method
response = await mlflow.save_model(input_data)
# Verify error handling
assert response['success'] is False
assert response['result'] is None
assert 'Reports directory does not exist' in response['error_message']
# Call the method - should raise exception
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
await mlflow.save_model(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
@@ -447,13 +439,9 @@ async def test_save_model_save_run_error(mlflow):
'train_result': train_result,
}
# Call the method
response = await mlflow.save_model(input_data)
# Verify error handling
assert response['success'] is False
assert response['result'] is None
assert 'One or more metrics (MSE, R2, MAE) are None' in response['error_message']
# Call the method - should raise exception
with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
await mlflow.save_model(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
@@ -491,10 +479,9 @@ async def test_save_model_missing_metadata(mlflow):
# Call the method
response = await mlflow.save_model(input_data)
# Verify it still works (metadata defaults to {})
assert response['success'] is True
assert response['result'] == train_result
assert response['error_message'] is None
# Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
assert response == train_result
assert response.run_name == 'test_experiment-1'
@mark.asyncio
@@ -540,10 +527,8 @@ async def test_save_model_complete_flow(mlflow):
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
# Verify response
assert response['success'] is True
assert response['result'] == updated_result
assert response['error_message'] is None
assert updated_result.run_name == 'production_model-5'
assert updated_result.run_dir is not None
assert updated_result.report_path is not None
# Verify response - returns TrainModelResult directly
assert response == updated_result
assert response.run_name == 'production_model-5'
assert response.run_dir is not None
assert response.report_path is not None

229
tests/activities/test_training.py
View File

@@ -3,6 +3,7 @@
from io import BytesIO
from unittest.mock import MagicMock, patch
import pytest
from pytest import mark
from model_manager.activities.training import Training
@@ -74,10 +75,11 @@ async def test_train_model_success(mock_training_repository_class):
# 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
# 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()
@@ -124,11 +126,10 @@ async def test_train_model_invalid_file_type(mock_training_repository_class):
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise ValueError
with pytest.raises(ValueError, match='uploaded_file must be BytesIO'):
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()
@@ -176,11 +177,10 @@ async def test_train_model_training_error(mock_training_repository_class):
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise ValueError
with pytest.raises(ValueError, match='Training data is empty'):
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()
@@ -227,16 +227,13 @@ async def test_train_model_sends_notification_on_error(mock_training_repository_
'train_params': train_params,
}
result = await training.train_model(input_data)
# 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()
# 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')
@@ -283,11 +280,10 @@ async def test_train_model_after_calculation_error(mock_training_repository_clas
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise Exception
with pytest.raises(Exception, match='Metric calculation failed'):
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()
@@ -316,11 +312,188 @@ async def test_train_model_invalid_train_params_type(mock_training_repository_cl
}, # This is a dict, not TrainModelParams
}
result = await training.train_model(input_data)
# 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)
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()
@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'

6
tests/utils/models/test_experiment_status.py
View File

@@ -15,8 +15,8 @@ def test_experiment_status_values():
def test_experiment_status_count():
"""Test that enum has exactly 7 status values."""
assert len(ExperimentStatus) == 7
"""Test that enum has exactly 8 status values."""
assert len(ExperimentStatus) == 8
def test_experiment_status_is_string():
@@ -40,7 +40,7 @@ def test_experiment_status_membership():
def test_experiment_status_iteration():
"""Test that enum can be iterated."""
statuses = list(ExperimentStatus)
assert len(statuses) == 7
assert len(statuses) == 8
assert ExperimentStatus.MAGE_WAITING_PROC in statuses
assert ExperimentStatus.TRAINING_SUCCESS in statuses
assert ExperimentStatus.TRAINING_ERROR in statuses

672
tests/workflows/test_train_model.py Normal file
View File

@@ -0,0 +1,672 @@
"""Unit tests for TrainModel workflow."""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pytest import fixture, mark
from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.workflows.train_model import TrainModel
@fixture
def train_model_workflow() -> TrainModel:
"""Fixture for TrainModel workflow instance."""
return TrainModel()
@fixture
def mock_train_params():
"""Fixture for mock TrainModelParams."""
return 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=[],
)
@fixture
def mock_train_result(mock_train_params):
"""Fixture for mock TrainModelResult."""
result = MagicMock(spec=TrainModelResult)
result.params = mock_train_params
result.run_name = 'test_experiment-1'
result.run_dir = 'test_run_dir' # Relative path instead of /tmp
return result
# ============================================================================
# Tests for run() - Complete workflow
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.shutil')
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_run_success_complete_flow(
workflow_mock: AsyncMock,
mock_shutil: MagicMock,
train_model_workflow: TrainModel,
mock_train_params,
mock_train_result,
):
"""Test successful complete workflow execution."""
input_data = {
'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': [],
}
# Mock activity responses
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update_experiment_run (MAGE_WAITING_PROC)
b'file_content', # fetch_file_from_minio
mock_train_result, # train_model
None, # update_experiment_run (TRAINING_SUCCESS)
mock_train_result, # save_model
None, # update_experiment_run (MLFLOW_SENT with run_name)
None, # delete_file_from_minio
None, # update_experiment_run (FILE_DELETED)
]
)
# Execute workflow
await train_model_workflow.run(input_data)
# Verify all activity calls
assert workflow_mock.execute_activity_method.call_count == 9
# Verify shutil.rmtree was called
mock_shutil.rmtree.assert_called_once_with('test_run_dir')
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_run_missing_experiment_run_id(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test workflow fails when experiment_run_id is missing."""
input_data = {
'target_variable': 'price',
'variable_columns': ['feature1'],
}
# Mock workflow.logger to avoid NotInWorkflowEventLoopError
workflow_mock.logger = MagicMock()
with pytest.raises(ValueError, match='experiment_run_id is required'):
await train_model_workflow.run(input_data)
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_run_invalid_experiment_run_id_type(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test workflow fails when experiment_run_id has invalid type."""
input_data = {
'experiment_run_id': 'invalid', # Should be int
'target_variable': 'price',
}
# Mock workflow.logger to avoid NotInWorkflowEventLoopError
workflow_mock.logger = MagicMock()
with pytest.raises(ValueError, match='experiment_run_id must be an integer'):
await train_model_workflow.run(input_data)
# ============================================================================
# Tests for _validate_experiment_run_id()
# ============================================================================
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_success(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test successful experiment_run_id validation."""
workflow_mock.logger = MagicMock()
input_data = {'experiment_run_id': 456}
result = train_model_workflow._validate_experiment_run_id(input_data)
assert result == 456
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_missing(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test validation fails when experiment_run_id is missing."""
workflow_mock.logger = MagicMock()
input_data = {}
with pytest.raises(ValueError, match='experiment_run_id is required'):
train_model_workflow._validate_experiment_run_id(input_data)
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_none(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test validation fails when experiment_run_id is None."""
workflow_mock.logger = MagicMock()
input_data = {'experiment_run_id': None}
with pytest.raises(ValueError, match='experiment_run_id is required'):
train_model_workflow._validate_experiment_run_id(input_data)
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_invalid_type(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test validation fails when experiment_run_id is not an integer."""
workflow_mock.logger = MagicMock()
input_data = {'experiment_run_id': 'not_an_int'}
with pytest.raises(ValueError, match='experiment_run_id must be an integer'):
train_model_workflow._validate_experiment_run_id(input_data)
# ============================================================================
# Tests for _validate_training_parameters()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_validate_training_parameters_success(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test successful parameter validation."""
input_data = {
'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',
'removed_intervals': [],
}
metadata = {
'metadata': {
'experiment_run_id': 123,
'workflow_name': 'train_model',
}
}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update_experiment_run
]
)
result = await train_model_workflow._validate_training_parameters(input_data, 123, metadata)
assert result == mock_train_params
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_validate_training_parameters_validation_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test parameter validation handles errors correctly."""
input_data = {'experiment_run_id': 123}
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
ValueError('Missing required field'), # validate_train_params fails
None, # update_experiment_run with error
]
)
with pytest.raises(ValueError, match='Missing required field'):
await train_model_workflow._validate_training_parameters(input_data, 123, metadata)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
# ============================================================================
# Tests for _download_and_train_model()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_success(
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_params,
mock_train_result,
):
"""Test successful download and training."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
b'file_content', # fetch_file_from_minio
mock_train_result, # train_model
None, # update_experiment_run
]
)
result = await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
assert result == mock_train_result
assert workflow_mock.execute_activity_method.call_count == 3
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_download_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test download error is handled correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
Exception('MinIO connection failed'), # fetch_file_from_minio fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='MinIO connection failed'):
await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_training_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test training error is handled correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
b'file_content', # fetch_file_from_minio succeeds
Exception('Training failed'), # train_model fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='Training failed'):
await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 3
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_closes_bytesio_on_success(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result
):
"""Test that BytesIO is closed in finally block on success."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Create a mock BytesIO with close method
mock_file = MagicMock()
mock_file.close = MagicMock()
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_file, # fetch_file_from_minio returns BytesIO
mock_train_result, # train_model succeeds
None, # update_experiment_run
]
)
result = await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify BytesIO.close() was called in finally block
mock_file.close.assert_called_once()
assert result == mock_train_result
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_closes_bytesio_on_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test that BytesIO is closed in finally block even on error."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Create a mock BytesIO with close method
mock_file = MagicMock()
mock_file.close = MagicMock()
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_file, # fetch_file_from_minio returns BytesIO
Exception('Training failed'), # train_model fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='Training failed'):
await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify BytesIO.close() was called in finally block even after exception
mock_file.close.assert_called_once()
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_handles_file_without_close(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result
):
"""Test that workflow handles file objects without close method gracefully."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Create a mock file without close method
mock_file = MagicMock(spec=[])
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_file, # fetch_file_from_minio returns object without close
mock_train_result, # train_model succeeds
None, # update_experiment_run
]
)
# Should not raise error even if file doesn't have close method
result = await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
assert result == mock_train_result
# ============================================================================
# Tests for _save_model_to_mlflow()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_save_model_to_mlflow_success(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result
):
"""Test successful model saving to MLFlow."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_train_result, # save_model
None, # update_experiment_run with MODEL_SAVED
]
)
result = await train_model_workflow._save_model_to_mlflow(
train_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
assert result == mock_train_result
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_save_model_to_mlflow_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result
):
"""Test MLFlow save error is handled correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
Exception('MLFlow connection failed'), # save_model fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='MLFlow connection failed'):
await train_model_workflow._save_model_to_mlflow(
train_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
# ============================================================================
# Tests for _cleanup_resources()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
@patch('model_manager.workflows.train_model.shutil')
async def test_cleanup_resources_success(
mock_shutil: MagicMock,
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_result,
):
"""Test successful resource cleanup."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
None, # delete_file_from_minio
None, # update_experiment_run
]
)
await train_model_workflow._cleanup_resources(
saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify shutil.rmtree was called
mock_shutil.rmtree.assert_called_once_with('test_run_dir')
# Verify activities were called
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
@patch('model_manager.workflows.train_model.shutil')
async def test_cleanup_resources_delete_error(
mock_shutil: MagicMock,
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_result,
):
"""Test cleanup handles delete errors correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
Exception('MinIO delete failed'), # delete_file_from_minio fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='MinIO delete failed'):
await train_model_workflow._cleanup_resources(
saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
@patch('model_manager.workflows.train_model.shutil')
async def test_cleanup_resources_without_run_dir(
mock_shutil: MagicMock,
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_result,
):
"""Test cleanup works when run_dir is not set."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Mock result without run_dir
mock_train_result.run_dir = None
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
None, # delete_file_from_minio
None, # update_experiment_run
]
)
await train_model_workflow._cleanup_resources(
saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify shutil.rmtree was NOT called
mock_shutil.rmtree.assert_not_called()
# Verify activities were still called
assert workflow_mock.execute_activity_method.call_count == 2
# ============================================================================
# Tests for _update_experiment_run()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_update_experiment_run_status_only(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test updating experiment run with status only."""
from model_manager.activities.experiment_tracking import UpdateType
metadata = {'metadata': {'experiment_run_id': 123}}
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
await train_model_workflow._update_experiment_run(
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.STATUS,
status=ExperimentStatus.TRAINING_SUCCESS,
)
workflow_mock.execute_activity_method.assert_called_once()
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_update_experiment_run_with_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test updating experiment run with error message."""
from model_manager.activities.experiment_tracking import UpdateType
metadata = {'metadata': {'experiment_run_id': 123}}
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
await train_model_workflow._update_experiment_run(
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.TRAINING_ERROR,
error_message='Training failed',
)
workflow_mock.execute_activity_method.assert_called_once()
call_args = workflow_mock.execute_activity_method.call_args[0][1]
assert call_args['error_message'] == 'Training failed'
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_update_experiment_run_with_run_name(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test updating experiment run with run_name."""
from model_manager.activities.experiment_tracking import UpdateType
metadata = {'metadata': {'experiment_run_id': 123}}
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
await train_model_workflow._update_experiment_run(
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.MLFLOW_SENT,
run_name='test_experiment-1',
)
workflow_mock.execute_activity_method.assert_called_once()
call_args = workflow_mock.execute_activity_method.call_args[0][1]
assert call_args['run_name'] == 'test_experiment-1'