diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py index 73f8f86..b09df5b 100644 --- a/orchestrator/activities/formatters.py +++ b/orchestrator/activities/formatters.py @@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through(): predictions_batch, scouter, drift, + simple_metrics, ) topic_separator = '\n ========== \n' @@ -140,6 +141,13 @@ class Formatters(SientiaMonitoring): 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ) ), } + elif pipeline['workflow_type'] == 'simple_metrics': + schedule_config[self.laborious_namespace][pipeline['schedule_name']] = { + **simple_metrics(pipeline), + 'updated_at': pipeline.get( + 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ) + ), + } self.info('Processed schedules', metadata=metadata) self.debug(json.dumps(schedule_config, indent=4, sort_keys=True), metadata=metadata) diff --git a/orchestrator/utils/orchestrator_functions.py b/orchestrator/utils/orchestrator_functions.py index 54c388b..65f1cc6 100644 --- a/orchestrator/utils/orchestrator_functions.py +++ b/orchestrator/utils/orchestrator_functions.py @@ -36,7 +36,6 @@ def drift(config: dict[str, Any]): """ Build drift configuration from pipeline config. """ - model = config['model'] return { **common_config(config), 'schema': 'sientia_data', @@ -51,6 +50,20 @@ def drift(config: dict[str, Any]): } +def simple_metrics(config: dict[str, Any]): + """ + Build simple metrics configuration from pipeline config. + """ + return { + **common_config(config), + 'schema': 'sientia_data', + 'predictions_table_name': 'predictions', + 'data_table_name': 'laborious_data', + 'interval_minutes': config.get('interval_minutes', 60), + 'metrics': config.get('metrics', ['rmse', 'mse', 'mae', 'r2']), + } + + def minimal_retrain(config: dict[str, Any]): """ Build minimal retrain configuration from pipeline config. diff --git a/tests/orchestrator/activities/test_formatters.py b/tests/orchestrator/activities/test_formatters.py index 12f6901..51b0e08 100644 --- a/tests/orchestrator/activities/test_formatters.py +++ b/tests/orchestrator/activities/test_formatters.py @@ -48,8 +48,17 @@ metadata = { 'orchestrator.activities.formatters.drift', return_value={'test_drift': 'test_drift'}, ) +@patch( + 'orchestrator.activities.formatters.simple_metrics', + return_value={'test_simple_metrics': 'test_simple_metrics'}, +) async def test_process_schedules( - mock_drift, mock_minimal_retrain, mock_predictions_batch, mock_scouter, formatters + mock_simple_metrics, + mock_drift, + mock_minimal_retrain, + mock_predictions_batch, + mock_scouter, + formatters, ): input_data = { 'pipelines': [ @@ -81,6 +90,13 @@ async def test_process_schedules( 'model_id': 'test_model_id', 'updated_at': '2021-01-04', }, + { + 'schedule_name': 'test_schedule_name5', + 'workflow_type': 'simple_metrics', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', + 'updated_at': '2021-01-05', + }, ] } @@ -103,6 +119,10 @@ async def test_process_schedules( 'test_drift': 'test_drift', 'updated_at': '2021-01-04', }, + 'test_schedule_name5': { + 'test_simple_metrics': 'test_simple_metrics', + 'updated_at': '2021-01-05', + }, }, } @@ -110,6 +130,7 @@ async def test_process_schedules( mock_predictions_batch.assert_called_once_with(input_data['pipelines'][1]) mock_minimal_retrain.assert_called_once_with(input_data['pipelines'][2]) mock_drift.assert_called_once_with(input_data['pipelines'][3]) + mock_simple_metrics.assert_called_once_with(input_data['pipelines'][4]) @mark.asyncio diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py index e3add38..f5a7c13 100644 --- a/tests/orchestrator/activities/test_temporal_manager.py +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -18,8 +18,7 @@ metadata = { @fixture -@patch('orchestrator.activities.temporal_manager.Client.connect') -def temporal_manager(connect_mock): +def temporal_manager(): temporal_manager = TemporalManager( host='localhost:7233', scouter_namespace='scouter', diff --git a/tests/orchestrator/utils/test_orchestrator_functions.py b/tests/orchestrator/utils/test_orchestrator_functions.py index 7c8f5da..de0eddd 100644 --- a/tests/orchestrator/utils/test_orchestrator_functions.py +++ b/tests/orchestrator/utils/test_orchestrator_functions.py @@ -10,6 +10,7 @@ from orchestrator.utils.orchestrator_functions import ( predictions_batch, process_path_priority, scouter, + simple_metrics, ) @@ -66,6 +67,36 @@ def test_drift(): assert result == expected +def test_simple_metrics(): + config = { + 'workflow_type': 'simple_metrics', + 'schedule_name': 'test_schedule', + 'model_id': 'test_model_id', + 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}}, + 'interval_minutes': 120, + 'metrics': ['rmse', 'mse'], + } + result = simple_metrics(config) + expected = { + 'workflow_type': 'simple_metrics', + 'schedule_name': 'test_schedule', + 'frequency': '1m', + 'offset': '0m', + 'max_retry_policy': 1, + 'model_id': 'test_model_id', + 'model_name': 'test_model_name', + 'model_config': {'test_config': 'test_config'}, + 'schema': 'sientia_data', + 'predictions_table_name': 'predictions', + 'data_table_name': 'laborious_data', + 'interval_minutes': 120, + 'metrics': ['rmse', 'mse'], + 'execution_timeout_seconds': 300, + 'task_timeout_seconds': 300, + } + assert result == expected + + def test_minimal_retrain(): config = { 'workflow_type': 'minimal_retrain',