feat: update training workflow and repository management
- Replaced synchronous MinIO repository calls with asynchronous counterparts in the Training class for improved performance. - Enhanced logging throughout the training process to provide better insights into model metadata loading, parameter validation, and training execution. - Updated the train_test_split function to enforce DataFrame input type, ensuring consistency in data handling. - Removed the deprecated model_repository.py file to streamline the codebase. - Adjusted cleanup schedule logic to improve error handling and logging during schedule reconciliation. - Updated tests to reflect changes in the training workflow and repository interactions.
This commit is contained in:
vitor-aignosi
@@ -18,7 +18,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
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_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
@@ -78,14 +78,18 @@ class Training(SientiaMonitoring):
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info(f'Loading model metadata for {input_data}', metadata)
|
||||
|
||||
try:
|
||||
train_params = TrainModelParams.from_dict(input_data)
|
||||
model_metadata = self.plugin_store.get_model_index(
|
||||
model_name=train_params.model_name,
|
||||
model_type=train_params.model_type,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
train_params.model_metadata = model_metadata
|
||||
self.info(f'Model metadata loaded successfully for {input_data}', metadata)
|
||||
self.debug(f'Model metadata: {model_metadata}', metadata)
|
||||
return train_params.to_dict()
|
||||
except Exception as exc:
|
||||
trace = traceback.format_exc()
|
||||
@@ -120,6 +124,9 @@ class Training(SientiaMonitoring):
|
||||
Exception: If validation fails (after sending notification)
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info(f'Validating training parameters for {input_data}', metadata)
|
||||
|
||||
try:
|
||||
train_params = TrainModelParams.from_dict(input_data)
|
||||
|
||||
@@ -132,6 +139,10 @@ class Training(SientiaMonitoring):
|
||||
metadata,
|
||||
)
|
||||
|
||||
self.debug(
|
||||
f'Training parameters validated successfully: {train_params.to_dict()}', metadata
|
||||
)
|
||||
|
||||
return train_params.to_dict()
|
||||
except Exception as e:
|
||||
error_msg = f'Error validating training parameters: {str(e)}'
|
||||
@@ -173,9 +184,15 @@ class Training(SientiaMonitoring):
|
||||
metadata = input_data.get('metadata')
|
||||
train_params = TrainModelParams.from_dict(input_data['train_params'])
|
||||
|
||||
self.info('Starting train_model process', metadata)
|
||||
|
||||
try:
|
||||
# Download training file bytes from MinIO
|
||||
train_bytes = self.minio_repository.download_file_sync(
|
||||
|
||||
self.info(
|
||||
f'Downloading training file from MinIO for {train_params.file_name}', metadata
|
||||
)
|
||||
train_bytes = self.minio_repository.download_file(
|
||||
object_name=train_params.file_name,
|
||||
bucket=train_params.bucket_name,
|
||||
metadata=metadata,
|
||||
@@ -185,12 +202,14 @@ class Training(SientiaMonitoring):
|
||||
val_bytes: bytes | None = None
|
||||
validation_name = train_params.val_file_name
|
||||
if validation_name is not None:
|
||||
val_bytes = self.minio_repository.download_file_sync(
|
||||
self.info(f'Downloading validation file from MinIO for {validation_name}', metadata)
|
||||
val_bytes = self.minio_repository.download_file(
|
||||
object_name=validation_name,
|
||||
bucket=train_params.bucket_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.info(f'Preparing training data for {train_params.file_name}', metadata)
|
||||
train_result = self.data_manager_repository.prepare_training_data(
|
||||
train_file_bytes=train_bytes,
|
||||
validation_file_bytes=val_bytes,
|
||||
@@ -198,8 +217,9 @@ class Training(SientiaMonitoring):
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.info(f'Getting model wrapper for {train_params.model_type}', metadata)
|
||||
wrapper = self.plugin_store.get_model(
|
||||
model_name=train_params.model_name,
|
||||
model_type=train_params.model_type,
|
||||
force_download=False,
|
||||
opt_params=train_params.opt_params or {},
|
||||
model_kwargs=train_params.model_kwargs or {},
|
||||
@@ -207,6 +227,7 @@ class Training(SientiaMonitoring):
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.info(f'Training model for {train_params.model_type}', metadata)
|
||||
train_data = train_result.train_data
|
||||
val_data = train_result.val_data
|
||||
|
||||
@@ -216,6 +237,10 @@ class Training(SientiaMonitoring):
|
||||
target=train_params.target_variable,
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Generating predictions using the trained wrapper for {train_params.model_type}',
|
||||
metadata,
|
||||
)
|
||||
# Generate predictions using the trained wrapper
|
||||
transformed_train, _ = wrapper.transform(train_data)
|
||||
transformed_val, _ = wrapper.transform(val_data)
|
||||
@@ -229,11 +254,13 @@ class Training(SientiaMonitoring):
|
||||
train_result.y_train_pred = y_train_pred_df
|
||||
train_result.y_pred = y_val_pred_df
|
||||
|
||||
self.info(f'Computing regression metrics for {train_params.model_type}', metadata)
|
||||
train_result = self.data_manager_repository.compute_regression_metrics(
|
||||
train_result,
|
||||
wrapper,
|
||||
)
|
||||
|
||||
self.info(f'Starting MLflow run for {train_params.model_type}', metadata)
|
||||
with self.mlflow_repository.start_run(
|
||||
model_name=train_params.model_name,
|
||||
run_name=None,
|
||||
@@ -273,6 +300,7 @@ class Training(SientiaMonitoring):
|
||||
wrapper: Any,
|
||||
metadata: dict[str, Any] | None,
|
||||
) -> None:
|
||||
self.info(f'Generating report for {train_params.model_type}', metadata)
|
||||
train_result = self.data_manager_repository.generate_report(
|
||||
train_result,
|
||||
metadata=metadata,
|
||||
@@ -285,7 +313,9 @@ class Training(SientiaMonitoring):
|
||||
):
|
||||
raise ValueError('Report path, train data path, or test data path is not set')
|
||||
|
||||
self.info(f'Storing model for {train_params.model_type}', metadata)
|
||||
wrapper.store_model(name=train_params.model_name)
|
||||
self.info(f'Logging artifacts for {train_params.model_type}', metadata)
|
||||
mlflow.log_artifact(train_result.report_path)
|
||||
mlflow.log_artifact(train_result.train_data_path)
|
||||
mlflow.log_artifact(train_result.test_data_path)
|
||||
|
||||
@@ -116,7 +116,6 @@ async def create_cleanup_schedule(
|
||||
runtime = (os.getenv('RUNTIME') or 'single').strip()
|
||||
cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single')
|
||||
schedule_id = build_cleanup_schedule_id(runtime)
|
||||
created = False
|
||||
updated = False
|
||||
|
||||
if await schedule_exists(client, schedule_id, logger, metadata):
|
||||
@@ -147,15 +146,13 @@ async def create_cleanup_schedule(
|
||||
),
|
||||
),
|
||||
)
|
||||
created = not updated
|
||||
|
||||
if updated:
|
||||
logger.custom_info(
|
||||
f"Schedule '{schedule_id}' reconciled successfully. "
|
||||
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
|
||||
metadata,
|
||||
)
|
||||
elif created:
|
||||
else:
|
||||
logger.custom_info(
|
||||
f"Schedule '{schedule_id}' created successfully. "
|
||||
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
|
||||
|
||||
@@ -34,7 +34,7 @@ from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
def train_test_split(
|
||||
data: pd.DataFrame | pd.Series,
|
||||
data: pd.DataFrame,
|
||||
train_size: float,
|
||||
random_state: int | None = None,
|
||||
shuffle: bool = True,
|
||||
@@ -57,11 +57,8 @@ def train_test_split(
|
||||
train_indices = indices[:n_train]
|
||||
test_indices = indices[n_train:]
|
||||
|
||||
# 5. Retornar os dados fatiados (funciona para DataFrame ou Series)
|
||||
if isinstance(data, (pd.DataFrame, pd.Series)):
|
||||
return data.iloc[train_indices], data.iloc[test_indices]
|
||||
|
||||
return data[train_indices], data[test_indices]
|
||||
# 5. Retornar os dados fatiados
|
||||
return data.iloc[train_indices], data.iloc[test_indices]
|
||||
|
||||
|
||||
def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
|
||||
|
||||
@@ -1,471 +0,0 @@
|
||||
"""
|
||||
MLFlow Repository
|
||||
|
||||
This module contains the MLFlowRepository class, which is responsible for
|
||||
handling model training artifacts and MLFlow operations for the Model Manager system.
|
||||
|
||||
It includes methods for generating training reports, managing artifacts,
|
||||
and logging model runs to MLFlow.
|
||||
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import warnings
|
||||
from datetime import datetime
|
||||
from os import makedirs, path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
from model_manager.sientia.model_serving import ModelServing # type: ignore[import-untyped]
|
||||
from model_manager.sientia.reports import Reports # type: ignore[import-untyped]
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
# Suppress sklearn FutureWarning about 'squared' deprecation without changing business logic
|
||||
warnings.filterwarnings('ignore', category=FutureWarning, message=".*'squared' is deprecated.*")
|
||||
|
||||
|
||||
class ModelRepository:
|
||||
def __init__(self, url, username, password, logger: Logger):
|
||||
self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
|
||||
self.logger = logger
|
||||
self.logger.info(f'MLFlow client initialized at {url}')
|
||||
|
||||
def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
|
||||
"""
|
||||
Save a trained ML model and its artifacts to MLflow.
|
||||
|
||||
This activity orchestrates the complete model saving pipeline:
|
||||
1. Generates the next run name for the experiment
|
||||
2. Creates and organizes artifacts (reports, data files)
|
||||
3. Logs model, parameters, metrics, and artifacts to MLflow
|
||||
|
||||
Args:
|
||||
input_data: Configuration for model saving operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- train_result (TrainModelResult): Training result with model and metrics
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated training result with run_name and artifacts
|
||||
|
||||
Raises:
|
||||
Exception: If model saving fails (after sending notification)
|
||||
"""
|
||||
experiment_name = train_result.params.experiment_name
|
||||
train_result.run_name = self._get_next_run_name(experiment_name)
|
||||
train_result = self._generate_artifacts(train_result)
|
||||
self._save_run(train_result)
|
||||
|
||||
self.logger.info(
|
||||
f'Model saved successfully - experiment run id: {train_result.params.experiment_run_id}, '
|
||||
f'experiment name: {experiment_name}, '
|
||||
f'run name: {train_result.run_name}'
|
||||
)
|
||||
|
||||
return train_result
|
||||
|
||||
def cleanup_run_directory(self, run_dir: str) -> None:
|
||||
"""
|
||||
Clean up temporary run directory after model training.
|
||||
|
||||
This activity deletes the temporary directory created during model training
|
||||
and artifact generation. It implements idempotent cleanup to handle cases
|
||||
where the directory may have already been deleted.
|
||||
|
||||
Args:
|
||||
run_dir (str): Path to the run directory to delete
|
||||
"""
|
||||
if not run_dir:
|
||||
self.logger.info('No run directory specified, skipping cleanup')
|
||||
return
|
||||
|
||||
if os.path.exists(run_dir):
|
||||
shutil.rmtree(run_dir)
|
||||
self.logger.info(f'Run directory deleted successfully: {run_dir}')
|
||||
else:
|
||||
self.logger.info(f'Run directory already deleted: {run_dir}')
|
||||
|
||||
def _get_next_run_name(self, experiment_name: str) -> str:
|
||||
"""
|
||||
Generates the next run name for a given experiment.
|
||||
|
||||
Args:
|
||||
experiment_name (str): The name of the experiment for which the next run name is being generated.
|
||||
|
||||
Returns:
|
||||
str: A unique run name in the format "<experiment_name>-<next_run_number>".
|
||||
"""
|
||||
runs = self.model_serving.search_runs_by_name(
|
||||
experiment_names=[experiment_name], order_by=['start_time desc']
|
||||
)
|
||||
|
||||
next_run_number = len(runs) + 1
|
||||
return f'{experiment_name}-{next_run_number}'
|
||||
|
||||
def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
|
||||
"""
|
||||
Generates and organizes artifacts related to the training process, such as reports and data files.
|
||||
|
||||
Args:
|
||||
data: The training model result containing the datasets, model, and parameters.
|
||||
|
||||
Returns:
|
||||
The updated result object with paths to the generated artifacts.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the reports directory or header.html file does not exist.
|
||||
ValueError: If run_name is not set.
|
||||
"""
|
||||
# Validate that run_name is set
|
||||
if not data.run_name:
|
||||
error_msg = 'run_name must be set before generating artifacts'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
reference_data, current_data = self._init_artifacts_data(data)
|
||||
base_path = self._get_reports_directory()
|
||||
|
||||
# Validate that reports directory exists
|
||||
if not path.exists(base_path):
|
||||
error_msg = f'Reports directory does not exist: {base_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise FileNotFoundError(error_msg)
|
||||
|
||||
data.run_dir = self._create_run_directory(base_path, data.run_name)
|
||||
header_file_path = path.join(base_path, 'header.html')
|
||||
|
||||
# Validate that header.html exists
|
||||
if not path.exists(header_file_path):
|
||||
error_msg = f'Header file does not exist: {header_file_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise FileNotFoundError(error_msg)
|
||||
|
||||
self._setup_run_directory(data.run_dir, header_file_path)
|
||||
return self._generate_report(reference_data, current_data, data)
|
||||
|
||||
def _save_run(self, data: TrainModelResult):
|
||||
"""
|
||||
Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
|
||||
to the Sientia tracking system.
|
||||
|
||||
Args:
|
||||
data: The training model result containing the datasets, model, parameters,
|
||||
and evaluation metrics.
|
||||
|
||||
Raises:
|
||||
ValueError: If required metrics or artifacts are missing.
|
||||
Exception: If MLflow logging fails for any reason.
|
||||
"""
|
||||
# Validate that required artifacts exist before attempting to log
|
||||
if not data.report_path or not path.exists(data.report_path):
|
||||
error_msg = f'Report file does not exist: {data.report_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if not data.train_data_path or not path.exists(data.train_data_path):
|
||||
error_msg = f'Training data file does not exist: {data.train_data_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if not data.test_data_path or not path.exists(data.test_data_path):
|
||||
error_msg = f'Test data file does not exist: {data.test_data_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Validate that metrics are present
|
||||
if data.mse_val is None or data.r2_val is None or data.mae_val is None:
|
||||
error_msg = 'One or more metrics (MSE, R2, MAE) are None'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Prepare parameters
|
||||
interval_strs = [
|
||||
(str(interval[0]), str(interval[1]))
|
||||
for interval in (data.params.removed_intervals or [])
|
||||
]
|
||||
|
||||
# Set experiment and create run
|
||||
self.model_serving.set_experiment(data.params.experiment_name)
|
||||
|
||||
with self.model_serving.save_experiment(
|
||||
run_name=data.run_name, description=data.params.experiment_name
|
||||
):
|
||||
# Log model parameters
|
||||
self.model_serving.log_param('model_name', data.params.model_name)
|
||||
self.model_serving.log_param(
|
||||
'models_params',
|
||||
{'degree': data.params.degree, 'interaction_only': data.params.interaction_only},
|
||||
)
|
||||
self.model_serving.log_param('target_variable', data.params.target_variable)
|
||||
self.model_serving.log_param('input_variables', data.params.variable_columns)
|
||||
self.model_serving.log_param('nan_treatment', data.params.nan_treatment)
|
||||
self.model_serving.log_param('lag_train', data.params.lag_train)
|
||||
self.model_serving.log_param('lag_transform', data.params.lag_val)
|
||||
static_threshold_value = None
|
||||
if data.params.rem_static_win:
|
||||
static_threshold_value = (
|
||||
data.params.static_threshold if data.params.static_threshold is not None else 1
|
||||
)
|
||||
self.model_serving.log_param('static_threshold', static_threshold_value)
|
||||
self.model_serving.log_param('lower_limits', data.params.low_lim)
|
||||
self.model_serving.log_param('upper_limits', data.params.upp_lim)
|
||||
self.model_serving.log_param('scaler_name', data.params.scaler_name)
|
||||
self.model_serving.log_param('scaler_params', data.scaler_dict)
|
||||
self.model_serving.log_param('include_ar', data.params.include_ar)
|
||||
self.model_serving.log_param('train_size', round(data.params.train_size / 100, 2))
|
||||
self.model_serving.log_param('test_size', round(1 - (data.params.train_size / 100), 2))
|
||||
self.model_serving.log_param('start_date', data.params.start_date)
|
||||
self.model_serving.log_param('end_date', data.params.end_date)
|
||||
self.model_serving.log_param('removed_intervals', interval_strs)
|
||||
self.model_serving.log_param('retrain', False)
|
||||
self.model_serving.log_param('support_filters', data.params.support_filters)
|
||||
|
||||
# Log evaluation metrics
|
||||
self.model_serving.log_metric('MSE', data.mse_val)
|
||||
self.model_serving.log_metric('R2', data.r2_val)
|
||||
self.model_serving.log_metric('MAE', data.mae_val)
|
||||
|
||||
# Log models
|
||||
self.model_serving.log_model(data.process_data, 'data_model')
|
||||
self.model_serving.log_model(data.regr, 'prediction_model')
|
||||
|
||||
# Log artifacts
|
||||
self.model_serving.log_artifact(data.report_path)
|
||||
self.model_serving.log_artifact(data.train_data_path)
|
||||
self.model_serving.log_artifact(data.test_data_path)
|
||||
|
||||
# Log equation artifact if available
|
||||
if data.equation_path and path.exists(data.equation_path):
|
||||
self.model_serving.log_artifact(data.equation_path)
|
||||
|
||||
def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""
|
||||
Prepares the reference and current datasets for artifact generation.
|
||||
|
||||
Args:
|
||||
data: The training model result containing the datasets and model.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- reference_data: The training dataset with predictions added.
|
||||
- current_data: The testing dataset with predictions added.
|
||||
|
||||
Raises:
|
||||
ValueError: If training or test datasets are empty or invalid.
|
||||
AttributeError: If required attributes are missing from the data object.
|
||||
"""
|
||||
# Validate that required DataFrames are not empty
|
||||
# Note: x_train, y_train, x_test, y_test, and regr are required fields in TrainModelResult
|
||||
# so we only check if they are empty, not None
|
||||
if data.x_train.empty:
|
||||
error_msg = 'Training features (x_train) are empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.y_train.empty:
|
||||
error_msg = 'Training target (y_train) is empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.x_test.empty:
|
||||
error_msg = 'Test features (x_test) are empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.y_test.empty:
|
||||
error_msg = 'Test target (y_test) is empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Validate that predictions exist (y_pred is optional, so check for None)
|
||||
if data.y_pred is None:
|
||||
error_msg = 'Test predictions (y_pred) are None'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.y_train_pred is None:
|
||||
error_msg = 'Training predictions (y_train_pred) are None'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Prepare reference data (training set)
|
||||
reference_data = pd.concat([data.x_train, data.y_train], axis=1)
|
||||
reference_data = reference_data.rename(columns={data.params.target_variable: 'target'})
|
||||
# Use pre-calculated predictions (calculated before denormalization to avoid overflow)
|
||||
reference_data['prediction'] = data.y_train_pred
|
||||
|
||||
# Prepare current data (test set)
|
||||
current_data = pd.concat([data.x_test, data.y_test], axis=1)
|
||||
current_data = current_data.rename(columns={data.params.target_variable: 'target'})
|
||||
current_data['prediction'] = data.y_pred
|
||||
|
||||
return reference_data, current_data
|
||||
|
||||
def _create_run_directory(self, base_path: str, run_name: str) -> str:
|
||||
"""
|
||||
Creates a directory inside the 'reports' folder with the run name and a timestamp.
|
||||
|
||||
Uses microsecond precision in timestamp to minimize collision probability
|
||||
in high-concurrency scenarios.
|
||||
|
||||
Args:
|
||||
base_path (str): The path to the 'reports' folder.
|
||||
run_name (str): The name of the run.
|
||||
|
||||
Returns:
|
||||
str: The path to the created directory.
|
||||
|
||||
Raises:
|
||||
PermissionError: If there are insufficient permissions to create the directory.
|
||||
OSError: If directory creation fails for any other reason.
|
||||
"""
|
||||
# Use microsecond precision to reduce collision probability
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
|
||||
run_dir = path.join(base_path, 'temp', f'{run_name}_{timestamp}')
|
||||
|
||||
try:
|
||||
makedirs(run_dir, exist_ok=True)
|
||||
return run_dir
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when creating directory: {run_dir}'
|
||||
self.logger.error(error_msg)
|
||||
raise PermissionError(error_msg) from e
|
||||
except OSError as e:
|
||||
error_msg = f'Failed to create directory {run_dir}: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
def _setup_run_directory(self, run_dir: str, header_file_path: str):
|
||||
"""
|
||||
Creates empty files and copies a header file into the specified run directory.
|
||||
|
||||
Note: Lock removed as each run has its own unique directory, so no synchronization
|
||||
is needed between different runs. File operations within the same directory are
|
||||
atomic at the OS level.
|
||||
|
||||
Args:
|
||||
run_dir (str): The path to the run directory where the files will be created.
|
||||
header_file_path (str): The path to the header.html file to be copied.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the header file does not exist.
|
||||
PermissionError: If there are insufficient permissions to create files.
|
||||
OSError: If file creation or copying fails for any other reason.
|
||||
"""
|
||||
empty_files = ['data_drift.html', 'data_quality.html', 'regression.html']
|
||||
|
||||
try:
|
||||
# Create empty placeholder files
|
||||
for file_name in empty_files:
|
||||
file_path = path.join(run_dir, file_name)
|
||||
with open(file_path, 'w'):
|
||||
pass # Create empty file
|
||||
|
||||
# Copy header file to run directory
|
||||
header_dest = path.join(run_dir, 'header.html')
|
||||
shutil.copy(header_file_path, header_dest)
|
||||
except FileNotFoundError as e:
|
||||
error_msg = f'Header file not found: {header_file_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise FileNotFoundError(error_msg) from e
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when setting up directory: {run_dir}'
|
||||
self.logger.error(error_msg)
|
||||
raise PermissionError(error_msg) from e
|
||||
except OSError as e:
|
||||
error_msg = f'Failed to setup run directory {run_dir}: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
def _generate_report(
|
||||
self, reference_data: pd.DataFrame, current_data: pd.DataFrame, data: TrainModelResult
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Generates a comprehensive report summarizing data quality, data drift, and regression analysis.
|
||||
|
||||
Args:
|
||||
reference_data (pd.DataFrame): The training dataset with predictions added.
|
||||
current_data (pd.DataFrame): The testing dataset with predictions added.
|
||||
data: The training model result containing the datasets, model, and parameters.
|
||||
|
||||
Returns:
|
||||
The updated result object with paths to the generated report and data files.
|
||||
|
||||
Raises:
|
||||
ValueError: If data conversion to float64 fails or DataFrames are invalid.
|
||||
PermissionError: If there are insufficient permissions to write files.
|
||||
OSError: If file writing fails for any other reason.
|
||||
"""
|
||||
try:
|
||||
# Convert data to float64 for report generation
|
||||
# This may raise ValueError if data contains non-numeric values
|
||||
reference_data_float = reference_data.astype(np.float64)
|
||||
current_data_float = current_data.astype(np.float64)
|
||||
|
||||
# Initialize report generator
|
||||
report = Reports(
|
||||
reference_data=reference_data_float,
|
||||
current_data=current_data_float,
|
||||
base_path=data.run_dir,
|
||||
)
|
||||
|
||||
# Generate report sections
|
||||
report.add_data_quality_section(columns=data.params.variable_columns + ['target'])
|
||||
report.add_data_drift_section(columns=data.params.variable_columns + ['target'])
|
||||
report.add_regression_section()
|
||||
|
||||
# Validate that run_dir is set (should be set by _create_run_directory)
|
||||
if not data.run_dir:
|
||||
error_msg = 'run_dir is not set after directory creation'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Save HTML report
|
||||
data.report_path = path.join(data.run_dir, 'report.html')
|
||||
report.save_all_sections_html(data.report_path)
|
||||
|
||||
# Save training data CSV
|
||||
data.train_data_path = path.join(data.run_dir, 'train_data.csv')
|
||||
reference_data.to_csv(data.train_data_path, index=False)
|
||||
|
||||
# Save test data CSV
|
||||
data.test_data_path = path.join(data.run_dir, 'test_data.csv')
|
||||
current_data.to_csv(data.test_data_path, index=False)
|
||||
|
||||
# Save equation as JSON
|
||||
if data.equation is not None:
|
||||
data.equation_path = path.join(data.run_dir, 'model_equation.json')
|
||||
with open(data.equation_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(data.equation, f, indent=2, ensure_ascii=False)
|
||||
|
||||
return data
|
||||
except ValueError as e:
|
||||
error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg) from e
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when writing report files to: {data.run_dir}'
|
||||
self.logger.error(error_msg)
|
||||
raise PermissionError(error_msg) from e
|
||||
except OSError as e:
|
||||
error_msg = f'Failed to generate report in {data.run_dir}: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
def _get_reports_directory(self) -> str:
|
||||
"""
|
||||
Get the absolute path to the reports directory.
|
||||
|
||||
Returns:
|
||||
str: Absolute path to model_manager/reports directory.
|
||||
"""
|
||||
# Get the directory where this file is located (model_manager/utils/repository/)
|
||||
current_file_dir = path.dirname(path.abspath(__file__))
|
||||
# Navigate up to model_manager/ and then to reports/
|
||||
model_manager_dir = path.dirname(path.dirname(current_file_dir))
|
||||
reports_dir = path.join(model_manager_dir, 'reports')
|
||||
return reports_dir
|
||||
@@ -161,14 +161,15 @@ async def main():
|
||||
|
||||
logger.custom_info(f'SDK metrics server initialized on port {SDK_METRICS_PORT}', metadata)
|
||||
|
||||
namespace = os.getenv('TEMPORAL_NAMESPACE', 'model-manager')
|
||||
temporal_client = await client.Client.connect(
|
||||
target_host=host,
|
||||
namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
|
||||
namespace=namespace,
|
||||
runtime=new_runtime,
|
||||
tls=use_tls,
|
||||
)
|
||||
|
||||
logger.custom_info(f'Temporal client initialized at {host}', metadata)
|
||||
logger.custom_info(f'Temporal client initialized at {host}/{namespace}', metadata)
|
||||
|
||||
# Create cleanup schedule (idempotent - only creates if doesn't exist)
|
||||
try:
|
||||
|
||||
@@ -93,6 +93,8 @@ class TrainModel:
|
||||
Raises:
|
||||
ValueError: If experiment_run_id is missing or invalid
|
||||
"""
|
||||
workflow.logger.info(f'Starting train_model workflow for {input_data}')
|
||||
|
||||
experiment_run_id = self._validate_experiment_run_id(input_data)
|
||||
input_data = {**input_data, 'experiment_run_id': experiment_run_id}
|
||||
|
||||
|
||||
@@ -17,3 +17,4 @@ pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
|
||||
# Development Tools
|
||||
ipython>=8.12.0 # Enhanced Python shell
|
||||
ipdb>=0.13.13 # IPython debugger
|
||||
ipykernel<=6.26.0 # IPython kernel
|
||||
|
||||
@@ -32,6 +32,9 @@ import psycopg2
|
||||
from dotenv import load_dotenv
|
||||
from psycopg2.extras import Json
|
||||
from temporalio import client
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# %%
|
||||
# --- configuration (edit here or use `.env` at repo root) ---
|
||||
@@ -54,9 +57,14 @@ PG = {
|
||||
|
||||
TEMPORAL_HOST = os.getenv('TEMPORAL_HOST')
|
||||
TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE')
|
||||
TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE')
|
||||
TRAIN_TASK_QUEUE = "train_model-single-queue"
|
||||
TEMPORAL_TLS = os.getenv('TEMPORAL_USE_TLS', 'false').lower() in ('1', 'true', 'yes')
|
||||
|
||||
|
||||
print(MINIO_MC_ALIAS, MINIO_BUCKET, OBJECT_NAME, LOCAL_CSV)
|
||||
print(PG)
|
||||
print(TEMPORAL_HOST, TEMPORAL_NAMESPACE, TRAIN_TASK_QUEUE, TEMPORAL_TLS)
|
||||
|
||||
# %%
|
||||
# --- 1) database: delete previous row (same id), then insert `experiment_run` ---
|
||||
# Primary key column is `id` (see `experiment_tracking` updates). `request_data` matches the SQL sample in `input-sample.md`.
|
||||
@@ -97,7 +105,7 @@ with psycopg2.connect(**PG) as conn:
|
||||
EXPERIMENT_RUN_ID,
|
||||
'test-experiment-name',
|
||||
'test-run-name',
|
||||
'test-username',
|
||||
'vitor.santos@aignosi.com.br',
|
||||
'ORCHESTRATOR_WAITING_PROC',
|
||||
None,
|
||||
now,
|
||||
@@ -142,15 +150,16 @@ _workflow_input = {
|
||||
'opt_params': {},
|
||||
}
|
||||
|
||||
# %%
|
||||
|
||||
c = await client.Client.connect(
|
||||
target_host=TH,
|
||||
namespace=TN,
|
||||
tls=TEMPORAL_TLS,
|
||||
)
|
||||
|
||||
# %%
|
||||
|
||||
wid = f'train-model-test-{uuid.uuid4()}'
|
||||
await c.execute_workflow( # type: ignore[call-overload]
|
||||
result = await c.execute_workflow( # type: ignore[call-overload]
|
||||
'train_model',
|
||||
_workflow_input,
|
||||
id=wid,
|
||||
@@ -160,3 +169,6 @@ await c.execute_workflow( # type: ignore[call-overload]
|
||||
task_timeout=timedelta(minutes=5),
|
||||
)
|
||||
print(wid)
|
||||
print(result)
|
||||
|
||||
# %%
|
||||
|
||||
@@ -93,7 +93,7 @@ def test_train_model_download_fails_notifies(training):
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
)
|
||||
training.minio_repository.download_file_sync = MagicMock(side_effect=OSError('minio'))
|
||||
training.minio_repository.download_file = MagicMock(side_effect=OSError('minio'))
|
||||
training.send_notification = MagicMock()
|
||||
with pytest.raises(OSError, match='minio'):
|
||||
training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp.to_dict()})
|
||||
@@ -129,7 +129,7 @@ def test_train_model_success_serializes_result(mock_mlflow, training):
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(return_value=b'csv')
|
||||
training.minio_repository.download_file = MagicMock(return_value=b'csv')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: setattr(x, 'mse_val', 0.1) or x
|
||||
@@ -185,7 +185,7 @@ def test_train_model_train_params_as_dict(mock_mlflow, training):
|
||||
tp = TrainModelParams.from_dict(d)
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(return_value=b'csv')
|
||||
training.minio_repository.download_file = MagicMock(return_value=b'csv')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: x
|
||||
@@ -240,7 +240,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
|
||||
return b'val'
|
||||
raise AssertionError(object_name)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(side_effect=_dl)
|
||||
training.minio_repository.download_file = MagicMock(side_effect=_dl)
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: x
|
||||
@@ -272,7 +272,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
|
||||
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|
||||
assert training.minio_repository.download_file_sync.call_count == 2
|
||||
assert training.minio_repository.download_file.call_count == 2
|
||||
mock_mlflow.log_artifact.assert_called()
|
||||
|
||||
|
||||
@@ -288,7 +288,7 @@ def test_train_model_value_error_when_paths_missing_after_report(training):
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(return_value=b'x')
|
||||
training.minio_repository.download_file = MagicMock(return_value=b'x')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: x
|
||||
|
||||
@@ -125,6 +125,29 @@ async def test_schedule_exists_handles_exception(mock_temporal_client, mock_logg
|
||||
assert 'Error checking if schedule exists' in mock_logger.custom_error.call_args[0][0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_needs_schedule_reconcile_handles_describe_exception(mock_logger, metadata):
|
||||
"""Test _needs_schedule_reconcile returns True and logs when describe fails."""
|
||||
from model_manager.schedules.cleanup_schedule import _needs_schedule_reconcile
|
||||
|
||||
handle = AsyncMock()
|
||||
handle.describe = AsyncMock(side_effect=RuntimeError('describe failed'))
|
||||
|
||||
needs_reconcile = await _needs_schedule_reconcile(
|
||||
schedule_handle=handle,
|
||||
cleanup_task_queue='cleanup_files-model-manager-worker-queue',
|
||||
logger=mock_logger,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
assert needs_reconcile is True
|
||||
mock_logger.custom_error.assert_called_once()
|
||||
assert (
|
||||
'Error describing cleanup schedule for reconcile'
|
||||
in mock_logger.custom_error.call_args[0][0]
|
||||
)
|
||||
|
||||
|
||||
# --- create_cleanup_schedule Tests ---
|
||||
|
||||
|
||||
|
||||
9
tests/sientia/test_exceptions.py
Normal file
9
tests/sientia/test_exceptions.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""Unit tests for custom exception aliases."""
|
||||
|
||||
from mlflow.exceptions import MlflowException
|
||||
|
||||
from model_manager.sientia.exceptions import SientiaMlException
|
||||
|
||||
|
||||
def test_sientia_ml_exception_is_mlflow_exception_alias():
|
||||
assert SientiaMlException is MlflowException
|
||||
@@ -27,9 +27,11 @@ def test_train_test_split_dataframe_no_shuffle():
|
||||
assert list(tr['a']) == [0, 1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_train_test_split_ndarray():
|
||||
arr = np.arange(20).reshape(10, 2)
|
||||
tr, te = dmr.train_test_split(arr, train_size=0.5, shuffle=False, random_state=None)
|
||||
def test_train_test_split_dataframe_returns_dataframes():
|
||||
df = pd.DataFrame(np.arange(20).reshape(10, 2), columns=['a', 'b'])
|
||||
tr, te = dmr.train_test_split(df, train_size=0.5, shuffle=False, random_state=None)
|
||||
assert isinstance(tr, pd.DataFrame)
|
||||
assert isinstance(te, pd.DataFrame)
|
||||
assert tr.shape[0] == 5 and te.shape[0] == 5
|
||||
|
||||
|
||||
|
||||
@@ -3,6 +3,12 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def test_build_queue_name_without_runtime_uses_default_suffix():
|
||||
from model_manager.worker.prepare_worker import build_queue_name
|
||||
|
||||
assert build_queue_name('TrainModel') == 'train_model-queue'
|
||||
|
||||
|
||||
def test_prepare_worker_train_queue_uses_train_limits():
|
||||
from model_manager.worker.prepare_worker import prepare_worker
|
||||
from model_manager.workflows.train_model import TrainModel
|
||||
|
||||
@@ -290,9 +290,11 @@ async def test_run_training_failure_skips_cleanup_activity(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_missing_experiment_run_id():
|
||||
@patch('model_manager.workflows.train_model.workflow')
|
||||
async def test_run_missing_experiment_run_id(mock_wf):
|
||||
from model_manager.workflows.train_model import TrainModel
|
||||
|
||||
mock_wf.logger = Mock()
|
||||
with pytest.raises(ValueError, match='experiment_run_id is required'):
|
||||
await TrainModel().run({})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user