Code import - branch release/SIENTIAPDE-1645

This commit is contained in:
2026-08-05 13:53:37 +00:00
commit d481e0acff
116 changed files with 92848 additions and 0 deletions

View File

View File

View File

@@ -0,0 +1,166 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
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_sync import MinioRepository
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
from sientia_model.model_repository.plugin_store import PluginStore
from model_manager.activities.cleanup import Cleanup
from model_manager.activities.experiment_tracking import ExperimentTracking
from model_manager.activities.training import Training
class Activities(ExperimentTracking, Training, Cleanup):
"""
Main activities orchestrator for the Model Manager system.
This class combines functionality from multiple activity classes to provide
a unified interface for all workflow operations. It manages database connections,
MLFlow model interactions, MinIO storage operations, and cleanup operations.
The class implements multiple inheritance to combine specialized functionality:
- ExperimentTracking: ML experiment lifecycle tracking and database operations
- Training: ML model training operations with MLFlow and MinIO integration
- Cleanup: File and directory cleanup operations for MinIO and local filesystem
Attributes:
postgres_config (dict): PostgreSQL connection configuration
mlflow_config (dict): MLFlow server configuration
minio_config (dict): MinIO storage configuration
logger (Logger): Logging and observability instance
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(
self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
minio_config: dict[str, Any],
plugin_store: PluginStore,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
"""
Initialize the Activities orchestrator with all required configurations.
This constructor initializes all parent classes with their respective
configurations and sets up the foundation for all activity operations.
Args:
postgres_config: PostgreSQL connection configuration dictionary
Required keys: host, port, user, password, dbname, min_connections, max_connections
mlflow_config: MLFlow server configuration dictionary
Required keys: host, port, username, password
minio_config: MinIO storage configuration dictionary
Required keys: endpoint_url, access_key, secret_key, region, use_ssl, default_bucket
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If any parent class initialization fails
"""
ExperimentTracking.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
self.mlflow_repository = SientiaMLflowRepository(
host=mlflow_config['url'],
username=mlflow_config['username'],
password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
# MinIO repository used for all object storage operations
endpoint_url = minio_config['endpoint_url']
# MinioRepository expects the endpoint without scheme
if endpoint_url.startswith('http://'):
endpoint = endpoint_url.removeprefix('http://')
elif endpoint_url.startswith('https://'):
endpoint = endpoint_url.removeprefix('https://')
else:
endpoint = endpoint_url
self.minio_repository = MinioRepository(
endpoint=endpoint,
access_key=minio_config['access_key'],
secret_key=minio_config['secret_key'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
secure=minio_config['use_ssl'],
bucket=minio_config['default_bucket'],
)
Training.__init__(
self,
mlflow_repository=self.mlflow_repository,
plugin_store=plugin_store,
minio_repository=self.minio_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
Cleanup.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
def __del__(self):
"""
Destructor to safely handle cleanup during garbage collection.
This prevents AttributeError when the parent Postgres.__del__ tries to access
self.engine in objects with multiple inheritance. Only attempts cleanup if
the engine attribute exists.
"""
# Only call parent __del__ if engine attribute exists
# This prevents AttributeError in multiple inheritance scenarios
if hasattr(self, 'engine'):
try:
# Call parent class __del__ if it exists
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
def shutdown(self):
"""
Gracefully shutdown all activities and clean up resources.
This method ensures proper cleanup of all resources including:
- PostgreSQL connection pools (via ExperimentTracking)
- Any other resources that need explicit cleanup
The method should be called before the application terminates to ensure
proper resource cleanup and prevent resource leaks.
Prefer calling this method explicitly rather than relying on __del__.
"""
ExperimentTracking.close(self)
self.info('Postgres client closed')
SientiaMonitoring.shutdown(self)

View File

@@ -0,0 +1,177 @@
"""
Cleanup activities for removing stale files from local filesystem.
This module provides activities for cleaning up temporary files and directories
that are older than the configured retention period. It operates independently
of the database, using timestamps embedded in filenames.
"""
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import os
import re
import shutil
import traceback
from datetime import datetime, timedelta
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from model_manager.runtime_paths import REPORTS_TEMP_DIR
RETENTION_HOURS = int(os.getenv('CLEANUP_RETENTION_HOURS', '24'))
DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true'
class Cleanup(SientiaMonitoring):
"""
Activity for cleaning up stale files and directories.
This activity extends SientiaMonitoring and handles cleanup of:
- Local temporary directories with timestamp suffixes
"""
def __init__(
self,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
"""
Initialize Cleanup activity.
Args:
logger: Logger instance for observability
notification_handler: Handler for sending notifications
metrics_controller: Controller for metrics emission
"""
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
# Configuration from environment variables
self.retention_hours = RETENTION_HOURS
self.dry_run = DRY_RUN
# Regex patterns for timestamp extraction
self.dir_timestamp_pattern = re.compile(
r'^(.+)_(\d{8}_\d{6}_\d{6})$'
) # name_YYYYMMDD_HHMMSS_microseconds
@activity.defn(name='cleanup_temp_directories')
def cleanup_temp_directories(self, input_data: dict[str, Any]) -> None:
"""
Clean up stale temporary directories based on timestamp in directory name.
This activity scans the reports/temp directory for subdirectories following
the pattern '{name}_{timestamp}' where timestamp is in YYYYMMDD_HHMMSS_microseconds format.
Directories older than the retention period are deleted.
Args:
input_data: Cleanup configuration containing:
- metadata (dict): Workflow execution metadata
- temp_path (str): Path to temp directory (optional, defaults to reports/temp)
Returns:
None: Results are logged and tracked via metrics
Raises:
Exception: If cleanup fails (after sending notification)
"""
metadata = input_data.get('metadata', {})
temp_path = input_data.get('temp_path', REPORTS_TEMP_DIR)
cutoff_time = datetime.now() - timedelta(hours=self.retention_hours)
try:
self.info(
f'Starting local directory cleanup - Path: {temp_path}, '
f'Retention: {self.retention_hours}h, Dry run: {self.dry_run}',
metadata,
)
if not os.path.exists(temp_path):
self.warning(f'Temp directory does not exist: {temp_path}', metadata)
return
directories_scanned = 0
directories_deleted = 0
errors = []
for item_name in os.listdir(temp_path):
item_path = os.path.join(temp_path, item_name)
if not os.path.isdir(item_path):
continue
directories_scanned += 1
# Extract timestamp from directory name
match = self.dir_timestamp_pattern.match(item_name)
if not match:
self.debug(
f'Skipping directory without timestamp pattern: {item_name}', metadata
)
continue
timestamp_str = match.group(2)
try:
# Parse YYYYMMDD_HHMMSS_microseconds
dir_time = datetime.strptime(timestamp_str, '%Y%m%d_%H%M%S_%f')
if dir_time < cutoff_time:
age_hours = (datetime.now() - dir_time).total_seconds() / 3600
if self.dry_run:
self.info(
f'[DRY RUN] Would delete directory: {item_name} (age: {age_hours:.1f}h)',
metadata,
)
directories_deleted += 1
else:
try:
shutil.rmtree(item_path)
self.info(
f'Deleted stale directory: {item_name} (age: {age_hours:.1f}h)',
metadata,
)
directories_deleted += 1
except OSError as e:
error_msg = f'Failed to delete directory {item_name}: {str(e)}'
errors.append(error_msg)
self.error(error_msg, metadata)
else:
age_hours = (datetime.now() - dir_time).total_seconds() / 3600
self.debug(
f'Keeping recent directory: {item_name} (age: {age_hours:.1f}h)',
metadata,
)
except ValueError as e:
error_msg = f'Invalid timestamp format in directory {item_name}: {str(e)}'
errors.append(error_msg)
self.error(error_msg, metadata)
self.info(
f'Directory cleanup completed - Scanned: {directories_scanned}, '
f'Deleted: {directories_deleted}, Errors: {len(errors)}',
metadata,
)
except Exception as e:
error_msg = f'Error in directory cleanup: {str(e)}'
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='CLEANUP_DIRECTORIES_ERROR',
message=error_msg,
block='cleanup_temp_directories',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise

View File

@@ -0,0 +1,280 @@
"""
Experiment tracking activities for managing ML experiment lifecycle.
This module provides activities for tracking and updating experiment run status
in the PostgreSQL database, extending the synchronous Postgres client with specialized
methods for experiment management.
"""
import enum
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.temporal.activities.postgres_sync import Postgres
from sqlalchemy import text
class UpdateType(enum.StrEnum):
"""Types of experiment run updates."""
STATUS = 'status'
STATUS_WITH_ERROR = 'status_with_error'
MODEL_SAVED = 'model_saved'
class ExperimentTracking(Postgres):
"""
Activity for tracking ML experiment lifecycle and status updates.
This activity extends the Postgres activity to provide specialized methods
for managing experiment runs, including status updates, error tracking, and
model registration. It maintains the experiment lifecycle from initialization
through training, model saving, and cleanup.
The activity uses a SQLAlchemy engine / connection pool and adds experiment-specific
operations with proper error handling and notifications.
"""
def __init__(
self,
host: str,
port: int,
user: str,
password: str,
dbname: str,
min_connections: int,
max_connections: int,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
"""
Initialize ExperimentTracking activity with database configuration.
Args:
host: PostgreSQL server hostname
port: PostgreSQL server port
user: Database user
password: Database password
dbname: Database name
min_connections: Minimum connections in pool
max_connections: Maximum connections in pool
logger: Logger instance for observability
notification_handler: Notification handler for alerts
metrics_controller: Metrics controller for observability
Raises:
ConnectionError: If database connection cannot be established
"""
Postgres.__init__(
self,
host=host,
port=port,
user=user,
password=password,
dbname=dbname,
min_connections=min_connections,
max_connections=max_connections,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
self.info(f'Postgres client initialized at {host}:{port}')
def __del__(self):
"""
Destructor to safely handle cleanup during garbage collection.
This prevents AttributeError when used in multiple inheritance scenarios
where the parent Postgres.__del__ might be called on objects without
the engine attribute.
"""
# Only call parent __del__ if engine attribute exists
if hasattr(self, 'engine'):
try:
if hasattr(super(), '__del__'):
super().__del__()
except Exception: # noqa: S110, BLE001
# Silently ignore errors during garbage collection
pass
def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
"""
Execute an UPDATE SQL statement.
Args:
query: Parameterized SQL string to execute.
params: Mapping of parameters for the SQL query.
Returns:
dict: A dictionary containing the affected row count: {'rowcount': int}.
"""
with self.engine.begin() as connection:
result = connection.execute(text(query), params)
return {'rowcount': result.rowcount}
def _build_status_update_query(
self, status: str | None, experiment_run_id: int
) -> tuple[str, dict[str, Any]]:
"""Build SQL query for simple status update."""
if not isinstance(status, str) or not status:
raise ValueError('status is required for STATUS update type')
sql_query = """
UPDATE experiment_run
SET status = :status, updated_at = :updated_at
WHERE id = :experiment_run_id
"""
query_params = {
'status': status,
'updated_at': datetime.now(UTC),
'experiment_run_id': experiment_run_id,
}
return sql_query, query_params
def _build_status_with_error_query(
self, status: str | None, error_message: str | None, experiment_run_id: int
) -> tuple[str, dict[str, Any]]:
"""Build SQL query for status update with error message."""
if not isinstance(status, str) or not status:
raise ValueError('status is required for STATUS_WITH_ERROR update type')
if not isinstance(error_message, str) or not error_message:
raise ValueError('error_message is required for STATUS_WITH_ERROR update type')
# Truncate error message if too long
truncated_error = error_message[:1024] if len(error_message) > 1024 else error_message
sql_query = """
UPDATE experiment_run
SET status = :status, error_message = :error_message, updated_at = :updated_at
WHERE id = :experiment_run_id
"""
query_params = {
'status': status,
'error_message': truncated_error,
'updated_at': datetime.now(UTC),
'experiment_run_id': experiment_run_id,
}
return sql_query, query_params
def _build_model_saved_query(
self, run_name: str | None, status: str | None, experiment_run_id: int
) -> tuple[str, dict[str, Any]]:
"""Build SQL query for model saved update."""
if not isinstance(run_name, str) or not run_name:
raise ValueError('run_name is required for MODEL_SAVED update type')
if not isinstance(status, str) or not status:
raise ValueError('status is required for MODEL_SAVED update type')
sql_query = """
UPDATE experiment_run
SET run_name = :run_name, status = :status, updated_at = :updated_at
WHERE id = :experiment_run_id
"""
query_params = {
'run_name': run_name,
'status': status,
'updated_at': datetime.now(UTC),
'experiment_run_id': experiment_run_id,
}
return sql_query, query_params
def _get_update_query_and_params(
self, update_type: str, experiment_run_id: int, input_data: dict[str, Any]
) -> tuple[str, dict[str, Any]]:
"""Get SQL query and parameters based on update type."""
status = input_data.get('status')
error_message = input_data.get('error_message')
run_name = input_data.get('run_name')
if update_type == UpdateType.STATUS:
return self._build_status_update_query(status, experiment_run_id)
if update_type == UpdateType.STATUS_WITH_ERROR:
return self._build_status_with_error_query(status, error_message, experiment_run_id)
if update_type == UpdateType.MODEL_SAVED:
return self._build_model_saved_query(run_name, status, experiment_run_id)
raise ValueError(f'Invalid update_type: {update_type}')
@activity.defn(name='update_experiment_run')
def update_experiment_run(self, input_data: dict[str, Any]) -> None:
"""
Update experiment run with status, errors, or model information.
This activity provides a unified interface for all experiment run updates,
supporting different update types through a single method. It automatically
selects the appropriate SQL query based on the update type and parameters.
Args:
input_data: Configuration for experiment run update operation
Required keys:
- metadata (dict): Workflow execution metadata
- experiment_run_id (int): Unique identifier for the experiment run
- update_type (str): Type of update (status, status_with_error, model_saved)
Optional keys:
- status (str): New status for the experiment run
- error_message (str): Error message if update failed
- run_name (str): MLFlow run name if model was saved
Raises:
ValueError: If required parameters are missing for the update type
RuntimeError: If update operation fails
"""
metadata = input_data.get('metadata')
experiment_run_id = input_data['experiment_run_id']
update_type = input_data['update_type']
status = input_data.get('status')
try:
sql_query, query_params = self._get_update_query_and_params(
update_type, experiment_run_id, input_data
)
result = self._execute_update(sql_query, query_params)
if result.get('rowcount', 0) == 0:
error_msg = (
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)
self.info(
f'Successfully updated experiment run {experiment_run_id} with status {status}',
metadata,
)
except Exception as e: # noqa: BLE001
error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Status: {status}, Error: {str(e)}'
trace = traceback.format_exc()
self.send_notification(
metadata=metadata or {},
notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
message=error_msg,
block='update_experiment_run',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise RuntimeError(error_msg) from e

View File

@@ -0,0 +1,497 @@
"""
Training activities for ML model training operations.
This module provides activities for training machine learning models.
The activity extends BaseActivity and receives pre-downloaded files
and raises `ModelTrainingError` when training fails.
"""
from sientia_model.wrappers.sientia_model import SientiaModel
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import os
import time
import traceback
from typing import Any
import mlflow
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.minio_repository_sync import MinioRepository
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
from sientia_model.model_repository.plugin_store import PluginStore
from model_manager import metrics as mm_metrics
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.utils.repository.data_manager_repository import DataManagerRepository
class Training(SientiaMonitoring):
"""
Activity for ML model training operations.
This activity extends SientiaMonitoring and handles machine learning model
training with comprehensive error handling. It receives pre-downloaded
files from the workflow and raises `ModelTrainingError` on failure so the
workflow can map the correct experiment status.
"""
def __init__(
self,
mlflow_repository: SientiaMLflowRepository,
plugin_store: PluginStore,
minio_repository: MinioRepository,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
"""
Initialize Training activity.
Args:
logger: Logger instance for observability
notification_handler: Handler for sending notifications
"""
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.data_manager_repository = DataManagerRepository(logger)
self.mlflow_repository = mlflow_repository
self.plugin_store = plugin_store
self.minio_repository = minio_repository
@activity.defn(name='load_model_metadata')
def load_model_metadata(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Load model metadata/schemas from the model store.
This activity is responsible for fetching model metadata/schemas from the
model store index and extracting a serializable `model_metadata` dict that
`TrainModelParams.validate_business_rules()` depends on.
Args:
input_data: Workflow input at the same level as `validate_train_params`,
including at least `model_name` and the fields required by
`TrainModelParams.from_dict` to build wrapper kwargs.
Return:
dict[str, Any]: Updated `input_data` containing `input_data['model_metadata']`.
"""
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_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()
self.send_notification(
metadata=metadata,
notification_id='LOAD_MODEL_METADATA_ERROR',
message=f'Error loading model metadata: {str(exc)}',
block='load_model_metadata',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise
@activity.defn(name='validate_train_params')
def validate_train_params(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Validate and convert training parameters from dict to TrainModelParams.
This activity validates the input training parameters and converts them
to a TrainModelParams object.
Args:
input_data: Training parameters and metadata at the same level
Required keys:
- metadata (dict): Workflow execution metadata
- All TrainModelParams fields (experiment_run_id, target_variable, etc.)
Returns:
dict[str, Any]: Validated and converted training parameters as dictionary
Raises:
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)
train_params.validate_business_rules()
self.info(
f'Training parameters validated successfully - '
f'Target: {train_params.target_variable}, '
f'Experiment: {train_params.experiment_name}',
metadata,
)
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)}'
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
message=error_msg,
block='validate_train_params',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise
@activity.defn(name='train_model')
def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Train a machine learning model.
This activity orchestrates the ML training pipeline:
1. Validate input parameters.
2. Prepare data via DataManagerRepository.
3. Train the model and compute metrics.
Args:
input_data: Training configuration containing:
- metadata (dict): Workflow execution metadata.
- uploaded_file (BytesIO): Training data already downloaded from MinIO.
- train_params (dict): Training parameters.
Returns:
dict[str, Any]: Serializable summary (run identifiers, run_dir for cleanup, regression metrics).
Raises:
ValueError: If input validation fails.
Exception: If training fails (after sending notification).
"""
metadata = input_data.get('metadata')
train_params = TrainModelParams.from_dict(input_data['train_params'])
labels = self._get_training_labels(train_params)
self.info('Starting train_model process', metadata)
try:
# Download training file bytes from MinIO
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,
)
# Download optional validation file bytes from the same bucket
val_bytes: bytes | None = None
validation_name = train_params.val_file_name
if validation_name is not None:
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._prepare_data(train_bytes, val_bytes, train_params, metadata)
mm_metrics.SIENTIA_TRAINING_DATASET_TRAIN_ROWS.labels(**labels).set(
len(train_result.train_data)
)
mm_metrics.SIENTIA_TRAINING_DATASET_VAL_ROWS.labels(**labels).set(
len(train_result.val_data)
)
mm_metrics.SIENTIA_TRAINING_FEATURE_COUNT.labels(**labels).set(
len(train_params.variable_columns)
)
self.info(f'Getting model wrapper for {train_params.model_type}', metadata)
wrapper = self.plugin_store.get_model(
model_type=train_params.model_type,
force_download=False,
opt_params=train_params.opt_params or {},
model_kwargs=train_params.model_kwargs or {},
data_model_kwargs=train_params.data_model_kwargs or {},
metadata=metadata,
)
if self.logger is not None:
wrapper.logger = self.logger.base_logger
self.info(f'Training model for {train_params.model_type}', metadata)
train_result = self._fit_model(wrapper, train_result, train_params, metadata)
self.info(f'Computing regression metrics for {train_params.model_type}', metadata)
train_result = self.data_manager_repository.compute_regression_metrics(
train_result,
wrapper,
metadata=metadata,
)
if train_result.mse_val is not None:
mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels(**labels).set(
train_result.mse_val
)
if train_result.mae_val is not None:
mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels(**labels).set(
train_result.mae_val
)
if train_result.r2_val is not None:
mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels(**labels).set(
train_result.r2_val
)
mm_metrics.SIENTIA_TRAINING_INFO.labels(
pod_id=labels['pod_id'],
model_name=train_params.model_name,
model_type=train_params.model_type,
dataset_train_rows=str(len(train_result.train_data)),
dataset_val_rows=str(len(train_result.val_data)),
feature_count=str(len(train_params.variable_columns)),
mse=str(train_result.mse_val) if train_result.mse_val is not None else '',
mae=str(train_result.mae_val) if train_result.mae_val is not None else '',
r2=str(train_result.r2_val) if train_result.r2_val is not None else '',
).set(time.time() * 1000)
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=train_result.run_name,
experiment_name=train_result.experiment_name,
tags=None,
metadata=metadata,
) as run_info:
train_result.run_id = run_info.run_id
self._persist_training_artifacts(train_result, train_params, wrapper, metadata)
self.emit_metric_sync(
metric_object=mm_metrics.SIENTIA_TRAINING_MODEL_TRAINED_TOTAL,
tags=labels,
)
return {
'run_name': train_result.run_name,
'experiment_name': train_result.experiment_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)}'
trace = traceback.format_exc()
self.send_notification(
metadata=metadata or {},
notification_id='TRAIN_MODEL_ERROR',
message=error_msg,
block='train_model',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise e
def _get_training_labels(self, train_params: TrainModelParams) -> dict:
return {
'pod_id': os.getenv('HOSTNAME', 'localhost'),
'model_name': train_params.model_name,
'model_type': train_params.model_type,
}
def _prepare_data(
self,
train_bytes: bytes,
val_bytes: bytes | None,
train_params: TrainModelParams,
metadata: dict | None,
) -> TrainModelResult:
labels = self._get_training_labels(train_params)
start_time = time.time()
try:
return self.data_manager_repository.prepare_training_data(
train_file_bytes=train_bytes,
validation_file_bytes=val_bytes,
params=train_params,
metadata=metadata,
)
except Exception:
self.emit_metric_sync(
metric_object=mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_ERROR_COUNT_TOTAL,
tags=labels,
)
raise
finally:
self.observe_lag_sync(
start_time,
mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_LAG,
labels,
)
def _fit_model(
self,
wrapper: Any,
train_result: TrainModelResult,
train_params: TrainModelParams,
metadata: dict | None,
) -> TrainModelResult:
labels = self._get_training_labels(train_params)
train_data = train_result.train_data
val_data = train_result.val_data
self.debug(
f'train_model prepared data (head 10):\ntrain:\n{train_data.head(10).to_string()}'
f'\nval:\n{val_data.head(10).to_string()}',
metadata,
)
start_time = time.time()
try:
wrapper.train(
train_data=train_data,
val_data=val_data,
target=train_params.target_variable,
)
self.info(
f'Generating predictions using the trained wrapper for {train_params.model_type}',
metadata,
)
transformed_train, _ = wrapper.transform(train_data)
transformed_val, _ = wrapper.transform(val_data)
self.debug(
f'train_model transform (head 10):\ntrain:\n{transformed_train.head(10).to_string()}'
f'\nval:\n{transformed_val.head(10).to_string()}',
metadata,
)
y_train_pred_df, _ = wrapper.predict({}, transformed_train)
y_val_pred_df, _ = wrapper.predict({}, transformed_val)
self.debug(
f'train_model predict (head 10):\ntrain:\n{y_train_pred_df.head(10).to_string()}'
f'\nval:\n{y_val_pred_df.head(10).to_string()}',
metadata,
)
y_train_pred_df.sort_index(inplace=True, ascending=False)
y_val_pred_df.sort_index(inplace=True, ascending=False)
train_result.y_train_pred = y_train_pred_df
train_result.y_pred = y_val_pred_df
return train_result
except Exception:
self.emit_metric_sync(
metric_object=mm_metrics.SIENTIA_TRAINING_MODEL_FIT_ERROR_COUNT_TOTAL,
tags=labels,
)
raise
finally:
self.observe_lag_sync(
start_time,
mm_metrics.SIENTIA_TRAINING_MODEL_FIT_LAG,
labels,
)
def _persist_training_artifacts(
self,
train_result: TrainModelResult,
train_params: TrainModelParams,
wrapper: SientiaModel,
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,
)
if (
train_result.report_path is None
or train_result.train_data_path is None
or train_result.test_data_path is None
):
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._input_example = None
wrapper.store_model(name=train_params.model_name)
self._log_regression_metrics_as_params(train_result)
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)
if train_result.equation_path is not None:
mlflow.log_artifact(train_result.equation_path)
def _log_regression_metrics_as_params(self, train_result: TrainModelResult) -> None:
"""
Persist computed regression metrics as MLflow params.
Args:
train_result: Training output containing computed regression metrics.
"""
metric_params = {
'mse_val': train_result.mse_val,
'mae_val': train_result.mae_val,
'r2_val': train_result.r2_val,
}
for key, value in metric_params.items():
if value is not None:
mlflow.log_param(key, value)
@activity.defn(name='cleanup_resources')
def cleanup_resources(self, input_data: dict[str, Any]) -> None:
"""
Cleanup temporary resources created during training.
Args:
input_data: Cleanup configuration containing:
- metadata (dict): Workflow execution metadata.
- run_dir (str): Temporary directory to remove.
Raises:
Exception: If cleanup fails (after sending notification).
"""
metadata = input_data.get('metadata', {})
run_dir = input_data.get('run_dir', '')
try:
self.data_manager_repository.cleanup_run_directory(run_dir, metadata)
except Exception as e: # noqa: BLE001
error_msg = f'Error cleaning up resources - Run directory: {run_dir}, Error: {str(e)}'
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='CLEANUP_RESOURCES_ERROR',
message=error_msg,
block='cleanup_resources',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise

110
model_manager/metrics.py Normal file
View File

@@ -0,0 +1,110 @@
"""
Model Manager Metrics Module
This module defines all Prometheus metrics used by the Sientia DataOps Model Manager system
for monitoring and observability. The metrics provide insights into system performance,
training operations, and operational health.
The metrics are designed to be scraped by Prometheus and can be visualized in
Grafana or other monitoring dashboards to provide real-time visibility into
the system's operation.
Key Metric Categories:
- Application Health: Overall system status and availability
Metric Labels:
- pod_id: Kubernetes pod identifier for multi-instance deployments
"""
from prometheus_client import Counter, Gauge, Histogram
# Application health metric
APP_UP = Gauge(
'app_up',
'Indicates if the application is running (1) or shutting down (0)',
['pod_id'],
)
_TRAINING_LABELS = ['pod_id', 'model_name', 'model_type']
SIENTIA_TRAINING_INFO = Gauge(
'sientia_training_info',
'Metadata and execution timestamp (ms) of the last successful model training run',
[
'pod_id',
'model_name',
'model_type',
'dataset_train_rows',
'dataset_val_rows',
'feature_count',
'mse',
'mae',
'r2',
],
)
SIENTIA_TRAINING_MODEL_TRAINED_TOTAL = Counter(
'sientia_training_model_trained_total',
'Number of successfully completed model training runs',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_DATA_PREPARATION_LAG = Histogram(
'sientia_training_data_preparation_lag',
'Latency of prepare_training_data() in seconds',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_DATA_PREPARATION_ERROR_COUNT_TOTAL = Counter(
'sientia_training_data_preparation_error_count_total',
'Number of failures in prepare_training_data()',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_MODEL_FIT_LAG = Histogram(
'sientia_training_model_fit_lag',
'Latency of wrapper.train() (model fitting) in seconds',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_MODEL_FIT_ERROR_COUNT_TOTAL = Counter(
'sientia_training_model_fit_error_count_total',
'Number of failures in wrapper.train() (model fitting)',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_MODEL_QUALITY_MSE = Gauge(
'sientia_training_model_quality_mse',
'Mean Squared Error of the last successful model training',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_MODEL_QUALITY_MAE = Gauge(
'sientia_training_model_quality_mae',
'Mean Absolute Error of the last successful model training',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_MODEL_QUALITY_R2 = Gauge(
'sientia_training_model_quality_r2',
'R-squared of the last successful model training',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_DATASET_TRAIN_ROWS = Gauge(
'sientia_training_dataset_train_rows',
'Number of rows in the training dataset after preparation',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_DATASET_VAL_ROWS = Gauge(
'sientia_training_dataset_val_rows',
'Number of rows in the validation dataset after preparation',
_TRAINING_LABELS,
)
SIENTIA_TRAINING_FEATURE_COUNT = Gauge(
'sientia_training_feature_count',
'Number of input feature columns used for training',
_TRAINING_LABELS,
)

View File

@@ -0,0 +1,167 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<meta http-equiv="X-UA-Compatible" content="IE=edge" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<link rel="stylesheet" href="style.css" />
<title>Report</title>
<style>
* {
font-family: "Franklin Gothic Medium", "Arial Narrow", Arial, sans-serif;
margin: 0;
padding: 0;
box-sizing: border-box;
}
.logo {
padding-top: 70;
padding-bottom: 70;
position: absolute;
margin-left: -48px;
}
h1 {
color: #fff;
position: absolute;
margin-left: 45%;
}
html,
body {
scroll-behavior: smooth;
}
section {
padding-top: 90px;
width: 100%;
display: fixed;
justify-content: center;
align-items: center;
background-color: rgb(217, 217, 214, 0.7);
}
.material-symbols-outlined {
font-variation-settings: "FILL" 0, "wght" 400, "GRAD" 0, "opsz" 24;
color: #ffff;
}
/* Tooltip text */
.tooltiptext {
visibility: hidden;
background-color: rgb(0, 30, 96, 0.9);
padding: 10px;
margin-left: -90px;
font-size: 16px;
position: absolute;
top: 85px;
border-bottom-left-radius: 12px;
border-bottom-right-radius: 12px;
}
/* Show the tooltip text when you mouse over the tooltip container */
.material-symbols-outlined:hover .tooltiptext {
visibility: visible;
}
header {
position: fixed;
top: 0;
width: 100%;
height: 85px;
background: rgb(0, 30, 96, 0.95);
display: flex;
justify-content: space-between;
align-items: center;
padding: 10px 50px 0;
}
header nav {
display: absolute;
margin-left: 80%;
gap: 10px;
}
header nav a {
position: relative;
text-decoration: none;
padding: 12px 18px;
color: #fff;
font-weight: 500;
}
header nav a.active {
background-color: #001540;
position: relative;
border-radius: 12px;
}
</style>
<link
rel="stylesheet"
href="https://fonts.googleapis.com/css2?family=Material+Symbols+Outlined:opsz,wght,FILL,GRAD@20..48,100..700,0..1,-50..200"
/>
</head>
<body>
<main>
<header>
<a href="#" class="logo">
<img
src="https://aignosi.blob.core.windows.net/sientia/20231016-Aignosi_Logo_WHITE.png"
alt="Aignosi Logo"
width="247"
height="70"
/>
</a>
<h1>Report</h1>
<nav>
<a href="#data_quality" class="active"> Summary </a>
<a href="#data_drift"> Drift </a>
<a href="#regression"> Regression </a>
</nav>
<div class="material-symbols-outlined">
info
<p class="tooltiptext">
Note that "current" <br />
is related to the test <br />
set while "reference" <br />
refers to the training <br />
set
</p>
</div>
</header>
<div class="quality_div">
<section id="data_quality"></section>
</div>
<div class="data_drift_div">
<section id="data_drift"></section>
</div>
<div class="regression_div">
<section id="regression"></section>
</div>
</main>
<script>
let sec = document.querySelectorAll("section");
let links = document.querySelectorAll("nav a");
window.onscroll = () => {
sec.forEach((section) => {
let top = window.scrollY;
let offset = section.offsetTop;
let height = section.offsetHeight;
let id = section.getAttribute("id");
if (top >= offset && top < offset + height) {
links.forEach((link) => {
link.classList.remove("active");
document.querySelector("nav a[href*=" + id + "]").classList.add("active");
});
}
});
};
</script>
</body>
</html>

View File

View File

@@ -0,0 +1,29 @@
"""Filesystem layout for worker runtime data outside the application package tree."""
from os import makedirs
from os.path import join
# Root for all mutable runtime data (not under /app; avoids clashing with git clone under /app).
RUNTIME_DATA_ROOT = '/var/lib/model-manager'
# Training reports (HTML, CSV exports, etc.) and related outputs.
REPORTS_ROOT = join(RUNTIME_DATA_ROOT, 'reports')
# project base path
PROJECT_BASE_PATH = '/app/model_manager'
# Per-training run folders (name + timestamp); cleanup cron deletes stale entries here.
REPORTS_TEMP_DIR = join(REPORTS_ROOT, 'temp')
# Worker log files when file logging is wired; stdout remains primary until then.
LOGS_DIR = join(RUNTIME_DATA_ROOT, 'logs')
def ensure_runtime_directories() -> None:
"""Create runtime directories expected by the worker process."""
# REPORTS_ROOT: base directory for report artifacts; remove if all outputs move elsewhere.
makedirs(REPORTS_ROOT, exist_ok=True)
# REPORTS_TEMP_DIR: transient run subdirs; remove after retention/cleanup is centralized.
makedirs(REPORTS_TEMP_DIR, exist_ok=True)
# LOGS_DIR: on-disk logs; remove if logging stays stdout-only forever.
makedirs(LOGS_DIR, exist_ok=True)

View File

View File

@@ -0,0 +1,160 @@
"""Schedule configuration for cleanup workflow."""
import os
from datetime import timedelta
from typing import Any
from sientia_do.observability.logger import Logger as SientiaLogger
from temporalio.client import (
Client,
Schedule,
ScheduleActionStartWorkflow,
ScheduleSpec,
)
from model_manager.worker.prepare_worker import build_queue_name
# Schedule configuration from environment variables
CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 0 * * *') # Default: midnight UTC
CLEANUP_TIMEZONE = os.getenv('CLEANUP_TIMEZONE', 'UTC')
CLEANUP_EXECUTION_TIMEOUT_HOURS = int(os.getenv('CLEANUP_EXECUTION_TIMEOUT_HOURS', '1'))
def build_cleanup_schedule_id(runtime: str | None) -> str:
"""
Build cleanup schedule ID using runtime-derived naming.
Args:
runtime: Runtime suffix used by workers
Return:
str: Cleanup schedule ID
"""
normalized_runtime = runtime.strip() if runtime else ''
return f'cleanup-files-{normalized_runtime or "single"}-daily'
async def _needs_schedule_reconcile(
schedule_handle: Any,
cleanup_task_queue: str,
logger: SientiaLogger,
metadata: dict[str, str | None],
) -> bool:
"""
Compare configured cleanup schedule against current expected values.
Args:
schedule_handle: Temporal schedule handle for current schedule ID
cleanup_task_queue: Expected cleanup task queue
logger: Logger instance for errors
metadata: Metadata dictionary for logging context
Return:
bool: True when schedule should be recreated to apply current config
"""
try:
schedule_description = await schedule_handle.describe()
schedule = getattr(schedule_description, 'schedule', None)
action = getattr(schedule, 'action', None)
spec = getattr(schedule, 'spec', None)
current_task_queue = getattr(action, 'task_queue', None)
current_execution_timeout = getattr(action, 'execution_timeout', None)
current_cron = getattr(spec, 'cron_expressions', None)
current_timezone = getattr(spec, 'time_zone_name', None)
return (
current_task_queue != cleanup_task_queue
or current_execution_timeout != timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS)
or current_cron != [CLEANUP_CRON]
or current_timezone != CLEANUP_TIMEZONE
)
except Exception as e: # noqa: BLE001
logger.custom_error(f'Error describing cleanup schedule for reconcile: {e}', metadata)
return True
async def schedule_exists(
client: Client, schedule_id: str, logger: SientiaLogger, metadata: dict[str, str | None]
) -> bool:
"""
Check if a schedule already exists.
Args:
client: Temporal client instance
schedule_id: ID of the schedule to check
logger: Logger instance for error logging
metadata: Metadata dictionary for logging context
Returns:
True if schedule exists, False otherwise
"""
try:
async for schedule in await client.list_schedules():
if schedule.id == schedule_id:
return True
return False
except Exception as e: # noqa: BLE001
logger.custom_error(f'Error checking if schedule exists: {e}', metadata)
return False
async def create_cleanup_schedule(
client: Client, logger: SientiaLogger, metadata: dict[str, str | None]
) -> None:
"""
Create or update the cleanup files schedule.
This function is idempotent and can be called multiple times safely.
It will only create the schedule if it doesn't already exist.
Args:
client: Temporal client instance
logger: Logger instance for logging schedule operations
metadata: Metadata dictionary for logging context
"""
runtime = (os.getenv('RUNTIME') or 'single').strip()
cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single')
schedule_id = build_cleanup_schedule_id(runtime)
updated = False
if await schedule_exists(client, schedule_id, logger, metadata):
handle = client.get_schedule_handle(schedule_id)
if await _needs_schedule_reconcile(handle, cleanup_task_queue, logger, metadata):
await handle.delete()
updated = True
else:
logger.custom_info(
f"Schedule '{schedule_id}' is already up to date, no-op reconcile",
metadata,
)
return
await client.create_schedule(
schedule_id,
Schedule(
action=ScheduleActionStartWorkflow(
'cleanup_files',
{}, # Empty input, will use default bucket from environment
id=f'cleanup-files-scheduled-{schedule_id}',
task_queue=cleanup_task_queue,
execution_timeout=timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS),
),
spec=ScheduleSpec(
cron_expressions=[CLEANUP_CRON],
time_zone_name=CLEANUP_TIMEZONE,
),
),
)
if updated:
logger.custom_info(
f"Schedule '{schedule_id}' reconciled successfully. "
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
metadata,
)
else:
logger.custom_info(
f"Schedule '{schedule_id}' created successfully. "
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
metadata,
)

View File

View File

@@ -0,0 +1,3 @@
from mlflow.exceptions import MlflowException
SientiaMlException = MlflowException

View File

@@ -0,0 +1,148 @@
import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
def mse(real_data: pd.Series, predictions: pd.Series) -> float:
"""
Calculates the mean squared error between the real data and the predictions.
"""
return round(
mean_squared_error(real_data.astype(np.float64), predictions.astype(np.float64)), 2
)
def mae(real_data: pd.Series, predictions: pd.Series) -> float:
"""
Calculates the mean absolute error between the real data and the predictions.
"""
return round(
mean_absolute_error(real_data.astype(np.float64), predictions.astype(np.float64)), 2
)
def r2(real_data: pd.Series, predictions: pd.Series) -> float:
"""
Calculates the R2 score between the real data and the predictions.
"""
return round(r2_score(real_data.astype(np.float64), predictions.astype(np.float64)), 2)
def silverman_radius(data: np.ndarray) -> float:
"""
Calculate the Silverman bandwidth (radius) for a given dataset.
Args:
data (np.ndarray): Input data (1D array)
Returns:
float: Silverman bandwidth (radius)
"""
n = len(data)
sigma = np.std(data)
iqr = np.percentile(data, 75) - np.percentile(data, 25)
radius = 0.9 * min(sigma, iqr / 1.34) * n ** (-1 / 5)
return radius
def rce_train(training_set: pd.DataFrame, radius: float | None = None) -> pd.DataFrame:
"""
Get the Reduced Coulomb Energy (RCE) prototypes.
Args:
training_set (pd.DataFrame): The training set
radius (float | None): The radius of the RCE prototypes. If None, computed using Silverman's rule.
Returns:
pd.DataFrame: The RCE prototypes
"""
train_vectors = training_set.values
# Vectorized distance computation for the radius calculation
diff_vectors = train_vectors[:, np.newaxis] - train_vectors[np.newaxis, :]
distances = np.linalg.norm(diff_vectors, axis=-1)
# Non-parametric radius: Silverman Radius (compute if not provided)
effective_radius = radius if radius is not None else silverman_radius(distances.flatten())
# Initialize prototypes with the first vector
prototypes = [train_vectors[0]]
for vector in train_vectors[1:]:
# Vectorized distance check between current vector and all prototypes
distances_to_prototypes = np.linalg.norm(prototypes - vector, axis=1)
# If no prototype is close, add the current vector as a new prototype
if np.all(distances_to_prototypes > effective_radius):
prototypes.append(vector)
return pd.DataFrame(prototypes)
def rce_test(test_set: pd.DataFrame, prototypes: pd.DataFrame) -> pd.Series:
"""
Get the signed Reduced Coulomb Energy (RCE) predictions.
Args:
test_set (pd.DataFrame): The test set
prototypes (pd.DataFrame): The RCE prototypes
Returns:
pd.Series: The signed distances to the closest prototype for each test vector
"""
test_vectors = test_set.values
prototype_vectors = prototypes.values
# Vectorized computation of distances between test vectors and all prototypes
diff_vectors = test_vectors[:, np.newaxis] - prototype_vectors[np.newaxis, :]
distances = np.linalg.norm(diff_vectors, axis=-1)
# Find the closest prototype for each test vector
min_distances = np.min(distances, axis=1)
closest_prototypes = prototype_vectors[np.argmin(distances, axis=1)]
# Compute the signed distance for each test vector
signed_distances = np.sqrt(min_distances**2) * np.sign(
np.mean(test_vectors - closest_prototypes, axis=1)
)
return pd.Series(signed_distances)
def rce_drift(reference_data: pd.DataFrame, real_data: pd.DataFrame, column: str) -> pd.Series:
"""
Detect drift using the Reduced Coulomb Energy (RCE) method.
Args:
reference_data (pd.DataFrame): The reference data
real_data (pd.DataFrame): The real data
column (str): The target column to be analyzed. 'target' or 'prediction'
Returns:
pd.Series: Normalized drift distances
"""
common_columns = list(set(reference_data.columns).intersection(real_data.columns))
reference_data = reference_data[common_columns]
real_data = real_data[common_columns]
# Get prototypes
if column == 'target':
prototypes = rce_train(reference_data.drop(columns=['prediction']), 0.1)
else:
prototypes = rce_train(reference_data.drop(columns=['target']), 0.1)
# Distances to prototypes
if column == 'target':
distances_train = rce_test(reference_data.drop(columns=['prediction']), prototypes)
distances_test = rce_test(real_data.drop(columns=['prediction']), prototypes)
else:
distances_train = rce_test(reference_data.drop(columns=['target']), prototypes)
distances_test = rce_test(real_data.drop(columns=['target']), prototypes)
# Find the maximum absolute distance in the training set
max_abs_distance = max(abs(distances_train.max()), abs(distances_train.min()))
# Normalize while preserving sign
distances = distances_test / max_abs_distance
return distances

View File

@@ -0,0 +1,292 @@
import os
from collections.abc import Sequence
from typing import Any
from bs4 import BeautifulSoup, Tag
from evidently.metric_preset import DataDriftPreset
from evidently.metrics import (
ColumnSummaryMetric,
ConflictTargetMetric,
DatasetCorrelationsMetric,
DatasetSummaryMetric,
RegressionAbsPercentageErrorPlot,
RegressionDummyMetric,
RegressionErrorDistribution,
RegressionErrorPlot,
RegressionPerformanceMetrics,
RegressionPredictedVsActualPlot,
RegressionPredictedVsActualScatter,
)
from evidently.metrics.base_metric import generate_column_metrics
from evidently.options import ColorOptions
from evidently.pipeline.column_mapping import ColumnMapping
from evidently.report import Report
COLOR_DISCRETE_SEQUENCE = (
'#ed0400',
'#0a5f38',
'#6c3461',
'#71aa34',
'#d8dcd6',
'#6b8ba4',
)
def load_html_from_file(file_path):
with open(file_path, encoding='utf-8') as file:
return file.read()
def inject_content(main_html, section_id, content):
soup = BeautifulSoup(main_html, 'html.parser')
section = soup.find(id=section_id)
# Verifica se a seção foi encontrada E se ela é uma Tag (não uma string)
if section and isinstance(section, Tag):
section.clear()
# Converte o conteúdo para um fragmento de BeautifulSoup e anexa
new_content = BeautifulSoup(content, 'html.parser')
section.append(new_content)
return str(soup)
class Reports:
"""
Report generator using Evidently library.
Thread-safety: This class is NOT thread-safe. Multiple threads should not
call add_*_section() methods on the same instance simultaneously as they
modify shared state (self.metrics, self.sections, self.options).
For multi-threaded environments:
- Create separate Reports instances per thread
- Or synchronize access using locks
- After generation, instances are safe for read-only operations
I/O Note: This class relies on Evidently's report.save_html() method
for file operations. Ensure Evidently properly manages file handles.
"""
def __init__(
self,
reference_data: Any,
current_data: Any,
target_name: str,
base_path: str | None = None,
template_path: str | None = None,
) -> None:
"""
Initializes an instance of the AigReport class.
Args:
reference_data: The reference data for the report.
current_data: The current data for the report.
base_path: The base path for the report.
"""
self.metrics: list[Any] = []
self.options: list[Any] | None = None
self.sections: dict[str, Any] = {}
self.report: Any = None
self.ref_data = reference_data
self.cur_data = current_data
self.target_name = target_name
self.set_color_options(primary_color='#0F4C81', secondary_color='#001E60')
self.base_path = base_path
self.template_path = template_path
def add_data_quality_section(self, columns: list[str] | None = None, run: bool = True) -> None:
"""
Adds a data quality section to the report.
Args:
columns: The list of columns to include in the data quality section. If None, all columns will be included.
run: Indicates whether to run the report immediately after adding the section.
"""
metrics = [
DatasetSummaryMetric(),
generate_column_metrics(ColumnSummaryMetric, columns=columns, skip_id_column=True),
ConflictTargetMetric(),
DatasetCorrelationsMetric(),
]
self.metrics.extend(metrics)
if run:
mapping = ColumnMapping()
mapping.target = self.target_name
report = Report(metrics=metrics, options=self.options)
report.run(
reference_data=self.ref_data,
current_data=self.cur_data,
column_mapping=mapping,
)
self.sections['data_quality'] = report.as_dict()
if self.base_path:
# Note: Relies on Evidently's save_html() to properly manage file I/O
report.save_html(os.path.join(self.base_path, 'data_quality.html'))
def add_data_drift_section(self, columns: list[str] | None = None, run: bool = True) -> None:
"""
Adds a data drift section to the report.
Args:
columns: The list of columns to include in the data drift section. If None, all columns will be included.
run: Indicates whether to run the report immediately after adding the section.
"""
self.metrics.append(DataDriftPreset(columns=columns))
if run:
mapping = ColumnMapping()
mapping.target = self.target_name
report = Report(metrics=[DataDriftPreset(columns=columns)], options=self.options)
report.run(
reference_data=self.ref_data,
current_data=self.cur_data,
column_mapping=mapping,
)
self.sections['data_drift'] = report.as_dict()
if self.base_path:
# Note: Relies on Evidently's save_html() to properly manage file I/O
report.save_html(os.path.join(self.base_path, 'data_drift.html'))
def add_regression_section(self, run: bool = True) -> None:
"""
Adds a regression section to the report.
Args:
run: Indicates whether to run the report immediately after adding the section.
"""
metrics = [
RegressionPerformanceMetrics(),
RegressionDummyMetric(),
RegressionPredictedVsActualScatter(),
RegressionPredictedVsActualPlot(),
RegressionErrorPlot(),
RegressionAbsPercentageErrorPlot(),
RegressionErrorDistribution(),
]
self.metrics.extend(metrics)
if run:
mapping = ColumnMapping()
mapping.target = self.target_name
mapping.prediction = 'prediction'
report = Report(metrics=metrics, options=self.options)
report.run(
reference_data=self.ref_data,
current_data=self.cur_data,
column_mapping=mapping,
)
self.sections['regression'] = report.as_dict()
if self.base_path:
# Note: Relies on Evidently's save_html() to properly manage file I/O
report.save_html(os.path.join(self.base_path, 'regression.html'))
def set_color_options(
self,
primary_color: str = '#0F4C81',
secondary_color: str = '#001E60',
current_data_color: str | None = None,
reference_data_color: str | None = None,
additional_data_color: str = '#0a5f38',
color_sequence: Sequence[str] = COLOR_DISCRETE_SEQUENCE,
fill_color: str = 'LightGreen',
zero_line_color: str = 'green',
non_visible_color: str = 'white',
underestimation_color: str = '#6574f7',
overestimation_color: str = '#ee5540',
majority_color: str = '#1acc98',
vertical_lines: str = 'green',
heatmap: str = 'RdBu_r',
) -> None:
"""
Sets the color options for the report.
Args:
primary_color: The primary color for the report.
secondary_color: The secondary color for the report.
current_data_color: The color for the current data.
reference_data_color: The color for the reference data.
additional_data_color: The color for additional data.
color_sequence: The color sequence for discrete values.
fill_color: The fill color for visualizations.
zero_line_color: The color for the zero line.
non_visible_color: The color for non-visible elements.
underestimation_color: The color for underestimation.
overestimation_color: The color for overestimation.
majority_color: The color for majority elements.
vertical_lines: The color for vertical lines.
heatmap: The color map for heatmaps.
"""
color_scheme = ColorOptions(
primary_color=primary_color,
secondary_color=secondary_color,
current_data_color=current_data_color,
reference_data_color=reference_data_color,
additional_data_color=additional_data_color,
color_sequence=color_sequence,
fill_color=fill_color,
zero_line_color=zero_line_color,
non_visible_color=non_visible_color,
underestimation_color=underestimation_color,
overestimation_color=overestimation_color,
majority_color=majority_color,
vertical_lines=vertical_lines,
heatmap=heatmap,
)
if self.options is None:
self.options = [color_scheme]
else:
self.options.append(color_scheme)
def save_all_sections_html(self, report_path):
"""
Saves the report with all sections as HTML.
Args:
report_path: The path to save the report HTML file.
Raises:
ValueError: If base_path is not set
OSError: If directory creation or file writing fails
Note:
This method uses context manager (with open) to ensure file is properly closed.
Creates parent directories if they don't exist.
"""
if not self.base_path:
raise ValueError('base_path is required to save all sections HTML')
if not self.template_path:
raise ValueError('template_path is required to save all sections HTML')
# Ensure output directory exists
output_dir = os.path.dirname(report_path)
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
print(f'Output directory: {output_dir}')
print(f'Report path: {report_path}')
print(f'Base path: {self.base_path}')
print(f'Template path: {self.template_path}')
# Load main HTML template
main_html_path = os.path.join(self.template_path, 'header.html')
main_html = load_html_from_file(main_html_path)
# Load content from data_drift.html, data_quality.html, and regression.html
data_drift_content = load_html_from_file(os.path.join(self.base_path, 'data_drift.html'))
data_quality_content = load_html_from_file(
os.path.join(self.base_path, 'data_quality.html')
)
regression_content = load_html_from_file(os.path.join(self.base_path, 'regression.html'))
# Inject content into the main HTML template
main_html = inject_content(main_html, 'data_drift', data_drift_content)
main_html = inject_content(main_html, 'data_quality', data_quality_content)
main_html = inject_content(main_html, 'regression', regression_content)
# Save the final HTML to a new file (report.html)
# Context manager ensures file is properly closed even if an error occurs
with open(report_path, 'w', encoding='utf-8') as report_file:
report_file.write(main_html)

View File

View File

@@ -0,0 +1,166 @@
from os import getenv
from typing import Any
def build_postgres_config() -> dict[str, Any]:
"""
Build PostgreSQL database configuration from environment variables.
This function constructs a PostgreSQL configuration dictionary from
environment variables with sensible defaults for local development.
It handles connection pool configuration and security parameters.
Environment Variables:
POSTGRES_HOST: Database hostname (default: localhost)
POSTGRES_PORT: Database port (default: 5432)
POSTGRES_USER: Database username (default: sientia)
POSTGRES_PASSWORD: Database password (default: sientia)
POSTGRES_DBNAME: Database name (default: sientia)
POSTGRES_MIN_CONNECTIONS: Minimum connection pool size (default: 5)
POSTGRES_MAX_CONNECTIONS: Maximum connection pool size (default: 20)
Returns:
dict: PostgreSQL configuration dictionary with all required parameters
"""
return {
'host': getenv('POSTGRES_HOST', 'localhost'),
'port': int(getenv('POSTGRES_PORT', '5432')),
'user': getenv('POSTGRES_USER', 'sientia'),
'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
}
def build_mlflow_config() -> dict[str, Any]:
"""
Build MLFlow server configuration from environment variables.
This function constructs an MLFlow configuration dictionary from
environment variables with sensible defaults for local development.
It handles server connection and authentication parameters.
Environment Variables:
MLFLOW_URL: Full MLflow tracking URL including scheme, host, and port
(default: http://localhost:5080)
MLFLOW_USERNAME: MLFlow username (default: aignosi)
MLFLOW_PASSWORD: MLFlow password (default: aignosi)
Returns:
dict: MLFlow configuration dictionary with all required parameters
"""
return {
'url': getenv('MLFLOW_URL', 'http://localhost:5080'),
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
}
def build_mongodb_config() -> dict[str, Any]:
"""
Build MongoDB configuration from environment variables.
This function constructs a MongoDB configuration dictionary from
environment variables with sensible defaults for local development.
It handles connection string and database name configuration.
Environment Variables:
MONGODB_USERNAME: MongoDB username (default: root)
MONGODB_PASSWORD: MongoDB password (default: wKZDbMNU1c)
MONGODB_URL: MongoDB connection URI (default: localhost:27018)
MONGODB_DATABASE: MongoDB database name (default: sientia)
MONGODB_TTL_INDEX_HOURS: TTL index duration in hours (default: 1)
Returns:
dict: MongoDB configuration dictionary with connection parameters
"""
username = getenv('MONGODB_USERNAME', 'root')
password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c')
uri = getenv('MONGODB_URL', 'localhost:27018')
connection_string = f'mongodb://{username}:{password}@{uri}'
return {
'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
'uri': uri,
}
def build_minio_config() -> dict[str, Any]:
"""
Build MinIO (S3-compatible) configuration from environment variables.
This function constructs a MinIO configuration dictionary from
environment variables with sensible defaults for local development.
It handles endpoint URL, authentication, connection parameters, and retry policies.
Environment Variables:
MINIO_ENDPOINT_URL: MinIO server endpoint URL (default: http://localhost:9000)
MINIO_ACCESS_KEY: MinIO access key ID (default: minioadmin)
MINIO_SECRET_KEY: MinIO secret access key (default: minioadmin)
MINIO_REGION: MinIO region name (default: us-east-1)
MINIO_SECURE: Whether to use SSL/TLS (default: false)
MINIO_MAX_RETRY_ATTEMPTS: Maximum number of retry attempts (default: 3)
MINIO_RETRY_MODE: Retry mode - standard, legacy, or adaptive (default: adaptive)
MINIO_CONNECT_TIMEOUT: Connection timeout in seconds (default: 10)
MINIO_READ_TIMEOUT: Read timeout in seconds (default: 60)
MINIO_DEFAULT_BUCKET: Default S3 bucket for MinioRepository (default: model-training)
Returns:
dict: MinIO configuration dictionary with all required parameters
"""
return {
'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
'region': getenv('MINIO_REGION', 'us-east-1'),
'use_ssl': getenv('MINIO_SECURE', 'false').lower() == 'true',
'max_retry_attempts': int(getenv('MINIO_MAX_RETRY_ATTEMPTS', '3')),
'retry_mode': getenv('MINIO_RETRY_MODE', 'adaptive'),
'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')),
'read_timeout': int(getenv('MINIO_READ_TIMEOUT', '60')),
'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'model-training'),
}
def build_plugin_store_config() -> dict[str, Any]:
"""
Build PluginStore configuration from environment variables.
This function constructs a configuration dictionary for the PluginStore
client using environment variables with sensible defaults for local
development.
Environment Variables:
STORE_BASE_URL: Base URL of the PluginStore backing Git server
(default: http://localhost:3000)
STORE_OWNER: Repository owner/organization (default: sientia)
STORE_REPO: Repository name (default: model-library-store)
STORE_BRANCH: Optional branch name
STORE_USERNAME: Optional username for Git HTTP authentication
STORE_PASSWORD: Optional password/token for Git HTTP authentication
PYPI_SERVER: Optional custom PyPI index URL for runtime installation
(default: http://localhost:5000)
PYPI_USERNAME: Optional username for PyPI authentication
PYPI_PASSWORD: Optional password/token for PyPI authentication
Returns:
dict: PluginStore configuration dictionary with all connector parameters
"""
cache_ttl_seconds = getenv('STORE_CACHE_TTL_SECONDS')
return {
'base_url': getenv('STORE_BASE_URL', 'http://localhost:3000'),
'owner': getenv('STORE_OWNER', 'sientia'),
'repo': getenv('STORE_REPO', 'model-library-store'),
'username': getenv('STORE_USERNAME'),
'password': getenv('STORE_PASSWORD'),
'branch': getenv('STORE_BRANCH'),
'cache_ttl_seconds': int(cache_ttl_seconds) if cache_ttl_seconds else None,
'pypi_index_url': getenv('PYPI_SERVER', 'http://localhost:5000'),
'pypi_username': getenv('PYPI_USERNAME'),
'pypi_password': getenv('PYPI_PASSWORD'),
}

View File

@@ -0,0 +1,27 @@
"""
Logger helper to prevent duplicate logs caused by propagation.
This module provides a wrapper around sientia_do Logger to disable
log propagation and prevent duplicate log entries in the Model Manager.
"""
from sientia_do.observability.logger import Logger as SientiaLogger
def get_logger(name: str) -> SientiaLogger:
"""
Create a Logger instance with propagation disabled.
This prevents duplicate logs caused by hierarchical propagation
in Python's logging system.
Args:
name: Logger name (typically __name__ of the calling module).
Returns:
Logger: Configured logger instance with propagation disabled.
"""
logger = SientiaLogger(name)
# Disable propagation to prevent duplicate logs
logger.base_logger.propagate = False
return logger

View File

@@ -0,0 +1,16 @@
"""
Models and DTOs for the Model Manager system.
This module contains data transfer objects (DTOs) and model classes used
throughout the Model Manager workflows and activities.
"""
from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
__all__ = [
'ExperimentStatus',
'TrainModelParams',
'TrainModelResult',
]

View File

@@ -0,0 +1,26 @@
from enum import StrEnum
class ExperimentStatus(StrEnum):
"""
Status values for experiment run lifecycle.
This enum defines all possible status values that an experiment run can have
throughout its lifecycle, from initialization through training, model saving,
and cleanup. These statuses are used to track progress and identify failures
in the training pipeline.
The status values follow the naming convention from the original Mage pipeline
to maintain compatibility with existing database records and monitoring systems.
Attributes:
ORCHESTRATOR_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing.
ORCHESTRATOR_VALIDATION_ERROR: Error in the parameters validation.
TRAINING_SUCCESS: Training completed successfully with model and metrics calculated.
TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions.
"""
ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR'
ORCHESTRATOR_WAITING_PROC = 'ORCHESTRATOR_WAITING_PROC'
TRAINING_SUCCESS = 'TRAINING_SUCCESS'
TRAINING_ERROR = 'TRAINING_ERROR'

View File

@@ -0,0 +1,363 @@
from dataclasses import dataclass
from typing import Any
from jsonschema import Draft202012Validator, ValidationError # type: ignore[import-untyped]
# Allowed frontend date formats and their strftime equivalents (single source of truth)
FRONTEND_DATE_FORMAT_TO_STRFTIME = {
'dd/MM/yyyy HH:mm:ss': '%d/%m/%Y %H:%M:%S',
'MM/dd/yyyy HH:mm:ss': '%m/%d/%Y %H:%M:%S',
'yyyy/MM/dd HH:mm:ss': '%Y/%m/%d %H:%M:%S',
'dd-MM-yyyy HH:mm:ss': '%d-%m-%Y %H:%M:%S',
'MM-dd-yyyy HH:mm:ss': '%m-%d-%Y %H:%M:%S',
'yyyy-MM-dd HH:mm:ss': '%Y-%m-%d %H:%M:%S',
}
ALLOWED_FRONTEND_DATE_FORMATS = frozenset(FRONTEND_DATE_FORMAT_TO_STRFTIME.keys())
# When the client omits date_format (or sends null/blank), parsing uses this frontend format.
DEFAULT_TRAIN_DATE_FORMAT = 'yyyy-MM-dd HH:mm:ss'
def validate_frontend_date_format(fmt: str | None) -> None:
"""Raise ValueError if fmt is set and not one of the allowed frontend date formats."""
if not fmt or not fmt.strip():
return
if fmt not in ALLOWED_FRONTEND_DATE_FORMATS:
allowed = ', '.join(sorted(ALLOWED_FRONTEND_DATE_FORMATS))
raise ValueError(f'Invalid date_format "{fmt}". Allowed formats: {allowed}')
# Model name constants
MODEL_LINEAR_REGRESSION = 'Linear Regression'
MODEL_POLYNOMIAL_REGRESSION = 'Polynomial Regression'
@dataclass
class TrainModelParams:
"""
Parameters for machine learning model training.
This class encapsulates all configuration parameters required for the training
pipeline, including data processing settings, model configuration, and experiment
tracking information. All parameters are validated upon initialization to ensure
data integrity and prevent runtime errors.
Use the `from_dict()` class method to create instances from dictionaries with
automatic validation of all fields.
Attributes:
variable_columns (list[str]): List of variable column names to use as features.
target_variable (str): Name of the target variable to predict.
bucket_name (str): Name of the MinIO bucket containing training data.
file_name (str): Name of the file in the MinIO bucket.
line_separator (str): Line separator used in the CSV file.
decimal_separator (str): Decimal separator used in the CSV file.
date_column (str): Name of the date/time column in the dataset (required).
date_format (str): Format of the date column (allowed frontend strings). If omitted or blank
in the input dict, defaults to DEFAULT_TRAIN_DATE_FORMAT.
train_size (int): Percentage of data to use for training (0-100).
shuffle (bool): Whether to shuffle the data during train/test split.
experiment_run_id (int): Unique identifier for the experiment run.
model_name (str): Name of the model type ('Linear Regression' or 'Polynomial Regression').
val_file_name (str | None): Name of the validation file in the MinIO bucket.
data_model_kwargs (dict | None): Keyword arguments for the data model.
model_kwargs (dict | None): Keyword arguments for the model.
opt_params (dict | None): Keyword optimazation arguments for the wrapper.
model_type (str): Type of the model to use (ex.: 'Linear Regression', 'XGBoost').
"""
# Old Parameters (keep)
variable_columns: list[str]
target_variable: str
bucket_name: str
file_name: str
line_separator: str
decimal_separator: str
date_column: str
date_format: str
train_size: int
shuffle: bool
random_state: int
experiment_run_id: int
model_name: str
experiment_name: str
# New Parameters
val_file_name: str | None
data_model_kwargs: dict | None # Removed params used in DataPreprocessor here
model_kwargs: dict | None # Removed params used in Linear Regression Model here
opt_params: dict | None
model_type: str
model_id: str | None
# Context Parameters
model_metadata: dict | None
@classmethod
def from_dict(cls, data: dict[str, Any]) -> 'TrainModelParams':
"""
Create TrainModelParams from dictionary with validation.
This factory method creates a TrainModelParams instance from a dictionary,
applying validation to ensure all required fields are present and have
the correct types. This is the recommended way to create instances from
workflow input data.
Args:
data: Dictionary containing training parameters with keys matching the
attribute names (e.g. variable_columns, date_column, data_model_kwargs, model_kwargs, opt_params).
Unknown keys are ignored by from_dict; missing required snake_case keys raise.
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
Raises:
ValueError: If any required field is missing or None
TypeError: If any field has an incorrect type
KeyError: If any required key is missing from the dictionary
"""
# `from_dict()` should only build the "raw" object from the input dict.
# Semantic validation and defaults must be handled by `validate_business_rules()`
# (using `model_metadata` JSON Schemas).
model_name = cls._check_none(data.get('model_name'), str, 'model_name')
return cls(
variable_columns=cls._check_none(
data.get('variable_columns'), list, 'variable_columns'
),
target_variable=cls._check_none(data.get('target_variable'), str, 'target_variable'),
bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'),
file_name=cls._check_none(data.get('file_name'), str, 'file_name'),
line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'),
decimal_separator=cls._check_none(
data.get('decimal_separator'), str, 'decimal_separator'
),
date_column=cls._check_none(data.get('date_column'), str, 'date_column'),
date_format=cls._resolve_date_format(data.get('date_format')),
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._coerce_experiment_run_id(data.get('experiment_run_id')),
model_name=model_name,
experiment_name=model_name,
val_file_name=data.get('val_file_name'),
data_model_kwargs=cls._check_none(
data.get('data_model_kwargs'), dict, 'data_model_kwargs'
),
model_kwargs=cls._check_none(data.get('model_kwargs'), dict, 'model_kwargs'),
opt_params=cls._check_none(data.get('opt_params'), dict, 'opt_params'),
model_type=cls._check_none(data.get('model_type'), str, 'model_type'),
model_id=data.get('model_id'),
model_metadata=cls._parse_optional_model_metadata(data.get('model_metadata')),
)
@staticmethod
def _resolve_date_format(raw: Any) -> str:
"""
Resolve date_format from workflow input.
Omitted, null, or blank values use DEFAULT_TRAIN_DATE_FORMAT. Non-string types raise.
Args:
raw: Raw date_format from the payload, or None if absent.
Return:
str: Canonical frontend date format string.
"""
if raw is None:
return DEFAULT_TRAIN_DATE_FORMAT
if isinstance(raw, str) and not raw.strip():
return DEFAULT_TRAIN_DATE_FORMAT
if not isinstance(raw, str):
raise TypeError(
f'date_format must be a string or omitted, but got {type(raw).__name__}.'
)
return raw.strip()
def to_dict(self) -> dict[str, Any]:
"""
Convert TrainModelParams to a dictionary.
"""
return self.__dict__
@staticmethod
def _check_none(value: Any | None, expected_type: type, field_name: str) -> Any:
"""
Validate that a value is not None and check its type.
This method ensures that required parameters are provided and have the
correct type, raising descriptive errors if validation fails.
Args:
value (Any | None): The value to validate.
expected_type (type): The expected type of the value.
field_name (str): The name of the field being validated (for error messages).
Returns:
Any: The validated value if it is not None and matches the expected type.
Raises:
ValueError: If the value is None.
TypeError: If the value is not of the expected type.
"""
if value is None:
error = f'{field_name} is required and cannot be None.'
raise ValueError(error)
return TrainModelParams._check_type(value, expected_type, field_name)
@staticmethod
def _check_type(value: Any | None, expected_type: type, field_name: str) -> Any:
"""
Validate that a value matches the expected type.
This method checks type compatibility and raises a descriptive error
if the value does not match the expected type.
Args:
value (Any | None): The value to validate.
expected_type (type): The expected type of the value.
field_name (str): The name of the field being validated (for error messages).
Returns:
Any: The validated value if it matches the expected type.
Raises:
TypeError: If the value is not of the expected type.
"""
if value is not None and not isinstance(value, expected_type):
error = f'{field_name} must be of type {expected_type.__name__}, but got {type(value).__name__}.'
raise TypeError(error)
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.
This method performs additional validation beyond type checking to ensure
that parameter values are within acceptable ranges and logically consistent.
It implements defense-in-depth validation to catch configuration errors
early in the workflow.
Raises:
ValueError: If any business rule is violated
"""
self._validate_numeric_ranges()
self._validate_model_params()
self._validate_required_strings()
self._validate_date_format()
def _validate_numeric_ranges(self) -> None:
"""Validate numeric parameters are within acceptable ranges."""
if not 10 <= self.train_size <= 100:
raise ValueError(f'train_size must be between 10 and 100, got {self.train_size}')
if not self.variable_columns:
raise ValueError('variable_columns cannot be empty')
def _validate_model_params(self) -> None:
"""Validate model-related parameters."""
if not self.model_metadata:
raise ValueError('model_metadata is required')
schemas = self.model_metadata.get('schemas', {}).get('components', {}).get('schemas')
if not schemas:
return
data_model_schema = schemas.get('data_model')
model_schema = schemas.get('model')
opt_params_schema = schemas.get('opt_params')
if data_model_schema:
self._validate_model_param(data_model_schema, self.data_model_kwargs)
if model_schema:
self._validate_model_param(model_schema, self.model_kwargs)
if opt_params_schema:
self._validate_model_param(opt_params_schema, self.opt_params)
def _validate_model_param(self, schema: dict[str, Any], value: Any) -> None:
"""Validate model parameter against schema."""
try:
validator = Draft202012Validator(schema)
validator.validate(value)
except ValidationError as e:
raise ValueError(f'Model parameters validation failed: {e.message}') from e
def _validate_required_strings(self) -> None:
"""Validate required string fields are not empty."""
if not self.target_variable.strip():
raise ValueError('target_variable cannot be empty or whitespace')
if not self.bucket_name.strip():
raise ValueError('bucket_name cannot be empty or whitespace')
if not self.file_name.strip():
raise ValueError('file_name cannot be empty or whitespace')
if not self.model_name.strip():
raise ValueError('model_name cannot be empty or whitespace')
if not self.date_column.strip():
raise ValueError('date_column cannot be empty or whitespace')
def _validate_date_format(self) -> None:
"""Validate date_format is one of the allowed frontend formats."""
validate_frontend_date_format(self.date_format)

View File

@@ -0,0 +1,53 @@
from dataclasses import dataclass
import pandas as pd
from model_manager.utils.models.train_model_params import TrainModelParams
@dataclass
class TrainModelResult:
"""
A data container for storing the results of a machine learning training process.
This dataclass encapsulates all outputs from the training pipeline, including
the prepared datasets, evaluation metrics, and paths to generated artifacts.
It is used to pass results between activities in the training workflow.
Attributes:
params (TrainModelParams): The parameters used to train the model.
x_train (pd.DataFrame): The training dataset features.
x_test (pd.DataFrame): The testing dataset features.
y_train (pd.DataFrame): The training dataset target values.
y_test (pd.DataFrame): The testing dataset target values.
y_pred (pd.Series | None): The predicted target values for the testing dataset. Default is None.
y_train_pred (pd.Series | None): The predicted target values for the training dataset. Default is None.
mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None.
mae_val (float | None): The Mean Absolute Error (MAE) of the predictions. Default is None.
r2_val (float | None): The R-squared (R²) value of the predictions. Default is None.
equation (dict | None): The equation of the model. Default is None.
equation_path (str | None): The path to the equation file. Default is None.
run_name (str | None): The name of the MLFlow run. Default is None.
report_path (str | None): The path to the generated HTML report file. Default is None.
train_data_path (str | None): The path to the training dataset CSV file. Default is None.
test_data_path (str | None): The path to the testing dataset CSV file. Default is None.
"""
params: TrainModelParams
train_data: pd.DataFrame
val_data: pd.DataFrame
y_pred: pd.DataFrame | None = None
y_train_pred: pd.DataFrame | None = None
mse_val: float | None = None
mae_val: float | None = None
r2_val: float | None = None
equation: dict | None = None
equation_path: str | None = None
run_name: str | None = None
experiment_name: str | None = None
run_id: str | None = None
report_path: str | None = None
train_data_path: str | None = None
test_data_path: str | None = None
run_dir: str | None = None

View File

@@ -0,0 +1,636 @@
"""
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.
"""
import json
from datetime import datetime
from io import BytesIO
from os import makedirs, path
from shutil import rmtree
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_model.wrappers.sientia_model import SientiaModel
from model_manager.runtime_paths import PROJECT_BASE_PATH, REPORTS_ROOT
from model_manager.sientia.metrics import mae, mse, r2
from model_manager.sientia.reports import Reports # type: ignore[import-untyped]
from model_manager.utils.models.train_model_params import (
FRONTEND_DATE_FORMAT_TO_STRFTIME,
TrainModelParams,
)
from model_manager.utils.models.train_model_result import TrainModelResult
def train_test_split(
data: pd.DataFrame,
train_size: float,
random_state: int | None = None,
shuffle: bool = True,
) -> tuple[pd.DataFrame, pd.DataFrame]:
# 1. Definir a semente (seed) para reprodutibilidade
if random_state is not None:
np.random.seed(random_state)
# 2. Gerar índices e embaralhar se necessário
indices = np.arange(len(data))
if shuffle:
np.random.shuffle(indices)
# 3. Calcular o ponto de corte (split point)
# Cálculo: N_treino = tamanho_total * proporcao_treino
n_train = int(len(data) * train_size)
# 4. Dividir os índices
train_indices = indices[:n_train]
test_indices = indices[n_train:]
# 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:
"""
Parse params.date_column using the frontend date_format mapping only.
date_column must exist in ``data`` (callers validate before prepare). No broad
pandas inference or alternate timezone formats here—clients must send a supported
date_format or rely on the TrainModelParams default.
"""
if params.date_column not in data.columns:
raise ValueError(
f'date_column "{params.date_column}" not found in dataset columns: {list(data.columns)}'
)
data = data.copy()
col = data[params.date_column]
if params.date_format not in FRONTEND_DATE_FORMAT_TO_STRFTIME:
raise ValueError(
f'date_format "{params.date_format}" is not mapped to a strftime pattern '
'(must be one of the allowed frontend formats).'
)
strf = FRONTEND_DATE_FORMAT_TO_STRFTIME[params.date_format]
try:
parsed = pd.to_datetime(col, format=strf, 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
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 _drop_rows_with_missing_timestamp(
self,
df: pd.DataFrame,
params: TrainModelParams,
metadata: dict[str, Any] | None,
) -> pd.DataFrame:
"""
Remove rows where the configured date_column is missing (NaN/NaT/blank string).
Empty timestamp cells cannot be placed on a DatetimeIndex and break
downstream joins and metrics.
"""
if params.date_column not in df.columns:
return df
series = df[params.date_column]
mask = series.notna()
if series.dtype == object:
stripped = series.astype(str).str.strip()
mask &= stripped.ne('')
mask &= stripped.str.lower().ne('nan')
n_drop = int((~mask).sum())
if n_drop:
self.info(
f'Dropping {n_drop} row(s) with missing or blank timestamp column '
f'"{params.date_column}"',
metadata,
)
return df.loc[mask].copy()
def _coerce_non_timestamp_columns_to_numeric(
self,
df: pd.DataFrame,
params: TrainModelParams,
metadata: dict[str, Any] | None,
) -> pd.DataFrame:
"""
Coerce all non-timestamp columns to numeric dtype.
The timestamp column defined by params.date_column is excluded from coercion.
Non-numeric values are coerced to NaN.
"""
out = df.copy()
for col in out.columns:
if col == params.date_column:
continue
original_na = int(out[col].isna().sum())
out[col] = pd.to_numeric(out[col], errors='coerce')
new_na = int(out[col].isna().sum())
introduced_na = new_na - original_na
if introduced_na > 0:
self.warning(
f'Column "{col}" had {introduced_na} non-numeric value(s) coerced to NaN',
metadata,
)
return out
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 = self._drop_rows_with_missing_timestamp(train_df, params, metadata)
train_df = _ensure_date_column_parsed(train_df, params)
train_df = self._configure_datetime_index(train_df, params, metadata)
train_df = self._set_timezone_on_index(train_df, metadata)
train_df = self._coerce_non_timestamp_columns_to_numeric(train_df, params, metadata)
if len(train_df) <= 0:
raise ValueError('Training data view is empty after transformation')
# Explicit validation dataset path
train_data = pd.DataFrame(train_df[params.variable_columns + [params.target_variable]])
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 = self._drop_rows_with_missing_timestamp(val_df, params, metadata)
val_df = _ensure_date_column_parsed(val_df, params)
val_df = self._configure_datetime_index(val_df, params, metadata)
val_df = self._set_timezone_on_index(val_df, metadata)
val_df = self._coerce_non_timestamp_columns_to_numeric(val_df, params, metadata)
if len(val_df) <= 0:
raise ValueError('Validation data view is empty after transformation')
val_data = pd.DataFrame(val_df[params.variable_columns + [params.target_variable]])
else:
# Fallback path: derive validation via train/test split from a single dataset.
train_data, val_data = train_test_split(
train_data,
train_size=params.train_size / 100,
shuffle=params.shuffle,
random_state=params.random_state,
)
self.info(
f'Data preprocessed and split successfully - experiment run id: {params.experiment_run_id}',
metadata,
)
experiment_name = f'{params.experiment_name}'
run_name = f'{experiment_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
return TrainModelResult(
params=params,
train_data=train_data,
val_data=val_data,
run_name=run_name,
experiment_name=experiment_name,
)
def _as_series(self, pred: pd.DataFrame | pd.Series) -> pd.Series:
if isinstance(pred, pd.Series):
return pred
# If the wrapper returns a single-column DataFrame, take its first column.
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,
metadata: dict[str, Any] | None = None,
) -> TrainModelResult:
"""
Compute regression metrics for training results.
This helper mirrors the previous TrainingRepository.after_train_calculation
behavior, assuming that predictions (y_pred/y_train_pred) are already on the
correct scale for metric calculation (any scaling is handled inside the
model wrapper).
Args:
tmr: Training result containing:
- train_data/val_data DataFrames with a target column
- y_train_pred/y_pred populated (model predictions for train/val)
wrapper: Trained model wrapper (used for linear equation extraction).
metadata: Optional workflow metadata for debug logging.
Return:
TrainModelResult: Same object with mse_val, mae_val and r2_val set.
"""
if tmr.y_pred is None:
raise ValueError('y_pred must be set before computing regression metrics')
params = tmr.params
target = params.target_variable
# True values are expected to come from val_data.
y_true_val = tmr.val_data[target]
y_pred_val = self._as_series(tmr.y_pred).sort_index()
y_true_val = y_true_val.sort_index()
# Align by index to avoid metric calculation errors if ordering differs.
common_index = y_true_val.index.intersection(y_pred_val.index)
head = min(5, len(y_true_val), len(y_pred_val))
self.debug(
'compute_regression_metrics index alignment: '
f'val_n={len(y_true_val)} pred_n={len(y_pred_val)} common_n={len(common_index)}; '
f'val_index_dtype={y_true_val.index.dtype} '
f'pred_index_dtype={y_pred_val.index.dtype}; '
f'val_index_sample={list(y_true_val.index[:head])} '
f'pred_index_sample={list(y_pred_val.index[:head])}',
metadata,
)
if len(common_index) == 0:
raise ValueError(
'No overlapping indices between val_data and y_pred. '
f'val_n={len(y_true_val)} pred_n={len(y_pred_val)} '
f'val_index_sample={list(y_true_val.index[:head])} '
f'pred_index_sample={list(y_pred_val.index[:head])}'
)
y_true_val = y_true_val.loc[common_index]
y_pred_val = y_pred_val.loc[common_index]
# Metrics helpers already round to 2 decimals.
tmr.mse_val = mse(y_true_val, y_pred_val)
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':
inner = getattr(wrapper, 'model', None)
regr = getattr(inner, 'regr', None) if inner is not None else None
if regr is not None and hasattr(regr, 'coef_') and hasattr(regr, 'intercept_'):
tmr.equation = self._extract_model_equation(inner, params)
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.
Uses only params.date_column and assumes it was already parsed exactly once by
_ensure_date_column_parsed.
Args:
data: The DataFrame to configure the datetime index for.
params: The training parameters.
metadata: The metadata for the training run.
Returns:
The DataFrame with the datetime index configured.
"""
if data is None:
raise ValueError(
'Data is None after load_data. '
'Check file format, line separator and decimal separator.'
)
if params.date_column not in data.columns:
raise ValueError(
f'date_column "{params.date_column}" not found in dataset columns: {list(data.columns)}'
)
if not pd.api.types.is_datetime64_any_dtype(data[params.date_column]):
raise ValueError(
f'date_column "{params.date_column}" must be datetime before index configuration'
)
data = data.set_index(params.date_column)
data = data.sort_index()
self.info(f'Configured datetime index from column: {params.date_column}', metadata)
return data
def _set_timezone_on_index(
self, data: pd.DataFrame, metadata: dict[str, Any] | None = None
) -> pd.DataFrame:
"""
Check if the index has a timezone and if not, set it to UTC timezone.
Args:
data: The DataFrame to set the timezone on.
metadata: The metadata for the training run.
Returns:
The DataFrame with the timezone set.
"""
if isinstance(data.index, pd.DatetimeIndex):
if data.index.tz is None:
data.index = data.index.tz_localize('UTC')
else:
data.index = data.index.tz_convert('UTC')
else:
raise ValueError('Index is not a DatetimeIndex')
return data
def _get_reports_directory(self) -> str:
"""
Get the absolute path to the reports directory.
Returns:
str: Absolute path to the runtime reports root.
"""
return REPORTS_ROOT
def _create_run_directory(
self, base_path: str, run_name: str, metadata: dict[str, Any] | None = None
) -> 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
run_dir = path.join(base_path, 'temp', f'{run_name}')
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.error(error_msg, metadata)
raise PermissionError(error_msg) from e
except OSError as e:
error_msg = f'Failed to create directory {run_dir}: {str(e)}'
self.error(error_msg, metadata)
raise OSError(error_msg) from e
def generate_report(
self, data: TrainModelResult, metadata: dict[str, Any] | None = None
) -> 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.
"""
if data.run_name is None:
raise ValueError('run_name is not set, cannot generate report')
if data.y_train_pred is None or data.y_pred is None:
raise ValueError('y_train_pred or y_pred is not set, cannot generate report')
y_train_pred = data.y_train_pred.rename(columns={data.params.target_variable: 'prediction'})
y_val_pred = data.y_pred.rename(columns={data.params.target_variable: 'prediction'})
# Join the predictions to the data
reference_data = y_train_pred[['prediction']].join(data.train_data, how='inner')
reference_data_float = reference_data.astype(np.float64)
current_data = y_val_pred[['prediction']].join(data.val_data, how='inner')
current_data_float = current_data.astype(np.float64)
# Evidently's ConflictTargetMetric expects a literal `target` column name.
# Keep the original target column and provide this alias for report metrics.
target_col = data.params.target_variable
reference_data_float['target'] = reference_data_float[target_col]
current_data_float['target'] = current_data_float[target_col]
# Initialize report generator
base_path = self._get_reports_directory()
data.run_dir = self._create_run_directory(base_path, data.run_name)
# Template path is the code path of the model_manager package
template_path = path.join(PROJECT_BASE_PATH, 'reports')
report = Reports(
reference_data=reference_data_float,
current_data=current_data_float,
base_path=data.run_dir,
template_path=template_path,
target_name=data.params.target_variable,
)
# Generate report sections
feature_and_target_cols = data.params.variable_columns + [target_col]
report.add_data_quality_section(columns=feature_and_target_cols)
report.add_data_drift_section(columns=feature_and_target_cols)
report.add_regression_section()
# Save HTML report
data.report_path = path.join(data.run_dir, 'report.html')
report.save_all_sections_html(data.report_path)
# Save training / validation CSVs using the same frames as the report (includes
# literal `target` alias for Evidently, plus predictions and float-cast features).
data.train_data_path = path.join(data.run_dir, 'train_data.csv')
reference_data_float.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_float.to_csv(data.test_data_path, index=False)
# Save equation as JSON
if data.equation is not None and data.params.model_type == 'linear_regression':
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
def cleanup_run_directory(self, run_dir: str, metadata: dict[str, Any] | None = None) -> 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.info('No run directory specified, skipping cleanup')
return
if path.exists(run_dir):
rmtree(run_dir)
self.info(f'Run directory deleted successfully: {run_dir}')
else:
self.info(f'Run directory already deleted: {run_dir}')

View File

View File

@@ -0,0 +1,119 @@
import os
import re
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from typing import Any
from sientia_do.observability.logger import Logger
from temporalio.client import Client
from temporalio.worker import PollerBehaviorAutoscaling, Worker
# Worker configuration parameters with default values.
parameters = [
('MAX_CONCURRENT_WORKFLOW_TASKS', '200'),
('MAX_CONCURRENT_ACTIVITIES', '200'),
('MAX_CONCURRENT_LOCAL_ACTIVITIES', '200'),
('MAX_CACHED_WORKFLOWS', '200'),
('WORKFLOW_POLLER_BEHAVIOUR_MINIMUM', '10'),
('WORKFLOW_POLLER_BEHAVIOUR_INITIAL', '100'),
('WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM', '200'),
('ACTIVITY_POLLER_BEHAVIOUR_MINIMUM', '10'),
('ACTIVITY_POLLER_BEHAVIOUR_INITIAL', '100'),
('ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM', '200'),
('ACTIVITY_EXECUTOR_MAX_WORKERS', '200'),
]
def camel_to_snake(text: str) -> str:
"""
Convert a CamelCase or camelCase string into snake_case.
Args:
- text: str, original string in CamelCase or camelCase format
Return:
str: converted string in snake_case format
"""
text = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', text)
text = re.sub('([a-z0-9])([A-Z])', r'\1_\2', text)
return text.lower()
def build_queue_name(workflow_name: str, runtime: str | None = None) -> str:
"""
Build Temporal queue name from workflow name and runtime.
Args:
- workflow_name: str, workflow class name in CamelCase format
- runtime: str | None, runtime suffix for environment-specific queues
Return:
str: queue name in the format <workflow>-<runtime>-queue or <workflow>-queue
"""
snake_workflow_name = camel_to_snake(workflow_name)
if runtime:
return f'{snake_workflow_name}-{runtime}-queue'
return f'{snake_workflow_name}-queue'
def prepare_worker(
main_workflow: type,
other_workflows: Sequence[type],
activities: Sequence[Any],
temporal_client: Client,
logger: Logger,
runtime: str | None = None,
) -> Worker:
"""
Build and configure a Temporal worker for the given workflow and activities.
Args:
- main_workflow: type, main workflow class used as worker entry point
- other_workflows: Sequence[type], additional workflows in the same worker
- activities: Sequence[Any], activity callables registered in this worker
- temporal_client: Client, Temporal client used by the worker
- logger: Logger, logger instance used during worker preparation
- runtime: str | None, runtime suffix appended to queue name when present
Return:
Worker: fully configured Temporal worker instance ready to run
"""
main_workflow_name = main_workflow.__name__.upper()
queue_name = build_queue_name(main_workflow.__name__, runtime)
local_workflow_parameters: dict[str, int] = {}
for parameter_name, default_value in parameters:
local_workflow_parameters[parameter_name] = int(
os.getenv(f'{main_workflow_name}_{parameter_name}', default_value)
)
logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}')
activity_executor = ThreadPoolExecutor(
max_workers=local_workflow_parameters['ACTIVITY_EXECUTOR_MAX_WORKERS'],
thread_name_prefix=f'{queue_name}-activity',
)
return Worker(
temporal_client,
task_queue=queue_name,
workflows=[main_workflow, *other_workflows],
activities=[*activities],
activity_executor=activity_executor,
max_concurrent_workflow_tasks=local_workflow_parameters['MAX_CONCURRENT_WORKFLOW_TASKS'],
max_concurrent_activities=local_workflow_parameters['MAX_CONCURRENT_ACTIVITIES'],
max_concurrent_local_activities=local_workflow_parameters[
'MAX_CONCURRENT_LOCAL_ACTIVITIES'
],
max_cached_workflows=local_workflow_parameters['MAX_CACHED_WORKFLOWS'],
workflow_task_poller_behavior=PollerBehaviorAutoscaling(
minimum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MINIMUM'],
initial=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_INITIAL'],
maximum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM'],
),
activity_task_poller_behavior=PollerBehaviorAutoscaling(
minimum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MINIMUM'],
initial=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_INITIAL'],
maximum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM'],
),
)

View File

@@ -0,0 +1,259 @@
"""Model Manager Worker Module
This module provides the main worker implementation for the Sientia DataOps Model Manager system.
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
model training and cleanup workflows.
The worker supports two task queues:
- train_model-<runtime>-queue: For ML model training workflows
- cleanup_files-<runtime>-queue: For file cleanup workflows
Key Features:
- Automatic scaling with PollerBehaviorAutoscaling
- Prometheus metrics integration
- Comprehensive error handling and logging
- Graceful shutdown with cleanup
- ML model training pipeline orchestration
- Automated cleanup schedule management
Environment Variables:
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: model-manager)
- TEMPORAL_USE_TLS: Enable TLS for Temporal connection (default: false)
- RUNTIME: Runtime identifier used in queue naming (default: single)
- POD_ID: Kubernetes pod identifier for metrics
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
- PROJECT_NAME: Project name for notifications (default: model-manager)
"""
from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
with workflow.unsafe.imports_passed_through():
import asyncio
import os
import sys
from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger as SientiaLogger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_model.model_repository.plugin_store import PluginStore
from model_manager import metrics
from model_manager.activities.activities import Activities
from model_manager.runtime_paths import ensure_runtime_directories
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
from model_manager.utils.connectors_config import (
build_minio_config,
build_mlflow_config,
build_mongodb_config,
build_plugin_store_config,
build_postgres_config,
)
from model_manager.utils.logger_helper import get_logger
from model_manager.worker.prepare_worker import prepare_worker
from model_manager.workflows.cleanup_files import CleanupFiles
from model_manager.workflows.train_model import TrainModel
POD_ID = os.getenv('POD_ID')
RUNTIME = os.getenv('RUNTIME')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
def _get_runtime(runtime: str | None) -> str:
"""
Resolve runtime using fallback when missing.
Args:
- runtime: str | None, runtime value from environment
Return:
str: normalized runtime value
"""
normalized_runtime = runtime.strip() if runtime else ''
return normalized_runtime or 'single'
async def main():
"""
Main entry point for the Model Manager worker application.
This function initializes and starts all components of the worker:
1. Sets up logging and metadata
2. Starts Prometheus metrics server
3. Initializes notification handler
4. Creates and configures activities
5. Starts Temporal client and workers
6. Manages worker lifecycle and graceful shutdown
The function runs indefinitely until interrupted or an error occurs.
On error, it performs cleanup and exits with a non-zero status code.
Raises:
Exception: Any unhandled exception during worker execution
SystemExit: On graceful shutdown or error conditions
"""
runtime = _get_runtime(RUNTIME)
ensure_runtime_directories()
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
logger = get_logger(__name__)
metadata = {
'pod_id': POD_ID,
'runtime': runtime,
}
start_prometheus_server(logger, metadata)
mongo_config = build_mongodb_config()
notification_handler = NotificationHandler(
connection_string=mongo_config['connection_string'],
database=mongo_config['database_name'],
logger=logger,
project_name=os.getenv('PROJECT_NAME', 'model-manager'),
)
logger.custom_info(f'MongoDB client initialized at {mongo_config["uri"]}', metadata)
logger.custom_info('Initializing metrics controller', metadata)
metrics_controller = MetricsController(logger=logger)
logger.custom_info(f'Installing runtime {runtime}', metadata)
plugin_store_parameters = build_plugin_store_config()
plugin_store = PluginStore(
base_url=plugin_store_parameters['base_url'],
owner=plugin_store_parameters['owner'],
repo=plugin_store_parameters['repo'],
username=plugin_store_parameters['username'],
password=plugin_store_parameters['password'],
branch=plugin_store_parameters['branch'],
cache_ttl_seconds=plugin_store_parameters['cache_ttl_seconds'],
pypi_index_url=plugin_store_parameters['pypi_index_url'],
pypi_username=plugin_store_parameters['pypi_username'],
pypi_password=plugin_store_parameters['pypi_password'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
await plugin_store.install_runtime(runtime_name=runtime)
activities = Activities(
postgres_config=build_postgres_config(),
mlflow_config=build_mlflow_config(),
minio_config=build_minio_config(),
plugin_store=plugin_store,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
new_runtime = Runtime(
telemetry=TelemetryConfig(
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
)
)
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=namespace,
runtime=new_runtime,
tls=use_tls,
)
logger.custom_info(f'Temporal client initialized at {host}/{namespace}', metadata)
# Create cleanup schedule (idempotent - only creates if doesn't exist)
try:
await create_cleanup_schedule(temporal_client, logger, metadata)
except Exception as e: # noqa: BLE001
logger.custom_error(f'Failed to configure cleanup schedule: {e}', metadata)
# Don't fail the worker startup if schedule creation fails
# The schedule can be created manually if needed
workers = [
prepare_worker(
main_workflow=TrainModel,
other_workflows=[],
activities=[
activities.update_experiment_run,
activities.load_model_metadata,
activities.validate_train_params,
activities.train_model,
activities.cleanup_resources,
],
temporal_client=temporal_client,
logger=logger,
runtime=runtime,
),
prepare_worker(
main_workflow=CleanupFiles,
other_workflows=[],
activities=[
activities.cleanup_temp_directories,
],
temporal_client=temporal_client,
logger=logger,
runtime=runtime,
),
]
handlers = [w.run() for w in workers]
logger.custom_info('Model manager workers initialized', metadata)
try:
# This will run the workers and wait for them to complete.
# If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers)
except BaseException as e: # noqa: BLE001
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
finally:
notification_handler.shutdown()
logger.custom_info('MongoDB client closed', metadata)
activities.shutdown()
# Exit with a non-zero status code to indicate failure to Kubernetes
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
sys.exit(1)
def start_prometheus_server(logger: SientiaLogger, metadata: dict[str, str | None]):
"""
Starts the Prometheus metrics server for monitoring and observability.
This function initializes the Prometheus HTTP server on the configured port
and sets the application health metric to indicate the service is running.
The server exposes metrics that can be scraped by Prometheus for monitoring
the health and performance of the Model Manager worker.
Environment Variables:
HTTP_METRICS_PORT: Port for the metrics server (default: 9090)
POD_ID: Pod identifier for metrics labeling
Raises:
SystemExit: If the metrics server fails to start
"""
try:
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port)
logger.custom_info(f'Prometheus server initialized on port {port}.', metadata)
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
except Exception as e: # noqa: BLE001
logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)
os._exit(1)
if __name__ == '__main__':
asyncio.run(main())

View File

View File

@@ -0,0 +1,64 @@
"""
Cleanup workflow for removing local filesystem.
This module provides a Temporal cron workflow that runs daily to clean up
temporary files and directories older than the configured retention period.
"""
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import os
from datetime import timedelta
from typing import Any
from model_manager.activities.activities import Activities
from model_manager.runtime_paths import REPORTS_TEMP_DIR
from model_manager.workflows.train_model import no_retry_policy
TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120'))
POD_ID = os.getenv('POD_ID')
@workflow.defn(name='cleanup_files')
class CleanupFiles:
"""
Cleanup workflow for removing stale files.
This workflow cleans up:
- Local temporary directories with timestamp suffixes
The workflow is designed to be simple and robust, with error handling
delegated to the individual activities.
"""
@workflow.run
async def run(self, input_data: dict[str, Any] | None = None) -> None:
"""
Execute the cleanup workflow.
This method orchestrates the cleanup of local directories
in sequence. No exception handling is needed as activities handle their
own errors and notifications.
"""
payload = input_data or {}
temp_path = payload.get('temp_path') or REPORTS_TEMP_DIR
# Metadata for tracking
metadata = {
'metadata': {
'pod_id': POD_ID,
'workflow_name': 'cleanup_files',
}
}
# Execute local directory cleanup
await workflow.execute_activity_method(
Activities.cleanup_temp_directories,
{
**metadata,
'temp_path': temp_path,
},
retry_policy=no_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_LOCAL),
)

View File

@@ -0,0 +1,400 @@
"""
Train Model Workflow for ML model training pipeline.
This workflow orchestrates the complete ML model training process, including:
- Parameter validation and conversion
- Data download from MinIO
- Model training
- Model saving to MLFlow
- Experiment tracking and status updates
"""
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import os
from datetime import timedelta
from typing import Any
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError
from model_manager.activities.activities import Activities
from model_manager.activities.experiment_tracking import UpdateType
from model_manager.utils.models.experiment_status import ExperimentStatus
# 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'))
TIMEOUT_UPDATE_DATABASE = int(os.getenv('TIMEOUT_UPDATE_DATABASE', '30'))
# Retry Policies - Granular strategies for different operation types
# Fast retry for transient network errors (MinIO operations)
network_retry_policy = RetryPolicy(
initial_interval=timedelta(seconds=1),
maximum_interval=timedelta(seconds=10),
backoff_coefficient=2.0,
maximum_attempts=5,
)
# No retry for training - data errors are permanent
no_retry_policy = RetryPolicy(
maximum_attempts=1,
)
# Database retry with exponential backoff
database_retry_policy = RetryPolicy(
initial_interval=timedelta(seconds=2),
maximum_interval=timedelta(seconds=20),
backoff_coefficient=2.0,
maximum_attempts=5,
)
@workflow.defn(name='train_model')
class TrainModel:
"""
Complete ML model training workflow.
This workflow implements the full training pipeline from parameter validation
through model training and saving to MLFlow. It provides comprehensive error
handling with database status updates at each stage.
The workflow ensures:
- Proper parameter validation before training starts
- Status tracking in database for monitoring
- Error handling with detailed error messages
- Cleanup and proper resource management
"""
@workflow.run
async def run(self, input_data: dict[str, Any]) -> dict[str, Any] | None:
"""
Execute the complete model training workflow.
This method orchestrates all steps of the training pipeline:
1. Validate and convert training parameters
2. Download training data from MinIO
3. Train the model
4. Save model to MLFlow
5. Update experiment tracking status
Args:
input_data: Complete configuration for the training workflow
All TrainModelParams fields at the same level:
- experiment_run_id (int): Unique identifier for the experiment run (REQUIRED)
- target_variable (str): Target variable to predict
- variable_columns (list[str]): Feature columns
- train_size (int): Training data percentage
- ... (all other TrainModelParams fields)
Raises:
ValueError: If experiment_run_id is missing or invalid
"""
workflow.logger.info(f'Starting train_model workflow for {input_data}')
try:
experiment_run_id = self._validate_experiment_run_id(input_data)
except ValueError as exc:
# Prevent workflow-task retries on deterministic input contract violations.
raise ApplicationError(str(exc), non_retryable=True) from exc
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,
'workflow_name': 'train_model',
'model_name': model_name,
'model_id': model_id,
}
}
train_params = await self._validate_training_parameters(
input_data, experiment_run_id, metadata
)
training_succeeded = False
train_result: dict[str, Any] | None = None
try:
train_result = await self._train_model(
train_params=train_params,
experiment_run_id=experiment_run_id,
metadata=metadata,
)
training_succeeded = True
finally:
try:
if train_result is not None:
await self._cleanup_resources(
run_dir=train_result.get('run_dir'),
metadata=metadata,
)
else:
pass
except Exception: # noqa: BLE001
# 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:
"""
Validate experiment_run_id from input data.
This method ensures that experiment_run_id is present and valid.
Without a valid experiment_run_id, we cannot update database status,
so this validation must happen before any other operation.
Args:
input_data: Input data dictionary containing experiment_run_id
Returns:
int: Validated experiment_run_id
Raises:
ValueError: If experiment_run_id is missing or not an integer
"""
experiment_run_id = input_data.get('experiment_run_id')
if experiment_run_id is None:
raise ValueError('experiment_run_id is required but was not provided')
if isinstance(experiment_run_id, int):
return experiment_run_id
if isinstance(experiment_run_id, str) and experiment_run_id.strip().isdigit():
return int(experiment_run_id.strip())
raise ValueError(
f'experiment_run_id must be an integer or numeric string, got {type(experiment_run_id).__name__}'
)
async def _validate_training_parameters(
self,
input_data: dict[str, Any],
experiment_run_id: int,
metadata: dict[str, Any],
) -> dict[str, Any]:
"""
Validate and convert training parameters from dict to TrainModelParams.
This method calls the validate_train_params activity to convert and validate
the input parameters. On success, updates DB status to ORCHESTRATOR_WAITING_PROC.
On error, updates DB status to ORCHESTRATOR_VALIDATION_ERROR.
Args:
input_data: Input data dictionary containing all training parameters
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Returns:
dict[str, Any]: Validated training parameters
Raises:
Exception: If validation fails (after updating DB status)
"""
try:
input_data = await workflow.execute_activity_method(
Activities.load_model_metadata,
{
**input_data,
**metadata,
},
retry_policy=no_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
)
train_params = await workflow.execute_activity_method(
Activities.validate_train_params,
{
**input_data,
**metadata,
},
retry_policy=no_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS,
status=ExperimentStatus.ORCHESTRATOR_WAITING_PROC,
)
return train_params
except Exception as e:
try:
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
error_message=self._extract_error_message(e),
)
except Exception as secondary: # noqa: BLE001
workflow.logger.warning(
'Failed to persist ORCHESTRATOR_VALIDATION_ERROR to experiment_run: %s',
secondary,
)
raise
async def _train_model(
self,
train_params: dict[str, Any],
experiment_run_id: int,
metadata: dict[str, Any],
) -> dict[str, Any]:
"""
Download file from MinIO and train model.
This method orchestrates the download and training steps using proper
resource management with try/catch/finally. On success, updates DB status
to TRAINING_SUCCESS. On error, updates DB status to TRAINING_ERROR.
Args:
train_params: TrainModelParams object with training configuration
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Returns:
dict[str, Any]: Serializable training summary from the train_model activity
Raises:
Exception: If download or training fails (after updating DB status)
"""
try:
train_result = await workflow.execute_activity_method(
Activities.train_model,
{
**metadata,
'train_params': train_params,
},
retry_policy=no_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.TRAINING_SUCCESS,
run_name=train_result.get('run_name'),
)
return train_result
except Exception as e:
try:
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.TRAINING_ERROR,
error_message=self._extract_error_message(e),
)
except Exception as secondary: # noqa: BLE001
workflow.logger.warning(
'Failed to persist TRAINING_ERROR status to experiment_run: %s',
secondary,
)
raise
async def _cleanup_resources(
self,
run_dir: str | None,
metadata: dict[str, Any],
) -> None:
"""
Cleanup resources.
This method removes the temporary run directory via activity.
Args:
run_dir: Temporary directory to remove
metadata: Workflow execution metadata
"""
if run_dir is None:
return
await workflow.execute_activity_method(
Activities.cleanup_resources,
{
**metadata,
'run_dir': run_dir,
},
retry_policy=network_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
)
async def _update_experiment_run(
self,
metadata: dict[str, Any],
experiment_run_id: int,
update_type: UpdateType,
status: ExperimentStatus,
error_message: str | None = None,
run_name: str | None = None,
) -> None:
"""
Update experiment run status in the database.
This is a helper method to simplify calls to the update_experiment_run activity.
It handles both success and error status updates.
Args:
metadata: Workflow execution metadata
experiment_run_id: Unique identifier for the experiment run
update_type: Type of update (STATUS, STATUS_WITH_ERROR, or MODEL_SAVED)
status: Status to set in the database
error_message: Error message (required if update_type is STATUS_WITH_ERROR)
run_name: MLFlow run name (required if update_type is MODEL_SAVED)
"""
update_input = {
**metadata,
'experiment_run_id': experiment_run_id,
'update_type': update_type,
'status': status,
}
if error_message is not None:
update_input['error_message'] = error_message
if run_name is not None:
update_input['run_name'] = run_name
await workflow.execute_activity_method(
Activities.update_experiment_run,
update_input,
retry_policy=database_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE),
)
def _extract_error_message(self, exc: Exception) -> str:
message_parts: list[str] = []
seen: set[int] = set()
current: Exception | None = exc
while current and id(current) not in seen:
seen.add(id(current))
text = str(current).strip()
if text and text not in message_parts:
message_parts.append(text)
cause = getattr(current, '__cause__', None)
context = getattr(current, '__context__', None)
current = cause if isinstance(cause, Exception) else context
if not message_parts:
return repr(exc)
return ' | '.join(message_parts)