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

Merge pull request #7 from Aignosi/feature/SIENTIAPDE-1253

Browse Source 
SIENTIAPDE-1253: Improve Train Model Workflow with Enhanced Error Handling, Validation, and Cleanup
This commit is contained in:
Bruno Domingues
2025-10-15 17:04:49 -03:00
committed by GitHub
parent 014e1cd910 4ff796a56f
commit 9ee12373b4
12 changed files with 2339 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
.env.example
View File

@@ -33,4 +33,14 @@ MINIO_USE_SSL="false"
MINIO_MAX_RETRY_ATTEMPTS="3"
MINIO_RETRY_MODE="adaptive"
MINIO_CONNECT_TIMEOUT="10"
MINIO_READ_TIMEOUT="60"
MINIO_READ_TIMEOUT="60"
# Workflow Activity Timeouts (in seconds)
# These timeouts are designed to handle large files (up to 200MB)
TIMEOUT_VALIDATE_PARAMS="30" # Parameter validation (fast operation)
TIMEOUT_DOWNLOAD_FILE="600" # File download from MinIO (10 min for 200MB @ 1MB/s with 3x buffer)
TIMEOUT_TRAIN_MODEL="1800" # Model training (30 min for large datasets)
TIMEOUT_SAVE_MODEL="300" # Save model to MLFlow (5 min for artifacts upload)
TIMEOUT_CLEANUP_DIRECTORY="60" # Cleanup temporary directory (1 min)
TIMEOUT_DELETE_FILE="60" # Delete file from MinIO (1 min)
TIMEOUT_UPDATE_DATABASE="30" # Database update operations (30 sec)

113
README.md
View File

