SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.
This commit is contained in:
Bruno Domingues
@@ -7,12 +7,12 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
from model_manager.activities.experiment_tracking import ExperimentTracking
|
||||
from model_manager.activities.minio import MinIO
|
||||
from model_manager.activities.mlflow import MLFlow
|
||||
from model_manager.activities.training import Training
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
|
||||
class Activities(ExperimentTracking, MLFlow, MinIO, Training):
|
||||
class Activities(ExperimentTracking, Training):
|
||||
"""
|
||||
Main activities orchestrator for the Model Manager system.
|
||||
|
||||
@@ -75,18 +75,14 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Training):
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
MLFlow.__init__(
|
||||
self,
|
||||
mlflow_host=mlflow_config['host'],
|
||||
mlflow_port=mlflow_config['port'],
|
||||
mlflow_username=mlflow_config['username'],
|
||||
mlflow_password=mlflow_config['password'],
|
||||
self.model_repository = ModelRepository(
|
||||
url=mlflow_config['url'],
|
||||
username=mlflow_config['username'],
|
||||
password=mlflow_config['password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
MinIO.__init__(
|
||||
self,
|
||||
self.storage_repository = StorageRepository(
|
||||
endpoint_url=minio_config['endpoint_url'],
|
||||
access_key=minio_config['access_key'],
|
||||
secret_key=minio_config['secret_key'],
|
||||
@@ -97,10 +93,15 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Training):
|
||||
connect_timeout=minio_config['connect_timeout'],
|
||||
read_timeout=minio_config['read_timeout'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
Training.__init__(self, logger=logger, notification_handler=notification_handler)
|
||||
Training.__init__(
|
||||
self,
|
||||
model_repository=self.model_repository,
|
||||
storage_repository=self.storage_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
|
||||
@@ -9,7 +9,9 @@ methods for experiment management.
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import asyncio
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
@@ -18,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.temporal.activities.postgres import Postgres
|
||||
from sqlalchemy import text
|
||||
|
||||
|
||||
class UpdateType(str, Enum):
|
||||
@@ -85,6 +88,7 @@ class ExperimentTracking(Postgres):
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
self.logger = logger
|
||||
self.notification_handler = notification_handler
|
||||
|
||||
@@ -105,6 +109,14 @@ class ExperimentTracking(Postgres):
|
||||
# Silently ignore errors during garbage collection
|
||||
pass
|
||||
|
||||
async def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
|
||||
def _run() -> dict[str, Any]:
|
||||
with self.engine.begin() as connection:
|
||||
result = connection.execute(text(query), params)
|
||||
return {'rowcount': result.rowcount}
|
||||
|
||||
return await asyncio.to_thread(_run)
|
||||
|
||||
@activity.defn(name='update_experiment_run')
|
||||
async def update_experiment_run(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
@@ -128,32 +140,6 @@ class ExperimentTracking(Postgres):
|
||||
Raises:
|
||||
ValueError: If required parameters are missing for the update type
|
||||
RuntimeError: If update operation fails
|
||||
|
||||
Example:
|
||||
# Update status only
|
||||
await update_experiment_run({
|
||||
'metadata': {},
|
||||
'experiment_run_id': 123,
|
||||
'update_type': 'status',
|
||||
'status': 'TRAINING_SUCCESS'
|
||||
})
|
||||
|
||||
# Update with error
|
||||
await update_experiment_run({
|
||||
'metadata': {},
|
||||
'experiment_run_id': 123,
|
||||
'update_type': 'status_with_error',
|
||||
'status': 'TRAINING_ERROR',
|
||||
'error_message': 'Model training failed: insufficient data'
|
||||
})
|
||||
|
||||
# Update with model saved
|
||||
await update_experiment_run({
|
||||
'metadata': {},
|
||||
'experiment_run_id': 123,
|
||||
'update_type': 'model_saved',
|
||||
'run_name': 'experiment-model-5'
|
||||
})
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
experiment_run_id = input_data['experiment_run_id']
|
||||
@@ -164,71 +150,86 @@ class ExperimentTracking(Postgres):
|
||||
|
||||
try:
|
||||
self.info(
|
||||
f'Updating experiment run {experiment_run_id} with type: {update_type}', metadata
|
||||
f'Updating experiment run {experiment_run_id} with status: {status}', metadata
|
||||
)
|
||||
|
||||
# Validate parameters based on update type
|
||||
query_params: tuple[Any, ...]
|
||||
query_params: dict[str, Any]
|
||||
|
||||
if update_type == UpdateType.STATUS:
|
||||
if not status:
|
||||
if not isinstance(status, str) or not status:
|
||||
raise ValueError('status is required for STATUS update type')
|
||||
|
||||
sql_query = """
|
||||
UPDATE experiment_run
|
||||
SET status = %s, updated_at = %s
|
||||
WHERE id = %s
|
||||
SET status = :status, updated_at = :updated_at
|
||||
WHERE id = :experiment_run_id
|
||||
"""
|
||||
query_params = (status, datetime.now(UTC), experiment_run_id)
|
||||
|
||||
query_params = {
|
||||
'status': status,
|
||||
'updated_at': datetime.now(UTC),
|
||||
'experiment_run_id': experiment_run_id,
|
||||
}
|
||||
elif update_type == UpdateType.STATUS_WITH_ERROR:
|
||||
if not status or not error_message:
|
||||
raise ValueError(
|
||||
'status and error_message are required for STATUS_WITH_ERROR update type'
|
||||
)
|
||||
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 the error_message to 1024 characters if necessary
|
||||
if len(error_message) > 1024:
|
||||
error_message = error_message[:1024]
|
||||
|
||||
sql_query = """
|
||||
UPDATE experiment_run
|
||||
SET status = %s, error_message = %s, updated_at = %s
|
||||
WHERE id = %s
|
||||
SET status = :status, error_message = :error_message, updated_at = :updated_at
|
||||
WHERE id = :experiment_run_id
|
||||
"""
|
||||
query_params = (
|
||||
status,
|
||||
error_message,
|
||||
datetime.now(UTC),
|
||||
experiment_run_id,
|
||||
)
|
||||
|
||||
query_params = {
|
||||
'status': status,
|
||||
'error_message': error_message,
|
||||
'updated_at': datetime.now(UTC),
|
||||
'experiment_run_id': experiment_run_id,
|
||||
}
|
||||
elif update_type == UpdateType.MODEL_SAVED:
|
||||
if not run_name:
|
||||
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 = %s, status = %s, updated_at = %s
|
||||
WHERE id = %s
|
||||
SET run_name = :run_name, status = :status, updated_at = :updated_at
|
||||
WHERE id = :experiment_run_id
|
||||
"""
|
||||
query_params = (run_name, status, datetime.now(UTC), experiment_run_id)
|
||||
|
||||
query_params = {
|
||||
'run_name': run_name,
|
||||
'status': status,
|
||||
'updated_at': datetime.now(UTC),
|
||||
'experiment_run_id': experiment_run_id,
|
||||
}
|
||||
else:
|
||||
raise ValueError(f'Invalid update_type: {update_type}')
|
||||
|
||||
# Execute update query
|
||||
result = await self.execute_query(sql_query, query_params)
|
||||
result = await self._execute_update(sql_query, query_params)
|
||||
|
||||
if result.get('rowcount', 0) == 0:
|
||||
error_msg = f'No rows updated for experiment run {experiment_run_id}'
|
||||
error_msg = (
|
||||
f'No rows updated for experiment run {experiment_run_id} with status {status}'
|
||||
)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
self.info(
|
||||
f'Successfully updated experiment run {experiment_run_id} with type {update_type}',
|
||||
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}, Type: {update_type}, Error: {str(e)}'
|
||||
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,
|
||||
notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
|
||||
@@ -237,5 +238,6 @@ class ExperimentTracking(Postgres):
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
self.error(trace, metadata=metadata)
|
||||
raise RuntimeError(error_msg) from e
|
||||
|
||||
@@ -1,229 +0,0 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
import boto3 # type: ignore[import-untyped]
|
||||
from botocore.config import Config # type: ignore[import-untyped]
|
||||
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.temporal.activities.base import BaseActivity
|
||||
|
||||
|
||||
class MinIO(BaseActivity):
|
||||
"""
|
||||
MinIO (S3-compatible) storage activities for file operations.
|
||||
|
||||
This class provides activities for interacting with MinIO object storage,
|
||||
including file download and deletion operations. It handles authentication,
|
||||
connection management, and comprehensive error handling.
|
||||
|
||||
The class implements best practices for S3/MinIO operations:
|
||||
- Connection reuse (boto3 client is thread-safe)
|
||||
- Automatic retry with exponential backoff
|
||||
- Comprehensive error handling and logging
|
||||
- Notification integration for critical errors
|
||||
|
||||
Attributes:
|
||||
endpoint_url (str): MinIO server endpoint URL
|
||||
access_key (str): MinIO access key ID
|
||||
secret_key (str): MinIO secret access key
|
||||
region (str): MinIO region name
|
||||
use_ssl (bool): Whether to use SSL/TLS for connections
|
||||
minio_client: Boto3 S3 client configured for MinIO
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
endpoint_url: str,
|
||||
access_key: str,
|
||||
secret_key: str,
|
||||
region: str,
|
||||
use_ssl: bool,
|
||||
max_retry_attempts: int,
|
||||
retry_mode: str,
|
||||
connect_timeout: int,
|
||||
read_timeout: int,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
):
|
||||
"""
|
||||
Initialize MinIO activities with server configuration.
|
||||
|
||||
This constructor creates a persistent boto3 S3 client that will be
|
||||
reused across all activity calls. The client is thread-safe and
|
||||
includes automatic retry configuration.
|
||||
|
||||
Args:
|
||||
endpoint_url: MinIO server endpoint URL (e.g., http://localhost:9000)
|
||||
access_key: MinIO access key ID for authentication
|
||||
secret_key: MinIO secret access key for authentication
|
||||
region: MinIO region name (e.g., us-east-1)
|
||||
use_ssl: Whether to use SSL/TLS for connections
|
||||
max_retry_attempts: Maximum number of retry attempts (e.g., 3)
|
||||
retry_mode: Retry mode - standard, legacy, or adaptive (e.g., adaptive)
|
||||
connect_timeout: Connection timeout in seconds (e.g., 10)
|
||||
read_timeout: Read timeout in seconds (e.g., 60)
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
|
||||
Raises:
|
||||
ConnectionError: If boto3 client initialization fails
|
||||
"""
|
||||
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
|
||||
self.endpoint_url = endpoint_url
|
||||
self.access_key = access_key
|
||||
self.secret_key = secret_key
|
||||
self.region = region
|
||||
self.use_ssl = use_ssl
|
||||
self.max_retry_attempts = max_retry_attempts
|
||||
self.retry_mode = retry_mode
|
||||
self.connect_timeout = connect_timeout
|
||||
self.read_timeout = read_timeout
|
||||
|
||||
# Configure boto3 with retry strategy
|
||||
# This handles transient network errors and connection issues automatically
|
||||
boto_config = Config(
|
||||
region_name=region,
|
||||
retries={
|
||||
'max_attempts': max_retry_attempts,
|
||||
'mode': retry_mode,
|
||||
},
|
||||
connect_timeout=connect_timeout,
|
||||
read_timeout=read_timeout,
|
||||
)
|
||||
|
||||
try:
|
||||
self.minio_client = boto3.client(
|
||||
's3',
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=access_key,
|
||||
aws_secret_access_key=secret_key,
|
||||
config=boto_config,
|
||||
use_ssl=use_ssl,
|
||||
)
|
||||
self.info(f'MinIO client initialized successfully: {endpoint_url}')
|
||||
except Exception as e:
|
||||
error_msg = f'Failed to initialize MinIO client: {str(e)}'
|
||||
self.error(error_msg)
|
||||
raise ConnectionError(error_msg) from e
|
||||
|
||||
@activity.defn(name='fetch_file_from_minio')
|
||||
async def fetch_file_from_minio(self, input_data: dict[str, Any]) -> BytesIO:
|
||||
"""
|
||||
Fetch a file from MinIO and return its content as a BytesIO object.
|
||||
|
||||
This activity downloads a file from a MinIO bucket and returns the
|
||||
content as a BytesIO object, which is a file-like object that can be
|
||||
used directly with many Python libraries (pandas, PIL, etc.).
|
||||
|
||||
The operation includes:
|
||||
1. Input validation
|
||||
2. File download from MinIO
|
||||
3. Content reading and wrapping in BytesIO
|
||||
4. Comprehensive error handling and logging
|
||||
|
||||
Args:
|
||||
input_data: Configuration for file fetch operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- bucket_name (str): MinIO bucket name
|
||||
- file_name (str): File path/key in the bucket
|
||||
|
||||
Returns:
|
||||
BytesIO: File content as a file-like object
|
||||
|
||||
Raises:
|
||||
OSError: If file fetch fails due to network, permission, or other errors
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
bucket_name = input_data['bucket_name']
|
||||
file_name = input_data['file_name']
|
||||
|
||||
self.info(f'Fetching file from MinIO: {bucket_name}/{file_name}', metadata)
|
||||
|
||||
try:
|
||||
# Download file from MinIO
|
||||
response = self.minio_client.get_object(Bucket=bucket_name, Key=file_name)
|
||||
|
||||
# Read file content
|
||||
with response['Body'] as body:
|
||||
file_content = body.read()
|
||||
|
||||
file_size = len(file_content)
|
||||
self.info(
|
||||
f'File fetched successfully: {bucket_name}/{file_name} ({file_size} bytes)',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return BytesIO(file_content)
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Error fetching file from MinIO - Bucket: {bucket_name}, File: {file_name}, Error: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='FETCH_FILE_FROM_MINIO_ERROR',
|
||||
message=error_msg,
|
||||
block='fetch_file_from_minio',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
@activity.defn(name='delete_file_from_minio')
|
||||
async def delete_file_from_minio(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Delete a file from MinIO storage.
|
||||
|
||||
This activity removes a file from a MinIO bucket. The operation is
|
||||
idempotent - deleting a non-existent file is considered successful.
|
||||
|
||||
The operation includes:
|
||||
1. Input validation
|
||||
2. File deletion from MinIO
|
||||
3. Comprehensive error handling and logging
|
||||
|
||||
Args:
|
||||
input_data: Configuration for file deletion operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- bucket_name (str): MinIO bucket name
|
||||
- file_name (str): File path/key to delete
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
OSError: If file deletion fails due to permission or other errors
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
bucket_name = input_data['bucket_name']
|
||||
file_name = input_data['file_name']
|
||||
|
||||
self.info(f'Deleting file from MinIO: {bucket_name}/{file_name}', metadata)
|
||||
|
||||
try:
|
||||
# Delete file from MinIO
|
||||
# Note: delete_object is idempotent - no error if file doesn't exist
|
||||
self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
|
||||
|
||||
self.info(f'File deleted successfully: {bucket_name}/{file_name}', metadata)
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Error deleting file from MinIO - Bucket: {bucket_name}, File: {file_name}, Error: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='DELETE_FILE_FROM_MINIO_ERROR',
|
||||
message=error_msg,
|
||||
block='delete_file_from_minio',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise OSError(error_msg) from e
|
||||
@@ -1,214 +0,0 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
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.temporal.activities.base import BaseActivity
|
||||
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
from model_manager.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
class MLFlow(BaseActivity):
|
||||
"""
|
||||
MLFlow integration activities for model training operations.
|
||||
|
||||
This class provides activities for saving trained models and managing
|
||||
artifacts in MLFlow. It handles model persistence, artifact generation,
|
||||
and cleanup operations with comprehensive error handling.
|
||||
|
||||
The class implements robust error handling and logging for all
|
||||
MLFlow operations, ensuring reliable model management in production environments.
|
||||
|
||||
Attributes:
|
||||
mlflow_host (str): MLFlow server hostname
|
||||
mlflow_port (int): MLFlow server port
|
||||
mlflow_username (str): MLFlow authentication username
|
||||
mlflow_password (str): MLFlow authentication password
|
||||
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mlflow_host: str,
|
||||
mlflow_port: int,
|
||||
mlflow_username: str,
|
||||
mlflow_password: str,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
):
|
||||
"""
|
||||
Initialize MLFlow activities with server configuration.
|
||||
|
||||
Args:
|
||||
mlflow_host: MLFlow server hostname or IP address
|
||||
mlflow_port: MLFlow server port number
|
||||
mlflow_username: Username for MLFlow authentication
|
||||
mlflow_password: Password for MLFlow authentication
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
|
||||
Raises:
|
||||
Exception: If MLFlowRepository initialization fails
|
||||
"""
|
||||
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
|
||||
self.mlflow_host = mlflow_host
|
||||
self.mlflow_port = mlflow_port
|
||||
self.mlflow_username = mlflow_username
|
||||
self.mlflow_password = mlflow_password
|
||||
|
||||
self.model_monitoring_repository = MLFlowRepository(
|
||||
f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
|
||||
)
|
||||
|
||||
@activity.defn(name='save_model')
|
||||
async def save_model(self, input_data: dict[str, Any]) -> TrainModelResult:
|
||||
"""
|
||||
Save a trained ML model and its artifacts to MLflow.
|
||||
|
||||
This activity orchestrates the complete model saving pipeline:
|
||||
1. Generates the next run name for the experiment
|
||||
2. Creates and organizes artifacts (reports, data files)
|
||||
3. Logs model, parameters, metrics, and artifacts to MLflow
|
||||
|
||||
Args:
|
||||
input_data: Configuration for model saving operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- train_result (TrainModelResult): Training result with model and metrics
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated training result with run_name and artifacts
|
||||
|
||||
Raises:
|
||||
Exception: If model saving fails (after sending notification)
|
||||
|
||||
Example:
|
||||
result = await save_model({
|
||||
'metadata': {'workflow_id': 'save-123', 'experiment_run_id': 456},
|
||||
'train_result': TrainModelResult(...)
|
||||
})
|
||||
# Returns: TrainModelResult with run_name and artifacts
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
train_result = input_data['train_result']
|
||||
|
||||
try:
|
||||
experiment_name = train_result.params.experiment_name
|
||||
|
||||
self.info(
|
||||
f'Starting model save for experiment: {experiment_name}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Step 1: Generate next run name
|
||||
self.info('Generating run name', metadata)
|
||||
train_result.run_name = self.model_monitoring_repository.get_next_run_name(
|
||||
experiment_name
|
||||
)
|
||||
self.info(f'Generated run name: {train_result.run_name}', metadata)
|
||||
|
||||
# Step 2: Generate artifacts (reports, CSV files)
|
||||
self.info('Generating artifacts', metadata)
|
||||
train_result = self.model_monitoring_repository.generate_artifacts(train_result)
|
||||
self.info('Artifacts generated successfully', metadata)
|
||||
|
||||
# Step 3: Save run to MLflow
|
||||
self.info('Saving run to MLflow', metadata)
|
||||
self.model_monitoring_repository.save_run(train_result)
|
||||
|
||||
self.info(
|
||||
f'Model saved successfully - Run: {train_result.run_name}, '
|
||||
f'Experiment: {experiment_name}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return train_result
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Error saving model - Experiment: {train_result.params.experiment_name if train_result and train_result.params else "unknown"}, Error: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
|
||||
# Send notification (MongoDB)
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='SAVE_MODEL_ERROR',
|
||||
message=error_msg,
|
||||
block='save_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
# Log error with metadata
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise
|
||||
|
||||
@activity.defn(name='cleanup_run_directory')
|
||||
async def cleanup_run_directory(self, input_data: dict[str, Any]) -> 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:
|
||||
input_data: Configuration for cleanup operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- run_dir (str): Path to the run directory to delete
|
||||
|
||||
Raises:
|
||||
Exception: If cleanup fails for reasons other than directory not existing
|
||||
|
||||
Example:
|
||||
await cleanup_run_directory({
|
||||
'metadata': {'workflow_id': 'cleanup-123'},
|
||||
'run_dir': '/path/to/run_dir'
|
||||
})
|
||||
"""
|
||||
import os
|
||||
import shutil
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
run_dir = input_data.get('run_dir')
|
||||
|
||||
try:
|
||||
if not run_dir:
|
||||
self.info('No run directory specified, skipping cleanup', metadata)
|
||||
return
|
||||
|
||||
self.info(f'Cleaning up run directory: {run_dir}', metadata)
|
||||
|
||||
# Idempotent cleanup: check if directory exists before deleting
|
||||
if os.path.exists(run_dir):
|
||||
shutil.rmtree(run_dir)
|
||||
self.info(f'Run directory deleted successfully: {run_dir}', metadata)
|
||||
else:
|
||||
self.info(f'Run directory already deleted: {run_dir}', metadata)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f'Error cleaning up run directory {run_dir}: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
|
||||
# Send notification
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
|
||||
message=error_msg,
|
||||
block='cleanup_run_directory',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
# Log error
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
# Re-raise exception
|
||||
raise
|
||||
@@ -10,7 +10,6 @@ from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
@@ -18,8 +17,10 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
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.model_repository import ModelRepository
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
from model_manager.utils.repository.training_repository import TrainingRepository
|
||||
|
||||
|
||||
@@ -39,6 +40,8 @@ class Training(BaseActivity):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_repository: ModelRepository,
|
||||
storage_repository: StorageRepository,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
):
|
||||
@@ -49,8 +52,10 @@ class Training(BaseActivity):
|
||||
logger: Logger instance for observability
|
||||
notification_handler: Handler for sending notifications
|
||||
"""
|
||||
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
|
||||
super().__init__(logger, notification_handler, set_error_counter=True)
|
||||
self.training_repository = TrainingRepository(logger)
|
||||
self.model_repository = model_repository
|
||||
self.storage_repository = storage_repository
|
||||
|
||||
@activity.defn(name='validate_train_params')
|
||||
async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
|
||||
@@ -71,27 +76,12 @@ class Training(BaseActivity):
|
||||
|
||||
Raises:
|
||||
ValueError, TypeError, KeyError: If validation fails (after sending notification)
|
||||
|
||||
Example:
|
||||
result = await validate_train_params({
|
||||
'metadata': {'workflow_id': 'train-123'},
|
||||
'experiment_run_id': 456,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1', 'feature2'],
|
||||
'train_size': 80,
|
||||
# ... other required fields at same level
|
||||
})
|
||||
# Returns: TrainModelParams(...)
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
try:
|
||||
self.info('Validating training parameters', metadata)
|
||||
|
||||
# Step 1: Convert input_data to TrainModelParams (validates types and required fields)
|
||||
train_params = TrainModelParams.from_dict(input_data)
|
||||
|
||||
# Step 2: Validate business rules (ranges, consistency, etc.)
|
||||
train_params.validate_business_rules()
|
||||
|
||||
self.info(
|
||||
@@ -102,12 +92,10 @@ class Training(BaseActivity):
|
||||
)
|
||||
|
||||
return train_params
|
||||
|
||||
except (ValueError, TypeError, KeyError) as e:
|
||||
error_msg = f'Error validating training parameters: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
|
||||
# Send notification (MongoDB)
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
|
||||
@@ -117,14 +105,11 @@ class Training(BaseActivity):
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
# Log error with metadata
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise
|
||||
|
||||
@activity.defn(name='train_model')
|
||||
async def train_model(self, input_data: dict[str, Any]) -> TrainModelResult:
|
||||
async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]:
|
||||
"""
|
||||
Train a machine learning model.
|
||||
|
||||
@@ -156,51 +141,42 @@ class Training(BaseActivity):
|
||||
# Returns: TrainModelResult(...)
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
uploaded_file = input_data['uploaded_file']
|
||||
train_params = input_data['train_params']
|
||||
|
||||
if isinstance(train_params, dict):
|
||||
train_params = TrainModelParams.from_dict(train_params)
|
||||
# type: ignore[assignment]
|
||||
|
||||
model_trained = False
|
||||
model_saved = False
|
||||
|
||||
try:
|
||||
# Validate uploaded_file is BytesIO
|
||||
if not isinstance(uploaded_file, BytesIO):
|
||||
raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')
|
||||
with self.storage_repository.fetch_file(
|
||||
train_params.bucket_name, train_params.file_name
|
||||
) as uploaded_file:
|
||||
train_result = self.training_repository.train(uploaded_file, train_params)
|
||||
|
||||
# Validate train_params is TrainModelParams
|
||||
if not isinstance(train_params, TrainModelParams):
|
||||
raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}')
|
||||
train_result = self.training_repository.after_train_calculation(
|
||||
train_params, train_result
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Starting model training for target: {train_params.target_variable}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Step 1: Train the model
|
||||
self.info('Training model with TrainingRepository', metadata)
|
||||
train_result = self.training_repository.train(uploaded_file, train_params)
|
||||
|
||||
# Step 2: Perform post-training calculations
|
||||
self.info('Performing post-training calculations', metadata)
|
||||
final_result = self.training_repository.after_train_calculation(
|
||||
train_params, train_result
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Model training completed successfully - '
|
||||
f'MSE: {final_result.mse_val}, MAE: {final_result.mae_val}, R²: {final_result.r2_val}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return final_result
|
||||
model_trained = True
|
||||
train_result = self.model_repository.save_model(train_result)
|
||||
model_saved = True
|
||||
|
||||
return {
|
||||
'run_name': train_result.run_name,
|
||||
'run_dir': train_result.run_dir,
|
||||
}
|
||||
except Exception as e: # noqa: BLE001
|
||||
target = (
|
||||
train_params.target_variable
|
||||
if hasattr(train_params, 'target_variable')
|
||||
else 'unknown'
|
||||
error_msg = (
|
||||
'Error training model - '
|
||||
f'model_trained={model_trained}, model_saved={model_saved}, '
|
||||
f'error: {str(e)}'
|
||||
)
|
||||
error_msg = f'Error training model - Target: {target}, Error: {str(e)}'
|
||||
|
||||
trace = traceback.format_exc()
|
||||
|
||||
# Send notification (MongoDB)
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='TRAIN_MODEL_ERROR',
|
||||
@@ -210,8 +186,39 @@ class Training(BaseActivity):
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
# Log error with metadata
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise ModelTrainingError(
|
||||
model_trained=model_trained,
|
||||
model_saved=model_saved,
|
||||
) from e
|
||||
|
||||
@activity.defn(name='cleanup_resources')
|
||||
async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
|
||||
metadata = input_data.get('metadata', {})
|
||||
run_dir = input_data.get('run_dir', '')
|
||||
bucket_name = input_data.get('bucket_name', '')
|
||||
file_name = input_data.get('file_name', '')
|
||||
|
||||
try:
|
||||
self.model_repository.cleanup_run_directory(run_dir)
|
||||
self.storage_repository.delete_file(bucket_name, file_name)
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = (
|
||||
f'Error cleaning up resources - Run directory: {run_dir}, '
|
||||
f'File: {bucket_name}/{file_name}, 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,
|
||||
)
|
||||
|
||||
self.error(trace, metadata=metadata)
|
||||
raise
|
||||
|
||||
@@ -31,7 +31,6 @@ class ModelServing:
|
||||
tracking_uri: str,
|
||||
username: str | None = None,
|
||||
password: str | None = None,
|
||||
logger: Any | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize ModelServing client.
|
||||
@@ -54,9 +53,6 @@ class ModelServing:
|
||||
if password is not None:
|
||||
os.environ['MLFLOW_TRACKING_PASSWORD'] = password
|
||||
|
||||
# Note: logger parameter is accepted but not used
|
||||
# Consider removing if not needed, or implement logging
|
||||
|
||||
# Function to list runs for a given experiment
|
||||
def search_runs_by_name(
|
||||
self, experiment_names: list[str], order_by: None | list[str] = None
|
||||
|
||||
@@ -51,8 +51,7 @@ def build_mlflow_config() -> dict[str, Any]:
|
||||
dict: MLFlow configuration dictionary with all required parameters
|
||||
"""
|
||||
return {
|
||||
'host': getenv('MLFLOW_HOST', 'http://localhost'),
|
||||
'port': int(getenv('MLFLOW_PORT', '5080')),
|
||||
'url': getenv('MLFLOW_URL', 'http://localhost:5080'),
|
||||
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
|
||||
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
|
||||
}
|
||||
|
||||
12
model_manager/utils/exceptions.py
Normal file
12
model_manager/utils/exceptions.py
Normal file
@@ -0,0 +1,12 @@
|
||||
class ModelTrainingError(Exception):
|
||||
def __init__(self, model_trained: bool, model_saved: bool, message: str | None = None):
|
||||
self.model_trained = model_trained
|
||||
self.model_saved = model_saved
|
||||
|
||||
if message is None:
|
||||
message = (
|
||||
'Model training workflow failed '
|
||||
f'(model_trained={model_trained}, model_saved={model_saved})'
|
||||
)
|
||||
|
||||
super().__init__(message)
|
||||
@@ -184,22 +184,22 @@ class TrainModelParams:
|
||||
>>> params = TrainModelParams.from_dict(data)
|
||||
>>> params.validate_business_rules() # Raises ValueError if invalid
|
||||
"""
|
||||
# Validate train_size range (1-99%)
|
||||
if not 1 <= self.train_size <= 99:
|
||||
raise ValueError(f'train_size must be between 1 and 99, got {self.train_size}')
|
||||
# Validate train_size range (10-100%)
|
||||
if not 10 <= self.train_size <= 100:
|
||||
raise ValueError(f'train_size must be between 10 and 100, got {self.train_size}')
|
||||
|
||||
# Validate variable_columns is not empty
|
||||
if not self.variable_columns:
|
||||
raise ValueError('variable_columns cannot be empty')
|
||||
|
||||
# Validate positive integers
|
||||
if self.lag_train <= 0:
|
||||
if self.lag_train < 0:
|
||||
raise ValueError(f'lag_train must be positive, got {self.lag_train}')
|
||||
|
||||
if self.lag_val <= 0:
|
||||
if self.lag_val < 0:
|
||||
raise ValueError(f'lag_val must be positive, got {self.lag_val}')
|
||||
|
||||
if self.window <= 0:
|
||||
if self.window < 0:
|
||||
raise ValueError(f'window must be positive, got {self.window}')
|
||||
|
||||
# Validate low_lim and upp_lim consistency
|
||||
@@ -218,13 +218,6 @@ class TrainModelParams:
|
||||
f'Got low_lim={self.low_lim[var]}, upp_lim={self.upp_lim[var]}'
|
||||
)
|
||||
|
||||
# Validate target_variable is in variable_columns
|
||||
if self.target_variable not in self.variable_columns:
|
||||
raise ValueError(
|
||||
f'target_variable "{self.target_variable}" must be in variable_columns: '
|
||||
f'{self.variable_columns}'
|
||||
)
|
||||
|
||||
# Validate bucket_name and file_name are not empty
|
||||
if not self.bucket_name.strip():
|
||||
raise ValueError('bucket_name cannot be empty or whitespace')
|
||||
|
||||
@@ -9,6 +9,7 @@ and logging model runs to MLFlow.
|
||||
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
from os import makedirs, path
|
||||
@@ -22,14 +23,81 @@ from model_manager.sientia.reports import Reports # type: ignore[import-untyped
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
class MLFlowRepository:
|
||||
def __init__(self, host, username, password, logger: Logger):
|
||||
self.model_serving = ModelServing(
|
||||
tracking_uri=host, username=username, password=password, logger=logger
|
||||
)
|
||||
class ModelRepository:
|
||||
def __init__(self, url, username, password, logger: Logger):
|
||||
self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
|
||||
|
||||
self.logger = logger
|
||||
|
||||
def get_next_run_name(self, experiment_name: str) -> str:
|
||||
def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
|
||||
"""
|
||||
Save a trained ML model and its artifacts to MLflow.
|
||||
|
||||
This activity orchestrates the complete model saving pipeline:
|
||||
1. Generates the next run name for the experiment
|
||||
2. Creates and organizes artifacts (reports, data files)
|
||||
3. Logs model, parameters, metrics, and artifacts to MLflow
|
||||
|
||||
Args:
|
||||
input_data: Configuration for model saving operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- train_result (TrainModelResult): Training result with model and metrics
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated training result with run_name and artifacts
|
||||
|
||||
Raises:
|
||||
Exception: If model saving fails (after sending notification)
|
||||
"""
|
||||
experiment_name = train_result.params.experiment_name
|
||||
self.logger.info(f'Starting model save for experiment: {experiment_name}')
|
||||
|
||||
# Step 1: Generate next run name
|
||||
self.logger.info('Generating run name')
|
||||
train_result.run_name = self._get_next_run_name(experiment_name)
|
||||
self.logger.info(f'Generated run name: {train_result.run_name}')
|
||||
|
||||
# Step 2: Generate artifacts (reports, CSV files)
|
||||
self.logger.info('Generating artifacts')
|
||||
train_result = self._generate_artifacts(train_result)
|
||||
self.logger.info('Artifacts generated successfully')
|
||||
|
||||
# Step 3: Save run to MLflow
|
||||
self.logger.info('Saving run to MLflow')
|
||||
self._save_run(train_result)
|
||||
|
||||
self.logger.info(
|
||||
f'Model saved successfully - Run: {train_result.run_name}, '
|
||||
f'Experiment: {experiment_name}'
|
||||
)
|
||||
|
||||
return train_result
|
||||
|
||||
def cleanup_run_directory(self, run_dir: str) -> None:
|
||||
"""
|
||||
Clean up temporary run directory after model training.
|
||||
|
||||
This activity deletes the temporary directory created during model training
|
||||
and artifact generation. It implements idempotent cleanup to handle cases
|
||||
where the directory may have already been deleted.
|
||||
|
||||
Args:
|
||||
run_dir (str): Path to the run directory to delete
|
||||
"""
|
||||
if not run_dir:
|
||||
self.logger.info('No run directory specified, skipping cleanup')
|
||||
return
|
||||
|
||||
self.logger.info(f'Cleaning up run directory: {run_dir}')
|
||||
|
||||
if os.path.exists(run_dir):
|
||||
shutil.rmtree(run_dir)
|
||||
self.logger.info(f'Run directory deleted successfully: {run_dir}')
|
||||
else:
|
||||
self.logger.info(f'Run directory already deleted: {run_dir}')
|
||||
|
||||
def _get_next_run_name(self, experiment_name: str) -> str:
|
||||
"""
|
||||
Generates the next run name for a given experiment.
|
||||
|
||||
@@ -46,7 +114,7 @@ class MLFlowRepository:
|
||||
next_run_number = len(runs) + 1
|
||||
return f'{experiment_name}-{next_run_number}'
|
||||
|
||||
def generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
|
||||
def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
|
||||
"""
|
||||
Generates and organizes artifacts related to the training process, such as reports and data files.
|
||||
|
||||
@@ -87,7 +155,7 @@ class MLFlowRepository:
|
||||
self._setup_run_directory(data.run_dir, header_file_path)
|
||||
return self._generate_report(reference_data, current_data, data)
|
||||
|
||||
def save_run(self, data: TrainModelResult):
|
||||
def _save_run(self, data: TrainModelResult):
|
||||
"""
|
||||
Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
|
||||
to the Sientia tracking system.
|
||||
@@ -122,60 +190,48 @@ class MLFlowRepository:
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
try:
|
||||
# Prepare parameters
|
||||
train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
|
||||
interval_strs = [
|
||||
(str(interval[0]), str(interval[1]))
|
||||
for interval in (data.params.removed_intervals or [])
|
||||
]
|
||||
# Prepare parameters
|
||||
train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
|
||||
|
||||
# Set experiment and create run
|
||||
self.model_serving.set_experiment(data.params.experiment_name)
|
||||
self.logger.info(
|
||||
f"Logging run '{data.run_name}' to experiment '{data.params.experiment_name}'"
|
||||
)
|
||||
interval_strs = [
|
||||
(str(interval[0]), str(interval[1]))
|
||||
for interval in (data.params.removed_intervals or [])
|
||||
]
|
||||
|
||||
with self.model_serving.save_experiment(
|
||||
run_name=data.run_name, description=data.params.experiment_name
|
||||
):
|
||||
# Log model parameters
|
||||
self.model_serving.log_param('model_type', 'Linear Regression')
|
||||
self.model_serving.log_param('target_variable', data.params.target_variable)
|
||||
self.model_serving.log_param('input_variables', data.params.variable_columns)
|
||||
self.model_serving.log_param('lag_train', data.params.lag_train)
|
||||
self.model_serving.log_param('lag_val', data.params.lag_val)
|
||||
self.model_serving.log_param('ma', data.params.window)
|
||||
self.model_serving.log_param('low_lim', data.params.low_lim)
|
||||
self.model_serving.log_param('upp_lim', data.params.upp_lim)
|
||||
self.model_serving.log_param('normalized', data.scaler_dict)
|
||||
self.model_serving.log_param('ar', data.params.include_ar)
|
||||
self.model_serving.log_param('Train_test_split', train_test_split)
|
||||
self.model_serving.log_param('Removed_intervals', interval_strs)
|
||||
self.model_serving.log_param('Retrain', False)
|
||||
# Set experiment and create run
|
||||
self.model_serving.set_experiment(data.params.experiment_name)
|
||||
|
||||
# Log evaluation metrics
|
||||
self.model_serving.log_metric('MSE', data.mse_val)
|
||||
self.model_serving.log_metric('R2', data.r2_val)
|
||||
self.model_serving.log_metric('MAE', data.mae_val)
|
||||
with self.model_serving.save_experiment(
|
||||
run_name=data.run_name, description=data.params.experiment_name
|
||||
):
|
||||
# Log model parameters
|
||||
self.model_serving.log_param('model_type', 'Linear Regression')
|
||||
self.model_serving.log_param('target_variable', data.params.target_variable)
|
||||
self.model_serving.log_param('input_variables', data.params.variable_columns)
|
||||
self.model_serving.log_param('lag_train', data.params.lag_train)
|
||||
self.model_serving.log_param('lag_val', data.params.lag_val)
|
||||
self.model_serving.log_param('ma', data.params.window)
|
||||
self.model_serving.log_param('low_lim', data.params.low_lim)
|
||||
self.model_serving.log_param('upp_lim', data.params.upp_lim)
|
||||
self.model_serving.log_param('normalized', data.scaler_dict)
|
||||
self.model_serving.log_param('ar', data.params.include_ar)
|
||||
self.model_serving.log_param('Train_test_split', train_test_split)
|
||||
self.model_serving.log_param('Removed_intervals', interval_strs)
|
||||
self.model_serving.log_param('Retrain', False)
|
||||
|
||||
# Log models
|
||||
self.model_serving.log_model(data.process_data, 'data_model')
|
||||
self.model_serving.log_model(data.regr, 'prediction_model')
|
||||
# Log evaluation metrics
|
||||
self.model_serving.log_metric('MSE', data.mse_val)
|
||||
self.model_serving.log_metric('R2', data.r2_val)
|
||||
self.model_serving.log_metric('MAE', data.mae_val)
|
||||
|
||||
# Log artifacts
|
||||
self.model_serving.log_artifact(data.report_path)
|
||||
self.model_serving.log_artifact(data.train_data_path)
|
||||
self.model_serving.log_artifact(data.test_data_path)
|
||||
# Log models
|
||||
self.model_serving.log_model(data.process_data, 'data_model')
|
||||
self.model_serving.log_model(data.regr, 'prediction_model')
|
||||
|
||||
self.logger.info(
|
||||
f"Successfully logged run '{data.run_name}' with metrics: MSE={data.mse_val:.4f}, R2={data.r2_val:.4f}, MAE={data.mae_val:.4f}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to save run '{data.run_name}' to MLflow: {str(e)}"
|
||||
self.logger.error(error_msg)
|
||||
raise RuntimeError(error_msg) from e
|
||||
# Log artifacts
|
||||
self.model_serving.log_artifact(data.report_path)
|
||||
self.model_serving.log_artifact(data.train_data_path)
|
||||
self.model_serving.log_artifact(data.test_data_path)
|
||||
|
||||
def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""
|
||||
@@ -258,7 +314,6 @@ class MLFlowRepository:
|
||||
|
||||
try:
|
||||
makedirs(run_dir, exist_ok=True)
|
||||
self.logger.info(f'Created run directory: {run_dir}')
|
||||
return run_dir
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when creating directory: {run_dir}'
|
||||
@@ -298,9 +353,6 @@ class MLFlowRepository:
|
||||
# Copy header file to run directory
|
||||
header_dest = path.join(run_dir, 'header.html')
|
||||
shutil.copy(header_file_path, header_dest)
|
||||
|
||||
self.logger.info(f'Run directory setup completed successfully in: {run_dir}')
|
||||
|
||||
except FileNotFoundError as e:
|
||||
error_msg = f'Header file not found: {header_file_path}'
|
||||
self.logger.error(error_msg)
|
||||
@@ -360,20 +412,16 @@ class MLFlowRepository:
|
||||
# Save HTML report
|
||||
data.report_path = path.join(data.run_dir, 'report.html')
|
||||
report.save_all_sections_html(data.report_path)
|
||||
self.logger.info(f'Generated HTML report: {data.report_path}')
|
||||
|
||||
# Save training data CSV
|
||||
data.train_data_path = path.join(data.run_dir, 'train_data.csv')
|
||||
reference_data.to_csv(data.train_data_path, index=False)
|
||||
self.logger.info(f'Saved training data: {data.train_data_path}')
|
||||
|
||||
# Save test data CSV
|
||||
data.test_data_path = path.join(data.run_dir, 'test_data.csv')
|
||||
current_data.to_csv(data.test_data_path, index=False)
|
||||
self.logger.info(f'Saved test data: {data.test_data_path}')
|
||||
|
||||
return data
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
|
||||
129
model_manager/utils/repository/storage_repository.py
Normal file
129
model_manager/utils/repository/storage_repository.py
Normal file
@@ -0,0 +1,129 @@
|
||||
from io import BytesIO
|
||||
|
||||
import boto3 # type: ignore[import-untyped]
|
||||
from botocore.config import Config # type: ignore[import-untyped]
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
|
||||
class StorageRepository:
|
||||
"""
|
||||
MinIO (S3-compatible) storage activities for file operations.
|
||||
|
||||
This class provides activities for interacting with MinIO object storage,
|
||||
including file download and deletion operations. It handles authentication,
|
||||
connection management, and comprehensive error handling.
|
||||
|
||||
The class implements best practices for S3/MinIO operations:
|
||||
- Connection reuse (boto3 client is thread-safe)
|
||||
- Automatic retry with exponential backoff
|
||||
- Comprehensive error handling and logging
|
||||
- Notification integration for critical errors
|
||||
|
||||
Attributes:
|
||||
endpoint_url (str): MinIO server endpoint URL
|
||||
access_key (str): MinIO access key ID
|
||||
secret_key (str): MinIO secret access key
|
||||
region (str): MinIO region name
|
||||
use_ssl (bool): Whether to use SSL/TLS for connections
|
||||
minio_client: Boto3 S3 client configured for MinIO
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
endpoint_url: str,
|
||||
access_key: str,
|
||||
secret_key: str,
|
||||
region: str,
|
||||
use_ssl: bool,
|
||||
max_retry_attempts: int,
|
||||
retry_mode: str,
|
||||
connect_timeout: int,
|
||||
read_timeout: int,
|
||||
logger: Logger,
|
||||
):
|
||||
"""
|
||||
Initialize a reusable MinIO client with retry configuration.
|
||||
|
||||
Args:
|
||||
endpoint_url: MinIO server endpoint URL (e.g., http://localhost:9000).
|
||||
access_key: MinIO access key ID for authentication.
|
||||
secret_key: MinIO secret access key for authentication.
|
||||
region: MinIO region name (e.g., us-east-1).
|
||||
use_ssl: Whether to use SSL/TLS for connections.
|
||||
max_retry_attempts: Maximum number of retry attempts (e.g., 3).
|
||||
retry_mode: Retry policy to apply (standard, legacy, adaptive).
|
||||
connect_timeout: Connection timeout in seconds.
|
||||
read_timeout: Read timeout in seconds.
|
||||
logger: Logger used for observability.
|
||||
"""
|
||||
self.endpoint_url = endpoint_url
|
||||
self.access_key = access_key
|
||||
self.secret_key = secret_key
|
||||
self.region = region
|
||||
self.use_ssl = use_ssl
|
||||
self.max_retry_attempts = max_retry_attempts
|
||||
self.retry_mode = retry_mode
|
||||
self.connect_timeout = connect_timeout
|
||||
self.read_timeout = read_timeout
|
||||
self.logger = logger
|
||||
|
||||
boto_config = Config(
|
||||
region_name=region,
|
||||
retries={
|
||||
'max_attempts': max_retry_attempts,
|
||||
'mode': retry_mode,
|
||||
},
|
||||
connect_timeout=connect_timeout,
|
||||
read_timeout=read_timeout,
|
||||
)
|
||||
|
||||
self.minio_client = boto3.client(
|
||||
's3',
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=access_key,
|
||||
aws_secret_access_key=secret_key,
|
||||
config=boto_config,
|
||||
use_ssl=use_ssl,
|
||||
)
|
||||
|
||||
self.logger.info(f'MinIO client initialized successfully: {endpoint_url}')
|
||||
|
||||
def fetch_file(self, bucket_name: str, file_name: str) -> BytesIO:
|
||||
"""
|
||||
Fetch an object from MinIO and return its contents as `BytesIO`.
|
||||
|
||||
Args:
|
||||
bucket_name: MinIO bucket where the object resides.
|
||||
file_name: Object key to download inside the bucket.
|
||||
|
||||
Returns:
|
||||
BytesIO: File-like stream containing the downloaded bytes.
|
||||
|
||||
Raises:
|
||||
OSError: If the download fails (network, permissions, missing key, etc.).
|
||||
"""
|
||||
self.logger.info(f'Fetching file from MinIO: {bucket_name}/{file_name}')
|
||||
response = self.minio_client.get_object(Bucket=bucket_name, Key=file_name)
|
||||
|
||||
with response['Body'] as body:
|
||||
file_content = body.read()
|
||||
|
||||
file_size = len(file_content)
|
||||
|
||||
self.logger.info(
|
||||
f'File fetched successfully: {bucket_name}/{file_name} ({file_size} bytes)'
|
||||
)
|
||||
|
||||
return BytesIO(file_content)
|
||||
|
||||
def delete_file(self, bucket_name: str, file_name: str) -> None:
|
||||
"""
|
||||
Remove an object from MinIO storage.
|
||||
|
||||
Args:
|
||||
bucket_name: Bucket that contains the object.
|
||||
file_name: Object key to delete.
|
||||
"""
|
||||
self.logger.info(f'Deleting file from MinIO: {bucket_name}/{file_name}')
|
||||
self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
|
||||
self.logger.info(f'File deleted successfully: {bucket_name}/{file_name}')
|
||||
@@ -66,19 +66,17 @@ class TrainingRepository:
|
||||
ValueError: If transformed data is empty
|
||||
Exception: If data loading, preprocessing, or training fails
|
||||
"""
|
||||
# Load data from BytesIO
|
||||
self.logger.info('Loading data from BytesIO file')
|
||||
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
|
||||
|
||||
# Initialize and fit data preprocessor
|
||||
process_data = self.init_data_preprocessor(params)
|
||||
self.logger.info('Initializing and fitting data preprocessor')
|
||||
process_data = self._init_data_preprocessor(params)
|
||||
process_data.fit(data)
|
||||
data_view = process_data.transform(data)
|
||||
|
||||
# Validate transformed data
|
||||
if len(data_view) <= 0:
|
||||
raise ValueError('Data view is empty after transformation')
|
||||
|
||||
# Split data into train/test sets
|
||||
self.logger.info('Splitting data into train/test sets')
|
||||
x_train, x_test, y_train, y_test = split_train_test(
|
||||
data_view[params.variable_columns],
|
||||
data_view[params.target_variable],
|
||||
@@ -87,18 +85,18 @@ class TrainingRepository:
|
||||
random_state=42,
|
||||
)
|
||||
|
||||
# Prepare training data
|
||||
self.logger.info('Preparing training data')
|
||||
data_train = pd.concat([x_train, y_train], axis=1)
|
||||
scaler_dict = self.init_scaler_dict(process_data, params)
|
||||
scaler_dict = self._init_scaler_dict(process_data, params)
|
||||
|
||||
# Create and train linear regression model
|
||||
self.logger.info('Training linear regression model')
|
||||
regr = LinearRegressionModel(
|
||||
target_variable=params.target_variable,
|
||||
variable_columns=params.variable_columns,
|
||||
)
|
||||
|
||||
regr.fit(data_train)
|
||||
|
||||
# Return training result
|
||||
return TrainModelResult(
|
||||
params=params,
|
||||
process_data=process_data,
|
||||
@@ -110,7 +108,70 @@ class TrainingRepository:
|
||||
scaler_dict=scaler_dict,
|
||||
)
|
||||
|
||||
def init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
|
||||
def after_train_calculation(
|
||||
self, params: TrainModelParams, tmr: TrainModelResult
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Perform post-training calculations: predictions, denormalization, and metrics.
|
||||
|
||||
This method completes the training pipeline by:
|
||||
1. Making predictions on test set
|
||||
2. Denormalizing all data (if scaler was used)
|
||||
3. Reordering data by index
|
||||
4. Calculating evaluation metrics (MSE, MAE, R²)
|
||||
|
||||
Args:
|
||||
params: Training parameters used during model training
|
||||
tmr: Result object from training
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated result with predictions, denormalized data,
|
||||
and metrics (mse_val, mae_val, r2_val)
|
||||
"""
|
||||
self.logger.info('Making predictions on test set')
|
||||
y_pred_array = tmr.regr.predict(tmr.x_test)
|
||||
|
||||
if params.use_scaler:
|
||||
scaler = tmr.process_data.get_scaler()
|
||||
self.logger.info('Denormalizing features')
|
||||
|
||||
for col in params.variable_columns:
|
||||
tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
|
||||
tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
|
||||
|
||||
self.logger.info('Denormalizing target variable')
|
||||
tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
|
||||
tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
|
||||
y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
|
||||
|
||||
self.logger.info('Adding index to predictions')
|
||||
tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
|
||||
tmr.y_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
self.logger.info('Reordering all data by index')
|
||||
tmr.x_train = tmr.x_train.sort_index()
|
||||
tmr.x_test = tmr.x_test.sort_index()
|
||||
tmr.y_train = tmr.y_train.sort_index()
|
||||
tmr.y_test = tmr.y_test.sort_index()
|
||||
tmr.y_pred = tmr.y_pred.sort_index()
|
||||
|
||||
self.logger.info('Calculating evaluation metrics')
|
||||
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
||||
|
||||
tmr.mse_val = round(
|
||||
mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
tmr.mae_val = round(
|
||||
mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
|
||||
return tmr
|
||||
|
||||
def _init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
|
||||
"""
|
||||
Initialize dictionary containing scaling parameters for features and target.
|
||||
|
||||
@@ -152,69 +213,7 @@ class TrainingRepository:
|
||||
|
||||
return scaler_dict
|
||||
|
||||
def after_train_calculation(
|
||||
self, params: TrainModelParams, tmr: TrainModelResult
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Perform post-training calculations: predictions, denormalization, and metrics.
|
||||
|
||||
This method completes the training pipeline by:
|
||||
1. Making predictions on test set
|
||||
2. Denormalizing all data (if scaler was used)
|
||||
3. Reordering data by index
|
||||
4. Calculating evaluation metrics (MSE, MAE, R²)
|
||||
|
||||
Args:
|
||||
params: Training parameters used during model training
|
||||
tmr: Result object from training
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated result with predictions, denormalized data,
|
||||
and metrics (mse_val, mae_val, r2_val)
|
||||
"""
|
||||
# Make predictions on test set
|
||||
y_pred_array = tmr.regr.predict(tmr.x_test)
|
||||
|
||||
# Denormalize data if scaler was used
|
||||
if params.use_scaler:
|
||||
scaler = tmr.process_data.get_scaler()
|
||||
|
||||
# Denormalize features
|
||||
for col in params.variable_columns:
|
||||
tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
|
||||
tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
|
||||
|
||||
# Denormalize target variable
|
||||
tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
|
||||
tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
|
||||
y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
|
||||
|
||||
# Add index to predictions
|
||||
tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
|
||||
tmr.y_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
# Reorder all data by index
|
||||
tmr.x_train = tmr.x_train.sort_index()
|
||||
tmr.x_test = tmr.x_test.sort_index()
|
||||
tmr.y_train = tmr.y_train.sort_index()
|
||||
tmr.y_test = tmr.y_test.sort_index()
|
||||
tmr.y_pred = tmr.y_pred.sort_index()
|
||||
|
||||
# Calculate evaluation metrics
|
||||
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
||||
tmr.mse_val = round(
|
||||
mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
tmr.mae_val = round(
|
||||
mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
|
||||
|
||||
return tmr
|
||||
|
||||
def init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
|
||||
def _init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
|
||||
"""
|
||||
Initialize DataPreprocessor with training parameters.
|
||||
|
||||
|
||||
@@ -128,18 +128,10 @@ async def main():
|
||||
task_queue='train_model-queue',
|
||||
workflows=[TrainModel],
|
||||
activities=[
|
||||
# Training & Validation
|
||||
activities.update_experiment_run,
|
||||
activities.validate_train_params,
|
||||
activities.train_model,
|
||||
# MLFlow
|
||||
activities.save_model,
|
||||
# MinIO
|
||||
activities.fetch_file_from_minio,
|
||||
activities.delete_file_from_minio,
|
||||
# Filesystem
|
||||
activities.cleanup_run_directory,
|
||||
# Database
|
||||
activities.update_experiment_run,
|
||||
activities.cleanup_resources,
|
||||
],
|
||||
max_concurrent_workflow_tasks=50,
|
||||
max_concurrent_activities=50,
|
||||
|
||||
@@ -20,18 +20,15 @@ with workflow.unsafe.imports_passed_through():
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.activities.experiment_tracking import UpdateType
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
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
|
||||
|
||||
# Activity Timeouts (in seconds) - Configurable via environment variables
|
||||
# Defaults are designed to handle large files (up to 200MB)
|
||||
TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
|
||||
TIMEOUT_DOWNLOAD_FILE = int(os.getenv('TIMEOUT_DOWNLOAD_FILE', '600'))
|
||||
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '1800'))
|
||||
TIMEOUT_SAVE_MODEL = int(os.getenv('TIMEOUT_SAVE_MODEL', '300'))
|
||||
TIMEOUT_CLEANUP_DIRECTORY = int(os.getenv('TIMEOUT_CLEANUP_DIRECTORY', '60'))
|
||||
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '60'))
|
||||
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
|
||||
@@ -48,14 +45,6 @@ with workflow.unsafe.imports_passed_through():
|
||||
maximum_attempts=1,
|
||||
)
|
||||
|
||||
# Moderate retry with backoff for MLFlow operations
|
||||
mlflow_retry_policy = RetryPolicy(
|
||||
initial_interval=timedelta(seconds=5),
|
||||
maximum_interval=timedelta(seconds=30),
|
||||
backoff_coefficient=2.0,
|
||||
maximum_attempts=3,
|
||||
)
|
||||
|
||||
# Database retry with exponential backoff
|
||||
database_retry_policy = RetryPolicy(
|
||||
initial_interval=timedelta(seconds=2),
|
||||
@@ -64,14 +53,6 @@ with workflow.unsafe.imports_passed_through():
|
||||
maximum_attempts=5,
|
||||
)
|
||||
|
||||
# Filesystem retry for cleanup operations
|
||||
filesystem_retry_policy = RetryPolicy(
|
||||
initial_interval=timedelta(seconds=2),
|
||||
maximum_interval=timedelta(seconds=10),
|
||||
backoff_coefficient=1.5,
|
||||
maximum_attempts=3,
|
||||
)
|
||||
|
||||
|
||||
@workflow.defn(name='train_model')
|
||||
class TrainModel:
|
||||
@@ -113,8 +94,6 @@ class TrainModel:
|
||||
Raises:
|
||||
ValueError: If experiment_run_id is missing or invalid
|
||||
"""
|
||||
# CRITICAL: Validate experiment_run_id first
|
||||
# Without it, we cannot update database status, so fail immediately
|
||||
experiment_run_id = self._validate_experiment_run_id(input_data)
|
||||
|
||||
metadata = {
|
||||
@@ -124,33 +103,21 @@ class TrainModel:
|
||||
}
|
||||
}
|
||||
|
||||
# Step 1: Validate and convert training parameters
|
||||
train_params = await self._validate_training_parameters(
|
||||
input_data, experiment_run_id, metadata
|
||||
)
|
||||
|
||||
# Step 2 & 3: Download from MinIO and Train model
|
||||
# The _download_and_train_model method handles both steps:
|
||||
# - Downloads file from MinIO (returns BytesIO)
|
||||
# - Trains model with the downloaded file
|
||||
# Any error (download OR training) = TRAINING_ERROR
|
||||
train_result = await self._download_and_train_model(
|
||||
train_result = await self._train_model(
|
||||
train_params=train_params,
|
||||
experiment_run_id=experiment_run_id,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
# Step 4: Save model to MLFlow
|
||||
saved_result = await self._save_model_to_mlflow(
|
||||
train_result=train_result,
|
||||
experiment_run_id=experiment_run_id,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
# Step 5: Cleanup resources and delete file from MinIO
|
||||
await self._cleanup_resources(
|
||||
saved_result=saved_result,
|
||||
experiment_run_id=experiment_run_id,
|
||||
run_dir=(train_result.get('run_dir') or ''),
|
||||
bucket_name=train_params.bucket_name,
|
||||
file_name=train_params.file_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@@ -174,16 +141,12 @@ class TrainModel:
|
||||
experiment_run_id = input_data.get('experiment_run_id')
|
||||
|
||||
if experiment_run_id is None:
|
||||
error_msg = 'experiment_run_id is required but was not provided'
|
||||
workflow.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
raise ValueError('experiment_run_id is required but was not provided')
|
||||
|
||||
if not isinstance(experiment_run_id, int):
|
||||
error_msg = (
|
||||
raise ValueError(
|
||||
f'experiment_run_id must be an integer, got {type(experiment_run_id).__name__}'
|
||||
)
|
||||
workflow.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
return experiment_run_id
|
||||
|
||||
@@ -211,21 +174,17 @@ class TrainModel:
|
||||
Raises:
|
||||
Exception: If validation fails (after updating DB status)
|
||||
"""
|
||||
validation_input = {
|
||||
**metadata,
|
||||
**input_data, # All training params at same level
|
||||
}
|
||||
|
||||
try:
|
||||
# Execute validation activity
|
||||
train_params = await workflow.execute_activity_method(
|
||||
Activities.validate_train_params,
|
||||
validation_input,
|
||||
retry_policy=no_retry_policy, # Validation errors are permanent
|
||||
{
|
||||
**metadata,
|
||||
**input_data,
|
||||
},
|
||||
retry_policy=no_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
|
||||
)
|
||||
|
||||
# Validation succeeded: Update status
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
@@ -234,41 +193,23 @@ class TrainModel:
|
||||
)
|
||||
|
||||
return train_params
|
||||
|
||||
except Exception as e:
|
||||
# Validation failed: Update status with error
|
||||
error_message = str(e)
|
||||
error_type = type(e).__name__
|
||||
|
||||
# Log with rich context for debugging
|
||||
workflow.logger.error(
|
||||
f'[VALIDATION_ERROR] Experiment {experiment_run_id} validation failed',
|
||||
extra={
|
||||
'step': 'validate_training_parameters',
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'error_type': error_type,
|
||||
'error_message': error_message,
|
||||
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
|
||||
},
|
||||
)
|
||||
|
||||
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=error_message,
|
||||
error_message=self._extract_error_message(e),
|
||||
)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise
|
||||
|
||||
async def _download_and_train_model(
|
||||
async def _train_model(
|
||||
self,
|
||||
train_params: TrainModelParams,
|
||||
experiment_run_id: int,
|
||||
metadata: dict[str, Any],
|
||||
) -> TrainModelResult:
|
||||
) -> dict[str, str | None]:
|
||||
"""
|
||||
Download file from MinIO and train model.
|
||||
|
||||
@@ -287,167 +228,53 @@ class TrainModel:
|
||||
Raises:
|
||||
Exception: If download or training fails (after updating DB status)
|
||||
"""
|
||||
uploaded_file = None
|
||||
|
||||
try:
|
||||
# Step 1: Download file from MinIO
|
||||
download_input = {
|
||||
**metadata,
|
||||
'bucket_name': train_params.bucket_name,
|
||||
'file_name': train_params.file_name,
|
||||
}
|
||||
|
||||
uploaded_file = await workflow.execute_activity_method(
|
||||
Activities.fetch_file_from_minio,
|
||||
download_input,
|
||||
retry_policy=network_retry_policy, # Fast retry for network issues
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_DOWNLOAD_FILE),
|
||||
)
|
||||
|
||||
# Step 2: Train model with downloaded file
|
||||
train_input = {
|
||||
**metadata,
|
||||
'uploaded_file': uploaded_file,
|
||||
'train_params': train_params,
|
||||
}
|
||||
|
||||
train_result = await workflow.execute_activity_method(
|
||||
Activities.train_model,
|
||||
train_input,
|
||||
retry_policy=no_retry_policy, # Training errors are permanent (bad data)
|
||||
{
|
||||
**metadata,
|
||||
'train_params': train_params,
|
||||
},
|
||||
retry_policy=no_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
|
||||
)
|
||||
|
||||
# Training succeeded: Update status
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.STATUS,
|
||||
status=ExperimentStatus.TRAINING_SUCCESS,
|
||||
)
|
||||
|
||||
return train_result
|
||||
|
||||
except Exception as e:
|
||||
# Download or training failed: Update status with error
|
||||
error_message = str(e)
|
||||
error_type = type(e).__name__
|
||||
|
||||
# Log with rich context for debugging
|
||||
workflow.logger.error(
|
||||
f'[TRAINING_ERROR] Experiment {experiment_run_id} training failed',
|
||||
extra={
|
||||
'step': 'download_and_train_model',
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'experiment_name': train_params.experiment_name,
|
||||
'bucket_name': train_params.bucket_name,
|
||||
'file_name': train_params.file_name,
|
||||
'error_type': error_type,
|
||||
'error_message': error_message,
|
||||
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
|
||||
},
|
||||
)
|
||||
|
||||
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=error_message,
|
||||
)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Ensure BytesIO is closed
|
||||
if uploaded_file and hasattr(uploaded_file, 'close'):
|
||||
uploaded_file.close()
|
||||
|
||||
async def _save_model_to_mlflow(
|
||||
self,
|
||||
train_result: TrainModelResult,
|
||||
experiment_run_id: int,
|
||||
metadata: dict[str, Any],
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Save trained model to MLFlow.
|
||||
|
||||
This method calls the save_model activity to save the trained model and
|
||||
its artifacts to MLFlow. On success, updates DB status to MLFLOW_SENT
|
||||
with MODEL_SAVED type and run_name. On error, updates DB status to
|
||||
MLFLOW_SEND_ERROR.
|
||||
|
||||
Args:
|
||||
train_result: TrainModelResult from training step
|
||||
experiment_run_id: Validated experiment run ID
|
||||
metadata: Workflow execution metadata
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated training result with MLFlow run name
|
||||
|
||||
Raises:
|
||||
Exception: If model saving fails (after updating DB status)
|
||||
"""
|
||||
try:
|
||||
# Save model to MLFlow
|
||||
save_input = {
|
||||
**metadata,
|
||||
'train_result': train_result,
|
||||
}
|
||||
|
||||
saved_result = await workflow.execute_activity_method(
|
||||
Activities.save_model,
|
||||
save_input,
|
||||
retry_policy=mlflow_retry_policy, # Retry MLFlow with backoff
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_SAVE_MODEL),
|
||||
)
|
||||
|
||||
# Model saved successfully: Update status to MLFLOW_SENT with run_name
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.MODEL_SAVED,
|
||||
status=ExperimentStatus.MLFLOW_SENT,
|
||||
run_name=saved_result.run_name,
|
||||
run_name=train_result.get('run_name'),
|
||||
)
|
||||
|
||||
return saved_result
|
||||
|
||||
return train_result
|
||||
except Exception as e:
|
||||
# Model saving failed: Update status with error
|
||||
error_message = str(e)
|
||||
error_type = type(e).__name__
|
||||
# Mapear flags -> status
|
||||
# False/False: erro no treino
|
||||
# True/False: erro ao salvar (MLflow)
|
||||
# False/True: estado inconsistente, tratar como erro de treino
|
||||
# True/True: não deveria cair aqui; tratar como erro genérico de treino
|
||||
status = ExperimentStatus.TRAINING_ERROR
|
||||
|
||||
# Log with rich context for debugging
|
||||
workflow.logger.error(
|
||||
f'[MLFLOW_ERROR] Experiment {experiment_run_id} model save failed',
|
||||
extra={
|
||||
'step': 'save_model_to_mlflow',
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'experiment_name': train_result.params.experiment_name,
|
||||
'run_dir': train_result.run_dir if hasattr(train_result, 'run_dir') else None,
|
||||
'error_type': error_type,
|
||||
'error_message': error_message,
|
||||
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
|
||||
},
|
||||
)
|
||||
if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
|
||||
status = ExperimentStatus.MLFLOW_SEND_ERROR
|
||||
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.STATUS_WITH_ERROR,
|
||||
status=ExperimentStatus.MLFLOW_SEND_ERROR,
|
||||
error_message=error_message,
|
||||
status=status,
|
||||
error_message=self._extract_error_message(e),
|
||||
)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise
|
||||
|
||||
async def _cleanup_resources(
|
||||
self,
|
||||
saved_result: TrainModelResult,
|
||||
experiment_run_id: int,
|
||||
run_dir: str,
|
||||
bucket_name: str,
|
||||
file_name: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
@@ -466,75 +293,33 @@ class TrainModel:
|
||||
Exception: If cleanup fails (after updating DB status)
|
||||
"""
|
||||
try:
|
||||
# Step 1: Remove temporary run directory via activity (deterministic)
|
||||
if hasattr(saved_result, 'run_dir') and saved_result.run_dir:
|
||||
cleanup_input = {
|
||||
**metadata,
|
||||
'run_dir': saved_result.run_dir,
|
||||
}
|
||||
|
||||
await workflow.execute_activity_method(
|
||||
Activities.cleanup_run_directory,
|
||||
cleanup_input,
|
||||
retry_policy=filesystem_retry_policy, # Retry filesystem operations
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_DIRECTORY),
|
||||
)
|
||||
|
||||
workflow.logger.info(
|
||||
f'Run directory cleanup completed for experiment {experiment_run_id}'
|
||||
)
|
||||
|
||||
# Step 2: Delete file from MinIO
|
||||
delete_input = {
|
||||
**metadata,
|
||||
'bucket_name': saved_result.params.bucket_name,
|
||||
'file_name': saved_result.params.file_name,
|
||||
}
|
||||
|
||||
await workflow.execute_activity_method(
|
||||
Activities.delete_file_from_minio,
|
||||
delete_input,
|
||||
retry_policy=network_retry_policy, # Fast retry for network issues
|
||||
Activities.cleanup_resources,
|
||||
{
|
||||
**metadata,
|
||||
'run_dir': run_dir,
|
||||
'bucket_name': bucket_name,
|
||||
'file_name': file_name,
|
||||
},
|
||||
retry_policy=network_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
|
||||
)
|
||||
|
||||
# Cleanup succeeded: Update status to FILE_DELETED
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.STATUS,
|
||||
status=ExperimentStatus.FILE_DELETED,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# Cleanup failed: Update status with error
|
||||
error_message = str(e)
|
||||
error_type = type(e).__name__
|
||||
|
||||
# Log with rich context for debugging
|
||||
workflow.logger.error(
|
||||
f'[CLEANUP_ERROR] Experiment {experiment_run_id} cleanup failed',
|
||||
extra={
|
||||
'step': 'cleanup_resources',
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'bucket_name': saved_result.params.bucket_name,
|
||||
'file_name': saved_result.params.file_name,
|
||||
'run_dir': saved_result.run_dir if hasattr(saved_result, 'run_dir') else None,
|
||||
'error_type': error_type,
|
||||
'error_message': error_message,
|
||||
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
|
||||
},
|
||||
)
|
||||
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.STATUS_WITH_ERROR,
|
||||
status=ExperimentStatus.FILE_DELETE_ERROR,
|
||||
error_message=error_message,
|
||||
error_message=self._extract_error_message(e),
|
||||
)
|
||||
|
||||
# Re-raise exception to stop workflow
|
||||
raise
|
||||
|
||||
async def _update_experiment_run(
|
||||
@@ -567,17 +352,34 @@ class TrainModel:
|
||||
'status': status,
|
||||
}
|
||||
|
||||
# Add error_message only if provided
|
||||
if error_message is not None:
|
||||
update_input['error_message'] = error_message
|
||||
|
||||
# Add run_name only if provided
|
||||
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, # Retry DB with exponential backoff
|
||||
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)
|
||||
|
||||
current = getattr(current, 'cause', None)
|
||||
|
||||
if not message_parts:
|
||||
return repr(exc)
|
||||
|
||||
return ' | '.join(message_parts)
|
||||
|
||||
Reference in New Issue
Block a user