SIENTIAPDE-1251: Implement ML model training activity and repository
This commit introduces the 'Training' activity and 'TrainingRepository' for handling ML model training operations within the Model Manager system. - Added model_manager/activities/training.py for the Training activity, which extends BaseActivity and integrates with Temporal workflows. - Added model_manager/utils/repository/training_repository.py for the TrainingRepository, which encapsulates the core training logic. - Updated model_manager/activities/activities.py to include the Training activity in the main activities orchestrator. - Updated README.md to document the new 'Training' component. - Added unit tests for the new activity and repository.
This commit is contained in:
Bruno Domingues
@@ -10,9 +10,10 @@ with workflow.unsafe.imports_passed_through():
|
||||
from model_manager.activities.gates import Gates
|
||||
from model_manager.activities.minio import MinIO
|
||||
from model_manager.activities.mlflow import MLFlow
|
||||
from model_manager.activities.training import Training
|
||||
|
||||
|
||||
class Activities(ExperimentTracking, MLFlow, MinIO, Gates):
|
||||
class Activities(ExperimentTracking, MLFlow, MinIO, Gates, Training):
|
||||
"""
|
||||
Main activities orchestrator for the Model Manager system.
|
||||
|
||||
@@ -25,6 +26,7 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Gates):
|
||||
- MLFlow: Model inference and transformation operations
|
||||
- MinIO: Object storage operations (file upload/download/delete)
|
||||
- Gates: Data quality validation and filtering mechanisms
|
||||
- Training: ML model training operations (extends BaseActivity)
|
||||
|
||||
Attributes:
|
||||
postgres_config (dict): PostgreSQL connection configuration
|
||||
@@ -102,6 +104,8 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Gates):
|
||||
|
||||
Gates.__init__(self, logger=logger, notification_handler=notification_handler)
|
||||
|
||||
Training.__init__(self, logger=logger, notification_handler=notification_handler)
|
||||
|
||||
async def shutdown(self):
|
||||
"""
|
||||
Gracefully shutdown all activities and clean up resources.
|
||||
|
||||
166
model_manager/activities/training.py
Normal file
166
model_manager/activities/training.py
Normal file
@@ -0,0 +1,166 @@
|
||||
"""
|
||||
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.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='train_model')
|
||||
async def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Train a machine learning model with comprehensive error handling.
|
||||
|
||||
This activity orchestrates the complete ML training pipeline:
|
||||
1. Validates input parameters
|
||||
2. Trains the model using TrainingRepository
|
||||
3. Performs post-training calculations
|
||||
4. Returns success/failure status with results or error message
|
||||
|
||||
The activity does NOT raise exceptions on failure - it catches all errors,
|
||||
sends notifications, and returns a failure status. This allows the workflow
|
||||
to handle the error gracefully and update the database accordingly.
|
||||
|
||||
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 (dict): Training parameters (converted to TrainModelParams)
|
||||
|
||||
Returns:
|
||||
dict: Training result with the following structure:
|
||||
{
|
||||
'success': bool, # True if training succeeded, False otherwise
|
||||
'result': TrainModelResult | None, # Training result if success=True
|
||||
'error_message': str | None # Error message if success=False
|
||||
}
|
||||
|
||||
Example:
|
||||
# Successful training
|
||||
result = await train_model({
|
||||
'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
|
||||
'uploaded_file': BytesIO(csv_data),
|
||||
'train_params': {
|
||||
'experiment_run_id': 456,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1', 'feature2'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': True,
|
||||
# ... other TrainModelParams fields
|
||||
}
|
||||
})
|
||||
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
|
||||
|
||||
# Failed training
|
||||
# Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
uploaded_file = input_data['uploaded_file']
|
||||
train_params_dict = input_data['train_params']
|
||||
|
||||
try:
|
||||
self.info(
|
||||
f'Starting model training for target: {train_params_dict.get("target_variable")}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Convert dict to TrainModelParams
|
||||
train_params = TrainModelParams.from_dict(train_params_dict)
|
||||
|
||||
# Validate uploaded_file is BytesIO
|
||||
if not isinstance(uploaded_file, BytesIO):
|
||||
raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')
|
||||
|
||||
# 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 {
|
||||
'success': True,
|
||||
'result': final_result,
|
||||
'error_message': None,
|
||||
}
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Error training model - Target: {train_params_dict.get("target_variable", "unknown")}, 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)
|
||||
|
||||
# Return failure result (do NOT raise exception)
|
||||
# This allows workflow to update database with error status
|
||||
return {
|
||||
'success': False,
|
||||
'result': None,
|
||||
'error_message': str(e),
|
||||
}
|
||||
242
model_manager/utils/repository/training_repository.py
Normal file
242
model_manager/utils/repository/training_repository.py
Normal file
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
Training repository for ML model training operations.
|
||||
|
||||
This module provides the core training logic for machine learning models,
|
||||
including data preprocessing, model training, and post-training calculations.
|
||||
Migrated from laborious/utils/train_model_utils.py.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sientia.linear_models import LinearRegressionModel
|
||||
from sientia.metrics import mae, mse, r2
|
||||
from sientia.preprocessing import DataPreprocessor
|
||||
from sientia.utils import split_train_test
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.operations.df_preprocessor import load_data
|
||||
from sientia_do.operations.normalization import MinMaxScaler, Z_Scaler
|
||||
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
class TrainingRepository:
|
||||
"""
|
||||
Repository for machine learning model training operations.
|
||||
|
||||
This class encapsulates the core logic for training ML models, migrated from
|
||||
laborious/utils/train_model_utils.py. Follows the same pattern as MLFlowRepository
|
||||
with instance methods and logger integration.
|
||||
|
||||
Attributes:
|
||||
logger (Logger): Logger instance for observability and debugging
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Logger):
|
||||
"""
|
||||
Initialize TrainingRepository with logger.
|
||||
|
||||
Args:
|
||||
logger: Logger instance for observability
|
||||
"""
|
||||
self.logger = logger
|
||||
|
||||
def train(self, uploaded_file: BytesIO, params: TrainModelParams) -> TrainModelResult:
|
||||
"""
|
||||
Train a machine learning model using the provided file and parameters.
|
||||
|
||||
This method orchestrates the training pipeline:
|
||||
1. Load data from BytesIO file
|
||||
2. Initialize and fit data preprocessor
|
||||
3. Transform data and validate
|
||||
4. Split into train/test sets
|
||||
5. Initialize scaler dictionary
|
||||
6. Train LinearRegression model
|
||||
|
||||
Args:
|
||||
uploaded_file: BytesIO object containing training data (CSV format)
|
||||
params: Training parameters (TrainModelParams)
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Object containing trained model, processed data,
|
||||
train/test splits, and scaler dictionary
|
||||
|
||||
Raises:
|
||||
ValueError: If transformed data is empty
|
||||
Exception: If data loading, preprocessing, or training fails
|
||||
"""
|
||||
# Load data from BytesIO
|
||||
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
|
||||
|
||||
# Initialize and fit data preprocessor
|
||||
process_data = self.init_data_preprocessor(params)
|
||||
process_data.fit(data)
|
||||
data_view = process_data.transform(data)
|
||||
|
||||
# Validate transformed data
|
||||
if len(data_view) <= 0:
|
||||
raise ValueError('Data view is empty after transformation')
|
||||
|
||||
# Split data into train/test sets
|
||||
x_train, x_test, y_train, y_test = split_train_test(
|
||||
data_view[params.variable_columns],
|
||||
data_view[params.target_variable],
|
||||
train_size=params.train_size / 100,
|
||||
shuffle=params.shuffle,
|
||||
random_state=42,
|
||||
)
|
||||
|
||||
# Prepare training data
|
||||
data_train = pd.concat([x_train, y_train], axis=1)
|
||||
scaler_dict = self.init_scaler_dict(process_data, params)
|
||||
|
||||
# Create and train linear regression model
|
||||
regr = LinearRegressionModel(
|
||||
target_variable=params.target_variable,
|
||||
variable_columns=params.variable_columns,
|
||||
)
|
||||
regr.fit(data_train)
|
||||
|
||||
# Return training result
|
||||
return TrainModelResult(
|
||||
params=params,
|
||||
process_data=process_data,
|
||||
x_train=x_train,
|
||||
x_test=x_test,
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=regr,
|
||||
scaler_dict=scaler_dict,
|
||||
)
|
||||
|
||||
def init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
|
||||
"""
|
||||
Initialize dictionary containing scaling parameters for features and target.
|
||||
|
||||
This method extracts scaling parameters from the fitted scaler to enable
|
||||
denormalization of predictions and debugging of the normalization process.
|
||||
|
||||
Args:
|
||||
process_data: Fitted DataPreprocessor object with scaler
|
||||
params: Training parameters including scaler configuration
|
||||
|
||||
Returns:
|
||||
dict: Scaling parameters for each feature and target variable.
|
||||
Structure depends on scaler type:
|
||||
- MinMaxScaler: {'feature': {'min': float, 'max': float}, ...}
|
||||
- Z_Scaler: Dictionary from scaler.create_dict()
|
||||
- Empty dict: If no scaler is used
|
||||
|
||||
Raises:
|
||||
AttributeError: If scaler doesn't have expected attributes
|
||||
"""
|
||||
scaler_dict = {}
|
||||
|
||||
if params.use_scaler:
|
||||
scaler = process_data.get_scaler()
|
||||
|
||||
if isinstance(scaler, MinMaxScaler):
|
||||
# Extract min/max for each feature
|
||||
for i, col in enumerate(params.variable_columns):
|
||||
scaler_dict[col] = {'min': scaler.x_min[i], 'max': scaler.x_max[i]}
|
||||
|
||||
# Extract min/max for target variable
|
||||
scaler_dict[params.target_variable] = {
|
||||
'min': scaler.y_min,
|
||||
'max': scaler.y_max,
|
||||
}
|
||||
|
||||
elif isinstance(scaler, Z_Scaler):
|
||||
scaler_dict = scaler.create_dict()
|
||||
|
||||
return scaler_dict
|
||||
|
||||
def after_train_calculation(
|
||||
self, params: TrainModelParams, tmr: TrainModelResult
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Perform post-training calculations: predictions, denormalization, and metrics.
|
||||
|
||||
This method completes the training pipeline by:
|
||||
1. Making predictions on test set
|
||||
2. Denormalizing all data (if scaler was used)
|
||||
3. Reordering data by index
|
||||
4. Calculating evaluation metrics (MSE, MAE, R²)
|
||||
|
||||
Args:
|
||||
params: Training parameters used during model training
|
||||
tmr: Result object from training
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated result with predictions, denormalized data,
|
||||
and metrics (mse_val, mae_val, r2_val)
|
||||
"""
|
||||
# Make predictions on test set
|
||||
tmr.y_pred = tmr.regr.predict(tmr.x_test)
|
||||
|
||||
# Denormalize data if scaler was used
|
||||
if params.use_scaler:
|
||||
scaler = tmr.process_data.get_scaler()
|
||||
|
||||
# Denormalize features
|
||||
for col in params.variable_columns:
|
||||
tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
|
||||
tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
|
||||
|
||||
# Denormalize target variable
|
||||
tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
|
||||
tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
|
||||
tmr.y_pred = scaler.denormalize_predictions(tmr.y_pred, params.target_variable)
|
||||
|
||||
# Add index to predictions
|
||||
tmr.y_pred = pd.Series(tmr.y_pred, index=tmr.y_test.index)
|
||||
tmr.y_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
# Reorder all data by index
|
||||
tmr.x_train = tmr.x_train.sort_index()
|
||||
tmr.x_test = tmr.x_test.sort_index()
|
||||
tmr.y_train = tmr.y_train.sort_index()
|
||||
tmr.y_test = tmr.y_test.sort_index()
|
||||
tmr.y_pred = tmr.y_pred.sort_index()
|
||||
|
||||
# Calculate evaluation metrics
|
||||
tmr.mse_val = round(
|
||||
mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
tmr.mae_val = round(
|
||||
mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
|
||||
|
||||
return tmr
|
||||
|
||||
def init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
|
||||
"""
|
||||
Initialize DataPreprocessor with training parameters.
|
||||
|
||||
Args:
|
||||
params: Training parameters containing preprocessor configuration
|
||||
|
||||
Returns:
|
||||
DataPreprocessor: Configured preprocessor ready for fitting
|
||||
"""
|
||||
# Create lag dictionaries for each variable
|
||||
lag_train_dict = dict.fromkeys(params.variable_columns, params.lag_train)
|
||||
lag_val_dict = dict.fromkeys(params.variable_columns, params.lag_val)
|
||||
|
||||
return DataPreprocessor(
|
||||
target_variable=params.target_variable,
|
||||
input_columns=params.variable_columns,
|
||||
lag_train=lag_train_dict,
|
||||
lag_transform=lag_val_dict,
|
||||
static_threshold=1 if params.rem_static_win else None,
|
||||
low_lim=params.low_lim,
|
||||
upp_lim=params.upp_lim,
|
||||
window=params.window,
|
||||
scaler_name='Standard Scaler' if params.use_scaler else 'None',
|
||||
ar_var=params.target_variable if params.include_ar else None,
|
||||
)
|
||||
Reference in New Issue
Block a user