@@ -17,7 +17,8 @@ A comprehensive AI model management platform for the complete machine learning l
- [Predictions Batch Workflow](#1-predictions-batch-workflow-predictions_batchpy)
- [Prediction Process Workflow](#2-prediction-process-workflow-prediction_processpy)
- [Format and Export Prediction Workflow](#3-format-and-export-prediction-workflow-format_and_export_predictionpy)
- [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy)
- [Train Model Workflow](#4-train-model-workflow-train_modelpy)
- [Minimal Retrain Workflow](#5-minimal-retrain-workflow-minimal_retrainpy)
- [Installation & Setup](#installation--setup)
- [Prerequisites](#prerequisites)
- [Environment Setup](#environment-setup)
@@ -379,7 +380,99 @@ flowchart LR
C -.-> Prometheus[Prometheus]
```
### 4. Minimal Retrain Workflow (`minimal_retrain.py`)
### 4. Train Model Workflow (`train_model.py`)
The **TrainModel** workflow orchestrates the complete ML model training pipeline from parameter validation through model saving and cleanup.
#### Purpose
- **Model Training**: Complete ML model training pipeline
- **Parameter Validation**: Defense-in-depth validation with business rules
- **Resource Management**: Automatic cleanup of temporary resources
- **Status Tracking**: Comprehensive experiment tracking in database
- **Error Handling**: Robust error handling with detailed context logging
#### Execution Flow
1. **Validate Experiment Run ID**: Critical validation before any DB updates
2. **Validate Training Parameters**: Type checking + business rules validation
3. **Download Training Data**: Fetch file from MinIO storage
4. **Train Model**: Execute ML model training with validated parameters
5. **Save to MLFlow**: Save trained model and artifacts to MLFlow
6. **Cleanup Resources**: Delete temporary files and MinIO data
#### Key Features
- **Granular Retry Policies**: Different strategies for network, training, MLFlow, database, and filesystem operations
- **Configurable Timeouts**: Environment variable-based timeouts supporting files up to 200MB
- **Idempotent Cleanup**: Safe replay with Temporal workflow replay mechanism
- **Structured Logging**: Rich context in error messages for debugging
- **Business Validation**: 10 business rules including range checks, consistency validation, and data integrity
#### Input Parameters
```json
{
"experiment_run_id": 123,
"target_variable": "price",
"variable_columns": ["feature1", "feature2", "price"],
"train_size": 80,
"shuffle": true,
"use_scaler": true,
"include_ar": false,
"bucket_name": "ml-data",
"file_name": "training_data.csv",
"line_separator": "\n",
"decimal_separator": ".",
"lag_train": 5,
"lag_val": 3,
"rem_static_win": false,
"low_lim": {"feature1": 0.0, "feature2": 0.0, "price": 0.0},
"upp_lim": {"feature1": 100.0, "feature2": 100.0, "price": 1000.0},
"window": 10,
"experiment_name": "production_model_v1",
"removed_intervals": []
}
```
#### Architecture Diagram
```mermaid
flowchart TD
A[1. validate_experiment_run_id] --> B[2. validate_train_params]
B --> C[3. fetch_file_from_minio]
C --> D[4. train_model]
D --> E[5. save_model]
E --> F[6. cleanup_run_directory]
F --> G[7. delete_file_from_minio]
B -.-> DB[(PostgreSQL)]
C -.-> MinIO[MinIO Storage]
D -.-> Training[ML Training]
E -.-> MLFlow[MLFlow]
F -.-> FS[Filesystem]
G -.-> MinIO
```
#### Retry Strategies
The workflow implements 5 different retry policies optimized for each operation type:
| Operation Type | Initial Interval | Max Interval | Backoff | Max Attempts | Use Case |
|---------------|------------------|--------------|---------|--------------|----------|
| **Network** | 1s | 10s | 2.0x | 5 | MinIO operations (transient network errors) |
| **No Retry** | - | - | - | 1 | Training/Validation (permanent data errors) |
| **MLFlow** | 5s | 30s | 2.0x | 3 | MLFlow operations (API timeouts) |
| **Database** | 2s | 20s | 2.0x | 5 | PostgreSQL updates (lock contention) |
| **Filesystem** | 2s | 10s | 1.5x | 3 | Cleanup operations (busy resources) |
#### Business Validation Rules
The workflow validates 10 business rules beyond type checking:
1. **train_size**: Must be between 1-99%
2. **variable_columns**: Cannot be empty
3. **lag_train, lag_val, window**: Must be positive integers
4. **low_lim/upp_lim**: Must have same keys and low < upp for each variable
5. **target_variable**: Must be in variable_columns
6. **bucket_name, file_name, experiment_name**: Cannot be empty or whitespace
### 5. Minimal Retrain Workflow (`minimal_retrain.py`)
The **MinimalRetrain** workflow handles automated model retraining and production model updates.
@@ -892,6 +985,22 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa
| `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No |
| `POD_ID` | Kubernetes pod identifier | `None` | No |
#### Workflow Activity Timeouts
These timeouts control how long each activity in the training workflow can run before timing out. All values are in seconds and are designed to handle large files (up to 200MB).
| Variable | Description | Default | Calculation Basis |
|----------|-------------|---------|-------------------|
| `TIMEOUT_VALIDATE_PARAMS` | Parameter validation timeout | `30` | Fast operation, no I/O |
| `TIMEOUT_DOWNLOAD_FILE` | File download from MinIO timeout | `600` | 200MB @ 1MB/s with 3x buffer (10 min) |
| `TIMEOUT_TRAIN_MODEL` | Model training timeout | `1800` | Large dataset processing (30 min) |
| `TIMEOUT_SAVE_MODEL` | Save model to MLFlow timeout | `300` | Artifact upload and logging (5 min) |
| `TIMEOUT_CLEANUP_DIRECTORY` | Cleanup temporary directory timeout | `60` | Local filesystem operation (1 min) |
| `TIMEOUT_DELETE_FILE` | Delete file from MinIO timeout | `60` | MinIO delete operation (1 min) |
| `TIMEOUT_UPDATE_DATABASE` | Database update timeout | `30` | PostgreSQL update query (30 sec) |
**Note**: These timeouts can be adjusted based on your infrastructure performance and file sizes. If you're processing files larger than 200MB or have slower network/compute resources, increase these values accordingly.
### Workflow Configuration
MongoDB pipeline configuration:

105
model_manager/activities/mlflow.py
View File

@@ -13,6 +13,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.utils.repository.model_repository import MLFlowRepository
@@ -346,19 +347,14 @@ class MLFlow(BaseActivity):
raise e
@activity.defn(name='save_model')
async def save_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
async def save_model(self, input_data: dict[str, Any]) -> TrainModelResult:
"""
Save a trained ML model and its artifacts to MLflow with comprehensive error handling.
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
4. Returns success/failure status with results or error message
The activity does NOT raise exceptions on failure - it catches all errors,
sends notifications, and returns a failure status. This allows the workflow
to handle the error gracefully and update the database accordingly.
Args:
input_data: Configuration for model saving operation
@@ -367,23 +363,17 @@ class MLFlow(BaseActivity):
- train_result (TrainModelResult): Training result with model and metrics
Returns:
dict: Save result with the following structure:
{
'success': bool, # True if saving succeeded, False otherwise
'result': TrainModelResult | None, # Updated result if success=True
'error_message': str | None # Error message if success=False
}
TrainModelResult: Updated training result with run_name and artifacts
Raises:
Exception: If model saving fails (after sending notification)
Example:
# Successful save
result = await save_model({
'metadata': {'workflow_id': 'save-123', 'experiment_run_id': 456},
'train_result': TrainModelResult(...)
})
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
# Failed save
# Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
# Returns: TrainModelResult with run_name and artifacts
"""
metadata = input_data.get('metadata', {})
train_result = input_data['train_result']
@@ -418,11 +408,7 @@ class MLFlow(BaseActivity):
metadata,
)
return {
'success': True,
'result': train_result,
'error_message': None,
}
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)}'
@@ -441,10 +427,69 @@ class MLFlow(BaseActivity):
# Log error with metadata
self.error(trace, metadata=metadata)
# Return failure result (do NOT raise exception)
# This allows workflow to update database with error status
return {
'success': False,
'result': None,
'error_message': str(e),
}
# 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

153
model_manager/activities/training.py
View File

@@ -19,6 +19,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.base import BaseActivity
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.training_repository import TrainingRepository
@@ -51,73 +52,127 @@ class Training(BaseActivity):
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.training_repository = TrainingRepository(logger)
@activity.defn(name='train_model')
async def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
@activity.defn(name='validate_train_params')
async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
"""
Train a machine learning model with comprehensive error handling.
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:
TrainModelParams: Validated and converted training parameters
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(
f'Training parameters validated successfully - '
f'Target: {train_params.target_variable}, '
f'Experiment: {train_params.experiment_name}',
metadata,
)
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',
message=error_msg,
block='validate_train_params',
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='train_model')
async def train_model(self, input_data: dict[str, Any]) -> TrainModelResult:
"""
Train a machine learning model.
This activity orchestrates the complete ML training pipeline:
1. Validates input parameters
2. Trains the model using TrainingRepository
3. Performs post-training calculations
4. Returns success/failure status with results or error message
The activity does NOT raise exceptions on failure - it catches all errors,
sends notifications, and returns a failure status. This allows the workflow
to handle the error gracefully and update the database accordingly.
Args:
input_data: Configuration for model training operation
Required keys:
- metadata (dict): Workflow execution metadata
- uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
- train_params (dict): Training parameters (converted to TrainModelParams)
- train_params (TrainModelParams): Training parameters object
Returns:
dict: Training result with the following structure:
{
'success': bool, # True if training succeeded, False otherwise
'result': TrainModelResult | None, # Training result if success=True
'error_message': str | None # Error message if success=False
}
TrainModelResult: Training result with model, metrics, and data
Raises:
ValueError: If input validation fails
Exception: If training fails (after sending notification)
Example:
# Successful training
result = await train_model({
'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
'uploaded_file': BytesIO(csv_data),
'train_params': {
'experiment_run_id': 456,
'target_variable': 'price',
'variable_columns': ['feature1', 'feature2'],
'train_size': 80,
'shuffle': True,
'use_scaler': True,
# ... other TrainModelParams fields
}
'train_params': TrainModelParams(...)
})
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
# Failed training
# Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
# Returns: TrainModelResult(...)
"""
metadata = input_data.get('metadata', {})
uploaded_file = input_data['uploaded_file']
train_params_dict = input_data['train_params']
train_params = input_data['train_params']
try:
self.info(
f'Starting model training for target: {train_params_dict.get("target_variable")}',
metadata,
)
# Convert dict to TrainModelParams
train_params = TrainModelParams.from_dict(train_params_dict)
# Validate uploaded_file is BytesIO
if not isinstance(uploaded_file, BytesIO):
raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')
# Validate train_params is TrainModelParams
if not isinstance(train_params, TrainModelParams):
raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}')
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)
@@ -134,14 +189,15 @@ class Training(BaseActivity):
metadata,
)
return {
'success': True,
'result': final_result,
'error_message': None,
}
return final_result
except Exception as e: # noqa: BLE001
error_msg = f'Error training model - Target: {train_params_dict.get("target_variable", "unknown")}, Error: {str(e)}'
target = (
train_params.target_variable
if hasattr(train_params, 'target_variable')
else 'unknown'
)
error_msg = f'Error training model - Target: {target}, Error: {str(e)}'
trace = traceback.format_exc()
# Send notification (MongoDB)
@@ -157,10 +213,5 @@ class Training(BaseActivity):
# Log error with metadata
self.error(trace, metadata=metadata)
# Return failure result (do NOT raise exception)
# This allows workflow to update database with error status
return {
'success': False,
'result': None,
'error_message': str(e),
}
# Re-raise exception to stop workflow
raise

1
model_manager/utils/models/experiment_status.py
View File

@@ -23,6 +23,7 @@ class ExperimentStatus(str, Enum):
FILE_DELETE_ERROR: Cleanup failed due to file system or MinIO errors.
"""
ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR'
MAGE_WAITING_PROC = 'MAGE_WAITING_PROC'
TRAINING_SUCCESS = 'TRAINING_SUCCESS'
TRAINING_ERROR = 'TRAINING_ERROR'

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

@@ -167,3 +167,71 @@ class TrainModelParams:
raise TypeError(error)
return value
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
Example:
>>> 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 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:
raise ValueError(f'lag_train must be positive, got {self.lag_train}')
if self.lag_val <= 0:
raise ValueError(f'lag_val must be positive, got {self.lag_val}')
if self.window <= 0:
raise ValueError(f'window must be positive, got {self.window}')
# Validate low_lim and upp_lim consistency
if set(self.low_lim.keys()) != set(self.upp_lim.keys()):
raise ValueError(
f'low_lim and upp_lim must have the same keys. '
f'low_lim keys: {set(self.low_lim.keys())}, '
f'upp_lim keys: {set(self.upp_lim.keys())}'
)
# Validate that low_lim < upp_lim for each variable
for var in self.low_lim:
if self.low_lim[var] >= self.upp_lim[var]:
raise ValueError(
f'low_lim must be less than upp_lim for variable "{var}". '
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')
if not self.file_name.strip():
raise ValueError('file_name cannot be empty or whitespace')
# Validate experiment_name is not empty
if not self.experiment_name.strip():
raise ValueError('experiment_name cannot be empty or whitespace')

583
model_manager/workflows/train_model.py Normal file
View File

@@ -0,0 +1,583 @@
"""
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 model_manager.activities.activities import Activities
from model_manager.activities.experiment_tracking import UpdateType
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_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,
)
# 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),
maximum_interval=timedelta(seconds=20),
backoff_coefficient=2.0,
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:
"""
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]) -> 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
"""
# 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 = {
'metadata': {
'experiment_run_id': experiment_run_id,
'workflow_name': 'train_model',
}
}
# 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_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,
metadata=metadata,
)
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:
error_msg = 'experiment_run_id is required but was not provided'
workflow.logger.error(error_msg)
raise ValueError(error_msg)
if not isinstance(experiment_run_id, int):
error_msg = (
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
async def _validate_training_parameters(
self,
input_data: dict[str, Any],
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelParams:
"""
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 MAGE_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:
TrainModelParams: Validated training parameters object
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
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,
update_type=UpdateType.STATUS,
status=ExperimentStatus.MAGE_WAITING_PROC,
)
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,
)
# Re-raise exception to stop workflow
raise
async def _download_and_train_model(
self,
train_params: TrainModelParams,
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelResult:
"""
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:
TrainModelResult: Training result from train_model activity
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)
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,
)
return saved_result
except Exception as e:
# Model saving failed: Update status with error
error_message = str(e)
error_type = type(e).__name__
# 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'),
},
)
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,
)
# Re-raise exception to stop workflow
raise
async def _cleanup_resources(
self,
saved_result: TrainModelResult,
experiment_run_id: int,
metadata: dict[str, Any],
) -> None:
"""
Cleanup resources and delete file from MinIO.
This method removes the temporary run directory via activity and deletes
the training file from MinIO. On success, updates DB status to FILE_DELETED.
On error, updates DB status to FILE_DELETE_ERROR.
Args:
saved_result: TrainModelResult with run_dir and params information
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Raises:
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
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,
)
# Re-raise exception to stop workflow
raise
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,
}
# 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
start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE),
)

