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

SIENTIAPDE-1645: Add comprehensive model training observability metrics and Grafana dashboard.

Browse Source
This commit is contained in:
vitor-aignosi
2026-06-17 10:34:56 -03:00
parent 445fe643fe
commit 76f926a8ab
12 changed files with 866 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
tests/activities/test_cleanup.py
View File

@@ -125,13 +125,11 @@ def test_cleanup_temp_directories_nonexistent_path(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.warning = MagicMock()
cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}})
cleanup.warning.assert_called_once()
cleanup._emit_metrics.assert_called_once()
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
@@ -152,8 +150,6 @@ def test_cleanup_temp_directories_success_with_deletions(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
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)
@@ -166,7 +162,6 @@ def test_cleanup_temp_directories_success_with_deletions(
assert not os.path.exists(old_dir)
assert os.path.exists(recent_dir)
cleanup._emit_metrics.assert_called_once()
@patch.dict(os.environ, {'CLEANUP_DRY_RUN': 'true'})
@@ -187,8 +182,6 @@ def test_cleanup_temp_directories_dry_run(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
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)
@@ -196,7 +189,6 @@ def test_cleanup_temp_directories_dry_run(
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
assert os.path.exists(old_dir)
cleanup._emit_metrics.assert_called_once()
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
@@ -217,7 +209,6 @@ def test_cleanup_temp_directories_delete_error(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.error = MagicMock()
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
@@ -228,35 +219,9 @@ def test_cleanup_temp_directories_delete_error(
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
cleanup.error.assert_called_once()
cleanup._emit_metrics.assert_called_once()
# --- Metrics and Utility Tests ---
def test_emit_metrics(
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test that _emit_metrics calls the public emit_metric method."""
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup.emit_metric_sync = MagicMock()
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_sync.call_count == 2
# --- Utility Tests ---
def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
@@ -276,7 +241,6 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.debug = MagicMock()
# Create a file and a directory with a non-matching name
@@ -290,7 +254,6 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
cleanup.debug.assert_called_with(
'Skipping directory without timestamp pattern: a_directory_with_no_timestamp', {}
)
cleanup._emit_metrics.assert_called_once()
def test_cleanup_temp_directories_invalid_timestamp_format(
@@ -310,7 +273,6 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.error = MagicMock()
# Create a directory with a malformed timestamp that matches the regex but fails parsing
@@ -320,7 +282,6 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
cleanup.error.assert_called_once()
cleanup._emit_metrics.assert_called_once()
def test_cleanup_temp_directories_generic_exception(
@@ -340,7 +301,6 @@ def test_cleanup_temp_directories_generic_exception(
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = MagicMock() # type: ignore[method-assign]
cleanup.send_notification = MagicMock()
with patch('os.listdir', side_effect=Exception('Unexpected OS Error')):
@@ -348,29 +308,3 @@ def test_cleanup_temp_directories_generic_exception(
cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}})
cleanup.send_notification.assert_called_once()
cleanup._emit_metrics.assert_called_once()
def test_emit_metrics_activity_only(
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test that _emit_metrics can emit only the activity metric."""
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup.emit_metric_sync = MagicMock()
cleanup._emit_metrics(
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
metrics_status='success',
activity_name='test_activity',
emit_workflow_metric=False,
)
cleanup.emit_metric_sync.assert_called_once()

261
tests/activities/test_training.py
View File

@@ -330,6 +330,267 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
mock_mlflow.log_artifact.assert_called()
def test_prepare_data_observes_lag_on_success(training):
tp = TrainModelParams.from_dict(
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
)
train_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
training.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
from model_manager import metrics as mm_metrics
training._prepare_data(b'csv', None, tp, {})
training.observe_lag_sync.assert_called_once()
call_args = training.observe_lag_sync.call_args
assert call_args.args[1] is mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_LAG
training.emit_metric_sync.assert_not_called()
def test_prepare_data_increments_error_counter_and_still_observes_lag_on_failure(training):
tp = TrainModelParams.from_dict(
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
)
training.data_manager_repository.prepare_training_data = MagicMock(
side_effect=RuntimeError('prep-fail')
)
training.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
from model_manager import metrics as mm_metrics
with pytest.raises(RuntimeError, match='prep-fail'):
training._prepare_data(b'csv', None, tp, {})
training.observe_lag_sync.assert_called_once()
training.emit_metric_sync.assert_called_once()
call_args = training.emit_metric_sync.call_args
assert call_args.kwargs['metric_object'] is mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_ERROR_COUNT_TOTAL
def test_fit_model_observes_lag_on_success(training):
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)
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.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
from model_manager import metrics as mm_metrics
training._fit_model(wrapper, tmr, tp, {})
training.observe_lag_sync.assert_called_once()
call_args = training.observe_lag_sync.call_args
assert call_args.args[1] is mm_metrics.SIENTIA_TRAINING_MODEL_FIT_LAG
training.emit_metric_sync.assert_not_called()
def test_fit_model_increments_error_counter_on_failure(training):
tp = TrainModelParams.from_dict(
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
)
train_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
wrapper = MagicMock()
wrapper.train = MagicMock(side_effect=RuntimeError('fit-fail'))
training.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
from model_manager import metrics as mm_metrics
with pytest.raises(RuntimeError, match='fit-fail'):
training._fit_model(wrapper, tmr, tp, {})
training.observe_lag_sync.assert_called_once()
training.emit_metric_sync.assert_called_once()
call_args = training.emit_metric_sync.call_args
assert call_args.kwargs['metric_object'] is mm_metrics.SIENTIA_TRAINING_MODEL_FIT_ERROR_COUNT_TOTAL
@patch('model_manager.activities.training.mm_metrics')
@patch('model_manager.activities.training.mlflow')
def test_train_model_sets_quality_gauges_after_compute_metrics(mock_mlflow, mock_mm_metrics, training):
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.minio_repository.download_file = MagicMock(return_value=b'csv')
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
def _set_metrics(x, _w, **_kw):
x.mse_val = 0.5
x.mae_val = 0.3
x.r2_val = -0.1
return x
training.data_manager_repository.compute_regression_metrics = MagicMock(
side_effect=_set_metrics
)
def _fill_report(x, **_kw):
x.report_path = '/tmp/r.html'
x.train_data_path = '/tmp/tr.csv'
x.test_data_path = '/tmp/te.csv'
x.run_dir = '/tmp/run'
return x
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report)
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 = MagicMock(return_value=wrapper)
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_id = 'rid'
yield info
training.mlflow_repository.start_run = _run_ctx
training.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels.return_value.set.assert_called_once_with(0.5)
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels.return_value.set.assert_called_once_with(0.3)
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels.return_value.set.assert_called_once_with(-0.1)
@patch('model_manager.activities.training.mm_metrics')
@patch('model_manager.activities.training.mlflow')
def test_train_model_skips_quality_gauges_when_none(_mock_mlflow, mock_mm_metrics, training):
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.minio_repository.download_file = 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, **_kw: x
)
def _fill_report(x, **_kw):
x.report_path = '/tmp/r.html'
x.train_data_path = '/tmp/tr.csv'
x.test_data_path = '/tmp/te.csv'
x.run_dir = '/tmp/run'
return x
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report)
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 = MagicMock(return_value=wrapper)
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_id = 'rid'
yield info
training.mlflow_repository.start_run = _run_ctx
training.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels.return_value.set.assert_not_called()
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels.return_value.set.assert_not_called()
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels.return_value.set.assert_not_called()
@patch('model_manager.activities.training.mm_metrics')
@patch('model_manager.activities.training.mlflow')
def test_train_model_increments_trained_total_on_success(_mock_mlflow, mock_mm_metrics, training):
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.minio_repository.download_file = 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, **_kw: x
)
def _fill_report(x, **_kw):
x.report_path = '/tmp/r.html'
x.train_data_path = '/tmp/tr.csv'
x.test_data_path = '/tmp/te.csv'
x.run_dir = '/tmp/run'
return x
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report)
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 = MagicMock(return_value=wrapper)
@contextmanager
def _run_ctx(*_a, **_k):
info = MagicMock()
info.run_id = 'rid'
yield info
training.mlflow_repository.start_run = _run_ctx
training.observe_lag_sync = MagicMock()
training.emit_metric_sync = MagicMock()
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
training.emit_metric_sync.assert_called_once_with(
metric_object=mock_mm_metrics.SIENTIA_TRAINING_MODEL_TRAINED_TOTAL,
tags=training._get_training_labels(tp),
)
def test_train_model_does_not_increment_trained_total_on_failure(training):
tp = TrainModelParams.from_dict(
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
)
training.minio_repository.download_file = MagicMock(side_effect=RuntimeError('dl-fail'))
training.send_notification = MagicMock()
training.emit_metric_sync = MagicMock()
with pytest.raises(RuntimeError):
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
training.emit_metric_sync.assert_not_called()
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(