feat: integrate PluginStore and MinIO repository into model manager activities
- 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.
This commit is contained in:
vitor-aignosi
@@ -12,18 +12,20 @@ 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.model_repository import ModelRepository
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
from model_manager.utils.repository.training_repository import TrainingRepository
|
||||
from model_manager.utils.repository.data_manager_repository import DataManagerRepository
|
||||
|
||||
|
||||
class Training(SientiaMonitoring):
|
||||
@@ -38,8 +40,9 @@ class Training(SientiaMonitoring):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_repository: ModelRepository,
|
||||
storage_repository: StorageRepository,
|
||||
mlflow_repository: SientiaMLflowRepository,
|
||||
plugin_store: PluginStore,
|
||||
minio_repository: MinioRepository,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
@@ -52,9 +55,10 @@ class Training(SientiaMonitoring):
|
||||
notification_handler: Handler for sending notifications
|
||||
"""
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
self.training_repository = TrainingRepository(logger)
|
||||
self.model_repository = model_repository
|
||||
self.storage_repository = storage_repository
|
||||
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:
|
||||
@@ -120,8 +124,8 @@ class Training(SientiaMonitoring):
|
||||
|
||||
This activity orchestrates the ML training pipeline:
|
||||
1. Validate input parameters.
|
||||
2. Train the model via TrainingRepository.
|
||||
3. Perform post-training calculations.
|
||||
2. Prepare data via DataManagerRepository.
|
||||
3. Train the model and compute metrics.
|
||||
|
||||
Args:
|
||||
input_data: Training configuration containing:
|
||||
@@ -130,41 +134,98 @@ class Training(SientiaMonitoring):
|
||||
- train_params (TrainModelParams | dict): Training parameters.
|
||||
|
||||
Returns:
|
||||
dict: Keys `run_name` and `run_dir` when training and saving succeed.
|
||||
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', {})
|
||||
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
|
||||
metrics_status = 'success'
|
||||
|
||||
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)
|
||||
# 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,
|
||||
)
|
||||
|
||||
train_result = self.training_repository.after_train_calculation(
|
||||
train_params, train_result
|
||||
# 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,
|
||||
)
|
||||
|
||||
model_trained = True
|
||||
train_result = self.model_repository.save_model(train_result)
|
||||
model_saved = True
|
||||
train_result = self.data_manager_repository.prepare_training_data(
|
||||
train_file_bytes=train_bytes,
|
||||
validation_file_bytes=val_bytes,
|
||||
params=train_params,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
return {
|
||||
'run_name': train_result.run_name,
|
||||
'run_dir': train_result.run_dir,
|
||||
}
|
||||
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'
|
||||
|
||||
@@ -205,7 +266,6 @@ class Training(SientiaMonitoring):
|
||||
Args:
|
||||
input_data: Cleanup configuration containing:
|
||||
- metadata (dict): Workflow execution metadata.
|
||||
- run_dir (str): Temporary directory to remove.
|
||||
- bucket_name (str): MinIO bucket of the uploaded file.
|
||||
- file_name (str): MinIO object key to delete.
|
||||
|
||||
@@ -213,19 +273,21 @@ class Training(SientiaMonitoring):
|
||||
Exception: If cleanup fails (after sending notification).
|
||||
"""
|
||||
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', '')
|
||||
metrics_status = 'success'
|
||||
|
||||
try:
|
||||
self.model_repository.cleanup_run_directory(run_dir)
|
||||
self.storage_repository.delete_file(bucket_name, file_name)
|
||||
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 = (
|
||||
f'Error cleaning up resources - Run directory: {run_dir}, '
|
||||
'Error cleaning up resources - '
|
||||
f'File: {bucket_name}/{file_name}, Error: {str(e)}'
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user