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
381
model_manager/utils/repository/data_manager_repository.py
Normal file
381
model_manager/utils/repository/data_manager_repository.py
Normal file
@@ -0,0 +1,381 @@
|
||||
"""
|
||||
Data management repository for the training pipeline.
|
||||
|
||||
This module provides the core data loading and preprocessing logic for the
|
||||
training pipeline, including:
|
||||
- CSV loading from in-memory bytes
|
||||
- datetime parsing and index configuration
|
||||
- optional support filters
|
||||
- train/test split management (when no explicit validation dataset is provided)
|
||||
|
||||
It is intentionally decoupled from any specific model implementation or MLflow
|
||||
integration. Models are trained elsewhere (e.g., via SientiaModel wrappers),
|
||||
and this repository focuses solely on preparing data structures for them.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
from model_manager.sientia.metrics import mae, mse, r2
|
||||
from model_manager.sientia.utils import split_train_test
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
|
||||
"""
|
||||
If date_column is set, parse the column as timezone-aware
|
||||
datetime to avoid comparison errors downstream.
|
||||
|
||||
The values are expected to follow the global DATETIME_FORMAT_WITH_TZ
|
||||
pattern defined in sientia_do.temporal.constants.
|
||||
"""
|
||||
if not params.date_column or params.date_column not in data.columns:
|
||||
return data
|
||||
try:
|
||||
data = data.copy()
|
||||
parsed = pd.to_datetime(
|
||||
data[params.date_column],
|
||||
format=DATETIME_FORMAT_WITH_TZ,
|
||||
errors='raise',
|
||||
)
|
||||
data[params.date_column] = parsed
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f'Failed to parse date column "{params.date_column}" with format "{params.date_format}": {e}'
|
||||
) from e
|
||||
return data
|
||||
|
||||
|
||||
def _single_variable_support_mask(
|
||||
data_view: pd.DataFrame,
|
||||
var_col: str,
|
||||
target_variable: str,
|
||||
config: dict,
|
||||
) -> np.ndarray | None:
|
||||
"""Compute keep mask for one variable's support lines; None if config is invalid or skipped."""
|
||||
if var_col not in data_view.columns:
|
||||
return None
|
||||
upper = config.get('upper_line') or config.get('upperLine')
|
||||
lower = config.get('lower_line') or config.get('lowerLine')
|
||||
if not upper or not lower:
|
||||
return None
|
||||
|
||||
x_vals = data_view[var_col].astype(float).to_numpy()
|
||||
y_vals = data_view[target_variable].astype(float).to_numpy()
|
||||
xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
|
||||
ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
|
||||
x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
|
||||
y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
|
||||
scale_ratio = y_range / x_range
|
||||
|
||||
b1 = float(upper.get('intercept', 0))
|
||||
deg1 = float(upper.get('angle', 0))
|
||||
b2 = float(lower.get('intercept', 0))
|
||||
deg2 = float(lower.get('angle', 0))
|
||||
m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
|
||||
m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
|
||||
y1 = m1 * x_vals + b1
|
||||
y2 = m2 * x_vals + b2
|
||||
lower_bound = np.minimum(y1, y2)
|
||||
upper_bound = np.maximum(y1, y2)
|
||||
return (y_vals >= lower_bound) & (y_vals <= upper_bound)
|
||||
|
||||
|
||||
def _apply_support_filters(
|
||||
data_view: pd.DataFrame,
|
||||
target_variable: str,
|
||||
support_filters: dict,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Keep only rows where (var, target) lies between the two guide lines for each variable.
|
||||
|
||||
For each variable in support_filters, the condition is lower(x_var) <= target <= upper(x_var),
|
||||
where lower/upper are the two lines (intercept + slope from angle, scaled by y_range/x_range).
|
||||
Global mask is AND across all variables. Matches DEMO logic in template_01.py.
|
||||
|
||||
Args:
|
||||
data_view: DataFrame after preprocessor transform.
|
||||
target_variable: Name of the target column (y axis).
|
||||
support_filters: Per-variable config with upper_line/lower_line, each {intercept, angle}.
|
||||
|
||||
Returns:
|
||||
data_view filtered to rows satisfying all variable conditions; unchanged if support_filters empty.
|
||||
"""
|
||||
if not support_filters or target_variable not in data_view.columns:
|
||||
return data_view
|
||||
|
||||
combined_keep_mask = np.ones(len(data_view), dtype=bool)
|
||||
n = len(data_view)
|
||||
for var_col, config in support_filters.items():
|
||||
keep_mask = _single_variable_support_mask(data_view, var_col, target_variable, config)
|
||||
if keep_mask is not None and len(keep_mask) == n:
|
||||
combined_keep_mask &= keep_mask
|
||||
|
||||
return data_view.loc[combined_keep_mask]
|
||||
|
||||
|
||||
class DataManagerRepository(SientiaMonitoring):
|
||||
"""
|
||||
Repository for data preparation in the training pipeline.
|
||||
|
||||
This class encapsulates the core logic for preparing ML training data:
|
||||
loading CSV bytes, applying date/index configuration, support filters, and
|
||||
constructing train/test splits (or using an explicit validation dataset).
|
||||
|
||||
Attributes:
|
||||
logger (Logger): Logger instance for observability and debugging
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Logger):
|
||||
"""
|
||||
Initialize DataManagerRepository with logger.
|
||||
|
||||
Args:
|
||||
logger: Logger instance for observability
|
||||
"""
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=None,
|
||||
metrics_controller=None,
|
||||
)
|
||||
|
||||
def prepare_training_data(
|
||||
self,
|
||||
train_file_bytes: bytes,
|
||||
validation_file_bytes: bytes | None,
|
||||
params: TrainModelParams,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Build TrainModelResult from raw CSV bytes for train (and optional validation) data.
|
||||
|
||||
This method orchestrates the data pipeline:
|
||||
1. Load training data from in-memory bytes
|
||||
2. Optionally load validation data from in-memory bytes
|
||||
3. Parse and configure datetime index
|
||||
4. Apply optional support filters
|
||||
5. Split into train/test sets when no explicit validation dataset is provided
|
||||
|
||||
Args:
|
||||
train_file_bytes: Raw bytes of the training CSV.
|
||||
validation_file_bytes: Raw bytes of the validation CSV, or None when
|
||||
validation should be derived via train/test split.
|
||||
params: Training parameters (TrainModelParams).
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Object containing processed data, train/test splits,
|
||||
and scaler dictionary.
|
||||
|
||||
Raises:
|
||||
ValueError: If transformed data is empty.
|
||||
Exception: If data loading or preprocessing fails.
|
||||
"""
|
||||
try:
|
||||
train_df = pd.read_csv(
|
||||
BytesIO(train_file_bytes),
|
||||
sep=params.line_separator,
|
||||
decimal=params.decimal_separator,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise ValueError(
|
||||
'Failed to load training CSV data from MinIO object. '
|
||||
'Check file encoding, line separator and decimal separator.'
|
||||
) from exc
|
||||
|
||||
train_df = _ensure_date_column_parsed(train_df, params)
|
||||
train_df = self._configure_datetime_index(train_df, params, metadata)
|
||||
|
||||
if params.support_filters:
|
||||
train_df = _apply_support_filters(
|
||||
train_df,
|
||||
params.target_variable,
|
||||
params.support_filters,
|
||||
)
|
||||
|
||||
if len(train_df) <= 0:
|
||||
raise ValueError('Training data view is empty after transformation')
|
||||
|
||||
# Explicit validation dataset path
|
||||
if validation_file_bytes is not None:
|
||||
try:
|
||||
val_df = pd.read_csv(
|
||||
BytesIO(validation_file_bytes),
|
||||
sep=params.line_separator,
|
||||
decimal=params.decimal_separator,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise ValueError(
|
||||
'Failed to load validation CSV data from MinIO object. '
|
||||
'Check file encoding, line separator and decimal separator.'
|
||||
) from exc
|
||||
|
||||
val_df = _ensure_date_column_parsed(val_df, params)
|
||||
val_df = self._configure_datetime_index(val_df, params, metadata)
|
||||
|
||||
if params.support_filters:
|
||||
val_df = _apply_support_filters(
|
||||
val_df,
|
||||
params.target_variable,
|
||||
params.support_filters,
|
||||
)
|
||||
|
||||
if len(val_df) <= 0:
|
||||
raise ValueError('Validation data view is empty after transformation')
|
||||
|
||||
x_train = pd.DataFrame(train_df[params.variable_columns])
|
||||
y_train = pd.Series(train_df[params.target_variable])
|
||||
x_test = pd.DataFrame(val_df[params.variable_columns])
|
||||
y_test = pd.Series(val_df[params.target_variable])
|
||||
else:
|
||||
# Fallback path: derive validation via train/test split from a single dataset.
|
||||
x_train, x_test, y_train, y_test = split_train_test(
|
||||
pd.DataFrame(train_df[params.variable_columns]),
|
||||
pd.Series(train_df[params.target_variable]),
|
||||
train_size=params.train_size / 100,
|
||||
shuffle=params.shuffle,
|
||||
random_state=42,
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Data preprocessed and split successfully - experiment run id: {params.experiment_run_id}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return TrainModelResult(
|
||||
params=params,
|
||||
x_train=x_train,
|
||||
x_test=x_test,
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
)
|
||||
|
||||
def compute_regression_metrics(
|
||||
self,
|
||||
params: TrainModelParams,
|
||||
tmr: TrainModelResult,
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Compute regression metrics for training results.
|
||||
|
||||
This helper mirrors the previous TrainingRepository.after_train_calculation
|
||||
behavior, assuming that y_pred/y_train_pred are already on the correct scale
|
||||
for metric calculation (any scaling is handled inside the model wrapper).
|
||||
|
||||
Args:
|
||||
params: Training parameters used during model training.
|
||||
tmr: Training result with y_train, y_test, y_train_pred and y_pred populated.
|
||||
|
||||
Return:
|
||||
Updated TrainModelResult with mse_val, mae_val and r2_val fields populated.
|
||||
"""
|
||||
del params # unused for now, kept for possible future extensions
|
||||
|
||||
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()
|
||||
|
||||
if tmr.y_pred is not None:
|
||||
tmr.y_pred = tmr.y_pred.sort_index()
|
||||
if tmr.y_train_pred is not None:
|
||||
tmr.y_train_pred = tmr.y_train_pred.sort_index()
|
||||
|
||||
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
||||
|
||||
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 _configure_datetime_index(
|
||||
self,
|
||||
data: pd.DataFrame | None,
|
||||
params: TrainModelParams,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Configure datetime index for the DataFrame.
|
||||
|
||||
Guards against None to avoid 'NoneType' object has no attribute 'index' downstream.
|
||||
Prefers params.date_column when set; otherwise looks for common timestamp column names.
|
||||
"""
|
||||
if data is None:
|
||||
raise ValueError(
|
||||
'Data is None after load_data. '
|
||||
'Check file format, line separator and decimal separator.'
|
||||
)
|
||||
if not isinstance(data, pd.DataFrame):
|
||||
raise TypeError(f'Expected DataFrame, got {type(data).__name__}')
|
||||
|
||||
if isinstance(data.index, pd.DatetimeIndex):
|
||||
self.info('DataFrame already has DatetimeIndex', metadata)
|
||||
return data.sort_index()
|
||||
|
||||
common_timestamp_columns = [
|
||||
'timestamp',
|
||||
'Timestamp',
|
||||
'TIMESTAMP',
|
||||
'date',
|
||||
'Date',
|
||||
'DATE',
|
||||
'DATA',
|
||||
'datetime',
|
||||
'DateTime',
|
||||
]
|
||||
timestamp_columns = ([params.date_column] if params.date_column else []) + [
|
||||
c for c in common_timestamp_columns if c != params.date_column
|
||||
]
|
||||
|
||||
for col in timestamp_columns:
|
||||
if col in data.columns:
|
||||
try:
|
||||
data[col] = pd.to_datetime(data[col])
|
||||
data = data.set_index(col)
|
||||
data = data.sort_index()
|
||||
self.info(f'Configured datetime index from column: {col}', metadata)
|
||||
return data
|
||||
except (ValueError, TypeError) as e:
|
||||
self.warning(f'Failed to convert column {col} to datetime: {e}', metadata)
|
||||
continue
|
||||
|
||||
# If no timestamp column found, check if first column looks like a timestamp
|
||||
first_col = data.columns[0]
|
||||
try:
|
||||
# Try to parse first column as datetime
|
||||
test_values = data[first_col].head(10).dropna()
|
||||
if len(test_values) > 0:
|
||||
pd.to_datetime(test_values)
|
||||
data[first_col] = pd.to_datetime(data[first_col])
|
||||
data = data.set_index(first_col)
|
||||
data = data.sort_index()
|
||||
self.info(f'Configured datetime index from first column: {first_col}', metadata)
|
||||
return data
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
self.warning(
|
||||
'No timestamp column found - some features may not work correctly',
|
||||
metadata,
|
||||
)
|
||||
return data
|
||||
|
||||
@@ -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
|
||||
@@ -1,164 +0,0 @@
|
||||
from io import BytesIO
|
||||
|
||||
import boto3 # type: ignore[import-untyped]
|
||||
from botocore.config import Config # type: ignore[import-untyped]
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
|
||||
class StorageRepository:
|
||||
"""
|
||||
MinIO (S3-compatible) storage activities for file operations.
|
||||
|
||||
This class provides activities for interacting with MinIO object storage,
|
||||
including file download and deletion operations. It handles authentication,
|
||||
connection management, and comprehensive error handling.
|
||||
|
||||
The class implements best practices for S3/MinIO operations:
|
||||
- Connection reuse (boto3 client is thread-safe)
|
||||
- Automatic retry with exponential backoff
|
||||
- Comprehensive error handling and logging
|
||||
- Notification integration for critical errors
|
||||
|
||||
Attributes:
|
||||
endpoint_url (str): MinIO server endpoint URL
|
||||
access_key (str): MinIO access key ID
|
||||
secret_key (str): MinIO secret access key
|
||||
region (str): MinIO region name
|
||||
use_ssl (bool): Whether to use SSL/TLS for connections
|
||||
minio_client: Boto3 S3 client configured for MinIO
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
endpoint_url: str,
|
||||
access_key: str,
|
||||
secret_key: str,
|
||||
region: str,
|
||||
use_ssl: bool,
|
||||
max_retry_attempts: int,
|
||||
retry_mode: str,
|
||||
connect_timeout: int,
|
||||
read_timeout: int,
|
||||
logger: Logger,
|
||||
):
|
||||
"""
|
||||
Initialize a reusable MinIO client with retry configuration.
|
||||
|
||||
Args:
|
||||
endpoint_url: MinIO server endpoint URL (e.g., http://localhost:9000).
|
||||
access_key: MinIO access key ID for authentication.
|
||||
secret_key: MinIO secret access key for authentication.
|
||||
region: MinIO region name (e.g., us-east-1).
|
||||
use_ssl: Whether to use SSL/TLS for connections.
|
||||
max_retry_attempts: Maximum number of retry attempts (e.g., 3).
|
||||
retry_mode: Retry policy to apply (standard, legacy, adaptive).
|
||||
connect_timeout: Connection timeout in seconds.
|
||||
read_timeout: Read timeout in seconds.
|
||||
logger: Logger used for observability.
|
||||
"""
|
||||
self.endpoint_url = endpoint_url
|
||||
self.access_key = access_key
|
||||
self.secret_key = secret_key
|
||||
self.region = region
|
||||
self.use_ssl = use_ssl
|
||||
self.max_retry_attempts = max_retry_attempts
|
||||
self.retry_mode = retry_mode
|
||||
self.connect_timeout = connect_timeout
|
||||
self.read_timeout = read_timeout
|
||||
self.logger = logger
|
||||
|
||||
boto_config = Config(
|
||||
region_name=region,
|
||||
retries={
|
||||
'max_attempts': max_retry_attempts,
|
||||
'mode': retry_mode,
|
||||
},
|
||||
connect_timeout=connect_timeout,
|
||||
read_timeout=read_timeout,
|
||||
)
|
||||
|
||||
self.minio_client = boto3.client(
|
||||
's3',
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=access_key,
|
||||
aws_secret_access_key=secret_key,
|
||||
config=boto_config,
|
||||
use_ssl=use_ssl,
|
||||
)
|
||||
|
||||
self.logger.info(f'MinIO client initialized at {endpoint_url}')
|
||||
|
||||
def close(self) -> None:
|
||||
self.minio_client.close()
|
||||
self.logger.info('MinIO client closed')
|
||||
|
||||
def fetch_file(self, bucket_name: str, file_name: str) -> BytesIO:
|
||||
"""
|
||||
Fetch an object from MinIO and return its contents as `BytesIO`.
|
||||
|
||||
Args:
|
||||
bucket_name: MinIO bucket where the object resides.
|
||||
file_name: Object key to download inside the bucket.
|
||||
|
||||
Returns:
|
||||
BytesIO: File-like stream containing the downloaded bytes.
|
||||
|
||||
Raises:
|
||||
OSError: If the download fails (network, permissions, missing key, etc.).
|
||||
"""
|
||||
response = self.minio_client.get_object(Bucket=bucket_name, Key=file_name)
|
||||
|
||||
with response['Body'] as body:
|
||||
file_content = body.read()
|
||||
|
||||
file_size = len(file_content)
|
||||
|
||||
self.logger.info(
|
||||
f'File fetched successfully: {bucket_name}/{file_name} ({file_size} bytes)'
|
||||
)
|
||||
|
||||
return BytesIO(file_content)
|
||||
|
||||
def delete_file(self, bucket_name: str, file_name: str) -> None:
|
||||
"""
|
||||
Remove an object from MinIO storage.
|
||||
|
||||
Args:
|
||||
bucket_name: Bucket that contains the object.
|
||||
file_name: Object key to delete.
|
||||
"""
|
||||
self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
|
||||
self.logger.info(f'File deleted successfully: {bucket_name}/{file_name}')
|
||||
|
||||
def list_bucket_objects(self, bucket_name: str, max_keys: int = 1000) -> list[str]:
|
||||
"""
|
||||
List objects in a MinIO bucket.
|
||||
|
||||
This method uses the MinIO/S3 list_objects_v2 API to retrieve objects
|
||||
from the specified bucket. This is optimized for cleanup operations
|
||||
by using configurable pagination.
|
||||
|
||||
Args:
|
||||
bucket_name: Name of the bucket to list objects from.
|
||||
max_keys: Maximum number of keys per page (default: 1000).
|
||||
|
||||
Returns:
|
||||
List[str]: List of object keys (file names).
|
||||
"""
|
||||
# Use list_objects_v2 for efficient pagination
|
||||
paginator = self.minio_client.get_paginator('list_objects_v2')
|
||||
|
||||
pages = paginator.paginate(Bucket=bucket_name, MaxKeys=max_keys)
|
||||
|
||||
objects = []
|
||||
total_count = 0
|
||||
|
||||
for page in pages:
|
||||
if 'Contents' in page:
|
||||
for obj in page['Contents']:
|
||||
objects.append(obj['Key'])
|
||||
total_count += 1
|
||||
|
||||
self.logger.info(f'Listed {total_count} objects from bucket {bucket_name}')
|
||||
|
||||
return objects
|
||||
@@ -1,519 +0,0 @@
|
||||
"""
|
||||
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_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.sientia.metrics import mae, mse, r2
|
||||
from model_manager.sientia.models import (
|
||||
DataPreprocessor,
|
||||
LinearRegressionModel,
|
||||
)
|
||||
from model_manager.sientia.models import (
|
||||
_frontend_date_format_to_strftime as _frontend_format_to_strftime,
|
||||
)
|
||||
from model_manager.sientia.utils import split_train_test
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
|
||||
"""If date_column and date_format are set, parse the column as datetime to avoid comparison errors downstream."""
|
||||
if not params.date_column or not params.date_format or params.date_column not in data.columns:
|
||||
return data
|
||||
try:
|
||||
python_fmt = _frontend_format_to_strftime(params.date_format)
|
||||
data = data.copy()
|
||||
data[params.date_column] = pd.to_datetime(
|
||||
data[params.date_column], format=python_fmt, errors='coerce'
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f'Failed to parse date column "{params.date_column}" with format "{params.date_format}": {e}'
|
||||
) from e
|
||||
return data
|
||||
|
||||
|
||||
def _single_variable_support_mask(
|
||||
data_view: pd.DataFrame,
|
||||
var_col: str,
|
||||
target_variable: str,
|
||||
config: dict,
|
||||
) -> np.ndarray | None:
|
||||
"""Compute keep mask for one variable's support lines; None if config is invalid or skipped."""
|
||||
if var_col not in data_view.columns:
|
||||
return None
|
||||
upper = config.get('upper_line') or config.get('upperLine')
|
||||
lower = config.get('lower_line') or config.get('lowerLine')
|
||||
if not upper or not lower:
|
||||
return None
|
||||
|
||||
x_vals = data_view[var_col].astype(float).to_numpy()
|
||||
y_vals = data_view[target_variable].astype(float).to_numpy()
|
||||
xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
|
||||
ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
|
||||
x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
|
||||
y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
|
||||
scale_ratio = y_range / x_range
|
||||
|
||||
b1 = float(upper.get('intercept', 0))
|
||||
deg1 = float(upper.get('angle', 0))
|
||||
b2 = float(lower.get('intercept', 0))
|
||||
deg2 = float(lower.get('angle', 0))
|
||||
m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
|
||||
m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
|
||||
y1 = m1 * x_vals + b1
|
||||
y2 = m2 * x_vals + b2
|
||||
lower_bound = np.minimum(y1, y2)
|
||||
upper_bound = np.maximum(y1, y2)
|
||||
return (y_vals >= lower_bound) & (y_vals <= upper_bound)
|
||||
|
||||
|
||||
def _apply_support_filters(
|
||||
data_view: pd.DataFrame,
|
||||
target_variable: str,
|
||||
support_filters: dict,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Keep only rows where (var, target) lies between the two guide lines for each variable.
|
||||
|
||||
For each variable in support_filters, the condition is lower(x_var) <= target <= upper(x_var),
|
||||
where lower/upper are the two lines (intercept + slope from angle, scaled by y_range/x_range).
|
||||
Global mask is AND across all variables. Matches DEMO logic in template_01.py.
|
||||
|
||||
Args:
|
||||
data_view: DataFrame after preprocessor transform.
|
||||
target_variable: Name of the target column (y axis).
|
||||
support_filters: Per-variable config with upper_line/lower_line, each {intercept, angle}.
|
||||
|
||||
Returns:
|
||||
data_view filtered to rows satisfying all variable conditions; unchanged if support_filters empty.
|
||||
"""
|
||||
if not support_filters or target_variable not in data_view.columns:
|
||||
return data_view
|
||||
|
||||
combined_keep_mask = np.ones(len(data_view), dtype=bool)
|
||||
n = len(data_view)
|
||||
for var_col, config in support_filters.items():
|
||||
keep_mask = _single_variable_support_mask(data_view, var_col, target_variable, config)
|
||||
if keep_mask is not None and len(keep_mask) == n:
|
||||
combined_keep_mask &= keep_mask
|
||||
|
||||
return data_view.loc[combined_keep_mask]
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
|
||||
if data is None:
|
||||
raise ValueError(
|
||||
'Failed to load CSV data: load_data returned None. '
|
||||
'Check file encoding, line separator and decimal separator.'
|
||||
)
|
||||
data = _ensure_date_column_parsed(data, params)
|
||||
data = self._configure_datetime_index(data, params)
|
||||
|
||||
process_data = self._init_data_preprocessor(params)
|
||||
process_data.fit(data)
|
||||
data_view = process_data.transform(data)
|
||||
|
||||
if params.support_filters:
|
||||
data_view = _apply_support_filters(
|
||||
data_view,
|
||||
params.target_variable,
|
||||
params.support_filters,
|
||||
)
|
||||
|
||||
if len(data_view) <= 0:
|
||||
raise ValueError('Data view is empty after transformation')
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
data_train = pd.concat([x_train, y_train], axis=1)
|
||||
scaler_dict = self._init_scaler_dict(process_data, params)
|
||||
|
||||
regr = LinearRegressionModel(
|
||||
target_variable=params.target_variable,
|
||||
variable_columns=params.variable_columns,
|
||||
degree=params.degree,
|
||||
interaction_only=params.interaction_only,
|
||||
)
|
||||
|
||||
regr.fit(data_train)
|
||||
self.logger.info(
|
||||
f'Model trained successfully - experiment run id: {params.experiment_run_id}'
|
||||
)
|
||||
|
||||
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 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)
|
||||
"""
|
||||
# Calculate predictions BEFORE denormalization (important for polynomial models)
|
||||
y_pred_array = tmr.regr.predict(tmr.x_test)
|
||||
y_train_pred_array = tmr.regr.predict(tmr.x_train)
|
||||
|
||||
if params.use_scaler:
|
||||
scaler = tmr.process_data.get_scaler()
|
||||
|
||||
# If using custom scaler with denormalize_* helpers
|
||||
if hasattr(scaler, 'denormalize_single_input'):
|
||||
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)
|
||||
|
||||
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)
|
||||
y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
|
||||
y_train_pred_array = scaler.denormalize_predictions(
|
||||
y_train_pred_array, params.target_variable
|
||||
)
|
||||
else:
|
||||
# Fallback for sklearn StandardScaler: only inverse-transform features
|
||||
feature_cols = getattr(
|
||||
tmr.process_data, 'feature_names_order', params.variable_columns
|
||||
)
|
||||
# Ensure columns are in the same order used during fit
|
||||
x_train_features = tmr.x_train[feature_cols]
|
||||
x_test_features = tmr.x_test[feature_cols]
|
||||
|
||||
tmr.x_train[feature_cols] = scaler.inverse_transform(x_train_features)
|
||||
tmr.x_test[feature_cols] = scaler.inverse_transform(x_test_features)
|
||||
# Target was not scaled with StandardScaler in preprocessing; leave y as-is
|
||||
|
||||
tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
|
||||
tmr.y_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
tmr.y_train_pred = pd.Series(y_train_pred_array, index=tmr.y_train.index)
|
||||
tmr.y_train_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
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()
|
||||
tmr.y_train_pred = tmr.y_train_pred.sort_index()
|
||||
|
||||
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
||||
|
||||
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)
|
||||
|
||||
# Extract model equation
|
||||
tmr.equation = self._extract_model_equation(tmr.regr, params)
|
||||
|
||||
self.logger.info(
|
||||
f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}'
|
||||
)
|
||||
return tmr
|
||||
|
||||
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 _get_static_threshold(self, params: TrainModelParams) -> int | None:
|
||||
"""
|
||||
Get the static threshold value based on parameters.
|
||||
|
||||
Args:
|
||||
params: Training parameters containing static window configuration
|
||||
|
||||
Returns:
|
||||
int | None: Static threshold value (1-1000) if rem_static_win is True, None otherwise
|
||||
"""
|
||||
if not params.rem_static_win:
|
||||
return None
|
||||
return params.static_threshold if params.static_threshold is not None else 1
|
||||
|
||||
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
|
||||
"""
|
||||
# Convert removed_intervals to list of tuples if needed
|
||||
removed_intervals = None
|
||||
if params.removed_intervals:
|
||||
removed_intervals = [
|
||||
(interval[0], interval[1]) if isinstance(interval, (list, tuple)) else interval
|
||||
for interval in params.removed_intervals
|
||||
]
|
||||
|
||||
return DataPreprocessor(
|
||||
target_variable=params.target_variable,
|
||||
input_columns=params.variable_columns,
|
||||
nan_treatment=params.nan_treatment,
|
||||
lag_train=params.lag_train,
|
||||
lag_transform=params.lag_val,
|
||||
start_date=params.start_date,
|
||||
end_date=params.end_date,
|
||||
date_format=params.date_format,
|
||||
removed_intervals=removed_intervals,
|
||||
static_threshold=self._get_static_threshold(params),
|
||||
low_lim=params.low_lim,
|
||||
upp_lim=params.upp_lim,
|
||||
scaler_name=params.scaler_name,
|
||||
scaler_params={} if params.use_scaler else None,
|
||||
ar_var=params.target_variable if params.include_ar else None,
|
||||
)
|
||||
|
||||
def _extract_model_equation(
|
||||
self, regr: LinearRegressionModel, params: TrainModelParams
|
||||
) -> dict:
|
||||
"""
|
||||
Extract the linear regression equation coefficients and create equation metadata.
|
||||
|
||||
This method extracts the coefficients and intercept from the trained model
|
||||
and creates a structured dictionary containing the equation information
|
||||
for serialization as JSON artifact.
|
||||
|
||||
Args:
|
||||
regr: Trained LinearRegressionModel object
|
||||
params: Training parameters containing variable information
|
||||
|
||||
Returns:
|
||||
dict: Equation metadata containing:
|
||||
- target_variable: Name of the target variable
|
||||
- coefficients: Dictionary mapping variable names to coefficients
|
||||
- intercept: Model intercept value
|
||||
- equation_string: Human-readable equation string
|
||||
- latex_equation: LaTeX formatted equation
|
||||
"""
|
||||
coefficients = regr.regr.coef_
|
||||
intercept = regr.regr.intercept_
|
||||
|
||||
# Get feature names - for polynomial models, use poly_feature_names
|
||||
if params.degree > 1 and regr.poly_feature_names:
|
||||
feature_names = regr.poly_feature_names
|
||||
else:
|
||||
feature_names = params.variable_columns
|
||||
|
||||
# Create coefficients dictionary
|
||||
coefficients_dict = {}
|
||||
for i, var in enumerate(feature_names):
|
||||
if i < len(coefficients):
|
||||
coefficients_dict[var] = float(coefficients[i])
|
||||
|
||||
# Create equation string
|
||||
equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()]
|
||||
equation_string = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(
|
||||
equation_parts
|
||||
)
|
||||
|
||||
# Create LaTeX equation
|
||||
latex_parts = [f'{coef:.6f} \\cdot {var}' for var, coef in coefficients_dict.items()]
|
||||
latex_equation = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(latex_parts)
|
||||
|
||||
return {
|
||||
'target_variable': params.target_variable,
|
||||
'coefficients': coefficients_dict,
|
||||
'intercept': float(intercept),
|
||||
'equation_string': equation_string,
|
||||
'latex_equation': latex_equation,
|
||||
'model_type': params.model_name,
|
||||
'degree': params.degree,
|
||||
'interaction_only': params.interaction_only,
|
||||
'original_features': params.variable_columns,
|
||||
}
|
||||
|
||||
def _configure_datetime_index(
|
||||
self, data: pd.DataFrame | None, params: TrainModelParams
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Configure datetime index for the DataFrame.
|
||||
|
||||
Guards against None to avoid 'NoneType' object has no attribute 'index' downstream.
|
||||
Prefers params.date_column when set; otherwise looks for common timestamp column names.
|
||||
"""
|
||||
if data is None:
|
||||
raise ValueError(
|
||||
'Data is None after load_data. '
|
||||
'Check file format, line separator and decimal separator.'
|
||||
)
|
||||
if not isinstance(data, pd.DataFrame):
|
||||
raise TypeError(f'Expected DataFrame, got {type(data).__name__}')
|
||||
|
||||
if isinstance(data.index, pd.DatetimeIndex):
|
||||
self.logger.info('DataFrame already has DatetimeIndex')
|
||||
return data.sort_index()
|
||||
|
||||
common_timestamp_columns = [
|
||||
'timestamp',
|
||||
'Timestamp',
|
||||
'TIMESTAMP',
|
||||
'date',
|
||||
'Date',
|
||||
'DATE',
|
||||
'DATA',
|
||||
'datetime',
|
||||
'DateTime',
|
||||
]
|
||||
timestamp_columns = ([params.date_column] if params.date_column else []) + [
|
||||
c for c in common_timestamp_columns if c != params.date_column
|
||||
]
|
||||
|
||||
for col in timestamp_columns:
|
||||
if col in data.columns:
|
||||
try:
|
||||
data[col] = pd.to_datetime(data[col])
|
||||
data = data.set_index(col)
|
||||
data = data.sort_index()
|
||||
self.logger.info(f'Configured datetime index from column: {col}')
|
||||
return data
|
||||
except (ValueError, TypeError) as e:
|
||||
self.logger.warning(f'Failed to convert column {col} to datetime: {e}')
|
||||
continue
|
||||
|
||||
# If no timestamp column found, check if first column looks like a timestamp
|
||||
first_col = data.columns[0]
|
||||
try:
|
||||
# Try to parse first column as datetime
|
||||
test_values = data[first_col].head(10).dropna()
|
||||
if len(test_values) > 0:
|
||||
pd.to_datetime(test_values)
|
||||
data[first_col] = pd.to_datetime(data[first_col])
|
||||
data = data.set_index(first_col)
|
||||
data = data.sort_index()
|
||||
self.logger.info(f'Configured datetime index from first column: {first_col}')
|
||||
return data
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
self.logger.warning('No timestamp column found - some features may not work correctly')
|
||||
return data
|
||||
Reference in New Issue
Block a user