225 lines
8.2 KiB
Python
225 lines
8.2 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 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.exceptions import ModelTrainingError
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.repository.model_repository import ModelRepository
|
|
from model_manager.utils.repository.storage_repository import StorageRepository
|
|
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,
|
|
model_repository: ModelRepository,
|
|
storage_repository: StorageRepository,
|
|
logger: Logger,
|
|
notification_handler: NotificationHandler,
|
|
):
|
|
"""
|
|
Initialize Training activity.
|
|
|
|
Args:
|
|
logger: Logger instance for observability
|
|
notification_handler: Handler for sending notifications
|
|
"""
|
|
super().__init__(logger, notification_handler, set_error_counter=True)
|
|
self.training_repository = TrainingRepository(logger)
|
|
self.model_repository = model_repository
|
|
self.storage_repository = storage_repository
|
|
|
|
@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)
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
|
|
try:
|
|
self.info('Validating training parameters', metadata)
|
|
train_params = TrainModelParams.from_dict(input_data)
|
|
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()
|
|
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
|
|
message=error_msg,
|
|
block='validate_train_params',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
|
|
self.error(trace, metadata=metadata)
|
|
raise
|
|
|
|
@activity.defn(name='train_model')
|
|
async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]:
|
|
"""
|
|
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', {})
|
|
train_params = input_data['train_params']
|
|
|
|
if isinstance(train_params, dict):
|
|
train_params = TrainModelParams.from_dict(train_params)
|
|
# type: ignore[assignment]
|
|
|
|
model_trained = False
|
|
model_saved = False
|
|
|
|
try:
|
|
with self.storage_repository.fetch_file(
|
|
train_params.bucket_name, train_params.file_name
|
|
) as uploaded_file:
|
|
train_result = self.training_repository.train(uploaded_file, train_params)
|
|
|
|
train_result = self.training_repository.after_train_calculation(
|
|
train_params, train_result
|
|
)
|
|
|
|
model_trained = True
|
|
train_result = self.model_repository.save_model(train_result)
|
|
model_saved = True
|
|
|
|
return {
|
|
'run_name': train_result.run_name,
|
|
'run_dir': train_result.run_dir,
|
|
}
|
|
except Exception as e: # noqa: BLE001
|
|
error_msg = (
|
|
'Error training model - '
|
|
f'model_trained={model_trained}, model_saved={model_saved}, '
|
|
f'error: {str(e)}'
|
|
)
|
|
|
|
trace = traceback.format_exc()
|
|
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='TRAIN_MODEL_ERROR',
|
|
message=error_msg,
|
|
block='train_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
|
|
self.error(trace, metadata=metadata)
|
|
|
|
raise ModelTrainingError(
|
|
model_trained=model_trained,
|
|
model_saved=model_saved,
|
|
) from e
|
|
|
|
@activity.defn(name='cleanup_resources')
|
|
async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
|
|
metadata = input_data.get('metadata', {})
|
|
run_dir = input_data.get('run_dir', '')
|
|
bucket_name = input_data.get('bucket_name', '')
|
|
file_name = input_data.get('file_name', '')
|
|
|
|
try:
|
|
self.model_repository.cleanup_run_directory(run_dir)
|
|
self.storage_repository.delete_file(bucket_name, file_name)
|
|
except Exception as e: # noqa: BLE001
|
|
error_msg = (
|
|
f'Error cleaning up resources - Run directory: {run_dir}, '
|
|
f'File: {bucket_name}/{file_name}, Error: {str(e)}'
|
|
)
|
|
|
|
trace = traceback.format_exc()
|
|
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='CLEANUP_RESOURCES_ERROR',
|
|
message=error_msg,
|
|
block='cleanup_resources',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
|
|
self.error(trace, metadata=metadata)
|
|
raise
|