219
tests/activities/test_mlflow.py
View File

@@ -1,6 +1,7 @@
from unittest.mock import ANY, MagicMock, patch
import numpy as np
import pytest
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
@@ -333,10 +334,9 @@ async def test_save_model_success(mlflow):
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
# Verify response
assert response['success'] is True
assert response['result'] == train_result
assert response['error_message'] is None
# Verify response - now returns TrainModelResult directly
assert response == train_result
assert response.run_name == 'test_experiment-1'
assert train_result.run_name == 'test_experiment-1'
@@ -362,13 +362,9 @@ async def test_save_model_get_next_run_name_error(mlflow):
'train_result': train_result,
}
# Call the method
response = await mlflow.save_model(input_data)
# Verify error handling
assert response['success'] is False
assert response['result'] is None
assert 'MLflow connection error' in response['error_message']
# Call the method - should raise exception
with pytest.raises(Exception, match='MLflow connection error'):
await mlflow.save_model(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
@@ -404,13 +400,9 @@ async def test_save_model_generate_artifacts_error(mlflow):
'train_result': train_result,
}
# Call the method
response = await mlflow.save_model(input_data)
# Verify error handling
assert response['success'] is False
assert response['result'] is None
assert 'Reports directory does not exist' in response['error_message']
# Call the method - should raise exception
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
await mlflow.save_model(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
@@ -447,13 +439,9 @@ async def test_save_model_save_run_error(mlflow):
'train_result': train_result,
}
# Call the method
response = await mlflow.save_model(input_data)
# Verify error handling
assert response['success'] is False
assert response['result'] is None
assert 'One or more metrics (MSE, R2, MAE) are None' in response['error_message']
# Call the method - should raise exception
with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
await mlflow.save_model(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
@@ -491,10 +479,9 @@ async def test_save_model_missing_metadata(mlflow):
# Call the method
response = await mlflow.save_model(input_data)
# Verify it still works (metadata defaults to {})
assert response['success'] is True
assert response['result'] == train_result
assert response['error_message'] is None
# Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
assert response == train_result
assert response.run_name == 'test_experiment-1'
@mark.asyncio
@@ -540,10 +527,170 @@ async def test_save_model_complete_flow(mlflow):
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
# Verify response
assert response['success'] is True
assert response['result'] == updated_result
assert response['error_message'] is None
assert updated_result.run_name == 'production_model-5'
assert updated_result.run_dir is not None
assert updated_result.report_path is not None
# Verify response - returns TrainModelResult directly
assert response == updated_result
assert response.run_name == 'production_model-5'
assert response.run_dir is not None
assert response.report_path is not None
# ============================================================================
# Tests for cleanup_run_directory
# ============================================================================
@mark.asyncio
async def test_cleanup_run_directory_success(mlflow):
"""Test successful cleanup of run directory."""
input_data = {
**metadata,
'run_dir': 'test_run_dir',
}
# Patch os and shutil inside the activity method
with (
patch('os.path.exists', return_value=True) as mock_exists,
patch('shutil.rmtree') as mock_rmtree,
):
# Call the method
await mlflow.cleanup_run_directory(input_data)
# Verify directory existence was checked
mock_exists.assert_called_once_with('test_run_dir')
# Verify shutil.rmtree was called
mock_rmtree.assert_called_once_with('test_run_dir')
@mark.asyncio
async def test_cleanup_run_directory_already_deleted(mlflow):
"""Test cleanup when directory is already deleted (idempotent)."""
input_data = {
**metadata,
'run_dir': 'already_deleted_dir',
}
# Patch os and shutil inside the activity method
with (
patch('os.path.exists', return_value=False) as mock_exists,
patch('shutil.rmtree') as mock_rmtree,
):
# Call the method - should not raise error
await mlflow.cleanup_run_directory(input_data)
# Verify directory existence was checked
mock_exists.assert_called_once_with('already_deleted_dir')
# Verify shutil.rmtree was NOT called
mock_rmtree.assert_not_called()
@mark.asyncio
async def test_cleanup_run_directory_no_run_dir(mlflow):
"""Test cleanup when no run_dir is provided."""
input_data = {
**metadata,
# No run_dir key
}
# Call the method - should not raise error
await mlflow.cleanup_run_directory(input_data)
# Should complete without errors
@mark.asyncio
async def test_cleanup_run_directory_empty_run_dir(mlflow):
"""Test cleanup when run_dir is empty string."""
input_data = {
**metadata,
'run_dir': '',
}
# Call the method - should not raise error
await mlflow.cleanup_run_directory(input_data)
# Should complete without errors
@mark.asyncio
async def test_cleanup_run_directory_none_run_dir(mlflow):
"""Test cleanup when run_dir is None."""
input_data = {
**metadata,
'run_dir': None,
}
# Call the method - should not raise error
await mlflow.cleanup_run_directory(input_data)
# Should complete without errors
@mark.asyncio
async def test_cleanup_run_directory_error(mlflow):
"""Test cleanup handles errors correctly."""
input_data = {
**metadata,
'run_dir': 'error_dir',
}
# Patch os and shutil inside the activity method
with (
patch('os.path.exists', return_value=True),
patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
):
# Call the method - should raise exception
with pytest.raises(PermissionError, match='Permission denied'):
await mlflow.cleanup_run_directory(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
message=ANY,
block='cleanup_run_directory',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_cleanup_run_directory_missing_metadata(mlflow):
"""Test cleanup handles missing metadata gracefully."""
input_data = {
'run_dir': 'no_metadata_dir',
}
# Patch os and shutil inside the activity method
with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
# Call the method
await mlflow.cleanup_run_directory(input_data)
# Verify directory was deleted
mock_rmtree.assert_called_once_with('no_metadata_dir')
@mark.asyncio
async def test_cleanup_run_directory_oserror(mlflow):
"""Test cleanup handles OSError correctly."""
input_data = {
**metadata,
'run_dir': 'os_error_dir',
}
# Patch os and shutil inside the activity method
with (
patch('os.path.exists', return_value=True),
patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
):
# Call the method - should raise exception
with pytest.raises(OSError, match='Directory not empty'):
await mlflow.cleanup_run_directory(input_data)
# Verify notification was sent
mlflow.send_notification.assert_called_once()
call_args = mlflow.send_notification.call_args[1]
assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
assert call_args['level'] == NotificationLevel.ERROR
assert 'Directory not empty' in call_args['message']

472
tests/activities/test_training.py
View File

@@ -3,9 +3,11 @@
from io import BytesIO
from unittest.mock import MagicMock, patch
import pytest
from pytest import mark
from model_manager.activities.training import Training
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
@@ -42,41 +44,42 @@ async def test_train_model_success(mock_training_repository_class):
# Test data
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
train_params_dict = {
'experiment_run_id': 123,
'target_variable': 'price',
'variable_columns': ['feature1', 'feature2'],
'train_size': 80,
'shuffle': True,
'use_scaler': True,
'include_ar': False,
'bucket_name': 'test-bucket',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0, 'feature2': 0.0},
'upp_lim': {'feature1': 100.0, 'feature2': 100.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
}
train_params = TrainModelParams(
experiment_run_id=123,
target_variable='price',
variable_columns=['feature1', 'feature2'],
train_size=80,
shuffle=True,
use_scaler=True,
include_ar=False,
bucket_name='test-bucket',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
lag_train=1,
lag_val=1,
rem_static_win=False,
low_lim={'feature1': 0.0, 'feature2': 0.0},
upp_lim={'feature1': 100.0, 'feature2': 100.0},
window=10,
experiment_name='test_experiment',
removed_intervals=[],
)
input_data = {
'metadata': {'workflow_id': 'test-123'},
'uploaded_file': uploaded_file,
'train_params': train_params_dict,
'train_params': train_params,
}
# Execute
result = await training.train_model(input_data)
# Assertions
assert result['success'] is True
assert result['result'] == mock_final_result
assert result['error_message'] is None
# Assertions - now returns TrainModelResult directly
assert result == mock_final_result
assert result.mse_val == 0.5
assert result.mae_val == 0.3
assert result.r2_val == 0.95
# Verify repository calls
mock_repository.train.assert_called_once()
@@ -95,37 +98,38 @@ async def test_train_model_invalid_file_type(mock_training_repository_class):
training = Training(logger=logger, notification_handler=notification_handler)
# Invalid file type (string instead of BytesIO)
train_params = TrainModelParams(
experiment_run_id=123,
target_variable='price',
variable_columns=['feature1'],
train_size=80,
shuffle=True,
use_scaler=False,
include_ar=False,
bucket_name='test',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
lag_train=1,
lag_val=1,
rem_static_win=False,
low_lim={'feature1': 0.0},
upp_lim={'feature1': 100.0},
window=10,
experiment_name='test_experiment',
removed_intervals=[],
)
input_data = {
'metadata': {},
'uploaded_file': 'not_a_bytesio',
'train_params': {
'experiment_run_id': 123,
'target_variable': 'price',
'variable_columns': ['feature1'],
'train_size': 80,
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise ValueError
with pytest.raises(ValueError, match='uploaded_file must be BytesIO'):
await training.train_model(input_data)
assert result['success'] is False
assert result['result'] is None
assert 'uploaded_file must be BytesIO' in result['error_message']
# Verify notification was sent (via BaseActivity)
notification_handler.send_notification.assert_called_once()
@@ -145,37 +149,38 @@ async def test_train_model_training_error(mock_training_repository_class):
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n')
train_params = TrainModelParams(
experiment_run_id=456,
target_variable='price',
variable_columns=['feature1'],
train_size=80,
shuffle=True,
use_scaler=False,
include_ar=False,
bucket_name='test',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
lag_train=1,
lag_val=1,
rem_static_win=False,
low_lim={'feature1': 0.0},
upp_lim={'feature1': 100.0},
window=10,
experiment_name='test_experiment',
removed_intervals=[],
)
input_data = {
'metadata': {'workflow_id': 'test-456'},
'uploaded_file': uploaded_file,
'train_params': {
'experiment_run_id': 456,
'target_variable': 'price',
'variable_columns': ['feature1'],
'train_size': 80,
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise ValueError
with pytest.raises(ValueError, match='Training data is empty'):
await training.train_model(input_data)
assert result['success'] is False
assert result['result'] is None
assert 'Training data is empty' in result['error_message']
# Verify notification was sent (via BaseActivity)
notification_handler.send_notification.assert_called_once()
@@ -194,42 +199,41 @@ async def test_train_model_sends_notification_on_error(mock_training_repository_
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n1,2')
train_params = TrainModelParams(
experiment_run_id=789,
target_variable='price',
variable_columns=['feature1'],
train_size=80,
shuffle=True,
use_scaler=False,
include_ar=False,
bucket_name='test',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
lag_train=1,
lag_val=1,
rem_static_win=False,
low_lim={'feature1': 0.0},
upp_lim={'feature1': 100.0},
window=10,
experiment_name='test_experiment',
removed_intervals=[],
)
input_data = {
'metadata': {'workflow_id': 'test-789', 'experiment_run_id': 789},
'uploaded_file': uploaded_file,
'train_params': {
'experiment_run_id': 789,
'target_variable': 'price',
'variable_columns': ['feature1'],
'train_size': 80,
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise Exception
with pytest.raises(Exception, match='Database connection failed'):
await training.train_model(input_data)
# Verify notification was sent (via BaseActivity)
notification_handler.send_notification.assert_called_once()
# Verify result - the important part is that error was caught and returned
assert result['success'] is False
assert result['result'] is None
assert 'Database connection failed' in result['error_message']
@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
@@ -248,36 +252,248 @@ async def test_train_model_after_calculation_error(mock_training_repository_clas
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
train_params = TrainModelParams(
experiment_run_id=999,
target_variable='price',
variable_columns=['feature1'],
train_size=80,
shuffle=True,
use_scaler=False,
include_ar=False,
bucket_name='test',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
lag_train=1,
lag_val=1,
rem_static_win=False,
low_lim={'feature1': 0.0},
upp_lim={'feature1': 100.0},
window=10,
experiment_name='test_experiment',
removed_intervals=[],
)
input_data = {
'metadata': {},
'uploaded_file': uploaded_file,
'train_params': {
'experiment_run_id': 999,
'target_variable': 'price',
'variable_columns': ['feature1'],
'train_size': 80,
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
},
'train_params': train_params,
}
result = await training.train_model(input_data)
# Should raise Exception
with pytest.raises(Exception, match='Metric calculation failed'):
await training.train_model(input_data)
assert result['success'] is False
assert result['result'] is None
assert 'Metric calculation failed' in result['error_message']
# Verify notification was sent (via BaseActivity)
notification_handler.send_notification.assert_called_once()
@mark.asyncio
@patch('model_manager.activities.training.TrainingRepository')
async def test_train_model_invalid_train_params_type(mock_training_repository_class):
"""Test training with invalid train_params type (dict instead of TrainModelParams)."""
mock_repository = MagicMock()
mock_training_repository_class.return_value = mock_repository
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
# Invalid train_params type (dict instead of TrainModelParams object)
input_data = {
'metadata': {'workflow_id': 'test-invalid'},
'uploaded_file': uploaded_file,
'train_params': {
'experiment_run_id': 123,
'target_variable': 'price',
'variable_columns': ['feature1'],
}, # This is a dict, not TrainModelParams
}
# Should raise ValueError
with pytest.raises(ValueError, match='train_params must be TrainModelParams.*dict'):
await training.train_model(input_data)
# ============================================================================
# Tests for validate_train_params
# ============================================================================
@mark.asyncio
async def test_validate_train_params_success():
"""Test successful validation of training parameters."""
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
training.info = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
'experiment_run_id': 456,
'target_variable': 'price',
'variable_columns': ['feature1', 'feature2', 'price'], # target_variable must be in list
'train_size': 80,
'shuffle': True,
'use_scaler': True,
'include_ar': False,
'bucket_name': 'test-bucket',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0, 'feature2': 0.0, 'price': 0.0},
'upp_lim': {'feature1': 100.0, 'feature2': 100.0, 'price': 1000.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
}
result = await training.validate_train_params(input_data)
assert isinstance(result, TrainModelParams)
assert result.experiment_run_id == 456
assert result.target_variable == 'price'
assert result.variable_columns == ['feature1', 'feature2', 'price']
assert result.train_size == 80
assert result.experiment_name == 'test_experiment'
assert training.info.call_count == 2
@mark.asyncio
async def test_validate_train_params_missing_required_field():
"""Test validation fails when required field is missing."""
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
training.info = MagicMock()
training.error = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-123'},
'experiment_run_id': 456,
'variable_columns': ['feature1'],
'train_size': 80,
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test',
'removed_intervals': [],
}
with pytest.raises(ValueError, match='target_variable'):
await training.validate_train_params(input_data)
notification_handler.send_notification.assert_called_once()
@mark.asyncio
async def test_validate_train_params_invalid_type():
"""Test validation fails when field has invalid type."""
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
training.info = MagicMock()
training.error = MagicMock()
input_data = {
'metadata': {'workflow_id': 'test-invalid'},
'experiment_run_id': 456,
'target_variable': 'price',
'variable_columns': ['feature1'],
'train_size': 'invalid',
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test',
'removed_intervals': [],
}
with pytest.raises((ValueError, TypeError)):
await training.validate_train_params(input_data)
notification_handler.send_notification.assert_called_once()
@mark.asyncio
async def test_validate_train_params_empty_input():
"""Test validation fails with empty input."""
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
training.info = MagicMock()
training.error = MagicMock()
input_data = {'metadata': {}}
with pytest.raises(ValueError):
await training.validate_train_params(input_data)
notification_handler.send_notification.assert_called_once()
@mark.asyncio
async def test_validate_train_params_without_metadata():
"""Test validation works even without metadata key."""
logger = MagicMock()
notification_handler = MagicMock()
training = Training(logger=logger, notification_handler=notification_handler)
training.info = MagicMock()
input_data = {
'experiment_run_id': 789,
'target_variable': 'temperature',
'variable_columns': ['sensor1', 'temperature'], # target_variable must be in list
'train_size': 75,
'shuffle': False,
'use_scaler': True,
'include_ar': True,
'bucket_name': 'sensors',
'file_name': 'data.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 2,
'lag_val': 2,
'rem_static_win': True,
'low_lim': {'sensor1': -50.0, 'temperature': -50.0},
'upp_lim': {'sensor1': 150.0, 'temperature': 150.0},
'window': 20,
'experiment_name': 'sensor_experiment',
'removed_intervals': [],
}
result = await training.validate_train_params(input_data)
assert isinstance(result, TrainModelParams)
assert result.experiment_run_id == 789
assert result.target_variable == 'temperature'
assert result.experiment_name == 'sensor_experiment'

6
tests/utils/models/test_experiment_status.py
View File

@@ -15,8 +15,8 @@ def test_experiment_status_values():
def test_experiment_status_count():
"""Test that enum has exactly 7 status values."""
assert len(ExperimentStatus) == 7
"""Test that enum has exactly 8 status values."""
assert len(ExperimentStatus) == 8
def test_experiment_status_is_string():
@@ -40,7 +40,7 @@ def test_experiment_status_membership():
def test_experiment_status_iteration():
"""Test that enum can be iterated."""
statuses = list(ExperimentStatus)
assert len(statuses) == 7
assert len(statuses) == 8
assert ExperimentStatus.MAGE_WAITING_PROC in statuses
assert ExperimentStatus.TRAINING_SUCCESS in statuses
assert ExperimentStatus.TRAINING_ERROR in statuses

198
tests/utils/models/test_train_model_params.py
View File

@@ -297,3 +297,201 @@ def test_train_model_params_check_type_method():
# Test that _check_type is a private method
assert hasattr(params, '_check_type')
assert callable(params._check_type)
# ============================================================================
# Tests for validate_business_rules method
# ============================================================================
def test_validate_business_rules_success(valid_params_dict):
"""Test that valid params pass business rules validation."""
# Ensure target_variable is in variable_columns
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
# Should not raise any exception
params.validate_business_rules()
def test_validate_business_rules_train_size_too_low(valid_params_dict):
"""Test that train_size < 1 raises ValueError."""
valid_params_dict['train_size'] = 0
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='train_size must be between 1 and 99'):
params.validate_business_rules()
def test_validate_business_rules_train_size_too_high(valid_params_dict):
"""Test that train_size > 99 raises ValueError."""
valid_params_dict['train_size'] = 100
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='train_size must be between 1 and 99'):
params.validate_business_rules()
def test_validate_business_rules_empty_variable_columns(valid_params_dict):
"""Test that empty variable_columns raises ValueError."""
valid_params_dict['variable_columns'] = []
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='variable_columns cannot be empty'):
params.validate_business_rules()
def test_validate_business_rules_lag_train_zero(valid_params_dict):
"""Test that lag_train = 0 raises ValueError."""
valid_params_dict['lag_train'] = 0
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='lag_train must be positive'):
params.validate_business_rules()
def test_validate_business_rules_lag_train_negative(valid_params_dict):
"""Test that lag_train < 0 raises ValueError."""
valid_params_dict['lag_train'] = -1
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='lag_train must be positive'):
params.validate_business_rules()
def test_validate_business_rules_lag_val_zero(valid_params_dict):
"""Test that lag_val = 0 raises ValueError."""
valid_params_dict['lag_val'] = 0
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='lag_val must be positive'):
params.validate_business_rules()
def test_validate_business_rules_window_zero(valid_params_dict):
"""Test that window = 0 raises ValueError."""
valid_params_dict['window'] = 0
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='window must be positive'):
params.validate_business_rules()
def test_validate_business_rules_low_lim_upp_lim_keys_mismatch(valid_params_dict):
"""Test that mismatched keys in low_lim and upp_lim raises ValueError."""
valid_params_dict['low_lim'] = {'var1': 0.0, 'var2': 0.0}
valid_params_dict['upp_lim'] = {'var1': 100.0, 'var3': 100.0} # var3 instead of var2
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='low_lim and upp_lim must have the same keys'):
params.validate_business_rules()
def test_validate_business_rules_low_lim_greater_than_upp_lim(valid_params_dict):
"""Test that low_lim >= upp_lim raises ValueError."""
valid_params_dict['low_lim'] = {'var1': 100.0, 'var2': 0.0, 'var3': 0.0}
valid_params_dict['upp_lim'] = {'var1': 50.0, 'var2': 100.0, 'var3': 100.0}
valid_params_dict['target_variable'] = 'var2'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='low_lim must be less than upp_lim for variable "var1"'):
params.validate_business_rules()
def test_validate_business_rules_low_lim_equal_to_upp_lim(valid_params_dict):
"""Test that low_lim == upp_lim raises ValueError."""
valid_params_dict['low_lim'] = {'var1': 50.0, 'var2': 0.0, 'var3': 0.0}
valid_params_dict['upp_lim'] = {'var1': 50.0, 'var2': 100.0, 'var3': 100.0}
valid_params_dict['target_variable'] = 'var2'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='low_lim must be less than upp_lim for variable "var1"'):
params.validate_business_rules()
def test_validate_business_rules_target_not_in_variable_columns(valid_params_dict):
"""Test that target_variable not in variable_columns raises ValueError."""
valid_params_dict['target_variable'] = 'nonexistent_var'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(
ValueError, match='target_variable "nonexistent_var" must be in variable_columns'
):
params.validate_business_rules()
def test_validate_business_rules_empty_bucket_name(valid_params_dict):
"""Test that empty bucket_name raises ValueError."""
valid_params_dict['bucket_name'] = ' '
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='bucket_name cannot be empty or whitespace'):
params.validate_business_rules()
def test_validate_business_rules_empty_file_name(valid_params_dict):
"""Test that empty file_name raises ValueError."""
valid_params_dict['file_name'] = ''
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='file_name cannot be empty or whitespace'):
params.validate_business_rules()
def test_validate_business_rules_empty_experiment_name(valid_params_dict):
"""Test that empty experiment_name raises ValueError."""
valid_params_dict['experiment_name'] = ' '
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
with pytest.raises(ValueError, match='experiment_name cannot be empty or whitespace'):
params.validate_business_rules()
def test_validate_business_rules_all_valid_edge_cases(valid_params_dict):
"""Test that edge case valid values pass validation."""
valid_params_dict['train_size'] = 1 # Minimum valid
valid_params_dict['lag_train'] = 1 # Minimum valid
valid_params_dict['lag_val'] = 1 # Minimum valid
valid_params_dict['window'] = 1 # Minimum valid
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
# Should not raise any exception
params.validate_business_rules()
def test_validate_business_rules_train_size_99(valid_params_dict):
"""Test that train_size = 99 (maximum valid) passes validation."""
valid_params_dict['train_size'] = 99
valid_params_dict['target_variable'] = 'var1'
params = TrainModelParams.from_dict(valid_params_dict)
# Should not raise any exception
params.validate_business_rules()

