diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index 01230a5..512ee77 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -18,6 +18,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL from model_manager.utils.exceptions import ModelTrainingError from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.repository.model_repository import ModelRepository @@ -76,6 +77,7 @@ class Training(SientiaMonitoring): ValueError, TypeError, KeyError: If validation fails (after sending notification) """ metadata = input_data.get('metadata', {}) + metrics_status = 'success' try: train_params = TrainModelParams.from_dict(input_data) @@ -90,6 +92,7 @@ class Training(SientiaMonitoring): return train_params except (ValueError, TypeError, KeyError) as e: + metrics_status = 'error' error_msg = f'Error validating training parameters: {str(e)}' trace = traceback.format_exc() @@ -102,6 +105,13 @@ class Training(SientiaMonitoring): attachment_content=trace, ) raise + finally: + await self._emit_metrics( + metadata=metadata, + metrics_status=metrics_status, + activity_name='validate_train_params', + emit_workflow_metric=(metrics_status == 'error'), + ) @activity.defn(name='train_model') async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]: @@ -135,6 +145,7 @@ class Training(SientiaMonitoring): model_trained = False model_saved = False + metrics_status = 'success' try: with self.storage_repository.fetch_file( @@ -155,6 +166,8 @@ class Training(SientiaMonitoring): 'run_dir': train_result.run_dir, } except Exception as e: # noqa: BLE001 + metrics_status = 'error' + error_msg = ( 'Error training model - ' f'model_trained={model_trained}, model_saved={model_saved}, ' @@ -176,6 +189,13 @@ class Training(SientiaMonitoring): model_trained=model_trained, model_saved=model_saved, ) from e + finally: + await self._emit_metrics( + metadata=metadata, + metrics_status=metrics_status, + activity_name='train_model', + emit_workflow_metric=(metrics_status == 'error'), + ) @activity.defn(name='cleanup_resources') async def cleanup_resources(self, input_data: dict[str, Any]) -> None: @@ -196,11 +216,14 @@ class Training(SientiaMonitoring): run_dir = input_data.get('run_dir', '') bucket_name = input_data.get('bucket_name', '') file_name = input_data.get('file_name', '') + metrics_status = 'success' try: self.model_repository.cleanup_run_directory(run_dir) self.storage_repository.delete_file(bucket_name, file_name) except Exception as e: # noqa: BLE001 + metrics_status = 'error' + error_msg = ( f'Error cleaning up resources - Run directory: {run_dir}, ' f'File: {bucket_name}/{file_name}, Error: {str(e)}' @@ -216,4 +239,46 @@ class Training(SientiaMonitoring): level=NotificationLevel.ERROR, attachment_content=trace, ) + raise + finally: + await self._emit_metrics( + metadata=metadata, + metrics_status=metrics_status, + activity_name='cleanup_resources', + emit_workflow_metric=True, + ) + + async def _emit_metrics( + self, + metadata: dict[str, Any], + metrics_status: str, + activity_name: str, + emit_workflow_metric: bool, + ) -> None: + """ + Emit workflow and activity execution metrics. + + Args: + metadata: Activity metadata containing pod_id and workflow_name + metrics_status: Execution status ('success' or 'error') + activity_name: Name of the activity being executed + """ + if emit_workflow_metric: + await self.emit_metric( + metric_object=WORKFLOW_EXECUTION_TOTAL, + tags={ + 'pod_id': metadata.get('pod_id'), + 'workflow_name': metadata.get('workflow_name'), + 'status': metrics_status, + }, + ) + + await self.emit_metric( + metric_object=ACTIVITY_EXECUTION_TOTAL, + tags={ + 'pod_id': metadata.get('pod_id'), + 'activity_name': activity_name, + 'status': metrics_status, + }, + ) diff --git a/model_manager/metrics.py b/model_manager/metrics.py index 0372ebb..9533e9b 100644 --- a/model_manager/metrics.py +++ b/model_manager/metrics.py @@ -16,7 +16,7 @@ Metric Labels: - pod_id: Kubernetes pod identifier for multi-instance deployments """ -from prometheus_client import Gauge +from prometheus_client import Counter, Gauge # Application health metric APP_UP = Gauge( @@ -24,3 +24,15 @@ APP_UP = Gauge( 'Indicates if the application is running (1) or shutting down (0)', ['pod_id'], ) + +WORKFLOW_EXECUTION_TOTAL = Counter( + 'model_manager_workflow_executions_total', + 'Total number of workflow executions', + ['pod_id', 'workflow_name', 'status'], # status: success, error +) + +ACTIVITY_EXECUTION_TOTAL = Counter( + 'model_manager_activity_executions_total', + 'Total number of activity executions', + ['pod_id', 'activity_name', 'status'], # status: success, error +) diff --git a/tests/activities/test_training.py b/tests/activities/test_training.py index c9013c2..4d6eb70 100644 --- a/tests/activities/test_training.py +++ b/tests/activities/test_training.py @@ -28,7 +28,12 @@ def mock_metrics_controller(): async def mock_shutdown(): pass + # Make emit an async coroutine + async def mock_emit(*args, **kwargs): + pass + controller.shutdown = mock_shutdown + controller.emit = mock_emit return controller