""" 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