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

feat: enhance configuration and error handling in project setup

Browse Source 
- Added new ignore rule for Ruff to allow temporary paths in tests.
- Introduced MyPy overrides for specific modules to ignore errors.
- Refactored `Cleanup` and `ExperimentTracking` classes to remove async keywords from methods, improving consistency in method signatures.
- Updated `Training` class methods to handle synchronous operations, enhancing performance and clarity.
- Adjusted `requirements.txt` to remove unnecessary Git dependency, streamlining project setup.
This commit is contained in:
vitor-aignosi
2026-04-07 10:25:17 -03:00
parent 6b1df7c3a7
commit 09ee92f100
21 changed files with 500 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/activities/test_activities.py
View File

@@ -44,7 +44,10 @@ def _minio(endpoint_url: str):
)
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.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,
@@ -66,7 +69,10 @@ def test_activities_strips_minio_endpoint_scheme(endpoint, expected_endpoint):
def test_activities_shutdown_calls_parents():
with (
patch('model_manager.activities.activities.ExperimentTracking.__init__', Mock(return_value=None)),
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'),
@@ -91,7 +97,10 @@ def test_activities_shutdown_calls_parents():
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.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'),
@@ -112,7 +121,10 @@ def test_activities_del_with_engine_runs_without_error():
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.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'),

63
tests/activities/test_cleanup.py
View File

@@ -1,6 +1,5 @@
"""Unit tests for the Cleanup activity, ensuring 100% code coverage."""
import asyncio
import os
import shutil
import tempfile
@@ -126,12 +125,10 @@ def test_cleanup_temp_directories_nonexistent_path(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
cleanup.warning = MagicMock()
asyncio.run(
cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}})
)
cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}})
cleanup.warning.assert_called_once()
cleanup._emit_metrics.assert_called_once()
@@ -155,7 +152,7 @@ def test_cleanup_temp_directories_success_with_deletions(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
@@ -165,7 +162,7 @@ def test_cleanup_temp_directories_success_with_deletions(
recent_dir = os.path.join(temp_dir, f'recent_dir_{recent_time}')
os.makedirs(recent_dir)
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
assert not os.path.exists(old_dir)
assert os.path.exists(recent_dir)
@@ -190,13 +187,13 @@ def test_cleanup_temp_directories_dry_run(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
os.makedirs(old_dir)
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
assert os.path.exists(old_dir)
cleanup._emit_metrics.assert_called_once()
@@ -220,7 +217,7 @@ def test_cleanup_temp_directories_delete_error(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
cleanup.error = MagicMock()
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
@@ -228,7 +225,7 @@ def test_cleanup_temp_directories_delete_error(
os.makedirs(old_dir)
with patch('shutil.rmtree', side_effect=OSError('Permission Denied')):
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
cleanup.error.assert_called_once()
cleanup._emit_metrics.assert_called_once()
@@ -250,18 +247,16 @@ def test_emit_metrics(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup.emit_metric = AsyncMock()
cleanup.emit_metric_sync = MagicMock()
asyncio.run(
cleanup._emit_metrics(
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
metrics_status='success',
activity_name='test_activity',
emit_workflow_metric=True,
)
cleanup._emit_metrics(
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
metrics_status='success',
activity_name='test_activity',
emit_workflow_metric=True,
)
assert cleanup.emit_metric.call_count == 2
assert cleanup.emit_metric_sync.call_count == 2
def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
@@ -281,7 +276,7 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
cleanup.debug = MagicMock()
# Create a file and a directory with a non-matching name
@@ -289,7 +284,7 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
f.write('hello')
os.makedirs(os.path.join(temp_dir, 'a_directory_with_no_timestamp'))
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
# Ensure the debug message for skipping was called for the unmatched directory
cleanup.debug.assert_called_with(
@@ -315,14 +310,14 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
cleanup.error = MagicMock()
# Create a directory with a malformed timestamp that matches the regex but fails parsing
malformed_dir_name = 'dir_20239999_999999_999999'
os.makedirs(os.path.join(temp_dir, malformed_dir_name))
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
cleanup.error.assert_called_once()
cleanup._emit_metrics.assert_called_once()
@@ -345,12 +340,12 @@ def test_cleanup_temp_directories_generic_exception(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup._emit_metrics = MagicMock()
cleanup.send_notification = MagicMock()
with patch('os.listdir', side_effect=Exception('Unexpected OS Error')):
with pytest.raises(Exception, match='Unexpected OS Error'):
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
cleanup.send_notification.assert_called_once()
cleanup._emit_metrics.assert_called_once()
@@ -369,15 +364,13 @@ def test_emit_metrics_activity_only(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup.emit_metric = AsyncMock()
cleanup.emit_metric_sync = MagicMock()
asyncio.run(
cleanup._emit_metrics(
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
metrics_status='success',
activity_name='test_activity',
emit_workflow_metric=False,
)
cleanup._emit_metrics(
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
metrics_status='success',
activity_name='test_activity',
emit_workflow_metric=False,
)
cleanup.emit_metric.assert_called_once()
cleanup.emit_metric_sync.assert_called_once()

63
tests/activities/test_experiment_tracking.py
View File

@@ -1,6 +1,5 @@
"""Unit tests for ExperimentTracking class with 100% coverage."""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -185,7 +184,7 @@ def test_execute_update_success(
mock_engine.begin.return_value.__enter__.return_value = mock_connection
et.engine = mock_engine
result = asyncio.run(et._execute_update('UPDATE test SET x = :x', {'x': 1}))
result = et._execute_update('UPDATE test SET x = :x', {'x': 1})
assert result == {'rowcount': 1}
mock_connection.execute.assert_called_once()
@@ -212,7 +211,7 @@ def test_update_experiment_run_status_success(
mock_execute = MagicMock()
async def mock_execute_update(*args, **kwargs):
def mock_execute_update(*args, **kwargs):
mock_execute(*args, **kwargs)
return {'rowcount': 1}
@@ -226,7 +225,7 @@ def test_update_experiment_run_status_success(
'status': 'running',
}
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
mock_execute.assert_called_once()
call_args = mock_execute.call_args
@@ -255,7 +254,7 @@ def test_update_experiment_run_status_missing_status(
metrics_controller=mock_metrics_controller,
)
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -264,9 +263,9 @@ def test_update_experiment_run_status_missing_status(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_update_experiment_run_status_with_error_success(
@@ -290,7 +289,7 @@ def test_update_experiment_run_status_with_error_success(
mock_execute = MagicMock()
async def mock_execute_update(*args, **kwargs):
def mock_execute_update(*args, **kwargs):
mock_execute(*args, **kwargs)
return {'rowcount': 1}
@@ -305,7 +304,7 @@ def test_update_experiment_run_status_with_error_success(
'error_message': 'Test error',
}
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
mock_execute.assert_called_once()
call_args = mock_execute.call_args
@@ -336,7 +335,7 @@ def test_update_experiment_run_status_with_error_truncate_message(
mock_execute = MagicMock()
async def mock_execute_update(*args, **kwargs):
def mock_execute_update(*args, **kwargs):
mock_execute(*args, **kwargs)
return {'rowcount': 1}
@@ -352,7 +351,7 @@ def test_update_experiment_run_status_with_error_truncate_message(
'error_message': long_error,
}
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
call_args = mock_execute.call_args
assert len(call_args[0][1]['error_message']) == 1024
@@ -377,7 +376,7 @@ def test_update_experiment_run_status_with_error_missing_error_message(
metrics_controller=mock_metrics_controller,
)
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -387,9 +386,9 @@ def test_update_experiment_run_status_with_error_missing_error_message(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_update_experiment_run_model_saved_success(
@@ -413,7 +412,7 @@ def test_update_experiment_run_model_saved_success(
mock_execute = MagicMock()
async def mock_execute_update(*args, **kwargs):
def mock_execute_update(*args, **kwargs):
mock_execute(*args, **kwargs)
return {'rowcount': 1}
@@ -428,7 +427,7 @@ def test_update_experiment_run_model_saved_success(
'run_name': 'run_001',
}
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
mock_execute.assert_called_once()
call_args = mock_execute.call_args
@@ -457,7 +456,7 @@ def test_update_experiment_run_model_saved_missing_run_name(
metrics_controller=mock_metrics_controller,
)
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -467,9 +466,9 @@ def test_update_experiment_run_model_saved_missing_run_name(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_update_experiment_run_invalid_update_type(
@@ -491,7 +490,7 @@ def test_update_experiment_run_invalid_update_type(
metrics_controller=mock_metrics_controller,
)
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -500,9 +499,9 @@ def test_update_experiment_run_invalid_update_type(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_update_experiment_run_no_rows_updated(
@@ -524,11 +523,11 @@ def test_update_experiment_run_no_rows_updated(
metrics_controller=mock_metrics_controller,
)
async def mock_execute_update(*args, **kwargs):
def mock_execute_update(*args, **kwargs):
return {'rowcount': 0}
et._execute_update = mock_execute_update
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -538,9 +537,9 @@ def test_update_experiment_run_no_rows_updated(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_update_experiment_run_status_with_error_missing_status(
@@ -562,7 +561,7 @@ def test_update_experiment_run_status_with_error_missing_status(
metrics_controller=mock_metrics_controller,
)
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -572,9 +571,9 @@ def test_update_experiment_run_status_with_error_missing_status(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_update_experiment_run_model_saved_missing_status(
@@ -596,7 +595,7 @@ def test_update_experiment_run_model_saved_missing_status(
metrics_controller=mock_metrics_controller,
)
et.send_notification_async = AsyncMock()
et.send_notification = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
@@ -606,9 +605,9 @@ def test_update_experiment_run_model_saved_missing_status(
}
with pytest.raises(RuntimeError):
asyncio.run(et.update_experiment_run(input_data))
et.update_experiment_run(input_data)
et.send_notification_async.assert_awaited_once()
et.send_notification.assert_called_once()
def test_experiment_tracking_del_with_engine_no_super_del(

154
tests/activities/test_training.py
View File

@@ -1,7 +1,7 @@
"""Unit tests for Training activities."""
from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, MagicMock, patch
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
@@ -49,46 +49,43 @@ def training():
)
@pytest.mark.asyncio
async def test_load_model_metadata_success(training):
training.plugin_store.get_model_index = MagicMock(return_value={'schemas': {'components': {'schemas': {}}}})
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)
out = 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):
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()
training.send_notification = MagicMock()
inp = {**_minimal_params_dict(), 'metadata': {}}
with pytest.raises(RuntimeError, match='idx'):
await training.load_model_metadata(inp)
training.send_notification_async.assert_awaited()
training.load_model_metadata(inp)
training.send_notification.assert_called_once()
@pytest.mark.asyncio
async def test_validate_train_params_success(training):
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'
out = training.validate_train_params(inp)
assert isinstance(out, dict)
assert out['target_variable'] == 't'
@pytest.mark.asyncio
async def test_validate_train_params_notifies(training):
training.send_notification_async = AsyncMock()
def test_validate_train_params_notifies(training):
training.send_notification = MagicMock()
inp = {'metadata': {}, 'experiment_run_id': 1}
with pytest.raises(Exception):
await training.validate_train_params(inp)
training.send_notification_async.assert_awaited()
with pytest.raises((KeyError, ValueError, TypeError)):
training.validate_train_params(inp)
training.send_notification.assert_called_once()
@pytest.mark.asyncio
async def test_train_model_download_fails_notifies(training):
def test_train_model_download_fails_notifies(training):
"""train_model notifies and re-raises when MinIO download fails."""
tp = TrainModelParams.from_dict(
{
@@ -96,32 +93,31 @@ async def test_train_model_download_fails_notifies(training):
'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.minio_repository.download_file_sync = MagicMock(side_effect=OSError('minio'))
training.send_notification = MagicMock()
with pytest.raises(RuntimeError, match='rm'):
await training.cleanup_resources({'metadata': {'pod': 'p'}, 'run_dir': '/tmp/x'})
with pytest.raises(OSError, match='minio'):
training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp.to_dict()})
training.send_notification.assert_called_once()
def test_cleanup_resources(training):
training.data_manager_repository.cleanup_run_directory = MagicMock()
training.cleanup_resources({'metadata': {}, 'run_dir': '/tmp/x'})
training.data_manager_repository.cleanup_run_directory.assert_called_once_with('/tmp/x', {})
def test_cleanup_resources_notifies_on_error(training):
training.data_manager_repository.cleanup_run_directory = MagicMock(
side_effect=RuntimeError('rm')
)
training.send_notification = MagicMock()
with pytest.raises(RuntimeError, match='rm'):
training.cleanup_resources({'metadata': {'pod': 'p'}, 'run_dir': '/tmp/x'})
training.send_notification.assert_called_once()
@pytest.mark.asyncio
@patch('model_manager.activities.training.mlflow')
async def test_train_model_success_serializes_result(mock_mlflow, training):
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(
{
@@ -133,7 +129,7 @@ async def test_train_model_success_serializes_result(mock_mlflow, training):
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
training.minio_repository.download_file = AsyncMock(return_value=b'csv')
training.minio_repository.download_file_sync = MagicMock(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
@@ -159,10 +155,10 @@ async def test_train_model_success_serializes_result(mock_mlflow, training):
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)
training.plugin_store.get_model = MagicMock(return_value=wrapper)
@asynccontextmanager
async def _run_ctx(*_a, **_k):
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_name = 'run-n'
info.run_id = 'run-i'
@@ -170,16 +166,15 @@ async def test_train_model_success_serializes_result(mock_mlflow, training):
training.mlflow_repository.start_run = _run_ctx
out = await training.train_model({'metadata': {'pod': 'p'}, 'train_params': tp})
out = training.train_model({'metadata': {'pod': 'p'}, 'train_params': tp.to_dict()})
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):
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(),
@@ -190,9 +185,12 @@ async def test_train_model_train_params_as_dict(mock_mlflow, training):
tp = TrainModelParams.from_dict(d)
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
training.minio_repository.download_file = AsyncMock(return_value=b'csv')
training.minio_repository.download_file_sync = MagicMock(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)
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'
@@ -207,10 +205,10 @@ async def test_train_model_train_params_as_dict(mock_mlflow, training):
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.plugin_store.get_model = MagicMock(return_value=wrapper)
@asynccontextmanager
async def _run_ctx(*_a, **_k):
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_name = 'n'
info.run_id = 'i'
@@ -218,13 +216,12 @@ async def test_train_model_train_params_as_dict(mock_mlflow, training):
training.mlflow_repository.start_run = _run_ctx
await training.train_model({'metadata': {}, 'train_params': d})
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):
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(),
@@ -236,16 +233,18 @@ async def test_train_model_downloads_validation_file_when_set(mock_mlflow, train
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
async def _dl(object_name, **_kwargs):
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.minio_repository.download_file = AsyncMock(side_effect=_dl)
training.minio_repository.download_file_sync = MagicMock(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)
training.data_manager_repository.compute_regression_metrics = MagicMock(
side_effect=lambda x, _w: x
)
def _fill(x, **_kw):
x.report_path = '/tmp/report.html'
@@ -261,10 +260,10 @@ async def test_train_model_downloads_validation_file_when_set(mock_mlflow, train
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.plugin_store.get_model = MagicMock(return_value=wrapper)
@asynccontextmanager
async def _run_ctx(*_a, **_k):
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_name = 'n'
info.run_id = 'i'
@@ -272,13 +271,12 @@ async def test_train_model_downloads_validation_file_when_set(mock_mlflow, train
training.mlflow_repository.start_run = _run_ctx
await training.train_model({'metadata': {}, 'train_params': tp})
assert training.minio_repository.download_file.await_count == 2
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
assert training.minio_repository.download_file_sync.call_count == 2
mock_mlflow.log_artifact.assert_called()
@pytest.mark.asyncio
async def test_train_model_value_error_when_paths_missing_after_report(training):
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(
{
@@ -290,26 +288,28 @@ async def test_train_model_value_error_when_paths_missing_after_report(training)
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
training.minio_repository.download_file = AsyncMock(return_value=b'x')
training.minio_repository.download_file_sync = MagicMock(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.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)
training.plugin_store.get_model = MagicMock(return_value=wrapper)
@asynccontextmanager
async def _run_ctx(*_a, **_k):
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_name = 'n'
info.run_id = 'i'
yield info
training.mlflow_repository.start_run = _run_ctx
training.send_notification_async = AsyncMock()
training.send_notification = MagicMock()
with pytest.raises(ValueError, match='Report path'):
await training.train_model({'metadata': {}, 'train_params': tp})
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})

12
tests/utils/models/test_train_model_params.py
View File

@@ -162,7 +162,11 @@ def test_validate_model_param_schema_validation_error(valid_train_params_dict):
'schemas': {
'components': {
'schemas': {
'data_model': {'type': 'object', 'properties': {'x': {'type': 'integer'}}, 'required': ['x']},
'data_model': {
'type': 'object',
'properties': {'x': {'type': 'integer'}},
'required': ['x'],
},
}
}
}
@@ -187,7 +191,7 @@ def test_validate_model_param_unexpected_validator_error(valid_train_params_dict
p = TrainModelParams.from_dict(d)
with patch('model_manager.utils.models.train_model_params.Draft202012Validator') as m:
m.return_value.validate.side_effect = RuntimeError('boom')
with pytest.raises(ValueError, match='Unexpected error'):
with pytest.raises(RuntimeError, match='boom'):
p.validate_business_rules()
@@ -224,9 +228,7 @@ def test_validate_model_param_only_data_model_schema(valid_train_params_dict):
def test_validate_model_param_only_model_schema(valid_train_params_dict):
d = copy.deepcopy(valid_train_params_dict)
d['model_metadata'] = {
'schemas': {'components': {'schemas': {'model': {'type': 'object'}}}}
}
d['model_metadata'] = {'schemas': {'components': {'schemas': {'model': {'type': 'object'}}}}}
p = TrainModelParams.from_dict(d)
p.model_kwargs = {}
p.validate_business_rules()

22
tests/utils/repository/test_data_manager_repository.py
View File

@@ -3,6 +3,7 @@
from __future__ import annotations
import json
from typing import Any
from unittest.mock import MagicMock, Mock, patch
import numpy as np
@@ -33,7 +34,7 @@ def test_train_test_split_ndarray():
def _params(**kwargs) -> TrainModelParams:
base = {
base: dict[str, Any] = {
'variable_columns': ['v1'],
'target_variable': 't',
'bucket_name': 'b',
@@ -276,7 +277,13 @@ def test_configure_datetime_index_already_datetime_index():
def test_configure_datetime_index_from_common_column():
repo = dmr.DataManagerRepository(MagicMock())
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
df = pd.DataFrame({'timestamp': pd.date_range('2024-01-01', periods=3, freq='D'), 'v1': [1, 2, 3], 't': [1, 2, 3]})
df = pd.DataFrame(
{
'timestamp': pd.date_range('2024-01-01', periods=3, freq='D'),
'v1': [1, 2, 3],
't': [1, 2, 3],
}
)
out = repo._configure_datetime_index(df, p, {})
assert isinstance(out.index, pd.DatetimeIndex)
@@ -329,14 +336,20 @@ def test_configure_datetime_index_first_column_numeric_parsed_as_time():
def test_create_run_directory_permission_error():
repo = dmr.DataManagerRepository(MagicMock())
with patch('model_manager.utils.repository.data_manager_repository.makedirs', side_effect=PermissionError('no')):
with patch(
'model_manager.utils.repository.data_manager_repository.makedirs',
side_effect=PermissionError('no'),
):
with pytest.raises(PermissionError, match='Permission denied'):
repo._create_run_directory('/tmp', 'run', {})
def test_create_run_directory_os_error():
repo = dmr.DataManagerRepository(MagicMock())
with patch('model_manager.utils.repository.data_manager_repository.makedirs', side_effect=OSError('disk')):
with patch(
'model_manager.utils.repository.data_manager_repository.makedirs',
side_effect=OSError('disk'),
):
with pytest.raises(OSError, match='Failed to create directory'):
repo._create_run_directory('/tmp', 'run', {})
@@ -355,7 +368,6 @@ def test_generate_report_success(tmp_path):
patch.object(repo, '_get_reports_directory', return_value=str(tmp_path)),
patch('model_manager.utils.repository.data_manager_repository.Reports') as mrep,
):
inst = mrep.return_value
instance = mrep.return_value
instance.save_all_sections_html = Mock()
out = repo.generate_report(tmr, {})

76
tests/worker/test_prepare_worker.py Normal file
View File

@@ -0,0 +1,76 @@
"""Unit tests for local worker factory."""
from unittest.mock import MagicMock, patch
def test_prepare_worker_train_queue_uses_train_limits():
from model_manager.worker.prepare_worker import prepare_worker
from model_manager.workflows.train_model import TrainModel
fake_worker = MagicMock()
fake_client = MagicMock()
fake_logger = MagicMock()
with patch(
'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
) as worker_class:
with patch.dict(
'os.environ',
{
'TRAINMODEL_ACTIVITY_EXECUTOR_MAX_WORKERS': '3',
'TRAINMODEL_MAX_CONCURRENT_ACTIVITIES': '6',
'TRAINMODEL_MAX_CONCURRENT_WORKFLOW_TASKS': '10',
},
clear=False,
):
worker = prepare_worker(
main_workflow=TrainModel,
other_workflows=[],
activities=[],
temporal_client=fake_client,
logger=fake_logger,
)
assert worker is fake_worker
worker_class.assert_called_once()
kwargs = worker_class.call_args.kwargs
assert kwargs['task_queue'] == 'train_model-queue'
assert kwargs['max_concurrent_activities'] == 6
assert kwargs['max_concurrent_workflow_tasks'] == 10
assert kwargs['activity_executor']._max_workers == 3
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)
def test_prepare_worker_cleanup_queue_uses_cleanup_limits():
from model_manager.worker.prepare_worker import prepare_worker
from model_manager.workflows.cleanup_files import CleanupFiles
fake_worker = MagicMock()
fake_client = MagicMock()
fake_logger = MagicMock()
with patch(
'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
) as worker_class:
with patch.dict(
'os.environ',
{
'CLEANUPFILES_ACTIVITY_EXECUTOR_MAX_WORKERS': '5',
'CLEANUPFILES_MAX_CONCURRENT_ACTIVITIES': '7',
},
clear=False,
):
worker = prepare_worker(
main_workflow=CleanupFiles,
other_workflows=[],
activities=[],
temporal_client=fake_client,
logger=fake_logger,
)
assert worker is fake_worker
kwargs = worker_class.call_args.kwargs
assert kwargs['task_queue'] == 'cleanup_files-queue'
assert kwargs['max_concurrent_activities'] == 7
assert kwargs['activity_executor']._max_workers == 5
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)

30
tests/worker/test_worker.py
View File

@@ -125,7 +125,7 @@ def test_start_prometheus_server_success(
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
metadata = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
start_prometheus_server(mock_logger, metadata)
@@ -148,7 +148,7 @@ def test_start_prometheus_server_custom_port(mock_metrics, mock_start_http_serve
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
metadata = {'pod_id': 'custom-pod', 'workflow_name': 'train_model'}
metadata: dict[str, str | None] = {'pod_id': 'custom-pod', 'workflow_name': 'train_model'}
start_prometheus_server(mock_logger, metadata)
@@ -166,7 +166,7 @@ def test_start_prometheus_server_failure(
mock_start_http_server.side_effect = OSError('Port already in use')
metadata = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
start_prometheus_server(mock_logger, metadata)
@@ -252,9 +252,7 @@ async def test_main_successful_startup(
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
@@ -347,9 +345,7 @@ async def test_main_handles_exception(
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
@@ -486,9 +482,7 @@ async def test_main_temporal_client_configuration(
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
@@ -581,9 +575,7 @@ async def test_main_worker_configuration(
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
@@ -703,9 +695,7 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
@@ -844,7 +834,7 @@ def test_start_prometheus_server_prints_success(
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
metadata = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
start_prometheus_server(mock_logger, metadata)
@@ -863,7 +853,7 @@ def test_start_prometheus_server_prints_failure(
mock_start_http_server.side_effect = Exception('Test error')
metadata = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
metadata: dict[str, str | None] = {'pod_id': 'test-pod-123', 'workflow_name': 'train_model'}
start_prometheus_server(mock_logger, metadata)

9
tests/workflows/test_train_model.py
View File

@@ -243,7 +243,9 @@ async def test_run_validation_error(mock_wf, sample_input_data):
@pytest.mark.asyncio
@patch('model_manager.workflows.train_model.workflow')
async def test_run_cleanup_failure_does_not_fail_workflow(mock_wf, sample_input_data, mock_train_params):
async def test_run_cleanup_failure_does_not_fail_workflow(
mock_wf, sample_input_data, mock_train_params
):
"""After successful training, cleanup failure is logged, workflow still returns result."""
from model_manager.workflows.train_model import TrainModel
@@ -267,7 +269,9 @@ async def test_run_cleanup_failure_does_not_fail_workflow(mock_wf, sample_input_
@pytest.mark.asyncio
@patch('model_manager.workflows.train_model.workflow')
async def test_run_training_failure_skips_cleanup_activity(mock_wf, sample_input_data, mock_train_params):
async def test_run_training_failure_skips_cleanup_activity(
mock_wf, sample_input_data, mock_train_params
):
"""When train_model raises, train_result stays None and cleanup activity is not scheduled."""
from model_manager.workflows.train_model import TrainModel
@@ -297,7 +301,6 @@ def test_module_constants():
from model_manager.workflows.train_model import (
TIMEOUT_DELETE_FILE,
TIMEOUT_TRAIN_MODEL,
TIMEOUT_UPDATE_DATABASE,
TIMEOUT_VALIDATE_PARAMS,
database_retry_policy,
network_retry_policy,