gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: enhance training and experiment tracking functionality

Browse Source 
- Updated `Activities` class to improve garbage collection handling.
- Enhanced error messaging in `ExperimentTracking` for better clarity on update failures.
- Refactored `Training` class to streamline exception handling and improve type hints.
- Introduced new methods in `TrainModelParams` for better handling of experiment run IDs and model metadata.
- Added functionality to extract model equations in `DataManagerRepository` for linear regression models.
This commit is contained in:
Vitor Pimentel dos Santos
2026-04-06 15:05:57 -03:00
parent 1352d1ac8f
commit 6b1df7c3a7
22 changed files with 1751 additions and 2085 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
model_manager/activities/activities.py
View File

@@ -141,9 +141,9 @@ class Activities(ExperimentTracking, Training, Cleanup):
if hasattr(self, 'engine'):
try:
# Call parent class __del__ if it exists
if hasattr(super(), '__del__'):
super().__del__()
except Exception: # noqa: S110, BLE001
if hasattr(super(), '__del__'): # pragma: no cover
super().__del__() # pragma: no cover
except Exception: # noqa: S110, BLE001 # pragma: no cover
# Silently ignore errors during garbage collection
# Logging here could cause issues if logger is already destroyed
pass

3
model_manager/activities/experiment_tracking.py
View File

@@ -260,7 +260,8 @@ class ExperimentTracking(Postgres):
if result.get('rowcount', 0) == 0:
error_msg = (
f'No rows updated for experiment run {experiment_run_id} with status {status}'
f'No experiment_run row updated for id={experiment_run_id} '
f'(row missing or id mismatch). update_type={update_type!r}, status={status!r}.'
)
raise ValueError(error_msg)

20
model_manager/activities/training.py
View File

@@ -118,7 +118,7 @@ class Training(SientiaMonitoring):
TrainModelParams: Validated and converted training parameters
Raises:
ValueError, TypeError, KeyError: If validation fails (after sending notification)
Exception: If validation fails (after sending notification)
"""
metadata = input_data.get('metadata', {})
try:
@@ -134,7 +134,7 @@ class Training(SientiaMonitoring):
)
return train_params
except (ValueError, TypeError, KeyError) as e:
except Exception as e:
error_msg = f'Error validating training parameters: {str(e)}'
trace = traceback.format_exc()
@@ -165,7 +165,7 @@ class Training(SientiaMonitoring):
- train_params (TrainModelParams | dict): Training parameters.
Returns:
dict: Key `run_name` when training and saving succeed.
dict[str, Any]: Serializable summary (run identifiers, run_dir for cleanup, regression metrics).
Raises:
ValueError: If input validation fails.
@@ -177,9 +177,6 @@ class Training(SientiaMonitoring):
if isinstance(train_params, dict):
train_params = TrainModelParams.from_dict(train_params)
model_trained = False
model_saved = False
try:
# Download training file bytes from MinIO
train_bytes = await self.minio_repository.download_file(
@@ -238,10 +235,9 @@ class Training(SientiaMonitoring):
train_result = self.data_manager_repository.compute_regression_metrics(
train_result,
wrapper,
)
model_trained = True
async with self.mlflow_repository.start_run(
model_name=train_params.model_name,
run_name=None,
@@ -266,7 +262,11 @@ class Training(SientiaMonitoring):
mlflow.log_artifact(train_result.train_data_path)
mlflow.log_artifact(train_result.test_data_path)
return train_result.to_dict()
return {
'run_name': train_result.run_name,
'run_id': train_result.run_id,
'run_dir': train_result.run_dir
}
except Exception as e: # noqa: BLE001
error_msg = f'Error training model - error: {str(e)}'
@@ -284,7 +284,7 @@ class Training(SientiaMonitoring):
raise e
activity.defn(name='cleanup_resources')
@activity.defn(name='cleanup_resources')
async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
"""
Cleanup temporary resources created during training.

60
model_manager/utils/models/train_model_params.py
View File

