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

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()})