SIENTIAPDE-1255: Refactor Model Manager to focus on ML model training pipeline and simplify architecture. This includes removing prediction workflows, updating core functionality, and improving parameter validation and resource management.
README.md | 411 deletions(-) 117 insertions(+) 1 file changed, 117 insertions(+), 411 deletions(-)
This commit is contained in:
Bruno Domingues
368
README.md
368
README.md
@@ -1,6 +1,6 @@
|
||||
# Sientia DataOps Model Manager
|
||||
|
||||
A comprehensive AI model management platform for the complete machine learning lifecycle. Handles model training, versioning, deployment, monitoring, and governance. Streamlines MLOps workflows with centralized model registry, automated pipelines, performance tracking, and enterprise-grade compliance features.
|
||||
An enterprise-grade ML model training orchestration platform built on Temporal. Provides robust, scalable workflows for training machine learning models with comprehensive validation, experiment tracking, and automated resource management. Integrates seamlessly with MLFlow for model persistence and PostgreSQL for experiment tracking.
|
||||
|
||||
## 📑 Table of Contents
|
||||
|
||||
@@ -14,11 +14,7 @@ A comprehensive AI model management platform for the complete machine learning l
|
||||
- [Data Flow Architecture](#data-flow-architecture)
|
||||
- [Security Architecture](#security-architecture)
|
||||
- [Workflows](#workflows)
|
||||
- [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)
|
||||
- [Train Model Workflow](#4-train-model-workflow-train_modelpy)
|
||||
- [Minimal Retrain Workflow](#5-minimal-retrain-workflow-minimal_retrainpy)
|
||||
- [Train Model Workflow](#train-model-workflow-train_modelpy)
|
||||
- [Installation & Setup](#installation--setup)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Environment Setup](#environment-setup)
|
||||
@@ -68,19 +64,19 @@ A comprehensive AI model management platform for the complete machine learning l
|
||||
## Features
|
||||
|
||||
### Core Functionality
|
||||
- **Batch Prediction Processing**: High-throughput ML model inference using MLFlow models
|
||||
- **Temporal Workflow Orchestration**: Robust workflow management with automatic retry policies and fault tolerance
|
||||
- **Data Quality Gates**: Configurable filtering for data validation, MLFlow API responses, and custom validation rules
|
||||
- **Multi-Model Support**: Flexible ML model management with retention policies and versioning
|
||||
- **Real-time Data Export**: PostgreSQL persistence for data storage
|
||||
- **ML Model Training Pipeline**: Complete training workflow from validation to deployment using MLFlow
|
||||
- **Temporal Workflow Orchestration**: Robust workflow management with granular retry policies and fault tolerance
|
||||
- **Parameter Validation**: Defense-in-depth validation with business rules and type checking
|
||||
- **Experiment Tracking**: Comprehensive status tracking in PostgreSQL database
|
||||
- **Resource Management**: Automatic cleanup of temporary files and storage
|
||||
- **Comprehensive Monitoring**: Prometheus metrics and detailed logging for operational visibility
|
||||
|
||||
### Advanced Capabilities
|
||||
- **Incremental Data Processing**: Timestamp-based data loading to avoid reprocessing
|
||||
- **Configurable Data Retention**: Model retention policies with automatic cleanup
|
||||
- **Granular Retry Policies**: Different strategies for network, training, MLFlow, database, and filesystem operations
|
||||
- **Configurable Timeouts**: Environment variable-based timeouts supporting large training files (up to 200MB)
|
||||
- **Notification System**: Integrated alerting and notification management via MongoDB
|
||||
- **Scalable Architecture**: Kubernetes-ready deployment with horizontal scaling support
|
||||
- **Model Retraining**: Automated model retraining workflows with production model updates
|
||||
- **MLFlow Integration**: Seamless model and artifact persistence to MLFlow tracking server
|
||||
|
||||
### Development & Quality Assurance
|
||||
- **Code Quality Tools**: Ruff (linting/formatting), mypy (type checking), Bandit (security analysis)
|
||||
@@ -136,18 +132,16 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
|
||||
- Health check endpoints for Kubernetes liveness/readiness probes
|
||||
- Graceful shutdown with cleanup procedures
|
||||
- Multi-instance deployment support
|
||||
- Two dedicated task queues: `predictions_batch-queue` and `minimal_retrain-queue`
|
||||
- Dedicated task queue: `train_model-queue` for ML model training workflows
|
||||
|
||||
#### **Workflows (`model_manager/workflows/`)**
|
||||
- **PredictionsBatch**: Main entry point for batch prediction pipelines
|
||||
- **PredictionProcess**: Core prediction pipeline with MLFlow integration
|
||||
- **FormatAndExportPrediction**: Data formatting and export operations
|
||||
- **MinimalRetrain**: Automated model retraining and deployment
|
||||
- **TrainModel**: Complete ML model training pipeline from validation to deployment
|
||||
- **Key Features**:
|
||||
- Temporal workflow definitions with retry policies
|
||||
- Child workflow orchestration and delegation
|
||||
- Comprehensive error handling and recovery
|
||||
- Configurable timeout and retry strategies
|
||||
- Temporal workflow definitions with granular retry policies
|
||||
- Parameter validation with business rules
|
||||
- Comprehensive error handling and status tracking
|
||||
- Configurable timeouts for different operation types
|
||||
- Automatic resource cleanup and management
|
||||
|
||||
#### **Activities (`model_manager/activities/`)**
|
||||
- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance
|
||||
@@ -162,13 +156,12 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
|
||||
- Returns success/failure status with TrainModelResult or error message
|
||||
- No exception raising on failure - allows workflow to handle errors gracefully
|
||||
- Integration with TrainingRepository for business logic separation
|
||||
- **Gates**: Data quality validation and filtering mechanisms
|
||||
- **MLFlow**: Model transformation and prediction operations
|
||||
- **MinIO**: Object storage operations for file management
|
||||
- **MLFlow**: Model saving and artifact management operations
|
||||
- **MinIO**: Object storage operations for training data management
|
||||
- **Key Features**:
|
||||
- Multiple inheritance pattern for unified activity interface
|
||||
- Configurable filter policies and validation rules
|
||||
- MLFlow model serving integration with configurable flavors
|
||||
- Parameter validation with business rules
|
||||
- MLFlow integration for model persistence
|
||||
- Comprehensive error handling and notification integration
|
||||
- Experiment tracking with automatic status management
|
||||
|
||||
@@ -188,17 +181,18 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
|
||||
|
||||
### Data Flow Architecture
|
||||
|
||||
#### **1. Batch Prediction Pipeline**
|
||||
#### **Model Training Pipeline**
|
||||
```
|
||||
Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform →
|
||||
MLFlow Prediction → Response Validation → Export (PostgreSQL)
|
||||
Training Request → Parameter Validation → MinIO Data Download →
|
||||
Model Training → MLFlow Model Save → Resource Cleanup → Status Update
|
||||
```
|
||||
|
||||
#### **2. Model Retraining Pipeline**
|
||||
```
|
||||
Training Data → Model Retraining → Quality Validation →
|
||||
Production Update → Notification & Monitoring
|
||||
```
|
||||
**Key Stages:**
|
||||
1. **Validation**: Experiment run ID and training parameters validation
|
||||
2. **Data Acquisition**: Download training data from MinIO storage
|
||||
3. **Training**: Execute ML model training with validated parameters
|
||||
4. **Persistence**: Save trained model and artifacts to MLFlow
|
||||
5. **Cleanup**: Remove temporary files and update experiment status
|
||||
|
||||
### Security Architecture
|
||||
|
||||
@@ -222,166 +216,7 @@ Production Update → Notification & Monitoring
|
||||
|
||||
## Workflows
|
||||
|
||||
### 1. Predictions Batch Workflow (`predictions_batch.py`)
|
||||
|
||||
The **PredictionsBatch** workflow is the main entry point for batch prediction pipelines. It orchestrates the complete prediction process and implements a robust data loading and processing pattern.
|
||||
|
||||
#### Purpose
|
||||
- **Batch Prediction Orchestration**: Coordinates data loading and prediction processing
|
||||
- **Data Preparation**: Loads data using custom SQL queries with configurable schemas
|
||||
- **Workflow Delegation**: Delegates actual prediction processing to the PredictionProcess workflow
|
||||
- **Configuration Management**: Handles model configuration, filters, and retention policies
|
||||
|
||||
#### Execution Flow
|
||||
1. **Data Loading**: Executes custom SQL query to load data from PostgreSQL
|
||||
2. **Input Preparation**: Prepares prediction input with metadata and configuration
|
||||
3. **Workflow Delegation**: Spawns PredictionProcess child workflow for actual processing
|
||||
4. **Error Handling**: Implements comprehensive error handling with retry policies
|
||||
|
||||
#### Key Features
|
||||
- **Custom Query Support**: Flexible SQL-based data loading
|
||||
- **Schema Configuration**: Configurable data schema definitions
|
||||
- **Automatic Retry**: Implements Temporal retry policies for fault tolerance
|
||||
- **Timeout Management**: 60-second timeout for all activities
|
||||
- **Comprehensive Error Handling**: Detailed error reporting and notification integration
|
||||
|
||||
#### Input Parameters
|
||||
```json
|
||||
{
|
||||
"schedule_name": "hourly_predictions",
|
||||
"model_name": "temperature_prediction_model",
|
||||
"model_id": "temp_pred_001",
|
||||
"query": "SELECT * FROM sensor_data WHERE timestamp > NOW() - INTERVAL '1 hour'",
|
||||
"schema": {
|
||||
"timestamp": "datetime",
|
||||
"temperature": "float",
|
||||
"humidity": "float"
|
||||
},
|
||||
"table_name": "predictions",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {"POLICY": "STOP"}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {"POLICY": "STOP"}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {"POLICY": "STOP"}
|
||||
},
|
||||
"model_retention": 60,
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"]
|
||||
}
|
||||
```
|
||||
|
||||
#### Architecture Diagram
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. load_custom_query] --> B[2. prediction_process 🔃]
|
||||
|
||||
A -.-> Database[(Database)]
|
||||
```
|
||||
|
||||
### 2. Prediction Process Workflow (`prediction_process.py`)
|
||||
|
||||
The **PredictionProcess** workflow implements the core prediction pipeline for ML model inference. It handles data quality validation, MLFlow model interactions, and prediction processing.
|
||||
|
||||
#### Purpose
|
||||
- **Data Quality Validation**: Applies configurable filters for data integrity
|
||||
- **MLFlow Integration**: Manages model transformation and prediction requests
|
||||
- **Response Validation**: Filters MLFlow API responses for quality assurance
|
||||
- **Prediction Export**: Delegates prediction formatting and export operations
|
||||
|
||||
#### Execution Flow
|
||||
1. **Timestamp Retrieval**: Gets the last processed timestamp for incremental processing
|
||||
2. **Input Data Gate**: Applies configured filters for data quality validation
|
||||
3. **Path Decision**: Determines processing path based on filter results
|
||||
4. **MLFlow Transform**: Requests data transformation using MLFlow models
|
||||
5. **Response Validation**: Filters transform responses for quality assurance
|
||||
6. **MLFlow Prediction**: Executes prediction using transformed data
|
||||
7. **Content Validation**: Filters prediction responses for final quality check
|
||||
8. **Export Delegation**: Delegates to FormatAndExportPrediction workflow
|
||||
|
||||
#### Key Features
|
||||
- **Configurable Quality Gates**: Multiple filter types with policy-based configuration
|
||||
- **Flexible Path Handling**: Configurable decision paths (STOP, CONTINUE, REPEAT)
|
||||
- **MLFlow Integration**: Comprehensive model management and inference
|
||||
- **Incremental Processing**: Timestamp-based data processing optimization
|
||||
- **Comprehensive Monitoring**: Detailed metrics and error reporting
|
||||
|
||||
#### Input Parameters
|
||||
```json
|
||||
{
|
||||
"metadata": {
|
||||
"schedule_name": "hourly_predictions",
|
||||
"model_name": "temperature_prediction_model",
|
||||
"model_id": "temp_pred_001",
|
||||
"workflow_name": "predictions_batch"
|
||||
},
|
||||
"data": {...},
|
||||
"schema": {...},
|
||||
"table_name": "predictions",
|
||||
"model_id": "temp_pred_001",
|
||||
"model_name": "temperature_prediction_model",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {"POLICY": "STOP"},
|
||||
"SPECIFIC_VARIABLES_NULL_VALUES": {
|
||||
"POLICY": "STOP",
|
||||
"config": {"variables": ["temperature", "humidity"]}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {"POLICY": "STOP"}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {"POLICY": "STOP"},
|
||||
"NAN_VALUES": {"POLICY": "STOP"}
|
||||
},
|
||||
"model_retention": 60,
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"]
|
||||
}
|
||||
```
|
||||
|
||||
#### Architecture Diagram
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. get_last_timestamp] --> B[2. input_gate] --> C[3. request_transform] --> D[4. mlflow_response_gate] --> E[5. mlflow_content_gate] --> F[6. request_predict] --> G[7. mlflow_response_gate] --> H[8. format_and_export_prediction🔃]
|
||||
|
||||
C -.-> MLFlow[MLFlow]
|
||||
F -.-> MLFlow[MLFlow]
|
||||
G -.-> Filters[MLFlow Filters]
|
||||
```
|
||||
|
||||
### 3. Format and Export Prediction Workflow (`format_and_export_prediction.py`)
|
||||
|
||||
The **FormatAndExportPrediction** workflow handles prediction data formatting and export operations to multiple destinations.
|
||||
|
||||
#### Purpose
|
||||
- **Data Formatting**: Formats prediction data for database storage
|
||||
- **PostgreSQL Export**: Persists predictions to database with metrics
|
||||
- **Metrics Recording**: Tracks export operations and performance metrics
|
||||
|
||||
#### Execution Flow
|
||||
1. **Path Decision**: Determines formatting path based on configuration
|
||||
2. **Data Formatting**: Formats prediction data for specific output requirements
|
||||
3. **PostgreSQL Export**: Writes formatted predictions to database
|
||||
4. **Metrics Recording**: Records export performance and success metrics
|
||||
|
||||
#### Key Features
|
||||
- **Flexible Formatting**: Configurable output formats for different destinations
|
||||
- **Database Export**: PostgreSQL integration for data persistence
|
||||
- **Performance Monitoring**: Comprehensive metrics for export operations
|
||||
- **Error Handling**: Robust error handling with notification integration
|
||||
|
||||
#### Architecture Diagram
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. format_prediction/format_default_prediction] --> B[2. export_data_to_postgres] --> C[3. write_metrics]
|
||||
|
||||
A -.-> Format[Data Formatting]
|
||||
B -.-> PostgreSQL[(PostgreSQL)]
|
||||
C -.-> Prometheus[Prometheus]
|
||||
```
|
||||
|
||||
### 4. Train Model Workflow (`train_model.py`)
|
||||
### Train Model Workflow (`train_model.py`)
|
||||
|
||||
The **TrainModel** workflow orchestrates the complete ML model training pipeline from parameter validation through model saving and cleanup.
|
||||
|
||||
@@ -473,34 +308,6 @@ The workflow validates 10 business rules beyond type checking:
|
||||
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.
|
||||
|
||||
#### Purpose
|
||||
- **Model Retraining**: Automates ML model retraining processes
|
||||
- **Production Updates**: Manages production model version updates
|
||||
- **Data Export**: Exports training data for model development
|
||||
- **Quality Assurance**: Ensures model quality before production deployment
|
||||
|
||||
#### Execution Flow
|
||||
1. **Data Loading**: Loads training data using custom queries
|
||||
2. **Model Retraining**: Executes model retraining process
|
||||
3. **Quality Validation**: Validates retrained model performance
|
||||
4. **Production Update**: Updates production model if quality criteria met
|
||||
5. **Data Export**: Exports training data for analysis
|
||||
|
||||
#### Architecture Diagram
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. load_custom_query] --> B[2. retrain_model] --> C[3. update_production_model] --> D[4. export_data_to_postgres]
|
||||
|
||||
A -.-> Database[(Database)]
|
||||
B -.-> MLFlow[MLFlow]
|
||||
C -.-> MLFlow[MLFlow]
|
||||
D -.-> PostgreSQL[(PostgreSQL)]
|
||||
```
|
||||
|
||||
## Installation & Setup
|
||||
|
||||
### Prerequisites
|
||||
@@ -986,8 +793,8 @@ pip install pytest pytest-cov pytest-asyncio
|
||||
pytest --cov=model_manager --cov-report=html
|
||||
|
||||
# Run specific test modules
|
||||
pytest tests/activities/test_gates.py
|
||||
pytest tests/workflows/test_predictions_batch.py
|
||||
pytest tests/activities/test_training.py
|
||||
pytest tests/workflows/test_train_model.py
|
||||
```
|
||||
|
||||
## Monitoring and Metrics
|
||||
@@ -1067,94 +874,6 @@ These timeouts control how long each activity in the training workflow can run b
|
||||
|
||||
**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:
|
||||
|
||||
#### Predictions Batch Workflow configuration sample
|
||||
|
||||
This is the configuration for the Predictions Batch Workflow, to be inserted into the MongoDB pipeline collection.
|
||||
|
||||
```json
|
||||
{
|
||||
"schedule_name": "laborious-orchestrated-pipeline",
|
||||
"model_id": "1",
|
||||
"workflow_type": "predictions_batch",
|
||||
"frequency": "30s",
|
||||
"max_retry_policy": 1,
|
||||
"query": "select * from sientia_data.laborious_data where model_id = 1 and \"timestamp\" > NOW() - INTERVAL '5 minutes' order by \"timestamp\" desc limit 30;",
|
||||
"write_tags": [
|
||||
{
|
||||
"server_id": "server1",
|
||||
"type": "prediction",
|
||||
"addr": "ns=2;i=5",
|
||||
"data_type": "double"
|
||||
},
|
||||
{
|
||||
"server_id": "server1",
|
||||
"type": "confidence",
|
||||
"addr": "ns=2;i=6",
|
||||
"data_type": "double"
|
||||
}
|
||||
],
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {"POLICY": "STOP"},
|
||||
"SPECIFIC_VARIABLES_NULL_VALUES": {
|
||||
"POLICY": "CONTINUE",
|
||||
"config": {"variables": ["Counter"]}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {"POLICY": "REPEAT"},
|
||||
"NAN_VALUES": {"POLICY": "STOP"}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {"POLICY": "CONTINUE"}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"active": true,
|
||||
"updated_at": {
|
||||
"$date": "2025-09-16T10:00:00.000Z"
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"],
|
||||
"predictions_storage_policy": "lts:1"
|
||||
}
|
||||
```
|
||||
|
||||
This is the configuration created by the Orchestrator in Temporal.
|
||||
|
||||
```json
|
||||
{
|
||||
"datetime_columns":["timestamp","created_at"],
|
||||
"frequency":"15m",
|
||||
"input_filters":{"EMPTY_DATA":{"config":{},"policy":"STOP"}},
|
||||
"max_retry_policy":1,
|
||||
"mlflow_predict_filters":{"API_ERROR":{"config":{},"policy":"CONTINUE"}},
|
||||
"mlflow_transform_filters":{
|
||||
"API_ERROR":{"config":{},"policy":"CONTINUE"},
|
||||
"EMPTY_DATA":{"config":{},"policy":"STOP"}
|
||||
},
|
||||
"model_config":{
|
||||
"is_compressed":true,
|
||||
"predict_flavor":"pyfunc",
|
||||
"retention_minutes":60,
|
||||
"retention_target":"artifact",
|
||||
"transform_function_keyword":"transform"
|
||||
},
|
||||
"model_id":"352",
|
||||
"model_name":"courier",
|
||||
"path_priority":["STOP","CONTINUE","REPEAT"],
|
||||
"predictions_storage_policy":"lts:1",
|
||||
"query":"select * from sientia_data.laborious_data where model_id = 352 order by \"timestamp\" desc limit 300;",
|
||||
"retention_time":3600,
|
||||
"schedule_name":"laborious-courier",
|
||||
"schema":"sientia_data",
|
||||
"table_name":"predictions",
|
||||
"updated_at":"2025-09-12 19:35:01.600000+0000",
|
||||
"workflow_type":"predictions_batch"
|
||||
}
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
### Code Quality & Testing
|
||||
@@ -1229,30 +948,27 @@ model_manager/
|
||||
├── activities/ # Temporal activity implementations
|
||||
│ ├── __init__.py
|
||||
│ ├── activities.py # Main activities orchestrator (combines all activities)
|
||||
│ ├── gates.py # Data quality gates and filtering logic
|
||||
│ ├── experiment_tracking.py # Experiment status tracking and database operations
|
||||
│ ├── training.py # ML model training operations
|
||||
│ ├── minio.py # MinIO object storage operations
|
||||
│ └── mlflow.py # MLFlow model operations (predict/transform)
|
||||
│ └── mlflow.py # MLFlow model saving and artifact management
|
||||
├── workflows/ # Temporal workflow definitions
|
||||
│ ├── __init__.py
|
||||
│ ├── predictions_batch.py # Main batch prediction workflow entry point
|
||||
│ ├── minimal_retrain.py # Model retraining workflow
|
||||
│ └── sub_workflows/ # Sub-workflow implementations
|
||||
│ ├── __init__.py
|
||||
│ ├── prediction_process.py # Core prediction pipeline
|
||||
│ └── format_and_export_prediction.py # Data export workflow
|
||||
│ └── train_model.py # Complete ML model training workflow
|
||||
├── worker/ # Worker implementation
|
||||
│ ├── __init__.py
|
||||
│ └── worker.py # Main worker orchestrator (Temporal client setup)
|
||||
├── utils/ # Utility functions and helpers
|
||||
│ ├── __init__.py
|
||||
│ ├── connectors_config.py # Environment-based configuration builders
|
||||
│ ├── filters/ # Data quality validation filters
|
||||
│ ├── models/ # Data models and schemas
|
||||
│ │ ├── __init__.py
|
||||
│ │ ├── conditional_filters.py # Input data validation filters
|
||||
│ │ └── mlflow_filters.py # MLFlow response validation filters
|
||||
│ │ ├── train_model_params.py # Training parameters model
|
||||
│ │ ├── train_model_result.py # Training result model
|
||||
│ │ └── experiment_status.py # Experiment status enum
|
||||
│ └── repository/ # Data access layer
|
||||
│ ├── __init__.py
|
||||
│ └── model_repository.py # MLFlow model operations and retraining
|
||||
│ └── training_repository.py # Training business logic
|
||||
├── metrics.py # Prometheus metrics definitions
|
||||
└── __init__.py
|
||||
```
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='minimal_retrain')
|
||||
class MinimalRetrain:
|
||||
"""
|
||||
Automated model retraining workflow for the Model Manager system.
|
||||
|
||||
This workflow implements a complete model retraining pipeline that loads
|
||||
training data, executes model retraining, updates production models,
|
||||
and maintains comprehensive audit trails. It's designed for automated
|
||||
model lifecycle management with minimal manual intervention.
|
||||
|
||||
The workflow provides a robust retraining process with:
|
||||
- Automated data loading from configured data sources
|
||||
- MLFlow model retraining with quality validation
|
||||
- Production model updates with version control
|
||||
- Comprehensive reporting and audit trail maintenance
|
||||
- Error handling and notification integration
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the automated model retraining workflow.
|
||||
|
||||
This method orchestrates the complete model retraining process by:
|
||||
1. Loading training data using the provided custom SQL query
|
||||
2. Executing MLFlow model retraining with the loaded data
|
||||
3. Updating production models with newly trained versions
|
||||
4. Persisting comprehensive retraining reports to database
|
||||
|
||||
The method implements comprehensive error handling and ensures all
|
||||
required parameters are properly configured before proceeding.
|
||||
|
||||
Args:
|
||||
input_data: Complete configuration for the retraining workflow
|
||||
Required keys:
|
||||
- schedule_name (str): Schedule identifier for the retraining
|
||||
- model_name (str): Name of the ML model to retrain
|
||||
- model_id (int): Unique identifier for the model version
|
||||
- query (str): SQL query for training data loading
|
||||
- schema (str, optional): Database schema for report storage
|
||||
- table_name (str, optional): Target table for retraining reports
|
||||
- datetime_columns (list[str], optional): Columns to treat as datetime
|
||||
|
||||
Returns:
|
||||
None: The workflow completes successfully when all steps finish
|
||||
|
||||
Raises:
|
||||
Exception: If any required parameters are missing or if the workflow fails
|
||||
during data loading, retraining, or model update operations
|
||||
"""
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_id': input_data['model_id'],
|
||||
'workflow_name': 'minimal_retrain',
|
||||
}
|
||||
}
|
||||
|
||||
model_name = input_data['model_name']
|
||||
|
||||
data = await workflow.execute_local_activity_method(
|
||||
Activities.load_custom_query,
|
||||
{
|
||||
**metadata,
|
||||
'query': input_data['query'],
|
||||
'datetime_columns': input_data.get('datetime_columns', []),
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
experiment_response = await workflow.execute_activity_method(
|
||||
Activities.retrain_model,
|
||||
{**metadata, 'data': data, 'model_name': model_name},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
report = await workflow.execute_activity_method(
|
||||
Activities.update_production_model,
|
||||
{
|
||||
**metadata,
|
||||
'model_name': model_name,
|
||||
'model_id': input_data['model_id'],
|
||||
**experiment_response,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
await workflow.execute_activity_method(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**metadata,
|
||||
'data': report,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
@@ -1,116 +0,0 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='predictions_batch')
|
||||
class PredictionsBatch:
|
||||
"""
|
||||
Main batch prediction workflow for the Model Manager system.
|
||||
|
||||
This workflow orchestrates the complete batch prediction process, handling
|
||||
data loading, configuration management, and workflow delegation. It serves
|
||||
as the primary entry point for batch prediction operations and ensures
|
||||
proper data preparation before ML model inference.
|
||||
|
||||
The workflow implements a robust data processing pipeline with:
|
||||
- Custom SQL query execution for data loading
|
||||
- Comprehensive configuration management
|
||||
- Data quality filter application
|
||||
- MLFlow model integration
|
||||
- Workflow delegation to specialized sub-workflows
|
||||
|
||||
Workflow Execution:
|
||||
1. Data Loading: Executes custom SQL query to load prediction data
|
||||
2. Configuration Preparation: Sets up prediction parameters and filters
|
||||
3. Workflow Delegation: Spawns PredictionProcess child workflow
|
||||
4. Error Handling: Implements comprehensive error handling and retry policies
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the batch prediction workflow.
|
||||
|
||||
This method orchestrates the complete batch prediction process by:
|
||||
1. Loading data using the provided custom SQL query
|
||||
2. Preparing prediction configuration and filters
|
||||
3. Delegating to the PredictionProcess workflow for ML operations
|
||||
|
||||
The method implements comprehensive error handling and ensures all
|
||||
required parameters are properly configured before proceeding.
|
||||
|
||||
Args:
|
||||
input_data: Complete configuration for the batch prediction
|
||||
Required keys:
|
||||
- schedule_name (str): Schedule identifier for the prediction
|
||||
- model_name (str): Name of the ML model to use
|
||||
- model_id (int): Unique identifier for the model
|
||||
- query (str): SQL query for data loading
|
||||
- schema (dict, optional): Data schema definition
|
||||
- table_name (str, optional): Target table for predictions
|
||||
- input_filters (dict, optional): Data quality filters
|
||||
- mlflow_transform_filters (dict, optional): MLFlow transform filters
|
||||
- mlflow_predict_filters (dict, optional): MLFlow prediction filters
|
||||
- model_retention (int, optional): Model retention period in minutes
|
||||
- path_priority (list[str]): Decision path priority configuration
|
||||
|
||||
- datetime_columns (list[str], optional): Columns to treat as datetime
|
||||
|
||||
Returns:
|
||||
None: The workflow completes successfully when the child workflow finishes
|
||||
|
||||
Raises:
|
||||
Exception: If any required parameters are missing or if the workflow fails
|
||||
during data loading or workflow delegation
|
||||
"""
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_id': input_data['model_id'],
|
||||
'workflow_name': 'predictions_batch',
|
||||
}
|
||||
}
|
||||
|
||||
# Load data using custom query
|
||||
data = await workflow.execute_local_activity_method(
|
||||
Activities.load_custom_query,
|
||||
{
|
||||
**metadata,
|
||||
'query': input_data['query'],
|
||||
'datetime_columns': input_data.get('datetime_columns', []),
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=300),
|
||||
)
|
||||
|
||||
# Prepare input for prediction_process workflow
|
||||
prediction_input = {
|
||||
'metadata': metadata,
|
||||
'data': data,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'model_id': input_data['model_id'],
|
||||
'model_name': input_data['model_name'],
|
||||
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
|
||||
'mlflow_transform_filters': input_data.get(
|
||||
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
||||
),
|
||||
'mlflow_predict_filters': input_data.get(
|
||||
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
||||
),
|
||||
'model_config': input_data.get('model_config', {}),
|
||||
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
|
||||
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
|
||||
}
|
||||
|
||||
# Execute prediction process workflow
|
||||
await workflow.execute_child_workflow('prediction_process', prediction_input)
|
||||
@@ -1,120 +0,0 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='format_and_export_prediction')
|
||||
class FormatAndExportPrediction:
|
||||
"""
|
||||
Data formatting and export workflow for prediction results.
|
||||
|
||||
This workflow handles the final stages of the prediction pipeline, including
|
||||
data formatting, database persistence, and metrics recording.
|
||||
It implements flexible formatting based on prediction quality and provides
|
||||
comprehensive export capabilities.
|
||||
|
||||
The workflow supports two main prediction paths:
|
||||
1. Normal Prediction: Formats and exports successful prediction results
|
||||
2. Default Prediction: Creates fallback predictions for error conditions
|
||||
|
||||
Export Destinations:
|
||||
- PostgreSQL Database: Persistent storage with timestamp conversion
|
||||
- Prometheus Metrics: Performance monitoring and operational visibility
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the prediction formatting and export workflow.
|
||||
|
||||
This method orchestrates the complete data export process by:
|
||||
1. Determining the appropriate formatting strategy based on path_flag
|
||||
2. Formatting prediction data according to quality and requirements
|
||||
3. Persisting data to PostgreSQL database with comprehensive metadata
|
||||
4. Recording performance metrics for operational monitoring
|
||||
|
||||
The method implements flexible formatting strategies:
|
||||
- Normal predictions: Full data formatting with confidence scores
|
||||
- Error predictions: Default formatting with error indicators
|
||||
|
||||
Args:
|
||||
input_data: Complete configuration for the export workflow
|
||||
Required keys:
|
||||
- path_flag (str | None): Decision path flag for formatting strategy
|
||||
- data (dict[str, Any]): Prediction data to format and export
|
||||
- prediction_confidence (float): Confidence score for the prediction
|
||||
- timestamp (str): ISO-formatted timestamp for the prediction
|
||||
- model_id (int): Unique identifier for the ML model
|
||||
- model_name (str): Name of the ML model
|
||||
- model_retention (str): Model retention policy configuration
|
||||
- comment (str): Operational comment or error description
|
||||
- schema (str): Database schema for data storage
|
||||
- table_name (str): Target table for data persistence
|
||||
- prediction_store_policy (str, optional): Data retention policy
|
||||
|
||||
Returns:
|
||||
bool: True if the workflow completes successfully, False otherwise
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
path_flag = input_data['path_flag']
|
||||
data = input_data['data']
|
||||
prediction_confidence = input_data['prediction_confidence']
|
||||
|
||||
if path_flag is None:
|
||||
# proceed with formatting and exporting
|
||||
prediction = await workflow.execute_local_activity_method(
|
||||
Activities.format_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'data': data,
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': prediction_confidence,
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
else:
|
||||
# create default prediction
|
||||
prediction = await workflow.execute_local_activity_method(
|
||||
Activities.format_default_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': prediction_confidence,
|
||||
'comment': input_data['comment'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
# write to postgres
|
||||
await workflow.execute_activity_method(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**metadata,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'data': prediction,
|
||||
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
await workflow.execute_activity_method(
|
||||
Activities.write_metrics,
|
||||
{**metadata, 'prediction': prediction},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
@@ -1,295 +0,0 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='prediction_process')
|
||||
class PredictionProcess:
|
||||
"""
|
||||
Core prediction processing workflow for the Model Manager system.
|
||||
|
||||
This workflow implements the complete ML model inference pipeline, handling
|
||||
data quality validation, MLFlow model interactions, and prediction processing.
|
||||
It serves as the central orchestrator for all prediction operations and ensures
|
||||
data quality throughout the entire process.
|
||||
|
||||
The workflow implements a robust data processing pipeline with:
|
||||
- Data quality validation using configurable filters
|
||||
- MLFlow model transformation and prediction
|
||||
- Response validation and quality assurance
|
||||
- Flexible decision path handling
|
||||
- Comprehensive error handling and retry policies
|
||||
|
||||
Workflow Execution:
|
||||
1. Timestamp Retrieval: Gets last processed timestamp for incremental processing
|
||||
2. Input Data Gate: Applies data quality filters
|
||||
3. Path Decision: Determines processing path based on filter results
|
||||
4. MLFlow Transform: Requests data transformation using MLFlow models
|
||||
5. Response Validation: Filters transform responses for quality assurance
|
||||
6. MLFlow Prediction: Executes prediction using transformed data
|
||||
7. Content Validation: Filters prediction responses for final quality check
|
||||
8. Export Delegation: Delegates to FormatAndExportPrediction workflow
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the prediction process workflow.
|
||||
|
||||
This method orchestrates the complete prediction processing pipeline by:
|
||||
1. Retrieving the last processed timestamp for incremental processing
|
||||
2. Applying data quality filters to validate input data
|
||||
3. Executing MLFlow model transformation and prediction
|
||||
4. Validating all responses for quality assurance
|
||||
5. Delegating to export workflow for data persistence
|
||||
|
||||
The method implements comprehensive error handling and ensures all
|
||||
data quality requirements are met before proceeding with ML operations.
|
||||
|
||||
Args:
|
||||
input_data: Complete configuration for the prediction process
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict): Input data for prediction processing
|
||||
- schema (dict): Data schema definition
|
||||
- table_name (str): Target table for predictions
|
||||
- model_id (str): ML model identifier
|
||||
- model_name (str): ML model name
|
||||
- input_filters (dict): Data quality filters
|
||||
- mlflow_transform_filters (dict): MLFlow transform filters
|
||||
- mlflow_predict_filters (dict): MLFlow prediction filters
|
||||
- model_retention (int): Model retention period in minutes
|
||||
- path_priority (list[str]): Decision path priority configuration
|
||||
|
||||
|
||||
Returns:
|
||||
None: The workflow completes successfully when export workflow finishes
|
||||
|
||||
Raises:
|
||||
Exception: If any required parameters are missing or if the workflow fails
|
||||
during data processing, MLFlow operations, or workflow delegation
|
||||
|
||||
"""
|
||||
|
||||
metadata = input_data['metadata']
|
||||
data = input_data['data']
|
||||
model_id = input_data['model_id']
|
||||
model_name = input_data['model_name']
|
||||
model_config = input_data.get('model_config', {})
|
||||
|
||||
# Get last timestamp for incremental processing
|
||||
last_timestamp = await workflow.execute_local_activity_method(
|
||||
Activities.get_last_timestamp,
|
||||
{**metadata, 'data': data},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
|
||||
# Apply input data quality gates
|
||||
gate_input = {
|
||||
**metadata,
|
||||
'filters': input_data['input_filters'],
|
||||
'data': data,
|
||||
'path_priority': input_data['path_priority'],
|
||||
}
|
||||
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
Activities.input_gate,
|
||||
gate_input,
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
|
||||
# Handle path decision based on filter results
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
):
|
||||
return
|
||||
|
||||
# Request MLFlow model transformation
|
||||
response_data = await workflow.execute_local_activity_method(
|
||||
Activities.request_transform,
|
||||
{**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=5),
|
||||
)
|
||||
|
||||
# Validate MLFlow transform response
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': response_data,
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
|
||||
# Handle path decision based on transform validation
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
):
|
||||
return
|
||||
|
||||
transformed_data = response_data['content']
|
||||
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': transformed_data,
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
):
|
||||
return
|
||||
|
||||
response_data = await workflow.execute_local_activity_method(
|
||||
Activities.request_predict,
|
||||
{
|
||||
**metadata,
|
||||
'data': transformed_data,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=5),
|
||||
)
|
||||
|
||||
# Validate MLFlow prediction response
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_predict_filters'],
|
||||
'data': response_data,
|
||||
'type': 'predict',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
|
||||
# Handle path decision based on prediction validation
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
):
|
||||
return
|
||||
|
||||
# Delegate to export workflow for data persistence
|
||||
await workflow.execute_child_workflow(
|
||||
'format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': path_flag,
|
||||
'data': response_data['content'],
|
||||
'prediction_confidence': confidence,
|
||||
'timestamp': last_timestamp,
|
||||
'model_id': model_id,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'comment': comment,
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
},
|
||||
)
|
||||
|
||||
async def path_flag_handler(
|
||||
self,
|
||||
data: dict,
|
||||
path_flag: str,
|
||||
input_data: dict,
|
||||
confidence: int,
|
||||
last_timestamp: str,
|
||||
comment: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Handle path decisions based on filter results and confidence levels.
|
||||
|
||||
This method determines the appropriate action based on the path flag
|
||||
returned by data quality filters. It can stop processing, continue,
|
||||
or repeat operations based on the configured path priority.
|
||||
|
||||
Args:
|
||||
data: Input data for processing
|
||||
path_flag: Path decision from filter (STOP, CONTINUE, REPEAT)
|
||||
input_data: Complete workflow input configuration
|
||||
confidence: Confidence level from filter validation
|
||||
last_timestamp: Last processed timestamp
|
||||
comment: Additional information about the filter result
|
||||
|
||||
Returns:
|
||||
bool: True if processing should stop, False to continue
|
||||
|
||||
Path Handling:
|
||||
- STOP: Terminates workflow execution
|
||||
- CONTINUE: Proceeds with normal processing
|
||||
- REPEAT: Repeats last prediction if available
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
|
||||
schema = input_data['schema']
|
||||
table_name = input_data['table_name']
|
||||
model_id = input_data['model_id']
|
||||
model_name = input_data['model_name']
|
||||
model_config = input_data.get('model_config', {})
|
||||
|
||||
path_flag = path_flag.upper() if path_flag else ''
|
||||
|
||||
if path_flag == 'STOP':
|
||||
# Stop processing and exit workflow
|
||||
return True
|
||||
elif path_flag == 'REPEAT':
|
||||
# Repeat last prediction if available
|
||||
await workflow.execute_activity_method(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model': model_id,
|
||||
'last_timestamp': last_timestamp,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
return True
|
||||
elif path_flag == 'CONTINUE':
|
||||
# call write workflow
|
||||
await workflow.execute_child_workflow(
|
||||
'format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': path_flag,
|
||||
'data': data,
|
||||
'prediction_confidence': confidence,
|
||||
'timestamp': last_timestamp,
|
||||
'model_id': model_id,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'comment': comment,
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
},
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -1,147 +0,0 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.workflows.sub_workflows.format_and_export_prediction import (
|
||||
FormatAndExportPrediction,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def format_and_export_prediction():
|
||||
return FormatAndExportPrediction()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'path_flag': None,
|
||||
'data': {'test': 'data'},
|
||||
'timestamp': '2021-01-01',
|
||||
'model_id': 1,
|
||||
'prediction_confidence': 0,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'prediction_store_policy': 'erl:1',
|
||||
}
|
||||
|
||||
await format_and_export_prediction.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.format_prediction,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': input_data['prediction_confidence'],
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
**metadata,
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert workflow_mock.execute_activity_method.call_count == 2
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'path_flag': 'default',
|
||||
'data': {'test': 'data'},
|
||||
'timestamp': '2021-01-01',
|
||||
'model_id': 1,
|
||||
'prediction_confidence': 0,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'comment': 'test_comment',
|
||||
}
|
||||
|
||||
await format_and_export_prediction.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.format_default_prediction,
|
||||
{
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': input_data['prediction_confidence'],
|
||||
'comment': input_data['comment'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
**metadata,
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert workflow_mock.execute_activity_method.call_count == 2
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||
@@ -1,780 +0,0 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
|
||||
|
||||
@fixture
|
||||
def prediction_process():
|
||||
return PredictionProcess()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(return_value=False)
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'prediction_store_policy': 'lts:1',
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
'2024-01-01', # get_last_timestamp
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
|
||||
# mlflow_response_gate (predict)
|
||||
('continue', 0.95, 'Error'),
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 7
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.get_last_timestamp,
|
||||
{
|
||||
**metadata,
|
||||
'data': input_data['data'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
**metadata,
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': 'transformed_data',
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_predict,
|
||||
{
|
||||
**metadata,
|
||||
'data': 'transformed_data',
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_predict_filters'],
|
||||
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'predict',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': 'continue',
|
||||
'data': 'predicted_data',
|
||||
'prediction_confidence': 0.95,
|
||||
'timestamp': '2024-01-01',
|
||||
'model_id': 1,
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': input_data['model_config'],
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'comment': 'Error',
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(return_value=True)
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
'2024-01-01', # get_last_timestamp
|
||||
('stop', 0.95, 'Input data with bad quality'), # input_gate
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 2
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.get_last_timestamp,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
),
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
),
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
'2024-01-01', # get_last_timestamp
|
||||
('repeat', 0.95, 'Input data with bad quality'), # input_gate
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
('continue', 0.95, 'Error'), # mlflow_response_gate (transform)
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 4
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.get_last_timestamp,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
'2024-01-01', # get_last_timestamp
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 5
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.get_last_timestamp,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': 'transformed_data',
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
'2024-01-01', # get_last_timestamp
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
|
||||
('continue', 0.95, 'Error'), # mlflow_response_gate (predict)
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 7
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.get_last_timestamp,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': 'transformed_data',
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_predict,
|
||||
{
|
||||
'data': 'transformed_data',
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_predict_filters'],
|
||||
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'predict',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'STOP'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'repeat'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'CONTINUE'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
prediction_store_policy = 'erl:1'
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'prediction_store_policy': prediction_store_policy,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'Prediction Process',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': path_flag,
|
||||
'data': data,
|
||||
'prediction_confidence': confidence,
|
||||
'timestamp': last_timestamp,
|
||||
'model_id': model,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'comment': 'Prediction Process',
|
||||
'prediction_store_policy': prediction_store_policy,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'unknown'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
prediction_store_policy = 'erl:1'
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'prediction_store_policy': prediction_store_policy,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is False
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
@@ -1,106 +0,0 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.workflows.minimal_retrain import MinimalRetrain
|
||||
|
||||
|
||||
@fixture
|
||||
def minimal_retrain() -> MinimalRetrain:
|
||||
return MinimalRetrain()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'minimal_retrain',
|
||||
'schedule_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
|
||||
input_data = {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'minimal_retrain',
|
||||
'schedule_name': 'test_schedule',
|
||||
'query': 'test_query',
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
}
|
||||
|
||||
workflow_mock.execute_activity_method = AsyncMock(
|
||||
return_value={
|
||||
'data1': '1',
|
||||
'data2': '2',
|
||||
}
|
||||
)
|
||||
|
||||
await minimal_retrain.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.load_custom_query,
|
||||
{
|
||||
**metadata,
|
||||
'query': input_data['query'],
|
||||
'datetime_columns': input_data.get('datetime_columns', []),
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.retrain_model,
|
||||
{
|
||||
**metadata,
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
'model_name': input_data['model_name'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.update_production_model,
|
||||
{
|
||||
**metadata,
|
||||
'model_name': input_data['model_name'],
|
||||
'model_id': input_data['model_id'],
|
||||
**workflow_mock.execute_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**metadata,
|
||||
'data': workflow_mock.execute_activity_method.return_value,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
@@ -1,77 +0,0 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
|
||||
@fixture
|
||||
def predictions_batch() -> PredictionsBatch:
|
||||
return PredictionsBatch()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'predictions_batch',
|
||||
'schedule_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
|
||||
workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
|
||||
input_data = {
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'query': 'SELECT * FROM test',
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
'prediction_store_policy': 'erl:1',
|
||||
'model_config': {'retention': '30'},
|
||||
}
|
||||
|
||||
await predictions_batch.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.load_custom_query,
|
||||
{
|
||||
**metadata,
|
||||
'query': input_data['query'],
|
||||
'datetime_columns': input_data.get('datetime_columns', []),
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
prediction_input = {
|
||||
'metadata': metadata,
|
||||
'data': {'data': 'test_data'},
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'model_id': input_data['model_id'],
|
||||
'model_name': input_data['model_name'],
|
||||
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
|
||||
'mlflow_transform_filters': input_data.get(
|
||||
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
||||
),
|
||||
'mlflow_predict_filters': input_data.get(
|
||||
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
||||
),
|
||||
'model_config': input_data.get('model_config', {}),
|
||||
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
|
||||
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
|
||||
}
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[call('prediction_process', prediction_input)]
|
||||
)
|
||||
Reference in New Issue
Block a user