@@ -84,6 +84,8 @@ class TrainModelParams:
Args:
data: Dictionary containing training parameters with keys matching the
attribute names (e.g. variable_columns, data_model_kwargs, model_kwargs, opt_params).
model_metadata may be omitted or None until load_model_metadata fills it.
experiment_run_id may be an int or numeric string.
Returns:
TrainModelParams: Validated instance with all fields populated
@@ -111,7 +113,7 @@ class TrainModelParams:
train_size=cls._check_none(data.get('train_size'), int, 'train_size'),
shuffle=cls._check_none(data.get('shuffle'), bool, 'shuffle'),
random_state=cls._check_none(data.get('random_state', 42), int, 'random_state'),
experiment_run_id=cls._check_none(data.get('experiment_run_id'), int, 'experiment_run_id'),
experiment_run_id=cls._coerce_experiment_run_id(data.get('experiment_run_id')),
model_name=model_name,
experiment_name=model_name + '_experiment',
val_file_name=data.get('val_file_name'),
@@ -121,7 +123,7 @@ class TrainModelParams:
model_type=cls._check_none(data.get('model_type'), str, 'model_type'),
model_id=data.get('model_id'),
model_metadata=cls._check_none(data.get('model_metadata'), dict, 'model_metadata'),
model_metadata=cls._parse_optional_model_metadata(data.get('model_metadata')),
)
def to_dict(self) -> dict[str, Any]:
@@ -181,6 +183,60 @@ class TrainModelParams:
return value
@staticmethod
def _coerce_experiment_run_id(value: Any) -> int:
"""
Coerce experiment_run_id to int.
Workflow clients may send numeric strings; this keeps from_dict aligned with
workflow validation.
Args:
value: Raw experiment_run_id from the payload.
Returns:
int: Parsed experiment run id.
Raises:
ValueError: If the value is None.
TypeError: If the value cannot be coerced to a non-boolean integer.
"""
if value is None:
raise ValueError('experiment_run_id is required and cannot be None.')
if isinstance(value, bool):
raise TypeError('experiment_run_id must be an integer, got bool.')
if isinstance(value, int):
return value
if isinstance(value, str) and value.strip().isdigit():
return int(value.strip())
if isinstance(value, float) and value.is_integer():
return int(value)
raise TypeError(
f'experiment_run_id must be an integer or numeric string, but got {type(value).__name__}.'
)
@staticmethod
def _parse_optional_model_metadata(value: Any) -> dict | None:
"""
Parse model_metadata for from_dict before load_model_metadata fills the index.
Args:
value: model_metadata from the payload, or None if not sent yet.
Returns:
dict | None: Dict when provided; None when absent (filled later by load_model_metadata).
Raises:
TypeError: If value is neither None nor a dict.
"""
if value is None:
return None
if isinstance(value, dict):
return value
raise TypeError(
f'model_metadata must be a dict or None, but got {type(value).__name__}.'
)
def validate_business_rules(self) -> None:
"""
Validate business rules and constraints for training parameters.

6
model_manager/utils/models/train_model_result.py
View File

@@ -51,9 +51,3 @@ class TrainModelResult:
test_data_path: str | None = None
run_dir: str | None = None
def to_dict(self) -> dict[str, Any]:
"""
Convert TrainModelResult to a dictionary.
"""
return self.__dict__

69
model_manager/utils/repository/data_manager_repository.py
View File

@@ -25,6 +25,7 @@ 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 sientia_model.wrappers.sientia_model import SientiaModel
from model_manager.sientia.metrics import mae, mse, r2
from model_manager.utils.models.train_model_params import TrainModelParams
@@ -206,10 +207,75 @@ class DataManagerRepository(SientiaMonitoring):
if pred.shape[1] == 1:
return pred.iloc[:, 0]
raise ValueError('y_pred/y_train_pred must be a Series or single-column DataFrame')
def _extract_model_equation(
self, regr: Any, 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
model_kwargs = params.model_kwargs or {}
degree = model_kwargs.get('degree', 1)
poly_feature_names = model_kwargs.get('poly_feature_names', None)
if degree > 1 and poly_feature_names:
feature_names = 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': degree,
'interaction_only': model_kwargs.get('interaction_only', False),
'original_features': feature_names,
}
def compute_regression_metrics(
self,
tmr: TrainModelResult,
wrapper: SientiaModel,
) -> TrainModelResult:
"""
Compute regression metrics for training results.
@@ -252,6 +318,9 @@ class DataManagerRepository(SientiaMonitoring):
tmr.mae_val = mae(y_true_val, y_pred_val)
tmr.r2_val = r2(y_true_val, y_pred_val)
if params.model_type == 'linear_regression':
tmr.equation = self._extract_model_equation(wrapper.model, params)
return tmr
def _configure_datetime_index(

4
model_manager/workflows/cleanup_files.py
View File

@@ -13,7 +13,7 @@ with workflow.unsafe.imports_passed_through():
from typing import Any
from model_manager.activities.activities import Activities
from model_manager.workflows.train_model import POD_ID, no_retry_policy
from model_manager.workflows.train_model import no_retry_policy
TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120'))
@@ -45,7 +45,7 @@ class CleanupFiles:
# Metadata for tracking
metadata = {
'metadata': {
'pod_id': POD_ID,
'pod_id': os.getenv('POD_ID'),
'workflow_name': 'cleanup_files',
}
}

41
model_manager/workflows/train_model.py
View File

@@ -23,8 +23,9 @@ with workflow.unsafe.imports_passed_through():
from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams
# Activity Timeouts (in seconds) - Configurable via environment variables
# Defaults are designed to handle large files (up to 200MB)
# Activity timeouts (seconds). Tune per environment (large uploads, long training).
# Training uses no_retry_policy: extend TIMEOUT_TRAIN_MODEL instead of adding retries
# to avoid duplicate MLflow side effects. Cleanup/delete uses network_retry_policy.
TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '2700'))
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '120'))
@@ -71,7 +72,7 @@ class TrainModel:
"""
@workflow.run
async def run(self, input_data: dict[str, Any]) -> dict[str, str | None] | None:
async def run(self, input_data: dict[str, Any]) -> dict[str, Any] | None:
"""
Execute the complete model training workflow.
@@ -95,9 +96,11 @@ class TrainModel:
ValueError: If experiment_run_id is missing or invalid
"""
experiment_run_id = self._validate_experiment_run_id(input_data)
input_data = {**input_data, 'experiment_run_id': experiment_run_id}
model_name = input_data.get('model_name')
model_id = input_data.get('model_id')
metadata = {
'metadata': {
'experiment_run_id': experiment_run_id,
@@ -112,7 +115,7 @@ class TrainModel:
)
training_succeeded = False
train_result: dict[str, str | None] | None = None
train_result: dict[str, Any] | None = None
try:
train_result = await self._train_model(
@@ -128,9 +131,15 @@ class TrainModel:
run_dir=train_result.get('run_dir'),
metadata=metadata,
)
else:
pass
except Exception:
if training_succeeded:
raise
# If cleanup fails after training failed, there is nothing extra to log (DB not committed).
if training_succeeded: # pragma: no branch
workflow.logger.warning(
'cleanup_resources failed after successful training; model and DB status '
'are already committed. Temp files may remain until scheduled cleanup.',
)
return train_result
def _validate_experiment_run_id(self, input_data: dict[str, Any]) -> int:
@@ -227,8 +236,11 @@ class TrainModel:
status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
error_message=self._extract_error_message(e),
)
except Exception:
pass
except Exception as secondary:
workflow.logger.warning(
'Failed to persist ORCHESTRATOR_VALIDATION_ERROR to experiment_run: %s',
secondary,
)
raise
async def _train_model(
@@ -236,7 +248,7 @@ class TrainModel:
train_params: TrainModelParams,
experiment_run_id: int,
metadata: dict[str, Any],
) -> dict[str, str | None]:
) -> dict[str, Any]:
"""
Download file from MinIO and train model.
@@ -250,7 +262,7 @@ class TrainModel:
metadata: Workflow execution metadata
Returns:
TrainModelResult: Training result from train_model activity
dict[str, Any]: Serializable training summary from the train_model activity
Raises:
Exception: If download or training fails (after updating DB status)
@@ -285,8 +297,11 @@ class TrainModel:
status=ExperimentStatus.TRAINING_ERROR,
error_message=self._extract_error_message(e),
)
except Exception:
pass
except Exception as secondary:
workflow.logger.warning(
'Failed to persist TRAINING_ERROR status to experiment_run: %s',
secondary,
)
raise
async def _cleanup_resources(