Code import - branch release/SIENTIAPDE-1645
This commit is contained in:
280
model_manager/activities/experiment_tracking.py
Normal file
280
model_manager/activities/experiment_tracking.py
Normal file
@@ -0,0 +1,280 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user