218 lines
8.1 KiB
Python
218 lines
8.1 KiB
Python
"""
|
|
Training activities for ML model training operations.
|
|
|
|
This module provides activities for training machine learning models.
|
|
The activity extends BaseActivity and receives pre-downloaded files
|
|
to return success/failure status without raising exceptions.
|
|
"""
|
|
|
|
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
import traceback
|
|
from io import BytesIO
|
|
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.base import BaseActivity
|
|
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
from model_manager.utils.repository.training_repository import TrainingRepository
|
|
|
|
|
|
class Training(BaseActivity):
|
|
"""
|
|
Activity for ML model training operations.
|
|
|
|
This activity extends BaseActivity and handles machine learning model
|
|
training with comprehensive error handling. It receives pre-downloaded
|
|
files from the workflow and returns success/failure status without
|
|
raising exceptions.
|
|
|
|
Attributes:
|
|
logger (Logger): Logger instance for observability (inherited from BaseActivity)
|
|
notification_handler (NotificationHandler): Handler for sending notifications (inherited)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
logger: Logger,
|
|
notification_handler: NotificationHandler,
|
|
):
|
|
"""
|
|
Initialize Training activity.
|
|
|
|
Args:
|
|
logger: Logger instance for observability
|
|
notification_handler: Handler for sending notifications
|
|
"""
|
|
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
|
|
self.training_repository = TrainingRepository(logger)
|
|
|
|
@activity.defn(name='validate_train_params')
|
|
async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
|
|
"""
|
|
Validate and convert training parameters from dict to TrainModelParams.
|
|
|
|
This activity validates the input training parameters and converts them
|
|
to a TrainModelParams object.
|
|
|
|
Args:
|
|
input_data: Training parameters and metadata at the same level
|
|
Required keys:
|
|
- metadata (dict): Workflow execution metadata
|
|
- All TrainModelParams fields (experiment_run_id, target_variable, etc.)
|
|
|
|
Returns:
|
|
TrainModelParams: Validated and converted training parameters
|
|
|
|
Raises:
|
|
ValueError, TypeError, KeyError: If validation fails (after sending notification)
|
|
|
|
Example:
|
|
result = await validate_train_params({
|
|
'metadata': {'workflow_id': 'train-123'},
|
|
'experiment_run_id': 456,
|
|
'target_variable': 'price',
|
|
'variable_columns': ['feature1', 'feature2'],
|
|
'train_size': 80,
|
|
# ... other required fields at same level
|
|
})
|
|
# Returns: TrainModelParams(...)
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
|
|
try:
|
|
self.info('Validating training parameters', metadata)
|
|
|
|
# Step 1: Convert input_data to TrainModelParams (validates types and required fields)
|
|
train_params = TrainModelParams.from_dict(input_data)
|
|
|
|
# Step 2: Validate business rules (ranges, consistency, etc.)
|
|
train_params.validate_business_rules()
|
|
|
|
self.info(
|
|
f'Training parameters validated successfully - '
|
|
f'Target: {train_params.target_variable}, '
|
|
f'Experiment: {train_params.experiment_name}',
|
|
metadata,
|
|
)
|
|
|
|
return train_params
|
|
|
|
except (ValueError, TypeError, KeyError) as e:
|
|
error_msg = f'Error validating training parameters: {str(e)}'
|
|
trace = traceback.format_exc()
|
|
|
|
# Send notification (MongoDB)
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
|
|
message=error_msg,
|
|
block='validate_train_params',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
|
|
# Log error with metadata
|
|
self.error(trace, metadata=metadata)
|
|
|
|
# Re-raise exception to stop workflow
|
|
raise
|
|
|
|
@activity.defn(name='train_model')
|
|
async def train_model(self, input_data: dict[str, Any]) -> TrainModelResult:
|
|
"""
|
|
Train a machine learning model.
|
|
|
|
This activity orchestrates the complete ML training pipeline:
|
|
1. Validates input parameters
|
|
2. Trains the model using TrainingRepository
|
|
3. Performs post-training calculations
|
|
|
|
Args:
|
|
input_data: Configuration for model training operation
|
|
Required keys:
|
|
- metadata (dict): Workflow execution metadata
|
|
- uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
|
|
- train_params (TrainModelParams): Training parameters object
|
|
|
|
Returns:
|
|
TrainModelResult: Training result with model, metrics, and data
|
|
|
|
Raises:
|
|
ValueError: If input validation fails
|
|
Exception: If training fails (after sending notification)
|
|
|
|
Example:
|
|
result = await train_model({
|
|
'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
|
|
'uploaded_file': BytesIO(csv_data),
|
|
'train_params': TrainModelParams(...)
|
|
})
|
|
# Returns: TrainModelResult(...)
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
uploaded_file = input_data['uploaded_file']
|
|
train_params = input_data['train_params']
|
|
|
|
try:
|
|
# Validate uploaded_file is BytesIO
|
|
if not isinstance(uploaded_file, BytesIO):
|
|
raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')
|
|
|
|
# Validate train_params is TrainModelParams
|
|
if not isinstance(train_params, TrainModelParams):
|
|
raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}')
|
|
|
|
self.info(
|
|
f'Starting model training for target: {train_params.target_variable}',
|
|
metadata,
|
|
)
|
|
|
|
# Step 1: Train the model
|
|
self.info('Training model with TrainingRepository', metadata)
|
|
train_result = self.training_repository.train(uploaded_file, train_params)
|
|
|
|
# Step 2: Perform post-training calculations
|
|
self.info('Performing post-training calculations', metadata)
|
|
final_result = self.training_repository.after_train_calculation(
|
|
train_params, train_result
|
|
)
|
|
|
|
self.info(
|
|
f'Model training completed successfully - '
|
|
f'MSE: {final_result.mse_val}, MAE: {final_result.mae_val}, R²: {final_result.r2_val}',
|
|
metadata,
|
|
)
|
|
|
|
return final_result
|
|
|
|
except Exception as e: # noqa: BLE001
|
|
target = (
|
|
train_params.target_variable
|
|
if hasattr(train_params, 'target_variable')
|
|
else 'unknown'
|
|
)
|
|
error_msg = f'Error training model - Target: {target}, Error: {str(e)}'
|
|
trace = traceback.format_exc()
|
|
|
|
# Send notification (MongoDB)
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='TRAIN_MODEL_ERROR',
|
|
message=error_msg,
|
|
block='train_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
|
|
# Log error with metadata
|
|
self.error(trace, metadata=metadata)
|
|
|
|
# Re-raise exception to stop workflow
|
|
raise
|