cf5111e520
- Added PluginStore integration for model management. - Replaced StorageRepository with MinIORepository in Activities, Cleanup, and Training classes. - Updated training logic to handle validation files and improved data management. - Enhanced configuration for MinIO and PluginStore in connectors. - Removed deprecated model repository and storage repository files. - Updated environment variable handling for new configurations.
347 lines
12 KiB
Python
347 lines
12 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
|
|
|
|
import pandas as pd
|
|
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.observability.sientia_monitoring import SientiaMonitoring
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
|
from sientia_model.model_repository.plugin_store import PluginStore
|
|
|
|
from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
|
|
from model_manager.utils.exceptions import ModelTrainingError
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.repository.data_manager_repository import DataManagerRepository
|
|
|
|
|
|
class Training(SientiaMonitoring):
|
|
"""
|
|
Activity for ML model training operations.
|
|
|
|
This activity extends SientiaMonitoring 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.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
mlflow_repository: SientiaMLflowRepository,
|
|
plugin_store: PluginStore,
|
|
minio_repository: MinioRepository,
|
|
logger: Logger,
|
|
notification_handler: NotificationHandler,
|
|
metrics_controller: MetricsController,
|
|
):
|
|
"""
|
|
Initialize Training activity.
|
|
|
|
Args:
|
|
logger: Logger instance for observability
|
|
notification_handler: Handler for sending notifications
|
|
"""
|
|
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
|
self.data_manager_repository = DataManagerRepository(logger)
|
|
self.mlflow_repository = mlflow_repository
|
|
self.plugin_store = plugin_store
|
|
self.minio_repository = minio_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', {})
|
|
metrics_status = 'success'
|
|
|
|
try:
|
|
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:
|
|
metrics_status = 'error'
|
|
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,
|
|
)
|
|
raise
|
|
finally:
|
|
await self._emit_metrics(
|
|
metadata=metadata,
|
|
metrics_status=metrics_status,
|
|
activity_name='validate_train_params',
|
|
emit_workflow_metric=(metrics_status == 'error'),
|
|
)
|
|
|
|
@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 ML training pipeline:
|
|
1. Validate input parameters.
|
|
2. Prepare data via DataManagerRepository.
|
|
3. Train the model and compute metrics.
|
|
|
|
Args:
|
|
input_data: Training configuration containing:
|
|
- metadata (dict): Workflow execution metadata.
|
|
- uploaded_file (BytesIO): Training data already downloaded from MinIO.
|
|
- train_params (TrainModelParams | dict): Training parameters.
|
|
|
|
Returns:
|
|
dict: Key `run_name` when training and saving succeed.
|
|
|
|
Raises:
|
|
ValueError: If input validation fails.
|
|
Exception: If training fails (after sending notification).
|
|
"""
|
|
metadata = input_data.get('metadata')
|
|
train_params = input_data['train_params']
|
|
|
|
if isinstance(train_params, dict):
|
|
train_params = TrainModelParams.from_dict(train_params)
|
|
|
|
|
|
model_trained = False
|
|
model_saved = False
|
|
metrics_status = 'success'
|
|
|
|
try:
|
|
# Download training file bytes from MinIO
|
|
train_bytes = await self.minio_repository.download_file(
|
|
object_name=train_params.file_name,
|
|
bucket=train_params.bucket_name,
|
|
metadata=metadata,
|
|
)
|
|
|
|
# Download optional validation file bytes from the same bucket
|
|
val_bytes: bytes | None = None
|
|
validation_name = getattr(train_params, 'validation_file_name', None)
|
|
if validation_name is not None:
|
|
val_bytes = await self.minio_repository.download_file(
|
|
object_name=validation_name,
|
|
bucket=train_params.bucket_name,
|
|
metadata=metadata,
|
|
)
|
|
|
|
train_result = self.data_manager_repository.prepare_training_data(
|
|
train_file_bytes=train_bytes,
|
|
validation_file_bytes=val_bytes,
|
|
params=train_params,
|
|
metadata=metadata,
|
|
)
|
|
|
|
wrapper = await self.plugin_store.get_model(
|
|
model_name=train_params.model_name,
|
|
force_download=False,
|
|
opt_params={},
|
|
model_kwargs={},
|
|
data_model_kwargs={},
|
|
metadata=metadata,
|
|
)
|
|
|
|
train_df = pd.concat([train_result.x_train, train_result.y_train], axis=1)
|
|
val_df = pd.concat([train_result.x_test, train_result.y_test], axis=1)
|
|
|
|
wrapper.train(
|
|
train_data=train_df,
|
|
val_data=val_df,
|
|
target=train_params.target_variable,
|
|
)
|
|
|
|
# Generate predictions using the trained wrapper
|
|
transformed_train, _ = wrapper.transform(train_df)
|
|
transformed_val, _ = wrapper.transform(val_df)
|
|
|
|
y_train_pred_df, _ = wrapper.predict({}, transformed_train)
|
|
y_val_pred_df, _ = wrapper.predict({}, transformed_val)
|
|
|
|
# Use the first column of the prediction DataFrame as the target prediction
|
|
train_result.y_train_pred = y_train_pred_df.iloc[:, 0]
|
|
train_result.y_pred = y_val_pred_df.iloc[:, 0]
|
|
|
|
train_result = self.data_manager_repository.compute_regression_metrics(
|
|
train_params,
|
|
train_result,
|
|
)
|
|
|
|
model_trained = True
|
|
|
|
async with self.mlflow_repository.start_run(
|
|
model_name=train_params.model_name,
|
|
run_name=None,
|
|
experiment_name=train_params.experiment_name,
|
|
tags=None,
|
|
metadata=metadata,
|
|
) as run_info:
|
|
wrapper.store_model(name=train_params.model_name)
|
|
|
|
model_saved = True
|
|
|
|
return {
|
|
'run_name': run_info.run_name or run_info.run_id,
|
|
}
|
|
except Exception as e: # noqa: BLE001
|
|
metrics_status = 'error'
|
|
|
|
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,
|
|
)
|
|
|
|
raise ModelTrainingError(
|
|
model_trained=model_trained,
|
|
model_saved=model_saved,
|
|
) from e
|
|
finally:
|
|
await self._emit_metrics(
|
|
metadata=metadata,
|
|
metrics_status=metrics_status,
|
|
activity_name='train_model',
|
|
emit_workflow_metric=(metrics_status == 'error'),
|
|
)
|
|
|
|
@activity.defn(name='cleanup_resources')
|
|
async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
|
|
"""
|
|
Cleanup temporary resources created during training.
|
|
|
|
Args:
|
|
input_data: Cleanup configuration containing:
|
|
- metadata (dict): Workflow execution metadata.
|
|
- bucket_name (str): MinIO bucket of the uploaded file.
|
|
- file_name (str): MinIO object key to delete.
|
|
|
|
Raises:
|
|
Exception: If cleanup fails (after sending notification).
|
|
"""
|
|
metadata = input_data.get('metadata', {})
|
|
bucket_name = input_data.get('bucket_name', '')
|
|
file_name = input_data.get('file_name', '')
|
|
metrics_status = 'success'
|
|
|
|
try:
|
|
await self.minio_repository.delete_file(
|
|
object_name=file_name,
|
|
bucket=bucket_name,
|
|
metadata=metadata,
|
|
)
|
|
except Exception as e: # noqa: BLE001
|
|
metrics_status = 'error'
|
|
|
|
error_msg = (
|
|
'Error cleaning up resources - '
|
|
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,
|
|
)
|
|
|
|
raise
|
|
finally:
|
|
await self._emit_metrics(
|
|
metadata=metadata,
|
|
metrics_status=metrics_status,
|
|
activity_name='cleanup_resources',
|
|
emit_workflow_metric=True,
|
|
)
|
|
|
|
async def _emit_metrics(
|
|
self,
|
|
metadata: dict[str, Any],
|
|
metrics_status: str,
|
|
activity_name: str,
|
|
emit_workflow_metric: bool,
|
|
) -> None:
|
|
"""
|
|
Emit workflow and activity execution metrics.
|
|
|
|
Args:
|
|
metadata: Activity metadata containing pod_id and workflow_name
|
|
metrics_status: Execution status ('success' or 'error')
|
|
activity_name: Name of the activity being executed
|
|
"""
|
|
if emit_workflow_metric:
|
|
await self.emit_metric(
|
|
metric_object=WORKFLOW_EXECUTION_TOTAL,
|
|
tags={
|
|
'pod_id': metadata.get('pod_id'),
|
|
'workflow_name': metadata.get('workflow_name'),
|
|
'status': metrics_status,
|
|
},
|
|
)
|
|
|
|
await self.emit_metric(
|
|
metric_object=ACTIVITY_EXECUTION_TOTAL,
|
|
tags={
|
|
'pod_id': metadata.get('pod_id'),
|
|
'activity_name': activity_name,
|
|
'status': metrics_status,
|
|
},
|
|
)
|