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 Pimentel dos Santos
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()})