Code import - branch release/SIENTIAPDE-1645
This commit is contained in:
0
model_manager/__init__.py
Normal file
0
model_manager/__init__.py
Normal file
0
model_manager/activities/__init__.py
Normal file
0
model_manager/activities/__init__.py
Normal file
166
model_manager/activities/activities.py
Normal file
166
model_manager/activities/activities.py
Normal 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)
|
||||
177
model_manager/activities/cleanup.py
Normal file
177
model_manager/activities/cleanup.py
Normal 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
|
||||
280
model_manager/activities/experiment_tracking.py
Normal file
280
model_manager/activities/experiment_tracking.py
Normal 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
|
||||
497
model_manager/activities/training.py
Normal file
497
model_manager/activities/training.py
Normal 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
110
model_manager/metrics.py
Normal 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,
|
||||
)
|
||||
167
model_manager/reports/header.html
Normal file
167
model_manager/reports/header.html
Normal 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>
|
||||
0
model_manager/reports/temp/.gitkeep
Normal file
0
model_manager/reports/temp/.gitkeep
Normal file
29
model_manager/runtime_paths.py
Normal file
29
model_manager/runtime_paths.py
Normal 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)
|
||||
0
model_manager/schedules/__init__.py
Normal file
0
model_manager/schedules/__init__.py
Normal file
160
model_manager/schedules/cleanup_schedule.py
Normal file
160
model_manager/schedules/cleanup_schedule.py
Normal 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,
|
||||
)
|
||||
0
model_manager/sientia/__init__.py
Normal file
0
model_manager/sientia/__init__.py
Normal file
3
model_manager/sientia/exceptions.py
Normal file
3
model_manager/sientia/exceptions.py
Normal file
@@ -0,0 +1,3 @@
|
||||
from mlflow.exceptions import MlflowException
|
||||
|
||||
SientiaMlException = MlflowException
|
||||
148
model_manager/sientia/metrics.py
Normal file
148
model_manager/sientia/metrics.py
Normal 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
|
||||
292
model_manager/sientia/reports.py
Normal file
292
model_manager/sientia/reports.py
Normal 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)
|
||||
0
model_manager/utils/__init__.py
Normal file
0
model_manager/utils/__init__.py
Normal file
166
model_manager/utils/connectors_config.py
Normal file
166
model_manager/utils/connectors_config.py
Normal 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'),
|
||||
}
|
||||
27
model_manager/utils/logger_helper.py
Normal file
27
model_manager/utils/logger_helper.py
Normal 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
|
||||
16
model_manager/utils/models/__init__.py
Normal file
16
model_manager/utils/models/__init__.py
Normal 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',
|
||||
]
|
||||
26
model_manager/utils/models/experiment_status.py
Normal file
26
model_manager/utils/models/experiment_status.py
Normal 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'
|
||||
363
model_manager/utils/models/train_model_params.py
Normal file
363
model_manager/utils/models/train_model_params.py
Normal 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)
|
||||
53
model_manager/utils/models/train_model_result.py
Normal file
53
model_manager/utils/models/train_model_result.py
Normal 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
|
||||
636
model_manager/utils/repository/data_manager_repository.py
Normal file
636
model_manager/utils/repository/data_manager_repository.py
Normal 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}')
|
||||
0
model_manager/worker/__init__.py
Normal file
0
model_manager/worker/__init__.py
Normal file
119
model_manager/worker/prepare_worker.py
Normal file
119
model_manager/worker/prepare_worker.py
Normal 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'],
|
||||
),
|
||||
)
|
||||
259
model_manager/worker/worker.py
Normal file
259
model_manager/worker/worker.py
Normal 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())
|
||||
0
model_manager/workflows/__init__.py
Normal file
0
model_manager/workflows/__init__.py
Normal file
64
model_manager/workflows/cleanup_files.py
Normal file
64
model_manager/workflows/cleanup_files.py
Normal 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),
|
||||
)
|
||||
400
model_manager/workflows/train_model.py
Normal file
400
model_manager/workflows/train_model.py
Normal 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)
|
||||
Reference in New Issue
Block a user