SIENTIAPDE-1250: Implement experiment tracking activity with status updates, error handling, and model registration. This includes a unified update method, error message truncation, and integration into the main activities orchestrator. (236+, 2-)

This commit is contained in:
Bruno Domingues
2025-10-07 11:27:23 -03:00
parent 2628dc92d2
commit 7e1eace77f
5 changed files with 628 additions and 14 deletions

View File

@@ -149,13 +149,20 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
#### **Activities (`model_manager/activities/`)** #### **Activities (`model_manager/activities/`)**
- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance - **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 - **Gates**: Data quality validation and filtering mechanisms
- **MLFlow**: Model transformation and prediction operations - **MLFlow**: Model transformation and prediction operations
- **MinIO**: Object storage operations for file management
- **Key Features**: - **Key Features**:
- Multiple inheritance pattern for unified activity interface - Multiple inheritance pattern for unified activity interface
- Configurable filter policies and validation rules - Configurable filter policies and validation rules
- MLFlow model serving integration with configurable flavors - MLFlow model serving integration with configurable flavors
- Comprehensive error handling and notification integration - Comprehensive error handling and notification integration
- Experiment tracking with automatic status management
#### **Data Services (`model_manager/utils/`)** #### **Data Services (`model_manager/utils/`)**
- **Connectors Config**: Environment variable-based configuration management - **Connectors Config**: Environment variable-based configuration management

View File

@@ -5,14 +5,14 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger 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.gates import Gates
from model_manager.activities.minio import MinIO from model_manager.activities.minio import MinIO
from model_manager.activities.mlflow import MLFlow 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. 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. MLFlow model interactions, MinIO storage operations, and data quality validation.
The class implements multiple inheritance to combine specialized functionality: 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 - MLFlow: Model inference and transformation operations
- MinIO: Object storage operations (file upload/download/delete) - MinIO: Object storage operations (file upload/download/delete)
- Gates: Data quality validation and filtering mechanisms - Gates: Data quality validation and filtering mechanisms
@@ -62,7 +62,7 @@ class Activities(Postgres, MLFlow, MinIO, Gates):
Exception: If any parent class initialization fails Exception: If any parent class initialization fails
""" """
# Initialize parent classes # Initialize parent classes
Postgres.__init__( ExperimentTracking.__init__(
self, self,
host=postgres_config['host'], host=postgres_config['host'],
port=postgres_config['port'], port=postgres_config['port'],
@@ -107,10 +107,10 @@ class Activities(Postgres, MLFlow, MinIO, Gates):
Gracefully shutdown all activities and clean up resources. Gracefully shutdown all activities and clean up resources.
This method ensures proper cleanup of all resources including: This method ensures proper cleanup of all resources including:
- PostgreSQL connection pools - PostgreSQL connection pools (via ExperimentTracking)
- Any other resources that need explicit cleanup - Any other resources that need explicit cleanup
The method should be called before the application terminates to ensure The method should be called before the application terminates to ensure
proper resource cleanup and prevent resource leaks. proper resource cleanup and prevent resource leaks.
""" """
Postgres.close(self) ExperimentTracking.close(self)

View File

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

View File

@@ -1,18 +1,20 @@
from unittest.mock import ANY, MagicMock, patch from unittest.mock import ANY, MagicMock, patch
from pytest import mark from pytest import mark
from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.activities import Activities 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.gates import Gates
from model_manager.activities.mlflow import MLFlow 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.MLFlow.__init__')
@patch('model_manager.activities.activities.MinIO.__init__') @patch('model_manager.activities.activities.MinIO.__init__')
@patch('model_manager.activities.activities.Gates.__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 = { postgres_config = {
'host': 'localhost', 'host': 'localhost',
'port': 5432, '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, Activities)
assert isinstance(activities, Postgres) assert isinstance(activities, ExperimentTracking)
assert isinstance(activities, MLFlow) assert isinstance(activities, MLFlow)
assert isinstance(activities, Gates) assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with( mock_experiment_tracking_init.assert_called_once_with(
ANY, ANY,
host=postgres_config['host'], host=postgres_config['host'],
port=postgres_config['port'], port=postgres_config['port'],
@@ -97,9 +99,9 @@ def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postg
@mark.asyncio @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()) @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 = { postgres_config = {
'host': 'localhost', 'host': 'localhost',
'port': 5432, 'port': 5432,
@@ -136,4 +138,4 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
) )
await activities.shutdown() await activities.shutdown()
mock_postgres_init.close.assert_called_once() mock_experiment_tracking_init.close.assert_called_once()

View File

@@ -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'