255 lines
9.5 KiB
Python
255 lines
9.5 KiB
Python
"""
|
|
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 asyncio
|
|
import traceback
|
|
from collections.abc import Mapping
|
|
from datetime import UTC, datetime
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
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
|
|
from sqlalchemy import text
|
|
|
|
|
|
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
|
|
|
|
def __del__(self):
|
|
"""
|
|
Destructor to safely handle cleanup during garbage collection.
|
|
|
|
This prevents AttributeError when used in multiple inheritance scenarios
|
|
where the parent Postgres.__del__ might be called on objects without
|
|
the engine attribute.
|
|
"""
|
|
# Only call parent __del__ if engine attribute exists
|
|
if hasattr(self, 'engine'):
|
|
try:
|
|
if hasattr(super(), '__del__'):
|
|
super().__del__()
|
|
except Exception: # noqa: S110, BLE001
|
|
# Silently ignore errors during garbage collection
|
|
pass
|
|
|
|
async def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Execute an UPDATE SQL statement asynchronously.
|
|
|
|
Args:
|
|
query: Parameterized SQL string to execute.
|
|
params: Mapping of parameters for the SQL query.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the affected row count: {'rowcount': int}.
|
|
"""
|
|
|
|
def _run() -> dict[str, Any]:
|
|
with self.engine.begin() as connection:
|
|
result = connection.execute(text(query), params)
|
|
return {'rowcount': result.rowcount}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
@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
|
|
"""
|
|
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 status: {status}', metadata
|
|
)
|
|
|
|
query_params: dict[str, Any]
|
|
|
|
if update_type == UpdateType.STATUS:
|
|
if not isinstance(status, str) or not status:
|
|
raise ValueError('status is required for STATUS update type')
|
|
|
|
sql_query = """
|
|
UPDATE experiment_run
|
|
SET status = :status, updated_at = :updated_at
|
|
WHERE id = :experiment_run_id
|
|
"""
|
|
|
|
query_params = {
|
|
'status': status,
|
|
'updated_at': datetime.now(UTC),
|
|
'experiment_run_id': experiment_run_id,
|
|
}
|
|
elif update_type == UpdateType.STATUS_WITH_ERROR:
|
|
if not isinstance(status, str) or not status:
|
|
raise ValueError('status is required for STATUS_WITH_ERROR update type')
|
|
|
|
if not isinstance(error_message, str) or not error_message:
|
|
raise ValueError('error_message is required for STATUS_WITH_ERROR update type')
|
|
|
|
if len(error_message) > 1024:
|
|
error_message = error_message[:1024]
|
|
|
|
sql_query = """
|
|
UPDATE experiment_run
|
|
SET status = :status, error_message = :error_message, updated_at = :updated_at
|
|
WHERE id = :experiment_run_id
|
|
"""
|
|
|
|
query_params = {
|
|
'status': status,
|
|
'error_message': error_message,
|
|
'updated_at': datetime.now(UTC),
|
|
'experiment_run_id': experiment_run_id,
|
|
}
|
|
elif update_type == UpdateType.MODEL_SAVED:
|
|
if not isinstance(run_name, str) or not run_name:
|
|
raise ValueError('run_name is required for MODEL_SAVED update type')
|
|
|
|
if not isinstance(status, str) or not status:
|
|
raise ValueError('status is required for MODEL_SAVED update type')
|
|
|
|
sql_query = """
|
|
UPDATE experiment_run
|
|
SET run_name = :run_name, status = :status, updated_at = :updated_at
|
|
WHERE id = :experiment_run_id
|
|
"""
|
|
|
|
query_params = {
|
|
'run_name': run_name,
|
|
'status': status,
|
|
'updated_at': datetime.now(UTC),
|
|
'experiment_run_id': experiment_run_id,
|
|
}
|
|
else:
|
|
raise ValueError(f'Invalid update_type: {update_type}')
|
|
|
|
result = await self._execute_update(sql_query, query_params)
|
|
|
|
if result.get('rowcount', 0) == 0:
|
|
error_msg = (
|
|
f'No rows updated for experiment run {experiment_run_id} with status {status}'
|
|
)
|
|
raise ValueError(error_msg)
|
|
|
|
self.info(
|
|
f'Successfully updated experiment run {experiment_run_id} with status {status}',
|
|
metadata,
|
|
)
|
|
except Exception as e: # noqa: BLE001
|
|
error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Status: {status}, 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
|