feat: enhance training and experiment tracking functionality
- Updated `Activities` class to improve garbage collection handling. - Enhanced error messaging in `ExperimentTracking` for better clarity on update failures. - Refactored `Training` class to streamline exception handling and improve type hints. - Introduced new methods in `TrainModelParams` for better handling of experiment run IDs and model metadata. - Added functionality to extract model equations in `DataManagerRepository` for linear regression models.
This commit is contained in:
130
tests/activities/test_activities.py
Normal file
130
tests/activities/test_activities.py
Normal file
@@ -0,0 +1,130 @@
|
||||
"""Unit tests for Activities orchestrator (constructor, shutdown, destructor)."""
|
||||
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
|
||||
def _postgres():
|
||||
return {
|
||||
'host': 'h',
|
||||
'port': 5432,
|
||||
'user': 'u',
|
||||
'password': 'p',
|
||||
'dbname': 'db',
|
||||
'min_connections': 1,
|
||||
'max_connections': 2,
|
||||
}
|
||||
|
||||
|
||||
def _mlflow():
|
||||
return {'url': 'http://mlflow:5000', 'username': 'u', 'password': 'p'}
|
||||
|
||||
|
||||
def _minio(endpoint_url: str):
|
||||
return {
|
||||
'endpoint_url': endpoint_url,
|
||||
'access_key': 'a',
|
||||
'secret_key': 's',
|
||||
'region': 'r',
|
||||
'use_ssl': True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'endpoint,expected_endpoint',
|
||||
[
|
||||
('http://minio:9000', 'minio:9000'),
|
||||
('https://minio:9000', 'minio:9000'),
|
||||
('minio:9000', 'minio:9000'),
|
||||
],
|
||||
)
|
||||
def test_activities_strips_minio_endpoint_scheme(endpoint, expected_endpoint):
|
||||
with (
|
||||
patch('model_manager.activities.activities.ExperimentTracking.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.SientiaMLflowRepository') as m_mlflow,
|
||||
patch('model_manager.activities.activities.MinioRepository') as m_minio,
|
||||
):
|
||||
Activities(
|
||||
postgres_config=_postgres(),
|
||||
mlflow_config=_mlflow(),
|
||||
minio_config=_minio(endpoint),
|
||||
plugin_store=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
m_minio.assert_called_once()
|
||||
assert m_minio.call_args.kwargs['endpoint'] == expected_endpoint
|
||||
m_mlflow.assert_called_once()
|
||||
|
||||
|
||||
def test_activities_shutdown_calls_parents():
|
||||
with (
|
||||
patch('model_manager.activities.activities.ExperimentTracking.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.SientiaMLflowRepository'),
|
||||
patch('model_manager.activities.activities.MinioRepository'),
|
||||
patch('model_manager.activities.activities.ExperimentTracking.close') as m_close,
|
||||
patch('model_manager.activities.activities.SientiaMonitoring.shutdown') as m_mon,
|
||||
):
|
||||
a = Activities(
|
||||
postgres_config=_postgres(),
|
||||
mlflow_config=_mlflow(),
|
||||
minio_config=_minio('http://x:9000'),
|
||||
plugin_store=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
with patch.object(SientiaMonitoring, 'info', Mock()):
|
||||
a.shutdown()
|
||||
m_close.assert_called_once()
|
||||
m_mon.assert_called_once()
|
||||
|
||||
|
||||
def test_activities_del_with_engine_runs_without_error():
|
||||
with (
|
||||
patch('model_manager.activities.activities.ExperimentTracking.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.SientiaMLflowRepository'),
|
||||
patch('model_manager.activities.activities.MinioRepository'),
|
||||
):
|
||||
a = Activities(
|
||||
postgres_config=_postgres(),
|
||||
mlflow_config=_mlflow(),
|
||||
minio_config=_minio('http://x:9000'),
|
||||
plugin_store=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
a.engine = MagicMock()
|
||||
Activities.__del__(a)
|
||||
|
||||
|
||||
def test_activities_del_without_engine_runs_without_error():
|
||||
with (
|
||||
patch('model_manager.activities.activities.ExperimentTracking.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Training.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.Cleanup.__init__', Mock(return_value=None)),
|
||||
patch('model_manager.activities.activities.SientiaMLflowRepository'),
|
||||
patch('model_manager.activities.activities.MinioRepository'),
|
||||
):
|
||||
a = Activities(
|
||||
postgres_config=_postgres(),
|
||||
mlflow_config=_mlflow(),
|
||||
minio_config=_minio('http://x:9000'),
|
||||
plugin_store=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
Activities.__del__(a)
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Unit tests for ExperimentTracking class with 100% coverage."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -22,12 +22,8 @@ def mock_notification_handler():
|
||||
def mock_metrics_controller():
|
||||
"""Create a mock metrics controller."""
|
||||
controller = MagicMock()
|
||||
|
||||
# Make shutdown an async coroutine
|
||||
async def mock_shutdown():
|
||||
pass
|
||||
|
||||
controller.shutdown = mock_shutdown
|
||||
controller.shutdown = AsyncMock()
|
||||
controller.emit = AsyncMock()
|
||||
return controller
|
||||
|
||||
|
||||
@@ -259,7 +255,7 @@ def test_update_experiment_run_status_missing_status(
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -270,7 +266,7 @@ def test_update_experiment_run_status_missing_status(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_update_experiment_run_status_with_error_success(
|
||||
@@ -381,7 +377,7 @@ def test_update_experiment_run_status_with_error_missing_error_message(
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -393,7 +389,7 @@ def test_update_experiment_run_status_with_error_missing_error_message(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_update_experiment_run_model_saved_success(
|
||||
@@ -461,7 +457,7 @@ def test_update_experiment_run_model_saved_missing_run_name(
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -473,7 +469,7 @@ def test_update_experiment_run_model_saved_missing_run_name(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_update_experiment_run_invalid_update_type(
|
||||
@@ -495,7 +491,7 @@ def test_update_experiment_run_invalid_update_type(
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -506,7 +502,7 @@ def test_update_experiment_run_invalid_update_type(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_update_experiment_run_no_rows_updated(
|
||||
@@ -532,7 +528,7 @@ def test_update_experiment_run_no_rows_updated(
|
||||
return {'rowcount': 0}
|
||||
|
||||
et._execute_update = mock_execute_update
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -544,7 +540,7 @@ def test_update_experiment_run_no_rows_updated(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_update_experiment_run_status_with_error_missing_status(
|
||||
@@ -566,7 +562,7 @@ def test_update_experiment_run_status_with_error_missing_status(
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -578,7 +574,7 @@ def test_update_experiment_run_status_with_error_missing_status(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_update_experiment_run_model_saved_missing_status(
|
||||
@@ -600,7 +596,7 @@ def test_update_experiment_run_model_saved_missing_status(
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
et.send_notification = MagicMock()
|
||||
et.send_notification_async = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
@@ -612,7 +608,7 @@ def test_update_experiment_run_model_saved_missing_status(
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(et.update_experiment_run(input_data))
|
||||
|
||||
et.send_notification.assert_called_once()
|
||||
et.send_notification_async.assert_awaited_once()
|
||||
|
||||
|
||||
def test_experiment_tracking_del_with_engine_no_super_del(
|
||||
|
||||
@@ -1,505 +1,315 @@
|
||||
"""Unit tests for Training class with 100% coverage."""
|
||||
"""Unit tests for Training activities."""
|
||||
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
from contextlib import asynccontextmanager
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger."""
|
||||
return MagicMock()
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_notification_handler():
|
||||
"""Create a mock notification handler."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_metrics_controller():
|
||||
"""Create a mock metrics controller."""
|
||||
controller = MagicMock()
|
||||
|
||||
# Make shutdown an async coroutine
|
||||
async def mock_shutdown():
|
||||
pass
|
||||
|
||||
# Make emit an async coroutine
|
||||
async def mock_emit(*args, **kwargs):
|
||||
pass
|
||||
|
||||
controller.shutdown = mock_shutdown
|
||||
controller.emit = mock_emit
|
||||
return controller
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_repository():
|
||||
"""Create a mock model repository."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_storage_repository():
|
||||
"""Create a mock storage repository."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_train_params():
|
||||
"""Create a mock TrainModelParams."""
|
||||
params = MagicMock()
|
||||
params.experiment_run_id = 1
|
||||
params.target_variable = 'target'
|
||||
params.model_name = 'Linear Regression'
|
||||
params.bucket_name = 'test-bucket'
|
||||
params.file_name = 'test-file.csv'
|
||||
params.validate_business_rules = MagicMock()
|
||||
return params
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_training_init(
|
||||
mock_training_repo,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test Training initialization."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
mock_training_repo.assert_called_once_with(mock_logger)
|
||||
assert training.model_repository is mock_model_repository
|
||||
assert training.storage_repository is mock_storage_repository
|
||||
assert training.metrics_controller is mock_metrics_controller
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
@patch('model_manager.activities.training.TrainModelParams')
|
||||
def test_validate_train_params_success(
|
||||
mock_train_params_class,
|
||||
mock_training_repo,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_train_params,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test validate_train_params with valid parameters."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
training.info = MagicMock()
|
||||
mock_train_params_class.from_dict.return_value = mock_train_params
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
def _minimal_params_dict():
|
||||
return {
|
||||
'variable_columns': ['a'],
|
||||
'target_variable': 't',
|
||||
'bucket_name': 'b',
|
||||
'file_name': 'f.csv',
|
||||
'line_separator': '\n',
|
||||
'decimal_separator': '.',
|
||||
'date_column': None,
|
||||
'date_format': None,
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'random_state': 42,
|
||||
'experiment_run_id': 1,
|
||||
'target_variable': 'target',
|
||||
'model_name': 'Linear Regression',
|
||||
'val_file_name': None,
|
||||
'data_model_kwargs': {},
|
||||
'model_kwargs': {},
|
||||
'opt_params': {},
|
||||
'model_type': 'linear_regression',
|
||||
'model_id': None,
|
||||
'model_metadata': None,
|
||||
}
|
||||
|
||||
result = asyncio.run(training.validate_train_params(input_data))
|
||||
|
||||
assert result is mock_train_params
|
||||
mock_train_params_class.from_dict.assert_called_once_with(input_data)
|
||||
mock_train_params.validate_business_rules.assert_called_once()
|
||||
training.info.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
@patch('model_manager.activities.training.TrainModelParams')
|
||||
def test_validate_train_params_value_error(
|
||||
mock_train_params_class,
|
||||
mock_training_repo,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test validate_train_params with ValueError."""
|
||||
@pytest.fixture
|
||||
def training():
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
return Training(
|
||||
mlflow_repository=MagicMock(),
|
||||
plugin_store=MagicMock(),
|
||||
minio_repository=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_model_metadata_success(training):
|
||||
training.plugin_store.get_model_index = MagicMock(return_value={'schemas': {'components': {'schemas': {}}}})
|
||||
inp = {**_minimal_params_dict(), 'metadata': {'w': '1'}}
|
||||
out = await training.load_model_metadata(inp)
|
||||
assert 'model_metadata' in out
|
||||
assert out['model_metadata']['schemas']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_model_metadata_notifies_on_error(training):
|
||||
training.plugin_store.get_model_index = MagicMock(side_effect=RuntimeError('idx'))
|
||||
training.send_notification_async = AsyncMock()
|
||||
inp = {**_minimal_params_dict(), 'metadata': {}}
|
||||
with pytest.raises(RuntimeError, match='idx'):
|
||||
await training.load_model_metadata(inp)
|
||||
training.send_notification_async.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_train_params_success(training):
|
||||
pdict = _minimal_params_dict()
|
||||
pdict['model_metadata'] = {'schemas': {'components': {'schemas': {}}}}
|
||||
inp = {**pdict, 'metadata': {}}
|
||||
out = await training.validate_train_params(inp)
|
||||
assert isinstance(out, TrainModelParams)
|
||||
assert out.target_variable == 't'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_train_params_notifies(training):
|
||||
training.send_notification_async = AsyncMock()
|
||||
inp = {'metadata': {}, 'experiment_run_id': 1}
|
||||
with pytest.raises(Exception):
|
||||
await training.validate_train_params(inp)
|
||||
training.send_notification_async.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_train_model_download_fails_notifies(training):
|
||||
"""train_model notifies and re-raises when MinIO download fails."""
|
||||
tp = TrainModelParams.from_dict(
|
||||
{
|
||||
**_minimal_params_dict(),
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
)
|
||||
training.minio_repository.download_file = AsyncMock(side_effect=OSError('minio'))
|
||||
training.send_notification_async = AsyncMock()
|
||||
with pytest.raises(OSError, match='minio'):
|
||||
await training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp})
|
||||
training.send_notification_async.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_resources(training):
|
||||
training.data_manager_repository.cleanup_run_directory = MagicMock()
|
||||
await training.cleanup_resources({'metadata': {}, 'run_dir': '/tmp/x'})
|
||||
training.data_manager_repository.cleanup_run_directory.assert_called_once_with('/tmp/x', {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_resources_notifies_on_error(training):
|
||||
training.data_manager_repository.cleanup_run_directory = MagicMock(side_effect=RuntimeError('rm'))
|
||||
training.send_notification = MagicMock()
|
||||
mock_train_params_class.from_dict.side_effect = ValueError('Invalid parameter')
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'experiment_run_id': 1,
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
asyncio.run(training.validate_train_params(input_data))
|
||||
|
||||
with pytest.raises(RuntimeError, match='rm'):
|
||||
await training.cleanup_resources({'metadata': {'pod': 'p'}, 'run_dir': '/tmp/x'})
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
@patch('model_manager.activities.training.TrainModelParams')
|
||||
def test_validate_train_params_type_error(
|
||||
mock_train_params_class,
|
||||
mock_training_repo,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test validate_train_params with TypeError."""
|
||||
from model_manager.activities.training import Training
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.activities.training.mlflow')
|
||||
async def test_train_model_success_serializes_result(mock_mlflow, training):
|
||||
"""Exercise train_model happy path with mocks (MinIO, plugin wrapper, MLflow)."""
|
||||
tp = TrainModelParams.from_dict(
|
||||
{
|
||||
**_minimal_params_dict(),
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
)
|
||||
train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]})
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
training.minio_repository.download_file = AsyncMock(return_value=b'csv')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: setattr(x, 'mse_val', 0.1) or x
|
||||
)
|
||||
|
||||
training.send_notification = MagicMock()
|
||||
mock_train_params_class.from_dict.side_effect = TypeError('Type mismatch')
|
||||
def _fill_report(x, **_kw):
|
||||
x.report_path = '/tmp/report.html'
|
||||
x.train_data_path = '/tmp/train.csv'
|
||||
x.test_data_path = '/tmp/test.csv'
|
||||
x.run_dir = '/tmp/run'
|
||||
return x
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'experiment_run_id': 1,
|
||||
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report)
|
||||
|
||||
wrapper = MagicMock()
|
||||
wrapper.transform = MagicMock(
|
||||
side_effect=[
|
||||
(train_df, None),
|
||||
(val_df, None),
|
||||
]
|
||||
)
|
||||
pred_train = pd.DataFrame({'p': [1.0, 2.0]})
|
||||
pred_val = pd.DataFrame({'p': [1.0]})
|
||||
wrapper.predict = MagicMock(side_effect=[(pred_train, None), (pred_val, None)])
|
||||
wrapper.store_model = MagicMock()
|
||||
training.plugin_store.get_model = AsyncMock(return_value=wrapper)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _run_ctx(*_a, **_k):
|
||||
info = MagicMock()
|
||||
info.run_name = 'run-n'
|
||||
info.run_id = 'run-i'
|
||||
yield info
|
||||
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
|
||||
out = await training.train_model({'metadata': {'pod': 'p'}, 'train_params': tp})
|
||||
assert out['run_name'] == 'run-n'
|
||||
assert out['run_id'] == 'run-i'
|
||||
assert out['run_dir'] == '/tmp/run'
|
||||
mock_mlflow.log_artifact.assert_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.activities.training.mlflow')
|
||||
async def test_train_model_train_params_as_dict(mock_mlflow, training):
|
||||
"""train_params may arrive as dict and is coerced via TrainModelParams.from_dict."""
|
||||
d = {
|
||||
**_minimal_params_dict(),
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]})
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tp = TrainModelParams.from_dict(d)
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
asyncio.run(training.validate_train_params(input_data))
|
||||
training.minio_repository.download_file = AsyncMock(return_value=b'csv')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(side_effect=lambda x, _w: x)
|
||||
def _fill_report2(x, **_kw):
|
||||
x.report_path = '/tmp/report.html'
|
||||
x.train_data_path = '/tmp/train.csv'
|
||||
x.test_data_path = '/tmp/test.csv'
|
||||
x.run_dir = '/tmp/run'
|
||||
return x
|
||||
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
@patch('model_manager.activities.training.TrainModelParams')
|
||||
def test_validate_train_params_key_error(
|
||||
mock_train_params_class,
|
||||
mock_training_repo,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test validate_train_params with KeyError."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report2)
|
||||
wrapper = MagicMock()
|
||||
wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)])
|
||||
wrapper.predict = MagicMock(
|
||||
side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)]
|
||||
)
|
||||
wrapper.store_model = MagicMock()
|
||||
training.plugin_store.get_model = AsyncMock(return_value=wrapper)
|
||||
|
||||
training.send_notification = MagicMock()
|
||||
mock_train_params_class.from_dict.side_effect = KeyError('missing_key')
|
||||
@asynccontextmanager
|
||||
async def _run_ctx(*_a, **_k):
|
||||
info = MagicMock()
|
||||
info.run_name = 'n'
|
||||
info.run_id = 'i'
|
||||
yield info
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'experiment_run_id': 1,
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
|
||||
await training.train_model({'metadata': {}, 'train_params': d})
|
||||
mock_mlflow.log_artifact.assert_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.activities.training.mlflow')
|
||||
async def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
|
||||
"""Second MinIO download when val_file_name is set (covers val_bytes branch)."""
|
||||
d = {
|
||||
**_minimal_params_dict(),
|
||||
'val_file_name': 'val.csv',
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
tp = TrainModelParams.from_dict(d)
|
||||
train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]})
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
with pytest.raises(KeyError):
|
||||
asyncio.run(training.validate_train_params(input_data))
|
||||
async def _dl(object_name, **_kwargs):
|
||||
if object_name == tp.file_name:
|
||||
return b'train'
|
||||
if object_name == 'val.csv':
|
||||
return b'val'
|
||||
raise AssertionError(object_name)
|
||||
|
||||
training.send_notification.assert_called_once()
|
||||
training.minio_repository.download_file = AsyncMock(side_effect=_dl)
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(side_effect=lambda x, _w: x)
|
||||
|
||||
def _fill(x, **_kw):
|
||||
x.report_path = '/tmp/report.html'
|
||||
x.train_data_path = '/tmp/train.csv'
|
||||
x.test_data_path = '/tmp/test.csv'
|
||||
x.run_dir = '/tmp/run'
|
||||
return x
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_train_model_success_with_params_object(
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_train_params,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test train_model with TrainModelParams object."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill)
|
||||
wrapper = MagicMock()
|
||||
wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)])
|
||||
wrapper.predict = MagicMock(
|
||||
side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)]
|
||||
)
|
||||
wrapper.store_model = MagicMock()
|
||||
training.plugin_store.get_model = AsyncMock(return_value=wrapper)
|
||||
|
||||
mock_file = BytesIO(b'test data')
|
||||
mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
|
||||
@asynccontextmanager
|
||||
async def _run_ctx(*_a, **_k):
|
||||
info = MagicMock()
|
||||
info.run_name = 'n'
|
||||
info.run_id = 'i'
|
||||
yield info
|
||||
|
||||
mock_train_result = MagicMock()
|
||||
mock_train_result.run_name = 'run_001'
|
||||
mock_train_result.run_dir = '/tmp/run_001' # noqa: S108
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
|
||||
training.training_repository.train.return_value = mock_train_result
|
||||
training.training_repository.after_train_calculation.return_value = mock_train_result
|
||||
mock_model_repository.save_model.return_value = mock_train_result
|
||||
await training.train_model({'metadata': {}, 'train_params': tp})
|
||||
assert training.minio_repository.download_file.await_count == 2
|
||||
mock_mlflow.log_artifact.assert_called()
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'train_params': mock_train_params,
|
||||
}
|
||||
|
||||
result = asyncio.run(training.train_model(input_data))
|
||||
|
||||
assert result == {'run_name': 'run_001', 'run_dir': '/tmp/run_001'} # noqa: S108
|
||||
mock_storage_repository.fetch_file.assert_called_once_with('test-bucket', 'test-file.csv')
|
||||
training.training_repository.train.assert_called_once_with(mock_file, mock_train_params)
|
||||
training.training_repository.after_train_calculation.assert_called_once_with(
|
||||
mock_train_params, mock_train_result
|
||||
@pytest.mark.asyncio
|
||||
async def test_train_model_value_error_when_paths_missing_after_report(training):
|
||||
"""Raises ValueError when report paths are not populated after generate_report."""
|
||||
tp = TrainModelParams.from_dict(
|
||||
{
|
||||
**_minimal_params_dict(),
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
)
|
||||
mock_model_repository.save_model.assert_called_once_with(mock_train_result)
|
||||
train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]})
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
@patch('model_manager.activities.training.TrainModelParams')
|
||||
def test_train_model_success_with_params_dict(
|
||||
mock_train_params_class,
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_train_params,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test train_model with dict parameters."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
training.minio_repository.download_file = AsyncMock(return_value=b'x')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(side_effect=lambda x, _w: x)
|
||||
training.data_manager_repository.generate_report = MagicMock(return_value=tmr)
|
||||
wrapper = MagicMock()
|
||||
wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)])
|
||||
wrapper.predict = MagicMock(
|
||||
side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)]
|
||||
)
|
||||
training.plugin_store.get_model = AsyncMock(return_value=wrapper)
|
||||
|
||||
mock_train_params_class.from_dict.return_value = mock_train_params
|
||||
mock_file = BytesIO(b'test data')
|
||||
mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
|
||||
@asynccontextmanager
|
||||
async def _run_ctx(*_a, **_k):
|
||||
info = MagicMock()
|
||||
info.run_name = 'n'
|
||||
info.run_id = 'i'
|
||||
yield info
|
||||
|
||||
mock_train_result = MagicMock()
|
||||
mock_train_result.run_name = 'run_002'
|
||||
mock_train_result.run_dir = '/tmp/run_002' # noqa: S108
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
training.send_notification_async = AsyncMock()
|
||||
|
||||
training.training_repository.train.return_value = mock_train_result
|
||||
training.training_repository.after_train_calculation.return_value = mock_train_result
|
||||
mock_model_repository.save_model.return_value = mock_train_result
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'train_params': {'experiment_run_id': 1, 'target_variable': 'target'},
|
||||
}
|
||||
|
||||
result = asyncio.run(training.train_model(input_data))
|
||||
|
||||
assert result == {'run_name': 'run_002', 'run_dir': '/tmp/run_002'} # noqa: S108
|
||||
mock_train_params_class.from_dict.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_train_model_training_fails(
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_train_params,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test train_model when training fails."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
training.send_notification = MagicMock()
|
||||
mock_file = BytesIO(b'test data')
|
||||
mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
|
||||
training.training_repository.train.side_effect = RuntimeError('Training failed')
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'train_params': mock_train_params,
|
||||
}
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
asyncio.run(training.train_model(input_data))
|
||||
|
||||
assert str(exc_info.value) == 'Training failed'
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_train_model_save_fails(
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_train_params,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test train_model when model saving fails."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
training.send_notification = MagicMock()
|
||||
mock_file = BytesIO(b'test data')
|
||||
mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
|
||||
|
||||
mock_train_result = MagicMock()
|
||||
training.training_repository.train.return_value = mock_train_result
|
||||
training.training_repository.after_train_calculation.return_value = mock_train_result
|
||||
mock_model_repository.save_model.side_effect = RuntimeError('Save failed')
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'train_params': mock_train_params,
|
||||
}
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
asyncio.run(training.train_model(input_data))
|
||||
|
||||
assert str(exc_info.value) == 'Save failed'
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_cleanup_resources_success(
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test cleanup_resources successfully."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'run_dir': '/tmp/run_001', # noqa: S108
|
||||
}
|
||||
|
||||
asyncio.run(training.cleanup_resources(input_data))
|
||||
|
||||
mock_model_repository.cleanup_run_directory.assert_called_once_with('/tmp/run_001') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_cleanup_resources_cleanup_fails(
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test cleanup_resources when cleanup fails."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
training.send_notification = MagicMock()
|
||||
mock_model_repository.cleanup_run_directory.side_effect = RuntimeError('Cleanup failed')
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'run_dir': '/tmp/run_001', # noqa: S108
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.csv',
|
||||
}
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(training.cleanup_resources(input_data))
|
||||
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
def test_cleanup_resources_with_empty_values(
|
||||
mock_training_repo_class,
|
||||
mock_model_repository,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test cleanup_resources with empty values."""
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
}
|
||||
|
||||
asyncio.run(training.cleanup_resources(input_data))
|
||||
|
||||
mock_model_repository.cleanup_run_directory.assert_called_once_with('')
|
||||
with pytest.raises(ValueError, match='Report path'):
|
||||
await training.train_model({'metadata': {}, 'train_params': tp})
|
||||
|
||||
Reference in New Issue
Block a user