281 lines
10 KiB
Python
281 lines
10 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 synchronous Postgres client with specialized
|
|
methods for experiment management.
|
|
"""
|
|
|
|
import enum
|
|
|
|
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
import traceback
|
|
from collections.abc import Mapping
|
|
from datetime import UTC, datetime
|
|
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.observability.metrics_controller import MetricsController
|
|
from sientia_do.temporal.activities.postgres_sync import Postgres
|
|
from sqlalchemy import text
|
|
|
|
|
|
class UpdateType(enum.StrEnum):
|
|
"""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 a SQLAlchemy engine / connection pool and adds experiment-specific
|
|
operations with proper error handling and 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,
|
|
metrics_controller: MetricsController,
|
|
):
|
|
"""
|
|
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
|
|
metrics_controller: Metrics controller for observability
|
|
|
|
Raises:
|
|
ConnectionError: If database connection cannot be established
|
|
"""
|
|
Postgres.__init__(
|
|
self,
|
|
host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password,
|
|
dbname=dbname,
|
|
min_connections=min_connections,
|
|
max_connections=max_connections,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
self.info(f'Postgres client initialized at {host}:{port}')
|
|
|
|
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
|
|
|
|
def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Execute an UPDATE SQL statement.
|
|
|
|
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}.
|
|
"""
|
|
|
|
with self.engine.begin() as connection:
|
|
result = connection.execute(text(query), params)
|
|
return {'rowcount': result.rowcount}
|
|
|
|
def _build_status_update_query(
|
|
self, status: str | None, experiment_run_id: int
|
|
) -> tuple[str, dict[str, Any]]:
|
|
"""Build SQL query for simple status update."""
|
|
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,
|
|
}
|
|
|
|
return sql_query, query_params
|
|
|
|
def _build_status_with_error_query(
|
|
self, status: str | None, error_message: str | None, experiment_run_id: int
|
|
) -> tuple[str, dict[str, Any]]:
|
|
"""Build SQL query for status update with error message."""
|
|
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')
|
|
|
|
# Truncate error message if too long
|
|
truncated_error = error_message[:1024] if len(error_message) > 1024 else error_message
|
|
|
|
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': truncated_error,
|
|
'updated_at': datetime.now(UTC),
|
|
'experiment_run_id': experiment_run_id,
|
|
}
|
|
|
|
return sql_query, query_params
|
|
|
|
def _build_model_saved_query(
|
|
self, run_name: str | None, status: str | None, experiment_run_id: int
|
|
) -> tuple[str, dict[str, Any]]:
|
|
"""Build SQL query for model saved update."""
|
|
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,
|
|
}
|
|
|
|
return sql_query, query_params
|
|
|
|
def _get_update_query_and_params(
|
|
self, update_type: str, experiment_run_id: int, input_data: dict[str, Any]
|
|
) -> tuple[str, dict[str, Any]]:
|
|
"""Get SQL query and parameters based on update type."""
|
|
status = input_data.get('status')
|
|
error_message = input_data.get('error_message')
|
|
run_name = input_data.get('run_name')
|
|
|
|
if update_type == UpdateType.STATUS:
|
|
return self._build_status_update_query(status, experiment_run_id)
|
|
|
|
if update_type == UpdateType.STATUS_WITH_ERROR:
|
|
return self._build_status_with_error_query(status, error_message, experiment_run_id)
|
|
|
|
if update_type == UpdateType.MODEL_SAVED:
|
|
return self._build_model_saved_query(run_name, status, experiment_run_id)
|
|
|
|
raise ValueError(f'Invalid update_type: {update_type}')
|
|
|
|
@activity.defn(name='update_experiment_run')
|
|
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')
|
|
|
|
try:
|
|
sql_query, query_params = self._get_update_query_and_params(
|
|
update_type, experiment_run_id, input_data
|
|
)
|
|
|
|
result = self._execute_update(sql_query, query_params)
|
|
|
|
if result.get('rowcount', 0) == 0:
|
|
error_msg = (
|
|
f'No experiment_run row updated for id={experiment_run_id} '
|
|
f'(row missing or id mismatch). update_type={update_type!r}, status={status!r}.'
|
|
)
|
|
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 or {},
|
|
notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
|
|
message=error_msg,
|
|
block='update_experiment_run',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
raise RuntimeError(error_msg) from e
|