660
tests/workflows/test_train_model.py Normal file
View File

@@ -0,0 +1,660 @@
"""Unit tests for TrainModel workflow."""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pytest import fixture, mark
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
from model_manager.workflows.train_model import TrainModel
@fixture
def train_model_workflow() -> TrainModel:
"""Fixture for TrainModel workflow instance."""
return TrainModel()
@fixture
def mock_train_params():
"""Fixture for mock TrainModelParams."""
return TrainModelParams(
experiment_run_id=123,
target_variable='price',
variable_columns=['feature1', 'feature2'],
train_size=80,
shuffle=True,
use_scaler=True,
include_ar=False,
bucket_name='test-bucket',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
lag_train=1,
lag_val=1,
rem_static_win=False,
low_lim={'feature1': 0.0, 'feature2': 0.0},
upp_lim={'feature1': 100.0, 'feature2': 100.0},
window=10,
experiment_name='test_experiment',
removed_intervals=[],
)
@fixture
def mock_train_result(mock_train_params):
"""Fixture for mock TrainModelResult."""
result = MagicMock(spec=TrainModelResult)
result.params = mock_train_params
result.run_name = 'test_experiment-1'
result.run_dir = 'test_run_dir' # Relative path instead of /tmp
return result
# ============================================================================
# Tests for run() - Complete workflow
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_run_success_complete_flow(
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_params,
mock_train_result,
):
"""Test successful complete workflow execution."""
input_data = {
'experiment_run_id': 123,
'target_variable': 'price',
'variable_columns': ['feature1', 'feature2'],
'train_size': 80,
'shuffle': True,
'use_scaler': True,
'include_ar': False,
'bucket_name': 'test-bucket',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0, 'feature2': 0.0},
'upp_lim': {'feature1': 100.0, 'feature2': 100.0},
'window': 10,
'experiment_name': 'test_experiment',
'removed_intervals': [],
}
# Mock activity responses
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update_experiment_run (MAGE_WAITING_PROC)
b'file_content', # fetch_file_from_minio
mock_train_result, # train_model
None, # update_experiment_run (TRAINING_SUCCESS)
mock_train_result, # save_model
None, # update_experiment_run (MLFLOW_SENT with run_name)
None, # cleanup_run_directory
None, # delete_file_from_minio
None, # update_experiment_run (FILE_DELETED)
]
)
# Execute workflow
await train_model_workflow.run(input_data)
# Verify all activity calls (now 10 instead of 9 due to cleanup_run_directory)
assert workflow_mock.execute_activity_method.call_count == 10
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_run_missing_experiment_run_id(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test workflow fails when experiment_run_id is missing."""
input_data = {
'target_variable': 'price',
'variable_columns': ['feature1'],
}
# Mock workflow.logger to avoid NotInWorkflowEventLoopError
workflow_mock.logger = MagicMock()
with pytest.raises(ValueError, match='experiment_run_id is required'):
await train_model_workflow.run(input_data)
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_run_invalid_experiment_run_id_type(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test workflow fails when experiment_run_id has invalid type."""
input_data = {
'experiment_run_id': 'invalid', # Should be int
'target_variable': 'price',
}
# Mock workflow.logger to avoid NotInWorkflowEventLoopError
workflow_mock.logger = MagicMock()
with pytest.raises(ValueError, match='experiment_run_id must be an integer'):
await train_model_workflow.run(input_data)
# ============================================================================
# Tests for _validate_experiment_run_id()
# ============================================================================
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_success(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test successful experiment_run_id validation."""
workflow_mock.logger = MagicMock()
input_data = {'experiment_run_id': 456}
result = train_model_workflow._validate_experiment_run_id(input_data)
assert result == 456
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_missing(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test validation fails when experiment_run_id is missing."""
workflow_mock.logger = MagicMock()
input_data = {}
with pytest.raises(ValueError, match='experiment_run_id is required'):
train_model_workflow._validate_experiment_run_id(input_data)
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_none(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test validation fails when experiment_run_id is None."""
workflow_mock.logger = MagicMock()
input_data = {'experiment_run_id': None}
with pytest.raises(ValueError, match='experiment_run_id is required'):
train_model_workflow._validate_experiment_run_id(input_data)
@patch('model_manager.workflows.train_model.workflow')
def test_validate_experiment_run_id_invalid_type(
workflow_mock: MagicMock, train_model_workflow: TrainModel
):
"""Test validation fails when experiment_run_id is not an integer."""
workflow_mock.logger = MagicMock()
input_data = {'experiment_run_id': 'not_an_int'}
with pytest.raises(ValueError, match='experiment_run_id must be an integer'):
train_model_workflow._validate_experiment_run_id(input_data)
# ============================================================================
# Tests for _validate_training_parameters()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_validate_training_parameters_success(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test successful parameter validation."""
input_data = {
'experiment_run_id': 123,
'target_variable': 'price',
'variable_columns': ['feature1'],
'train_size': 80,
'shuffle': True,
'use_scaler': False,
'include_ar': False,
'bucket_name': 'test',
'file_name': 'test.csv',
'line_separator': '\n',
'decimal_separator': '.',
'lag_train': 1,
'lag_val': 1,
'rem_static_win': False,
'low_lim': {'feature1': 0.0},
'upp_lim': {'feature1': 100.0},
'window': 10,
'experiment_name': 'test',
'removed_intervals': [],
}
metadata = {
'metadata': {
'experiment_run_id': 123,
'workflow_name': 'train_model',
}
}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update_experiment_run
]
)
result = await train_model_workflow._validate_training_parameters(input_data, 123, metadata)
assert result == mock_train_params
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_validate_training_parameters_validation_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test parameter validation handles errors correctly."""
input_data = {'experiment_run_id': 123}
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
ValueError('Missing required field'), # validate_train_params fails
None, # update_experiment_run with error
]
)
with pytest.raises(ValueError, match='Missing required field'):
await train_model_workflow._validate_training_parameters(input_data, 123, metadata)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
# ============================================================================
# Tests for _download_and_train_model()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_success(
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_params,
mock_train_result,
):
"""Test successful download and training."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
b'file_content', # fetch_file_from_minio
mock_train_result, # train_model
None, # update_experiment_run
]
)
result = await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
assert result == mock_train_result
assert workflow_mock.execute_activity_method.call_count == 3
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_download_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test download error is handled correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
Exception('MinIO connection failed'), # fetch_file_from_minio fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='MinIO connection failed'):
await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_training_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test training error is handled correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
b'file_content', # fetch_file_from_minio succeeds
Exception('Training failed'), # train_model fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='Training failed'):
await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 3
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_closes_bytesio_on_success(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result
):
"""Test that BytesIO is closed in finally block on success."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Create a mock BytesIO with close method
mock_file = MagicMock()
mock_file.close = MagicMock()
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_file, # fetch_file_from_minio returns BytesIO
mock_train_result, # train_model succeeds
None, # update_experiment_run
]
)
result = await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify BytesIO.close() was called in finally block
mock_file.close.assert_called_once()
assert result == mock_train_result
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_closes_bytesio_on_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
):
"""Test that BytesIO is closed in finally block even on error."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Create a mock BytesIO with close method
mock_file = MagicMock()
mock_file.close = MagicMock()
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_file, # fetch_file_from_minio returns BytesIO
Exception('Training failed'), # train_model fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='Training failed'):
await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
# Verify BytesIO.close() was called in finally block even after exception
mock_file.close.assert_called_once()
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_download_and_train_model_handles_file_without_close(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result
):
"""Test that workflow handles file objects without close method gracefully."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Create a mock file without close method
mock_file = MagicMock(spec=[])
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_file, # fetch_file_from_minio returns object without close
mock_train_result, # train_model succeeds
None, # update_experiment_run
]
)
# Should not raise error even if file doesn't have close method
result = await train_model_workflow._download_and_train_model(
train_params=mock_train_params, experiment_run_id=123, metadata=metadata
)
assert result == mock_train_result
# ============================================================================
# Tests for _save_model_to_mlflow()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_save_model_to_mlflow_success(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result
):
"""Test successful model saving to MLFlow."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
mock_train_result, # save_model
None, # update_experiment_run with MODEL_SAVED
]
)
result = await train_model_workflow._save_model_to_mlflow(
train_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
assert result == mock_train_result
assert workflow_mock.execute_activity_method.call_count == 2
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_save_model_to_mlflow_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result
):
"""Test MLFlow save error is handled correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
Exception('MLFlow connection failed'), # save_model fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='MLFlow connection failed'):
await train_model_workflow._save_model_to_mlflow(
train_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated
assert workflow_mock.execute_activity_method.call_count == 2
# ============================================================================
# Tests for _cleanup_resources()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_cleanup_resources_success(
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_result,
):
"""Test successful resource cleanup."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
None, # cleanup_run_directory
None, # delete_file_from_minio
None, # update_experiment_run
]
)
await train_model_workflow._cleanup_resources(
saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify activities were called (cleanup_run_directory + delete_file_from_minio + update_experiment_run)
assert workflow_mock.execute_activity_method.call_count == 3
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_cleanup_resources_delete_error(
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_result,
):
"""Test cleanup handles delete errors correctly."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
None, # cleanup_run_directory succeeds
Exception('MinIO delete failed'), # delete_file_from_minio fails
None, # update_experiment_run with error
]
)
with pytest.raises(Exception, match='MinIO delete failed'):
await train_model_workflow._cleanup_resources(
saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify error status was updated (cleanup_run_directory + delete_file_from_minio + update_experiment_run)
assert workflow_mock.execute_activity_method.call_count == 3
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_cleanup_resources_without_run_dir(
workflow_mock: AsyncMock,
train_model_workflow: TrainModel,
mock_train_result,
):
"""Test cleanup works when run_dir is not set."""
metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
# Mock result without run_dir
mock_train_result.run_dir = None
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
None, # delete_file_from_minio
None, # update_experiment_run
]
)
await train_model_workflow._cleanup_resources(
saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
)
# Verify cleanup_run_directory was NOT called (no run_dir)
# Only delete_file_from_minio + update_experiment_run
assert workflow_mock.execute_activity_method.call_count == 2
# ============================================================================
# Tests for _update_experiment_run()
# ============================================================================
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_update_experiment_run_status_only(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test updating experiment run with status only."""
from model_manager.activities.experiment_tracking import UpdateType
metadata = {'metadata': {'experiment_run_id': 123}}
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
await train_model_workflow._update_experiment_run(
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.STATUS,
status=ExperimentStatus.TRAINING_SUCCESS,
)
workflow_mock.execute_activity_method.assert_called_once()
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_update_experiment_run_with_error(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test updating experiment run with error message."""
from model_manager.activities.experiment_tracking import UpdateType
metadata = {'metadata': {'experiment_run_id': 123}}
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
await train_model_workflow._update_experiment_run(
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.TRAINING_ERROR,
error_message='Training failed',
)
workflow_mock.execute_activity_method.assert_called_once()
call_args = workflow_mock.execute_activity_method.call_args[0][1]
assert call_args['error_message'] == 'Training failed'
@mark.asyncio
@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
async def test_update_experiment_run_with_run_name(
workflow_mock: AsyncMock, train_model_workflow: TrainModel
):
"""Test updating experiment run with run_name."""
from model_manager.activities.experiment_tracking import UpdateType
metadata = {'metadata': {'experiment_run_id': 123}}
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
await train_model_workflow._update_experiment_run(
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.MLFLOW_SENT,
run_name='test_experiment-1',
)
workflow_mock.execute_activity_method.assert_called_once()
call_args = workflow_mock.execute_activity_method.call_args[0][1]
assert call_args['run_name'] == 'test_experiment-1'