diff --git a/README.md b/README.md index 12d5f34..a2fa182 100644 --- a/README.md +++ b/README.md @@ -149,13 +149,20 @@ The Model Manager system uses a Temporal-based workflow architecture with clear #### **Activities (`model_manager/activities/`)** - **Activities**: Main activity orchestrator combining all functionality through multiple inheritance +- **ExperimentTracking**: ML experiment lifecycle tracking and database operations (extends Postgres) + - Unified `update_experiment_run()` method for all experiment status updates + - Support for three update types: STATUS, STATUS_WITH_ERROR, MODEL_SAVED + - Automatic error message truncation (1024 chars) + - Connection pooling and retry logic via Postgres base class - **Gates**: Data quality validation and filtering mechanisms - **MLFlow**: Model transformation and prediction operations +- **MinIO**: Object storage operations for file management - **Key Features**: - Multiple inheritance pattern for unified activity interface - Configurable filter policies and validation rules - MLFlow model serving integration with configurable flavors - Comprehensive error handling and notification integration + - Experiment tracking with automatic status management #### **Data Services (`model_manager/utils/`)** - **Connectors Config**: Environment variable-based configuration management diff --git a/model_manager/activities/activities.py b/model_manager/activities/activities.py index 79daef4..501e3cb 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -5,14 +5,14 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger - from sientia_do.temporal.activities.postgres import Postgres + from model_manager.activities.experiment_tracking import ExperimentTracking from model_manager.activities.gates import Gates from model_manager.activities.minio import MinIO from model_manager.activities.mlflow import MLFlow -class Activities(Postgres, MLFlow, MinIO, Gates): +class Activities(ExperimentTracking, MLFlow, MinIO, Gates): """ Main activities orchestrator for the Model Manager system. @@ -21,7 +21,7 @@ class Activities(Postgres, MLFlow, MinIO, Gates): MLFlow model interactions, MinIO storage operations, and data quality validation. The class implements multiple inheritance to combine specialized functionality: - - Postgres: Database operations and data persistence + - ExperimentTracking: ML experiment lifecycle tracking and database operations (extends Postgres) - MLFlow: Model inference and transformation operations - MinIO: Object storage operations (file upload/download/delete) - Gates: Data quality validation and filtering mechanisms @@ -62,7 +62,7 @@ class Activities(Postgres, MLFlow, MinIO, Gates): Exception: If any parent class initialization fails """ # Initialize parent classes - Postgres.__init__( + ExperimentTracking.__init__( self, host=postgres_config['host'], port=postgres_config['port'], @@ -107,10 +107,10 @@ class Activities(Postgres, MLFlow, MinIO, Gates): Gracefully shutdown all activities and clean up resources. This method ensures proper cleanup of all resources including: - - PostgreSQL connection pools + - PostgreSQL connection pools (via ExperimentTracking) - Any other resources that need explicit cleanup The method should be called before the application terminates to ensure proper resource cleanup and prevent resource leaks. """ - Postgres.close(self) + ExperimentTracking.close(self) diff --git a/model_manager/activities/experiment_tracking.py b/model_manager/activities/experiment_tracking.py new file mode 100644 index 0000000..23cbf1b --- /dev/null +++ b/model_manager/activities/experiment_tracking.py @@ -0,0 +1,220 @@ +""" +Experiment tracking activities for managing ML experiment lifecycle. + +This module provides activities for tracking and updating experiment run status +in the PostgreSQL database, extending the base Postgres activity with specialized +methods for experiment management. +""" + +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import traceback + from datetime import datetime + from enum import Enum + from typing import Any + + import pytz + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.temporal.activities.postgres import Postgres + + +class UpdateType(str, Enum): + """Types of experiment run updates.""" + + STATUS = 'status' + STATUS_WITH_ERROR = 'status_with_error' + MODEL_SAVED = 'model_saved' + + +class ExperimentTracking(Postgres): + """ + Activity for tracking ML experiment lifecycle and status updates. + + This activity extends the Postgres activity to provide specialized methods + for managing experiment runs, including status updates, error tracking, and + model registration. It maintains the experiment lifecycle from initialization + through training, model saving, and cleanup. + + The activity uses the existing Postgres connection pool and adds experiment-specific + operations with proper error handling and notifications. + + Attributes: + logger (Logger): Logger instance for observability + notification_handler (NotificationHandler): Handler for sending notifications + """ + + def __init__( + self, + host: str, + port: int, + user: str, + password: str, + dbname: str, + min_connections: int, + max_connections: int, + logger: Logger, + notification_handler: NotificationHandler, + ): + """ + Initialize ExperimentTracking activity with database configuration. + + Args: + host: PostgreSQL server hostname + port: PostgreSQL server port + user: Database user + password: Database password + dbname: Database name + min_connections: Minimum connections in pool + max_connections: Maximum connections in pool + logger: Logger instance for observability + notification_handler: Notification handler for alerts + + Raises: + ConnectionError: If database connection cannot be established + """ + super().__init__( + host=host, + port=port, + user=user, + password=password, + dbname=dbname, + min_connections=min_connections, + max_connections=max_connections, + logger=logger, + notification_handler=notification_handler, + ) + self.logger = logger + self.notification_handler = notification_handler + + @activity.defn(name='update_experiment_run') + async def update_experiment_run(self, input_data: dict[str, Any]) -> None: + """ + Update experiment run with status, errors, or model information. + + This activity provides a unified interface for all experiment run updates, + supporting different update types through a single method. It automatically + selects the appropriate SQL query based on the update type and parameters. + + Args: + input_data: Configuration for experiment run update operation + Required keys: + - metadata (dict): Workflow execution metadata + - experiment_run_id (int): Unique identifier for the experiment run + - update_type (str): Type of update (status, status_with_error, model_saved) + Optional keys: + - status (str): New status for the experiment run + - error_message (str): Error message if update failed + - run_name (str): MLFlow run name if model was saved + + Raises: + ValueError: If required parameters are missing for the update type + RuntimeError: If update operation fails + + Example: + # Update status only + await update_experiment_run({ + 'metadata': {}, + 'experiment_run_id': 123, + 'update_type': 'status', + 'status': 'TRAINING_SUCCESS' + }) + + # Update with error + await update_experiment_run({ + 'metadata': {}, + 'experiment_run_id': 123, + 'update_type': 'status_with_error', + 'status': 'TRAINING_ERROR', + 'error_message': 'Model training failed: insufficient data' + }) + + # Update with model saved + await update_experiment_run({ + 'metadata': {}, + 'experiment_run_id': 123, + 'update_type': 'model_saved', + 'run_name': 'experiment-model-5' + }) + """ + metadata = input_data.get('metadata', {}) + experiment_run_id = input_data['experiment_run_id'] + update_type = input_data['update_type'] + status = input_data.get('status') + error_message = input_data.get('error_message') + run_name = input_data.get('run_name') + + try: + self.info( + f'Updating experiment run {experiment_run_id} with type: {update_type}', metadata + ) + + # Validate parameters based on update type + query_params: tuple[Any, ...] + if update_type == UpdateType.STATUS: + if not status: + raise ValueError('status is required for STATUS update type') + sql_query = """ + UPDATE experiment_run + SET status = %s, updated_at = %s + WHERE id = %s + """ + query_params = (status, datetime.now(pytz.utc), experiment_run_id) + + elif update_type == UpdateType.STATUS_WITH_ERROR: + if not status or not error_message: + raise ValueError( + 'status and error_message are required for STATUS_WITH_ERROR update type' + ) + + # Truncate the error_message to 1024 characters if necessary + if len(error_message) > 1024: + error_message = error_message[:1024] + + sql_query = """ + UPDATE experiment_run + SET status = %s, error_message = %s, updated_at = %s + WHERE id = %s + """ + query_params = (status, error_message, datetime.now(pytz.utc), experiment_run_id) + + elif update_type == UpdateType.MODEL_SAVED: + if not run_name: + raise ValueError('run_name is required for MODEL_SAVED update type') + sql_query = """ + UPDATE experiment_run + SET run_name = %s, status = %s, updated_at = %s + WHERE id = %s + """ + query_params = (run_name, status, datetime.now(pytz.utc), experiment_run_id) + + else: + raise ValueError(f'Invalid update_type: {update_type}') + + # Execute update query + result = await self.execute_query(sql_query, query_params) + + if result.get('rowcount', 0) == 0: + error_msg = f'No rows updated for experiment run {experiment_run_id}' + raise ValueError(error_msg) + + self.info( + f'Successfully updated experiment run {experiment_run_id} with type {update_type}', + metadata, + ) + + except Exception as e: # noqa: BLE001 + error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Type: {update_type}, Error: {str(e)}' + trace = traceback.format_exc() + self.send_notification( + metadata=metadata, + notification_id='UPDATE_EXPERIMENT_RUN_ERROR', + message=error_msg, + block='update_experiment_run', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata=metadata) + raise RuntimeError(error_msg) from e diff --git a/tests/activities/test_activities.py b/tests/activities/test_activities.py index 99f6083..e40c038 100644 --- a/tests/activities/test_activities.py +++ b/tests/activities/test_activities.py @@ -1,18 +1,20 @@ from unittest.mock import ANY, MagicMock, patch from pytest import mark -from sientia_do.temporal.activities.postgres import Postgres from model_manager.activities.activities import Activities +from model_manager.activities.experiment_tracking import ExperimentTracking from model_manager.activities.gates import Gates from model_manager.activities.mlflow import MLFlow -@patch('model_manager.activities.activities.Postgres.__init__') +@patch('model_manager.activities.activities.ExperimentTracking.__init__') @patch('model_manager.activities.activities.MLFlow.__init__') @patch('model_manager.activities.activities.MinIO.__init__') @patch('model_manager.activities.activities.Gates.__init__') -def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postgres_init): +def test___init__( + mock_gates_init, mock_minio_init, mock_mlflow_init, mock_experiment_tracking_init +): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -49,11 +51,11 @@ def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postg ) assert isinstance(activities, Activities) - assert isinstance(activities, Postgres) + assert isinstance(activities, ExperimentTracking) assert isinstance(activities, MLFlow) assert isinstance(activities, Gates) - mock_postgres_init.assert_called_once_with( + mock_experiment_tracking_init.assert_called_once_with( ANY, host=postgres_config['host'], port=postgres_config['port'], @@ -97,9 +99,9 @@ def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postg @mark.asyncio -@patch('model_manager.activities.activities.Postgres', return_value=MagicMock()) +@patch('model_manager.activities.activities.ExperimentTracking', return_value=MagicMock()) @patch('model_manager.activities.activities.MLFlow', return_value=MagicMock()) -async def test_shutdown(_mock_mlflow_init, mock_postgres_init): +async def test_shutdown(_mock_mlflow_init, mock_experiment_tracking_init): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -136,4 +138,4 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init): ) await activities.shutdown() - mock_postgres_init.close.assert_called_once() + mock_experiment_tracking_init.close.assert_called_once() diff --git a/tests/activities/test_experiment_tracking.py b/tests/activities/test_experiment_tracking.py new file mode 100644 index 0000000..1be5aec --- /dev/null +++ b/tests/activities/test_experiment_tracking.py @@ -0,0 +1,385 @@ +"""Unit tests for ExperimentTracking activity.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +from pytest import mark, raises + +from model_manager.activities.experiment_tracking import ExperimentTracking, UpdateType + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_status_success(mock_postgres_init): + """Test successful status update.""" + mock_postgres_init.return_value = None + + # Create instance + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + # Mock methods + tracking.info = MagicMock() + tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) + + # Test data + input_data = { + 'metadata': {'workflow_id': 'test-123'}, + 'experiment_run_id': 456, + 'update_type': UpdateType.STATUS, + 'status': 'TRAINING_SUCCESS', + } + + # Execute + await tracking.update_experiment_run(input_data) + + # Assertions + tracking.info.assert_called() + tracking.execute_query.assert_called_once() + call_args = tracking.execute_query.call_args + assert 'UPDATE experiment_run' in call_args[0][0] + assert 'SET status = %s' in call_args[0][0] + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_status_with_error_success(mock_postgres_init): + """Test successful status update with error message.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.info = MagicMock() + tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) + + input_data = { + 'metadata': {'workflow_id': 'test-123'}, + 'experiment_run_id': 456, + 'update_type': UpdateType.STATUS_WITH_ERROR, + 'status': 'TRAINING_ERROR', + 'error_message': 'Model training failed due to insufficient data', + } + + await tracking.update_experiment_run(input_data) + + tracking.info.assert_called() + tracking.execute_query.assert_called_once() + call_args = tracking.execute_query.call_args + assert 'UPDATE experiment_run' in call_args[0][0] + assert 'SET status = %s, error_message = %s' in call_args[0][0] + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_error_message_truncation(mock_postgres_init): + """Test that error messages longer than 1024 chars are truncated.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.info = MagicMock() + tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) + + # Create error message longer than 1024 characters + long_error = 'A' * 2000 + + input_data = { + 'metadata': {}, + 'experiment_run_id': 456, + 'update_type': UpdateType.STATUS_WITH_ERROR, + 'status': 'TRAINING_ERROR', + 'error_message': long_error, + } + + await tracking.update_experiment_run(input_data) + + # Check that error message was truncated to 1024 chars + call_args = tracking.execute_query.call_args + query_params = call_args[0][1] + assert len(query_params[1]) == 1024 # error_message is second param + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_model_saved_success(mock_postgres_init): + """Test successful model saved update.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.info = MagicMock() + tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) + + input_data = { + 'metadata': {}, + 'experiment_run_id': 456, + 'update_type': UpdateType.MODEL_SAVED, + 'run_name': 'experiment-model-123', + 'status': 'MLFLOW_SENT', + } + + await tracking.update_experiment_run(input_data) + + tracking.info.assert_called() + tracking.execute_query.assert_called_once() + call_args = tracking.execute_query.call_args + assert 'UPDATE experiment_run' in call_args[0][0] + assert 'SET run_name = %s, status = %s' in call_args[0][0] + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_missing_status_raises_error(mock_postgres_init): + """Test that missing status parameter raises ValueError.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.send_notification = MagicMock() + tracking.error = MagicMock() + + input_data = { + 'metadata': {}, + 'experiment_run_id': 456, + 'update_type': UpdateType.STATUS, + # Missing 'status' parameter + } + + with raises(RuntimeError, match='Error updating experiment run'): + await tracking.update_experiment_run(input_data) + + tracking.send_notification.assert_called_once() + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_missing_error_message_raises_error(mock_postgres_init): + """Test that missing error_message parameter raises ValueError.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.send_notification = MagicMock() + tracking.error = MagicMock() + + input_data = { + 'metadata': {}, + 'experiment_run_id': 456, + 'update_type': UpdateType.STATUS_WITH_ERROR, + 'status': 'TRAINING_ERROR', + # Missing 'error_message' parameter + } + + with raises(RuntimeError, match='Error updating experiment run'): + await tracking.update_experiment_run(input_data) + + tracking.send_notification.assert_called_once() + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_missing_run_name_raises_error(mock_postgres_init): + """Test that missing run_name parameter raises ValueError.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.send_notification = MagicMock() + tracking.error = MagicMock() + + input_data = { + 'metadata': {}, + 'experiment_run_id': 456, + 'update_type': UpdateType.MODEL_SAVED, + # Missing 'run_name' parameter + } + + with raises(RuntimeError, match='Error updating experiment run'): + await tracking.update_experiment_run(input_data) + + tracking.send_notification.assert_called_once() + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_invalid_update_type_raises_error(mock_postgres_init): + """Test that invalid update_type raises ValueError.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.send_notification = MagicMock() + tracking.error = MagicMock() + + input_data = { + 'metadata': {}, + 'experiment_run_id': 456, + 'update_type': 'INVALID_TYPE', + 'status': 'TRAINING_SUCCESS', + } + + with raises(RuntimeError, match='Error updating experiment run'): + await tracking.update_experiment_run(input_data) + + tracking.send_notification.assert_called_once() + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_no_rows_updated_raises_error(mock_postgres_init): + """Test that zero rows updated raises ValueError.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.info = MagicMock() + tracking.execute_query = AsyncMock(return_value={'rowcount': 0}) + tracking.send_notification = MagicMock() + tracking.error = MagicMock() + + input_data = { + 'metadata': {}, + 'experiment_run_id': 999, # Non-existent ID + 'update_type': UpdateType.STATUS, + 'status': 'TRAINING_SUCCESS', + } + + with raises(RuntimeError, match='Error updating experiment run'): + await tracking.update_experiment_run(input_data) + + tracking.send_notification.assert_called_once() + + +@mark.asyncio +@patch('model_manager.activities.experiment_tracking.Postgres.__init__') +async def test_update_experiment_run_sends_notification_on_error(mock_postgres_init): + """Test that notification is sent when update fails.""" + mock_postgres_init.return_value = None + + tracking = ExperimentTracking( + host='localhost', + port=5432, + user='test', + password='test', + dbname='test', + min_connections=1, + max_connections=10, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + tracking.info = MagicMock() + tracking.execute_query = AsyncMock(side_effect=Exception('Database error')) + tracking.send_notification = MagicMock() + tracking.error = MagicMock() + + input_data = { + 'metadata': {'workflow_id': 'test-123'}, + 'experiment_run_id': 456, + 'update_type': UpdateType.STATUS, + 'status': 'TRAINING_SUCCESS', + } + + with raises(RuntimeError): + await tracking.update_experiment_run(input_data) + + # Verify notification was sent + tracking.send_notification.assert_called_once() + call_args = tracking.send_notification.call_args + assert call_args[1]['notification_id'] == 'UPDATE_EXPERIMENT_RUN_ERROR' + assert call_args[1]['metadata'] == {'workflow_id': 'test-123'} + + +def test_update_type_enum_values(): + """Test UpdateType enum has correct values.""" + assert UpdateType.STATUS == 'status' + assert UpdateType.STATUS_WITH_ERROR == 'status_with_error' + assert UpdateType.MODEL_SAVED == 'model_saved'