Merge pull request #1 from Aignosi/release/SIENTIAPDE-1243

SIENTIAPDE-1243: Model Manager Project - Initial Implementation, Refactoring, and Enhancements
This commit is contained in:
Bruno Domingues
2025-10-02 14:59:16 -03:00
committed by GitHub
50 changed files with 7168 additions and 3 deletions

26
.env.example Normal file
View File

@@ -0,0 +1,26 @@
POSTGRES_HOST="paradedb-rw.paradedb.svc.cluster.local"
POSTGRES_PORT="5432"
POSTGRES_USER="sientia"
POSTGRES_PASSWORD="password"
POSTGRES_DBNAME="sientia"
POSTGRES_MIN_CONNECTIONS="10"
POSTGRES_MAX_CONNECTIONS="30"
MLFLOW_HOST="http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
MLFLOW_PORT="80"
MLFLOW_USERNAME="aignosi"
MLFLOW_PASSWORD="mlflow_password"
LOG_LEVEL="DEBUG"
HTTP_METRICS_PORT="9090"
HTTP_SDK_METRICS_PORT="9091"
PROJECT_NAME="sientia-model-manager"
TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233"
TEMPORAL_NAMESPACE="model-manager"
MONGODB_USERNAME="mongo_user"
MONGODB_PASSWORD="mongo_db_password"
MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"
MONGODB_DATABASE="sientia"
MONGODB_TTL_INDEX_HOURS="1"

249
.github/workflows/quality-gate.yml vendored Normal file
View File

@@ -0,0 +1,249 @@
name: Quality gate
on:
# push:
# branches:
# - main
pull_request:
branches:
- main
types: [ opened, synchronize, reopened ]
jobs:
sonar:
name: SonarQube Analysis
runs-on: ubuntu-latest
permissions: write-all
steps:
- name: ⬇️ Checkout Code
uses: actions/checkout@v4
with:
fetch-depth: 0
persist-credentials: false
- name: Calculate Version
id: calculate-version
if: github.event_name == 'pull_request'
uses: actions/github-script@v7
with:
script: |
const fs = require('fs');
// Função para calcular nova versão baseada no branch
function calculateVersion(lastVersion, branchName) {
const parseVersion = (v) => {
const match = v.match(/^(\d+)\.(\d+)\.(\d+)(?:-rc(\d+))?$/);
if (!match) throw new Error(`Invalid version format: ${v}`);
return {
major: parseInt(match[1]),
minor: parseInt(match[2]),
patch: parseInt(match[3]),
rc: match[4] ? parseInt(match[4]) : null
};
};
const current = parseVersion(lastVersion);
if (branchName.startsWith('release/')) {
return `${current.major + 1}.0.0`;
} else if (branchName.startsWith('feature/')) {
return `${current.major}.${current.minor + 1}.0`;
} else if (branchName.startsWith('fix/')) {
return `${current.major}.${current.minor}.${current.patch + 1}`;
} else if (branchName.startsWith('rc/')) {
if (current.rc !== null) {
return `${current.major}.${current.minor}.${current.patch}-rc${current.rc + 1}`;
} else {
return `${current.major}.${current.minor}.${current.patch}-rc1`;
}
}
return null; // Não sugerir para outros tipos de branch
}
try {
const branchName = context.payload.pull_request.head.ref;
console.log(`Branch name: ${branchName}`);
// Validar se o branch segue os padrões aceitos
const validPrefixes = ['release/', 'feature/', 'fix/', 'rc/'];
const isValidBranch = validPrefixes.some(prefix => branchName.startsWith(prefix));
if (!isValidBranch) {
const errorMessage = `## 🚨 Erro: Nome do Branch Inválido\n\n` +
`O branch \`${branchName}\` não segue os padrões de nomenclatura aceitos.\n\n` +
`### 📝 Padrões Aceitos:\n` +
`- \`release/*\`: Para releases de major version (ex: release/v2.0.0)\n` +
`- \`feature/*\`: Para novas funcionalidades (ex: feature/nova-funcionalidade)\n` +
`- \`fix/*\`: Para correções de bugs (ex: fix/correcao-bug)\n` +
`- \`rc/*\`: Para release candidates (ex: rc/v1.2.0-rc1)\n\n` +
`### 🔧 Como corrigir:\n` +
`1. Renomeie o branch para seguir um dos padrões acima\n` +
`2. Ou crie um novo branch com o nome correto\n`;
const prNumber = context.issue.number;
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: prNumber,
body: errorMessage
});
core.setFailed(`Invalid branch name: ${branchName}. Must start with release/, feature/, fix/, or rc/`);
return;
}
// Obter a última tag de release
console.log('Fetching latest release tag...');
const { data: releases } = await github.rest.repos.listReleases({
owner: context.repo.owner,
repo: context.repo.repo,
per_page: 1
});
let lastReleaseVersion = null;
let newVersion = null;
if (releases.length > 0) {
lastReleaseVersion = releases[0].tag_name.replace(/^v/, '');
console.log(`Latest release version: ${lastReleaseVersion}`);
// Calcular a nova versão baseada na última release
newVersion = calculateVersion(lastReleaseVersion, branchName);
} else {
console.log('No releases found, starting from 0.0.0');
lastReleaseVersion = 'N/A';
// Primeira release: começar com 0.0.0 independente do tipo de branch
newVersion = '0.0.0';
}
console.log(`Calculated version: ${newVersion}`);
// Exportar a versão como output
core.setOutput('version', newVersion);
// Adicionar comentário informativo no PR
const prNumber = context.issue.number;
const infoMessage = `## ✅ Versão Calculada Automaticamente\n\n` +
`**Branch:** \`${branchName}\`\n` +
`**Última release:** \`${lastReleaseVersion}\`\n` +
`**Nova versão:** \`${newVersion}\`\n\n` +
`### 📝 Regras Aplicadas:\n` +
`- \`release/*\`: Aumenta major, zera minor e patch (ex: 2.0.0)\n` +
`- \`feature/*\`: Mantém major, aumenta minor, zera patch (ex: 1.2.0)\n` +
`- \`fix/*\`: Mantém major e minor, aumenta patch (ex: 1.1.3)\n` +
`- \`rc/*\`: Mantém versão base, aumenta RC (ex: 1.1.2-rc2)\n`;
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: prNumber,
body: infoMessage
});
} catch (error) {
console.error('Error during version calculation:', error);
const prNumber = context.issue.number;
const errorMessage = `## 🚨 Erro no Cálculo de Versão\n\n` +
`Ocorreu um erro durante o cálculo da versão:\n\n\`\`\`\n${error.message}\n\`\`\`\n\n` +
`Por favor, verifique se o nome do branch está correto e tente novamente.`;
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: prNumber,
body: errorMessage
});
core.setFailed(`Version calculation error: ${error.message}`);
}
- name: Generate App Token
id: generate-app-token
uses: actions/create-github-app-token@v1
with:
app-id: ${{ secrets.APP_ID }}
private-key: ${{ secrets.APP_PRIVATE_KEY }}
owner: 'Aignosi'
repositories: 'sientia-dataops-library,sientia-mlops-library'
- name: Prepare requirements.txt
id: prepare-requirements
run: |
sed -e "s|git+ssh://git@github.com/|git+https://github.com/|g" \
-e "s|git@github.com:|git+https://github.com/|g" \
requirements.txt > requirements_prepared.txt
echo "PROCESSED_REQUIREMENTS_FILE=requirements_prepared.txt" >> $GITHUB_OUTPUT
- name: Configure Git to use App Token
env:
GH_APP_TOKEN: ${{ steps.generate-app-token.outputs.token }}
run: |
git config --global url."https://oauth2:${GH_APP_TOKEN}@github.com/".insteadOf "https://github.com/"
- name: 🧹 Free Disk Space
run: |
echo "Disk space before cleanup:"
df -h
sudo rm -rf /usr/share/dotnet
sudo rm -rf /usr/local/lib/android
sudo rm -rf /opt/ghc
sudo rm -rf /opt/hostedtoolcache/CodeQL
sudo docker image prune --all --force
echo "Disk space after cleanup:"
df -h
- name: 🔧 Setup Python
uses: actions/setup-python@v4
with:
python-version: "3.11"
- name: 🗄️ Cache Python dependencies
uses: actions/cache@v3
with:
path: ~/.cache/pip
key: ${{ runner.os }}-pip-${{ hashFiles(steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE) }}
restore-keys: |
${{ runner.os }}-pip-
- name: 📦 Install Dependencies
run: |
python -m pip install --upgrade pip
pip install --no-cache-dir -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }}
pip install --no-cache-dir -r requirements-dev.txt
- name: 📝 Code Formatting Check (Ruff)
run: |
echo "Checking code formatting..."
ruff format --check model_manager/ tests/
continue-on-error: false
- name: 🔎 Code Linting (Ruff)
run: |
echo "Running linting checks..."
ruff check model_manager/ tests/
continue-on-error: false
- name: 🏷️ Type Checking (mypy)
run: |
echo "Running type checks..."
mypy model_manager/
continue-on-error: true
- name: 🔒 Security Analysis (Bandit)
run: |
echo "Running security analysis..."
bandit -r model_manager/ -ll -q
continue-on-error: true
- name: 🧪 Run Tests with Pytest
run: |
pytest tests --junitxml=pytest.xml --cov=model_manager --cov-report=xml --cov-report=term
- name: Run SonarQube Analysis
uses: SonarSource/sonarqube-scan-action@v6
env:
SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }}
SONAR_HOST_URL: ${{ secrets.SONAR_HOST_URL }}
with:
args: >
-Dsonar.projectVersion=${{ steps.calculate-version.outputs.version || '0.0.0' }}

37
.gitignore vendored
View File

@@ -51,6 +51,10 @@ coverage.xml
.pytest_cache/
cover/
# Code quality tools cache
.bandit/
validate.txt
# Translations
*.mo
*.pot
@@ -173,7 +177,7 @@ cython_debug/
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
.idea/
# Abstra
# Abstra is an AI-powered process automation framework.
@@ -186,7 +190,7 @@ cython_debug/
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
# and can be added to the global gitignore or merged into this file. However, if you prefer,
# you could uncomment the following to ignore the entire vscode folder
# .vscode/
.vscode/
# Ruff stuff:
.ruff_cache/
@@ -205,3 +209,32 @@ cython_debug/
marimo/_static/
marimo/_lsp/
__marimo__/
# Ignore Docker volumes
docker-compose.override.yml
**/db_data/
**/kafka-volume/
**/zookeeper-volume/
**/mage_data/
**/minio_data/
**/venv/
**/certs/*.pem
**/certs/*.der
**/certs/*.csr
**/deploy/*.yaml
scouter/.file_versions/
scouter/pipelines/**/triggers.yaml
**/postgres_data/**
# Ignore Python cache files
*.pyc
*.pyo
*.pyd
# Ignore temporary files
*.swp
# Miscellaneous
git_key*
git_log
tmp/

7
Makefile Normal file
View File

@@ -0,0 +1,7 @@
VERSION = 1.0.8
name = sientia-model-manager
# ENVIRONMENT = production
docker-hub:
@docker build --no-cache -t aignosi.azurecr.io/$(name):$(VERSION) .
@docker push aignosi.azurecr.io/$(name):$(VERSION)

972
README.md
View File

@@ -1,2 +1,972 @@
# sientia-dataops-model-manager
# 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.
## 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
- **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
- **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
### Development & Quality Assurance
- **Code Quality Tools**: Ruff (linting/formatting), mypy (type checking), Bandit (security analysis)
- **Automated Validation**: Pre-commit validation script and CI/CD integration
- **Comprehensive Testing**: pytest with async support and 70%+ code coverage
- **Type Safety**: Static type checking with mypy for improved code reliability
## Architecture
The Model Manager system uses a Temporal-based workflow architecture with clear separation of concerns and robust error handling. The architecture is designed for high availability, scalability, and operational excellence in production ML environments.
### Architecture Principles
#### 1. **Separation of Concerns**
- **Worker Layer**: Manages Temporal workers, task queues, and application lifecycle
- **Workflow Layer**: Orchestrates business logic and process coordination
- **Activity Layer**: Implements specific operations and external system interactions
- **Data Layer**: Handles data persistence, caching, and external service connections
#### 2. **Fault Tolerance & Resilience**
- **Automatic Retry Policies**: Configurable retry strategies for transient failures
- **Circuit Breaker Pattern**: Prevents cascading failures in external service calls
- **Graceful Degradation**: System continues operating with reduced functionality
- **Comprehensive Error Handling**: Detailed error reporting and notification integration
#### 3. **Scalability & Performance**
- **Horizontal Scaling**: Multiple worker instances for load distribution
- **Task Queue Isolation**: Separate queues for different workflow types
- **Connection Pooling**: Optimized database and external service connections
- **Asynchronous Processing**: Non-blocking operations for improved throughput
#### 4. **Observability & Monitoring**
- **Prometheus Metrics**: Comprehensive system and business metrics
- **Structured Logging**: Consistent log format with correlation IDs
- **Health Checks**: Endpoint health monitoring and alerting
- **Performance Tracing**: Request flow tracking and bottleneck identification
### Key Components
#### **Worker (`model_manager/worker/worker.py`)**
- **Purpose**: Main application orchestrator managing Temporal workers and task queues
- **Responsibilities**:
- Temporal client initialization and connection management
- Worker lifecycle management and graceful shutdown
- Task queue configuration and load balancing
- Prometheus metrics server initialization
- Notification handler setup and configuration
- **Key Features**:
- Automatic scaling with `PollerBehaviorAutoscaling`
- 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`
#### **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
- **Key Features**:
- Temporal workflow definitions with retry policies
- Child workflow orchestration and delegation
- Comprehensive error handling and recovery
- Configurable timeout and retry strategies
#### **Activities (`model_manager/activities/`)**
- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance
- **Gates**: Data quality validation and filtering mechanisms
- **MLFlow**: Model transformation and prediction operations
- **Key Features**:
- Multiple inheritance pattern for unified activity interface
- Configurable filter policies and validation rules
- MLFlow model serving integration with configurable flavors
- Comprehensive error handling and notification integration
#### **Data Services (`model_manager/utils/`)**
- **Connectors Config**: Environment variable-based configuration management
- **Repository**: Data access layer for MLFlow operations
- `model_repository.py`: MLFlow model operations and retraining
- **Filters**: Data quality validation and MLFlow response filtering
- `conditional_filters.py`: Input data validation filters
- `mlflow_filters.py`: MLFlow API response validation filters
- **Key Features**:
- Environment variable-based configuration with sensible defaults
- Connection pool management and optimization
- Security credential management
- Configuration validation and error handling
- Support for MLFlow model flavors
### Data Flow Architecture
#### **1. Batch Prediction Pipeline**
```
Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform →
MLFlow Prediction → Response Validation → Export (PostgreSQL)
```
#### **2. Model Retraining Pipeline**
```
Training Data → Model Retraining → Quality Validation →
Production Update → Notification & Monitoring
```
### Security Architecture
#### **Authentication & Authorization**
- **MLFlow API Authentication**: Username/password with secure transmission
- **Database Connection Security**: Encrypted connections with credential management
- **Kubernetes Secrets Integration**: Secure credential storage and access
#### **Network Security**
- **TLS/SSL Encryption**: Secure communication channels
- **Network Isolation**: Kubernetes network policies and service mesh
- **Firewall Rules**: Controlled access to external services
- **VPN Integration**: Secure remote access and management
#### **Data Security**
- **Data Encryption**: At-rest and in-transit encryption
- **Access Control**: Role-based access control (RBAC)
- **Audit Logging**: Comprehensive access and operation logging
- **Data Retention**: Configurable data lifecycle management
## 🔄 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. 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)]
```
## 📋 Prerequisites
- Python 3.11+
- Temporal server/cluster
- PostgreSQL database
- MLFlow server
- MongoDB server (for notifications)
**Note**: External dependencies must be available either through:
- Kubernetes cluster deployment
- Docker Compose setup
- Cloud-managed services
- Local installations
### Temporal Namespace Setup
The Model Manager requires a dedicated Temporal namespace to isolate workflows and maintain proper execution history. The namespace must be created **before** starting the application.
#### Why Create a Namespace?
- **Isolation**: Separates Model Manager workflows from other applications
- **Retention Control**: Configures workflow history retention (default: 7 days)
- **Multi-tenancy**: Enables multiple environments (dev, staging, prod) on same cluster
- **Security**: Allows namespace-level access control and permissions
#### When to Create?
-**Before first deployment** in any environment
-**Once per environment** (dev, staging, production)
-**After Temporal cluster setup** or upgrade
#### How to Create the Namespace
**Option 1: Using Temporal Admin Tools Pod (Recommended for Kubernetes)**
```bash
# 1. List Temporal pods
kubectl get pods -n temporal
# 2. Connect to admin tools pod
kubectl exec -it -n temporal <temporal-admin-tools-pod-name> -- bash
# 3. Create namespace
tctl --namespace model-manager namespace register \
--retention 7 \
--description "Model Manager - ML Model Orchestration Namespace"
# 4. Verify creation
tctl --namespace model-manager namespace describe
# 5. Exit pod
exit
```
**Option 2: Using Port Forward (Local Development)**
```bash
# 1. Port forward Temporal frontend
kubectl port-forward -n temporal svc/temporal-frontend 7233:7233
# 2. In another terminal, create namespace
tctl --address localhost:7233 \
--namespace model-manager \
namespace register \
--retention 7 \
--description "Model Manager - ML Model Orchestration Namespace"
# 3. Verify
tctl --address localhost:7233 --namespace model-manager namespace describe
```
**Option 3: Direct kubectl exec (One-liner)**
```bash
kubectl exec -n temporal <temporal-admin-tools-pod-name> -- \
tctl --namespace model-manager namespace register \
--retention 7 \
--description "Model Manager - ML Model Orchestration Namespace"
```
#### Namespace Configuration
| Parameter | Value | Description |
|-----------|-------|-------------|
| **Name** | `model-manager` | Namespace identifier (configurable via `TEMPORAL_NAMESPACE` env var) |
| **Retention** | `7 days` | Workflow history retention period |
| **Description** | `Model Manager - ML Model Orchestration Namespace` | Human-readable description |
#### Verification
To verify the namespace was created successfully:
```bash
# List all namespaces
kubectl exec -n temporal <temporal-admin-tools-pod-name> -- tctl namespace list
# Describe specific namespace
kubectl exec -n temporal <temporal-admin-tools-pod-name> -- \
tctl --namespace model-manager namespace describe
```
#### Troubleshooting
**Error: "namespace already exists"**
- ✅ This is fine! The namespace is already configured
- No action needed, proceed with application deployment
**Error: "connection refused"**
- ❌ Temporal server is not accessible
- Verify Temporal cluster is running: `kubectl get pods -n temporal`
- Check network connectivity and port forwarding
**Error: "permission denied"**
- ❌ Insufficient permissions to create namespace
- Contact cluster administrator for namespace creation
- Or request elevated permissions for your service account
## 🚀 Installation
### Local Development Setup
1. **Clone the repository**
```bash
git clone <repository-url>
cd sientia-dataops-model-manager
```
2. **Create virtual environment**
```bash
conda create -p ./venv python=3.11
conda activate ./venv
```
3. **Install dependencies**
1. **Install github cli**
```bash
sudo apt update
sudo apt install gh -y
```
2. **Authenticate with github**
```bash
gh auth login
```
3. **Run the install_dependencies.sh script**
```bash
chmod +x install_dependencies.sh
./install_dependencies.sh
```
4. **Install Python dependencies**
```bash
python -m pip install --upgrade pip
# Install production dependencies
pip install -r requirements.txt
# Install development and testing tools
pip install -r requirements-dev.txt
```
4. **Create environment configuration file**
```bash
cp .env.example .env
# Edit .env with your connection details
```
5. **Configure external dependencies**
You'll need to set up port forwarding or connections to external services. For example:
```bash
# Port forwarding from Kubernetes cluster
kubectl port-forward svc/postgresql 5432:5432
kubectl port-forward svc/mlflow 5000:5000
kubectl port-forward svc/mongodb 27017:27017
# Or connect to external services
# Ensure services are accessible on localhost with appropriate ports
```
## 📦 How to Run
### Running the Model Manager Application
Use the provided script to run the application locally:
```bash
# Make script executable (first time only)
chmod +x run_local.sh
# Run the application
./run_local.sh
```
The script will:
- Activate the virtual environment
- Load environment variables from `.env`
- Start the model-manager worker application
### Running Tests and Coverage
Use the provided script to run tests with coverage:
```bash
# Make script executable (first time only)
chmod +x run_coverage.sh
# Run tests with coverage
./run_coverage.sh
```
The script will:
- Activate the virtual environment
- Run pytest with coverage reporting
- Generate HTML coverage report
- Open the coverage report in your browser
### Manual Test Execution
You can also run tests manually:
```bash
# Activate virtual environment
source ./venv/bin/activate
# Run all tests
pytest
# Run with coverage
pytest --cov=model_manager --cov-report=html
# Run specific test categories
pytest tests/activities/
pytest tests/workflow/
```
### Manual Application Execution
For manual execution without scripts:
```bash
# Activate virtual environment
source ./venv/bin/activate
# Load environment variables (if using .env file)
if [ -f .env ]; then
export $(cat .env | grep -v '^#' | xargs)
fi
# Start the model-manager worker
python -m model_manager.worker.worker
```
## 🔍 Code Quality & Validation
### Overview
Como Python não é uma linguagem compilada, utilizamos um conjunto robusto de ferramentas para validar a qualidade, segurança e correção do código antes da execução. Estas ferramentas detectam erros, problemas de estilo, vulnerabilidades de segurança e garantem a consistência do código.
### Ferramentas de Validação
#### 1. **Ruff** - Linting e Formatação ⚡
Ferramenta moderna e extremamente rápida (escrita em Rust) que substitui múltiplas ferramentas:
- **Linting**: Detecta erros de código, problemas de estilo (PEP 8), bugs comuns
- **Formatação**: Formata código automaticamente de forma consistente
- **Velocidade**: 10-100x mais rápido que Flake8/Black
#### 2. **mypy** - Type Checking 🏷️
Verificador de tipos estáticos que analisa type hints:
- Detecta erros de tipo antes da execução
- Melhora a documentação do código
- Previne bugs relacionados a tipos incorretos
#### 3. **Bandit** - Análise de Segurança 🔒
Scanner de vulnerabilidades de segurança:
- Detecta padrões inseguros de código
- Identifica hardcoded passwords, SQL injection, etc.
- Garante conformidade com práticas de segurança
#### 4. **pytest** - Testes Automatizados 🧪
Framework de testes com cobertura de código:
- Executa testes unitários e de integração
- Mede cobertura de código
- Suporta testes assíncronos
### Instalação das Ferramentas
```bash
# Instalar dependências de desenvolvimento
pip install -r requirements-dev.txt
```
### Validação Completa
#### Opção 1: Script Automatizado (Recomendado)
```bash
# Executar todas as validações de uma vez
./validate.sh
```
O script `validate.sh` executa automaticamente:
1. ✅ Verificação de formatação (Ruff)
2. ✅ Linting de código (Ruff)
3. ✅ Type checking (mypy)
4. ✅ Análise de segurança (Bandit)
5. ✅ Testes unitários com cobertura (pytest)
#### Opção 2: Comandos Individuais
```bash
# 1. Verificar formatação
ruff format --check model_manager/ tests/
# 2. Verificar linting
ruff check model_manager/ tests/
# 3. Verificar tipos
mypy model_manager/
# 4. Análise de segurança
bandit -r model_manager/ -ll
# 5. Executar testes
pytest tests/ --cov=model_manager --cov-report=term-missing
```
### Correção Automática
Algumas ferramentas podem corrigir problemas automaticamente:
```bash
# Formatar código automaticamente
ruff format model_manager/ tests/
# Corrigir problemas de linting automaticamente
ruff check --fix model_manager/ tests/
```
### Configuração
Todas as ferramentas são configuradas no arquivo `pyproject.toml`:
- **Ruff**: Regras de linting, formatação, complexidade
- **mypy**: Configurações de type checking
- **pytest**: Opções de teste e cobertura
- **Bandit**: Regras de segurança
### Integração com CI/CD
O workflow `.github/workflows/quality-gate.yml` executa automaticamente todas as validações em cada push/PR:
- ✅ Formatação e linting bloqueiam merge se falharem
- ⚠️ Type checking e segurança geram avisos mas não bloqueiam
- ✅ Testes devem passar com cobertura mínima de 70%
### Boas Práticas
1. **Antes de Commit**: Execute `./validate.sh` para garantir qualidade
2. **Durante Desenvolvimento**: Use `ruff check --watch` para feedback em tempo real
3. **Type Hints**: Adicione type hints em funções novas para melhor validação
4. **Testes**: Mantenha cobertura acima de 70%
5. **Segurança**: Revise e corrija todos os avisos do Bandit
## 🧪 Testing
### Test Structure
```
tests/
├── activities/ # Activity implementation tests
├── workflow/ # Workflow orchestration tests
├── utils/ # Utility function tests
└── integration/ # End-to-end workflow tests
```
### Test Execution
```bash
# Install test dependencies
pip install pytest pytest-cov pytest-asyncio
# Run tests with coverage
pytest --cov=model_manager --cov-report=html
# Run specific test modules
pytest tests/activities/test_gates.py
pytest tests/workflow/test_predictions_batch.py
```
## 📊 Monitoring and Metrics
The Model Manager system exposes comprehensive Prometheus metrics for operational visibility and performance monitoring:
### Application Health Metrics
- `app_up`: Application health status (1=healthy, 0=unhealthy)
- Labels: `pod_id`
### Prediction Operation Metrics
- `model_manager_predictions_written_count`: Counter for successful prediction exports
- Labels: `pod_id`, `model_name`, `pipeline_name`
- `model_manager_prediction_confidence_monitor`: Gauge for current prediction confidence levels
- Labels: `pod_id`, `model_name`, `pipeline_name`
- `model_manager_prediction_response_time_monitor`: Histogram for prediction response times
- Labels: `pod_id`, `model_name`, `pipeline_name`
- Buckets: [0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
### Data Quality Metrics
- Filter pass/fail rates through notification system
- MLFlow API response validation metrics
- Data quality gate performance tracking
## ⚙️ Configuration
### Environment Variables
| Variable | Description | Default | Required |
|----------|-------------|---------|----------|
| `TEMPORAL_HOST` | Temporal server address | `localhost:7233` | Yes |
| `TEMPORAL_NAMESPACE` | Temporal namespace | `model-manager` | No |
| `POSTGRES_HOST` | PostgreSQL hostname | `localhost` | Yes |
| `POSTGRES_PORT` | PostgreSQL port | `5432` | Yes |
| `POSTGRES_USER` | PostgreSQL username | `sientia` | Yes |
| `POSTGRES_PASSWORD` | PostgreSQL password | `sientia` | Yes |
| `POSTGRES_DBNAME` | PostgreSQL database | `sientia` | Yes |
| `POSTGRES_MIN_CONNECTIONS` | Minimum PostgreSQL connections | `5` | No |
| `POSTGRES_MAX_CONNECTIONS` | Maximum PostgreSQL connections | `20` | No |
| `MLFLOW_HOST` | MLFlow server hostname | `http://localhost` | Yes |
| `MLFLOW_PORT` | MLFlow server port | `5080` | Yes |
| `MLFLOW_USERNAME` | MLFlow username | `aignosi` | Yes |
| `MLFLOW_PASSWORD` | MLFlow password | `aignosi` | Yes |
| `MONGODB_URL` | MongoDB connection URI | `localhost:27018` | Yes |
| `MONGODB_USERNAME` | MongoDB username | `root` | Yes |
| `MONGODB_PASSWORD` | MongoDB password | `wKZDbMNU1c` | Yes |
| `MONGODB_DATABASE_NAME` | MongoDB database name | `sientia` | Yes |
| `MONGODB_TTL_INDEX_HOURS` | MongoDB TTL index hours | `1` | No |
| `LOG_LEVEL` | Application log level | `INFO` | No |
| `PROJECT_NAME` | Project name for metrics | `model-manager` | No |
| `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No |
| `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No |
| `POD_ID` | Kubernetes pod identifier | `None` | No |
### 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
### Project Structure
```
model_manager/
├── activities/ # Temporal activity implementations
│ ├── activities.py # Main activities orchestrator
│ ├── gates.py # Data quality gates and filtering
│ └── mlflow.py # MLFlow model operations
├── workflows/ # Temporal workflow definitions
│ ├── predictions_batch.py # Main batch prediction workflow
│ ├── minimal_retrain.py # Model retraining workflow
│ └── sub_workflows/ # Sub-workflow implementations
│ ├── prediction_process.py # Core prediction workflow
│ └── format_and_export_prediction.py # Export workflow
├── worker/ # Worker implementation
│ └── worker.py # Main worker orchestrator
├── utils/ # Utility functions
│ ├── connectors_config.py # Database configuration
│ ├── filters/ # Data quality filters
│ │ ├── conditional_filters.py # Conditional data filters
│ │ └── mlflow_filters.py # MLFlow response filters
│ └── repository/ # Data access layer
│ └── model_repository.py # MLFlow model operations
├── metrics.py # Prometheus metrics definitions
└── __init__.py
```
### Adding New Features
1. **Follow Temporal patterns** for new workflows and activities
2. **Add comprehensive docstrings** for all public methods
3. **Include Prometheus metrics** for monitoring
4. **Add unit tests** for new functionality
5. **Update this README** with new features and configuration
## 🐛 Troubleshooting
### Common Issues
1. **Temporal Connection Failures**
- Verify Temporal server is running and accessible
- Check namespace configuration and permissions
- Review server logs for connection issues
2. **MLFlow Connection Issues**
- Verify MLFlow server is running and accessible
- Check authentication credentials and permissions
- Ensure model names and versions exist
3. **Database Connection Issues**
- Verify PostgreSQL service is running
- Check connection credentials and network access
- Ensure proper connection pool configuration
4. **Workflow Execution Failures**
- Review activity error logs and notifications
- Check data quality filter configurations
- Verify input data format and required fields
### Debug Mode
Enable debug logging by setting the log level:
```bash
export LOG_LEVEL=DEBUG
```
## ⚡ Performance Tuning
### Key Parameters
- **Worker Concurrency**: Adjust `max_concurrent_workflow_tasks` and `max_concurrent_activities`
- **Connection Pools**: Optimize database connection pool sizes
- **Model Retention**: Configure MLFlow model retention based on requirements
- **Batch Sizes**: Adjust data processing batch sizes for optimal throughput
### Scaling Considerations
- **Horizontal Scaling**: Deploy multiple worker instances
- **Task Queue Distribution**: Use multiple task queues for different workflow types
- **Database Performance**: Optimize indexes and connection pooling
- **MLFlow Performance**: Configure appropriate model serving resources
## 🤝 Contributing
1. Fork the repository
2. Create a feature branch
3. Make your changes with comprehensive testing
4. Update documentation and docstrings
5. Submit a pull request
### Code Quality Standards
- Follow PEP 8 style guidelines
- Include comprehensive docstrings for all public methods
- Maintain test coverage above 80%
- Use type hints where appropriate
- Follow Temporal.io best practices
## 📄 License
This project is licensed under the terms specified in the LICENSE file.
## 🆘 Support
For support and questions:
- Check the troubleshooting section above
- Review the metrics and logs for error patterns
- Open an issue in the project repository
- Contact the development team
---
**Note**: The Model Manager system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments.

View File

View File

View File

@@ -0,0 +1,95 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.gates import Gates
from model_manager.activities.mlflow import MLFlow
class Activities(Postgres, MLFlow, Gates):
"""
Main activities orchestrator for the Model Manager system.
This class combines functionality from multiple activity classes to provide
a unified interface for all workflow operations. It manages database connections,
MLFlow model interactions, and data quality validation.
The class implements multiple inheritance to combine specialized functionality:
- Postgres: Database operations and data persistence
- MLFlow: Model inference and transformation operations
- Gates: Data quality validation and filtering mechanisms
Attributes:
postgres_config (dict): PostgreSQL connection configuration
mlflow_config (dict): MLFlow server configuration
logger (Logger): Logging and observability instance
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(
self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler,
):
"""
Initialize the Activities orchestrator with all required configurations.
This constructor initializes all parent classes with their respective
configurations and sets up the foundation for all activity operations.
Args:
postgres_config: PostgreSQL connection configuration dictionary
Required keys: host, port, user, password, dbname, min_connections, max_connections
mlflow_config: MLFlow server configuration dictionary
Required keys: host, port, username, password
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If any parent class initialization fails
"""
# Initialize parent classes
Postgres.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler,
)
MLFlow.__init__(
self,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler,
)
Gates.__init__(self, logger=logger, notification_handler=notification_handler)
async def shutdown(self):
"""
Gracefully shutdown all activities and clean up resources.
This method ensures proper cleanup of all resources including:
- PostgreSQL connection pools
- Any other resources that need explicit cleanup
The method should be called before the application terminates to ensure
proper resource cleanup and prevent resource leaks.
"""
Postgres.close(self)

View File

@@ -0,0 +1,583 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
from pandas import DataFrame
from sientia_do.formatters import create_sample_dict
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
from model_manager import metrics
from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
)
from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
# Input filter function mappings
input_filter_functions = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {'STOP': -1, 'CONTINUE': 2, 'REPEAT': -1},
}
# MLFlow response filter function mappings
mlflow_response_filter_functions = {
'API_ERROR': api_error_filter,
'path_confidence': {'STOP': -1, 'CONTINUE': 10, 'REPEAT': -1},
}
# MLFlow content filter function mappings
mlflow_content_filter_functions = {
'NAN_VALUES': nan_values_filter,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1},
}
class Gates(BaseActivity):
"""
Data quality gates and filtering activities for the Model Manager system.
This class implements comprehensive data quality validation and filtering
mechanisms that can be applied at different stages of the prediction pipeline.
It provides configurable filters with policy-based decision making to ensure
data integrity and quality throughout the ML workflow.
The class supports multiple filter types and implements a flexible policy
system that can be configured for different validation requirements. Each
filter returns a path decision (STOP, CONTINUE, REPEAT) along with confidence
scores and detailed comments for monitoring and debugging.
Attributes:
input_filter_functions (dict): Mapping of input filter names to functions
mlflow_response_filter_functions (dict): Mapping of MLFlow response filter names to functions
mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
"""
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
"""
Initialize data quality gates with logging and notification capabilities.
Args:
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If BaseActivity initialization fails
"""
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
@activity.defn(name='input_gate')
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Apply input data quality filters and validation.
This activity validates input data quality using configurable filters
before proceeding with ML operations. It applies multiple filter types
and returns a path decision based on the filter results and configured
policies.
The method implements a comprehensive filtering system that:
1. Applies configured filters to input data
2. Evaluates filter results against policy configurations
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
4. Provides confidence scores and detailed comments
5. Handles errors gracefully with notification integration
Args:
input_data: Configuration and data for input validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Filter configuration and policies
- data (dict): Input data to validate
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If filter execution fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
path_priority = input_data['path_priority']
filter_output = []
self.debug(f'Input data: {data.head(5).to_string()}', metadata)
self.debug(f'Filters: {filters}', metadata)
# Apply each configured filter
for fil, config in filters.items():
if fil not in input_filter_functions:
self.error(f'Filter {fil} not found', metadata)
continue
try:
if input_filter_functions[fil](data, config['config']): # type: ignore[operator]
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
filter_output.append(config['policy']) # type: ignore[index]
except Exception as e: # noqa: BLE001
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f'INTPUT_GATE_ERROR__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Input gate result: {path_flag}', metadata)
return (
path_flag,
input_filter_functions['path_confidence'][path_flag], # type: ignore[index]
'Input data with bad quality',
)
self.info('Nothing was filtered by the input gate', metadata)
return None, 0, ''
@activity.defn(name='mlflow_response_gate')
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow API response quality and integrity.
This activity validates MLFlow API responses to ensure they meet quality
standards before proceeding with further processing. It applies response-specific
filters and determines appropriate path decisions based on response quality.
The method implements response validation that:
1. Applies MLFlow response-specific filters
2. Evaluates API response quality and integrity
3. Determines path decisions based on response validation results
4. Provides confidence scores and detailed validation comments
5. Handles API errors and response validation failures
Args:
input_data: Configuration and data for response validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Response filter configuration and policies
- data (dict): MLFlow API response data to validate
- type (str): Type of MLFlow operation (transform, predict)
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If response validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing mlflow response gate...', metadata)
filters = input_data['filters']
data = input_data['data']
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}', metadata)
self.debug(f'Filters: {filters}', metadata)
comments = []
for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
self.error(f'Filter {fil} not found', metadata)
continue
try:
if mlflow_response_filter_functions[fil](data, config): # type: ignore[operator]
filter_output.append(config['policy']) # type: ignore[index]
comments.append(data['content']['message'])
self.send_notification(
metadata=metadata,
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
message=data['content']['message'],
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=data['content']['traceback'],
)
except Exception as e: # noqa: BLE001
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Mlflow response gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_response_filter_functions['path_confidence'][path_flag], # type: ignore[index]
', '.join(comments),
)
self.info('Nothing was filtered by the mlflow response gate', metadata)
return None, 0, ''
@activity.defn(name='mlflow_content_gate')
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow prediction content quality and integrity.
This activity validates the content of MLFlow predictions to ensure they
meet quality standards before export and persistence. It applies content-specific
filters and determines appropriate path decisions based on content quality.
The method implements content validation that:
1. Applies MLFlow content-specific filters
2. Evaluates prediction content quality and integrity
3. Determines path decisions based on content validation results
4. Provides confidence scores and detailed validation comments
5. Handles content validation failures and quality issues
Args:
input_data: Configuration and data for content validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Content filter configuration and policies
- data (dict): MLFlow prediction content to validate
- type (str): Type of MLFlow operation (transform, predict)
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If content validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.debug(f'Input data:\n {data.head(5).to_string()}', metadata)
self.debug(f'Filters: \n {create_sample_dict(filters)}', metadata)
for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
continue
try:
if mlflow_content_filter_functions[fil](data, config): # type: ignore[operator]
filter_output.append(config['policy']) # type: ignore[index]
self.send_notification(
metadata=metadata,
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
message=f'Data not passed the content filter {fil}:{config}',
block='mlflow_gate',
level=NotificationLevel.WARNING,
attachment_content=data.to_string(),
)
except Exception as e: # noqa: BLE001
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Mlflow content gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_content_filter_functions['path_confidence'][path_flag], # type: ignore[index]
'Transformed data not passed the content filter',
)
self.info('Nothing was filtered by the mlflow content gate', metadata)
return None, 0, ''
def get_prediction_store_policy(
self, prediction_store_policy: str, metadata: dict[str, Any]
) -> tuple[str, int]:
"""
Parse and validate prediction store policy configuration.
This method parses prediction store policy strings in the format 'type:value'
and validates them against allowed policy types and values. It provides
sensible defaults for invalid configurations and logs policy validation
failures for operational monitoring.
Supported Policy Types:
- 'lts': Latest timestamp - sorts data by timestamp descending
- 'erl': Earliest timestamp - sorts data by timestamp ascending
Args:
prediction_store_policy (str): Policy string in format 'type:value'
metadata (dict[str, Any]): Context metadata for logging and notifications
Returns:
tuple[str, int]: (policy_type, policy_value)
- policy_type (str): Validated policy type ('lts' or 'erl')
- policy_value (int): Number of rows to retain
"""
policy_elements = prediction_store_policy.split(':')
if len(policy_elements) < 2:
self.error(
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
policy_type = policy_elements[0]
policy_value = policy_elements[1]
# If the policy_type is not lts or erl, we use the default policy
# If the policty_value is not a number or 0, we use the default policy
if (
policy_type not in ['lts', 'erl']
or not policy_value.isdigit()
or int(policy_value) == 0
):
self.error(
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
return policy_type, int(policy_value)
@activity.defn(name='format_prediction')
async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Format prediction data according to configured storage policies.
This method formats prediction data for storage and export operations.
It applies timestamp-based sorting policies, adds metadata fields,
and ensures data consistency before persistence. The method supports
multiple storage policies for flexible data retention strategies.
Storage Policies:
- 'lts:N': Latest timestamp - retains N most recent predictions
- 'erl:N': Earliest timestamp - retains N oldest predictions
Args:
input_data (dict): Input data containing:
- data (dict[str, Any]): Raw prediction data to format
- timestamp (str): Default timestamp if data lacks timestamp column
- model_id (str): Unique identifier for the ML model
- prediction_confidence (float): Confidence score for the prediction
- prediction_store_policy (str): Storage policy in format 'type:value'
Returns:
dict: Formatted prediction data ready for storage and export
"""
metadata = input_data['metadata']
prediction_store_policy = input_data['prediction_store_policy']
self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data'])
# Create timestamp column from index and reset index
data['timestamp'] = data.index
data = data.reset_index(drop=True)
self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
policy_type, policy_value = self.get_prediction_store_policy(
prediction_store_policy, metadata
)
# If data has no timestamp, we use the default timestamp and not sort the data
self.info(
f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
)
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
if policy_type == 'lts':
self.debug('Sorting data by timestamp descending', metadata)
data = data.sort_values(by='timestamp', ascending=False)
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
elif policy_type == 'erl':
self.debug('Sorting data by timestamp ascending', metadata)
data = data.sort_values(by='timestamp', ascending=True)
else:
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
raise ValueError(f'Invalid policy type: {policy_type}')
data = data.head(int(policy_value))
data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good'
data['comments'] = ''
data = data.sort_values(by='timestamp', ascending=False)
data = data.reset_index(drop=True)
self.info(f'Prediction formatted: {len(data)} rows', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
return data.to_dict()
@activity.defn(name='format_default_prediction')
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Create and format default prediction data for error conditions.
This method generates default prediction data when the main prediction
pipeline encounters errors or quality issues. It creates a standardized
data structure with zero values for predictions and useful metadata
for operational monitoring and debugging.
The default prediction serves as a fallback mechanism to:
1. Maintain data pipeline continuity during failures
2. Provide operational visibility into prediction quality issues
3. Enable downstream systems to handle error conditions gracefully
4. Support debugging and troubleshooting efforts
Args:
input_data (dict): Input data containing:
- timestamp (str): Timestamp for the default prediction
- model_id (str): Unique identifier for the ML model
- prediction_confidence (float): Confidence score (typically low for errors)
- comment (str): Error description or operational comment
Returns:
dict: Formatted default prediction data with error indicators
"""
metadata = input_data['metadata']
self.debug('Formatting default prediction...', metadata)
data = DataFrame(
{
'prediction': [0],
'response_time': [0],
'timestamp': [input_data['timestamp']],
'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'],
'comments': [input_data['comment']],
}
)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
@activity.defn(name='get_last_timestamp')
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
"""
Extract the most recent timestamp from prediction data.
This method analyzes prediction data to find the latest timestamp,
enabling incremental processing and data continuity tracking.
It handles empty datasets gracefully by returning the current time
as a fallback timestamp.
The method is essential for:
1. Incremental data processing workflows
2. Data continuity validation
3. Timestamp-based data loading optimization
4. Workflow execution tracking
Args:
input_data (dict): Input data containing:
- data (dict[str, Any]): Prediction data to analyze
Returns:
str: Formatted timestamp string in UTC with timezone
"""
metadata = input_data['metadata']
self.info('Getting last timestamp...', metadata)
data = DataFrame(input_data['data'])
self.debug(f'Input data: {data.head(5).to_string()}', metadata)
if data.empty:
return now().strftime(DATETIME_FORMAT_WITH_TZ)
max_timestamp = max(data['timestamp'].values.tolist())
self.info(f'Last timestamp: {max_timestamp}', metadata)
return max_timestamp
@activity.defn(name='write_metrics')
async def write_metrics(self, input_data: dict[str, Any]):
"""
Write prediction performance metrics to Prometheus monitoring system.
This method records comprehensive metrics for prediction operations,
enabling operational monitoring, performance analysis, and alerting.
It tracks prediction counts, confidence levels, and response times
for each model and pipeline combination.
Metrics Recorded:
1. Prediction Count: Incremental counter for successful predictions
2. Confidence Monitor: Current confidence level for predictions
3. Response Time Monitor: Histogram of prediction response times
Args:
input_data (dict): Input data containing:
- metadata (dict[str, Any]): Workflow execution metadata
- prediction (dict[str, Any]): Prediction data with metrics
Raises:
Exception: If metrics writing fails or configuration is invalid
"""
metadata = input_data['metadata']
prediction = DataFrame(input_data['prediction'])
prediction_confidence = prediction['prediction_confidence'].values[0]
response_time = prediction['response_time'].values[0]
self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
metrics.PREDICTIONS_WRITTEN_COUNT.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
).inc()
metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
).set(prediction_confidence)
metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
).observe(response_time)
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

View File

@@ -0,0 +1,346 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
import numpy as np
from pandas import DataFrame, to_datetime
from sientia_do.formatters import create_sample_dict
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.utils.repository.model_repository import MLFlowRepository
class MLFlow(BaseActivity):
"""
MLFlow integration activities for model inference operations.
This class provides activities for interacting with MLFlow models, including
data transformation and prediction operations. It handles authentication,
data preprocessing, and model management with configurable retention policies.
The class implements comprehensive error handling and logging for all
MLFlow operations, ensuring reliable model inference in production environments.
Attributes:
mlflow_host (str): MLFlow server hostname
mlflow_port (int): MLFlow server port
mlflow_username (str): MLFlow authentication username
mlflow_password (str): MLFlow authentication password
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
"""
def __init__(
self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
mlflow_password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
"""
Initialize MLFlow activities with server configuration.
Args:
mlflow_host: MLFlow server hostname or IP address
mlflow_port: MLFlow server port number
mlflow_username: Username for MLFlow authentication
mlflow_password: Password for MLFlow authentication
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If MLFlowRepository initialization fails
"""
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
)
@activity.defn(name='request_transform')
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Transform input data using MLFlow models.
This activity processes input data through MLFlow model transformation,
including data preprocessing, format conversion, and validation. It handles
data deduplication, pivoting, and cleanup to ensure optimal model performance.
The transformation process includes:
1. Data deduplication based on variable and timestamp
2. Data pivoting for model input format
3. Null value handling and cleanup
4. MLFlow model transformation request
5. Response validation and logging
Args:
input_data: Configuration and data for transformation
Required keys:
- metadata (dict): Workflow execution metadata
- data (dict): Input data for transformation
- model_name (str): Name of the MLFlow model to use
- model_retention (int): Model retention period in minutes
Returns:
dict: Transformed data from MLFlow model
Raises:
Exception: If transformation fails or MLFlow model is unavailable
"""
metadata = input_data['metadata']
self.info('Transforming data...', metadata)
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug('Raw input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
data = data.sort_values('created_at', ascending=False).drop_duplicates(
subset=['variable', 'timestamp'], keep='first'
)
# Pivot data for model input format
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True)
data.columns.name = None
self.debug('Processed input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Request transformation from MLFlow model
response_data = self.model_monitoring_repository.transform(
model_name, data, model_config, metadata
)
self.debug(
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.debug(
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info('Data transformed successfully', metadata)
return response_data
@activity.defn(name='request_predict')
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Execute predictions using MLFlow models.
This activity performs ML model inference using MLFlow models with the
transformed data. It handles data format conversion, null value processing,
and model prediction requests with comprehensive error handling.
The prediction process includes:
1. Data format validation and cleanup
2. Null value handling for model compatibility
3. MLFlow model prediction request
4. Response validation and logging
5. Performance monitoring and metrics
Args:
input_data: Configuration and data for prediction
Required keys:
- metadata (dict): Workflow execution metadata
- data (dict): Transformed data for prediction
- model_name (str): Name of the MLFlow model to use
- model_retention (int): Model retention period in minutes
Returns:
dict: Prediction results from MLFlow model
Raises:
Exception: If prediction fails or MLFlow model is unavailable
"""
metadata = input_data['metadata']
self.info('Predicting data...', metadata)
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
# Convert numpy.nan to None for model compatibility
data.replace(np.nan, None, inplace=True)
data['timestamp'] = data.index
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
# Request prediction from MLFlow model
response_data = self.model_monitoring_repository.predict(
model_name, data, model_config, metadata
)
self.debug(
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info('Data predicted successfully', metadata)
return response_data
@activity.defn(name='retrain_model')
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Retrain MLFlow models with updated training data.
This activity orchestrates the complete model retraining process,
including data preparation, model retraining execution, and result
validation. It handles data preprocessing, column cleanup, and
comprehensive error handling for production model management.
The retraining process includes:
1. Data timestamp extraction and validation
2. Column cleanup and data preparation
3. Data pivoting for model input format
4. MLFlow model retraining execution
5. Result validation and error handling
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- data (dict[str, Any]): Training data for model retraining
- model_name (str): Name of the MLFlow model to retrain
Returns:
dict: Retraining results containing:
- status (str): Retraining operation status
- timestamp (str): Timestamp of the retraining operation
- experiment (str): MLFlow experiment identifier
Raises:
Exception: If retraining fails or encounters critical errors
"""
metadata = input_data['metadata']
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
self.info(f'Retraining model {model_name}...', metadata)
timestamp = data['timestamp'].max()
self.debug(f'Timestamp: {timestamp}', metadata)
data.drop(columns=['model_id'], inplace=True, errors='ignore')
data.drop(columns=['created_at'], inplace=True, errors='ignore')
data = data.pivot(index='timestamp', columns='variable', values='value')
data.sort_index(inplace=True)
data.reset_index(inplace=True)
data = data.dropna()
data.columns.name = None
try:
retrain_output, experiment = self.model_monitoring_repository.retrain_model(
data=data, model_name=model_name
)
return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment}
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='RETRAIN_MODEL_ERROR',
message=f'Error retraining model {model_name}: {e}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name='update_production_model')
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Update production model with newly trained model version.
This activity manages the critical process of updating production
models with newly trained versions. It handles model deployment,
status tracking, and comprehensive reporting for operational
visibility and audit trails.
The update process includes:
1. Production model update execution
2. Status and metadata tracking
3. Comprehensive reporting and logging
4. Error handling and notification
5. Audit trail maintenance
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- model_name (str): Name of the MLFlow model to update
- experiment (str): MLFlow experiment identifier
- model_id (str): Unique identifier for the model version
- timestamp (str): Timestamp of the update operation
- status (str): Current status of the model update
Returns:
dict[Any, Any]: Comprehensive update report containing:
- model_id (str): Model version identifier
- model_name (str): Name of the updated model
- timestamp (str): Update operation timestamp
- status (str): Update operation status
- Additional MLFlow response metadata
Raises:
Exception: If production model update fails
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
model_id = input_data['model_id']
experiment = input_data['experiment']
timestamp = input_data['timestamp']
status = input_data['status']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata
)
try:
response = self.model_monitoring_repository.update_production_model(
experiment=experiment, model_name=model_name
)
report = DataFrame([response])
report['model_id'] = model_id
report['model_name'] = model_name
report['timestamp'] = timestamp
report['status'] = status
self.info(f'Production model {model_name} updated successfully', metadata)
return report.to_dict() # type: ignore[no-any-return]
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
message=f'Error updating production model {model_name}: {e}',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e

57
model_manager/metrics.py Normal file
View File

@@ -0,0 +1,57 @@
"""
Model Manager Metrics Module
This module defines all Prometheus metrics used by the Sientia DataOps Model Manager system
for monitoring and observability. The metrics provide insights into system performance,
prediction quality, and operational health.
The metrics are designed to be scraped by Prometheus and can be visualized in
Grafana or other monitoring dashboards to provide real-time visibility into
the system's operation.
Key Metric Categories:
- Application Health: Overall system status and availability
- Prediction Operations: Count and performance of prediction operations
- Data Quality: Confidence levels and validation results
- Export Operations: Database export performance
- Response Times: Performance monitoring for various operations
Metric Labels:
- pod_id: Kubernetes pod identifier for multi-instance deployments
- model_name: Name of the ML model being used
- pipeline_name: Name of the prediction pipeline
"""
from prometheus_client import Counter, Gauge, Histogram
# Application health metric
APP_UP = Gauge(
'app_up',
'Indicates if the application is running (1) or shutting down (0)',
['pod_id'],
)
# Core labels used across multiple metrics
CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
# Prediction operation metrics
PREDICTIONS_WRITTEN_COUNT = Counter(
'model_manager_predictions_written_count',
'Number of predictions written to the database table predictions',
CORE_LABELS,
)
# Prediction quality metrics
PREDICTION_CONFIDENCE_MONITOR = Gauge(
'model_manager_prediction_confidence_monitor',
'Current confidence of each prediction',
CORE_LABELS,
)
# Performance monitoring metrics
PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
'model_manager_prediction_response_time_monitor',
'Current response time of each prediction',
CORE_LABELS,
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)

View File

View File

@@ -0,0 +1,89 @@
from os import getenv
from typing import Any
def build_postgres_config() -> dict[str, Any]:
"""
Build PostgreSQL database configuration from environment variables.
This function constructs a PostgreSQL configuration dictionary from
environment variables with sensible defaults for local development.
It handles connection pool configuration and security parameters.
Environment Variables:
POSTGRES_HOST: Database hostname (default: localhost)
POSTGRES_PORT: Database port (default: 5432)
POSTGRES_USER: Database username (default: sientia)
POSTGRES_PASSWORD: Database password (default: sientia)
POSTGRES_DBNAME: Database name (default: sientia)
POSTGRES_MIN_CONNECTIONS: Minimum connection pool size (default: 5)
POSTGRES_MAX_CONNECTIONS: Maximum connection pool size (default: 20)
Returns:
dict: PostgreSQL configuration dictionary with all required parameters
"""
return {
'host': getenv('POSTGRES_HOST', 'localhost'),
'port': int(getenv('POSTGRES_PORT', '5432')),
'user': getenv('POSTGRES_USER', 'sientia'),
'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
}
def build_mlflow_config() -> dict[str, Any]:
"""
Build MLFlow server configuration from environment variables.
This function constructs an MLFlow configuration dictionary from
environment variables with sensible defaults for local development.
It handles server connection and authentication parameters.
Environment Variables:
MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
MLFLOW_PORT: MLFlow server port (default: 5080)
MLFLOW_USERNAME: MLFlow username (default: aignosi)
MLFLOW_PASSWORD: MLFlow password (default: aignosi)
Returns:
dict: MLFlow configuration dictionary with all required parameters
"""
return {
'host': getenv('MLFLOW_HOST', 'http://localhost'),
'port': int(getenv('MLFLOW_PORT', '5080')),
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
}
def build_mongodb_config() -> dict[str, Any]:
"""
Build MongoDB configuration from environment variables.
This function constructs a MongoDB configuration dictionary from
environment variables with sensible defaults for local development.
It handles connection string and database name configuration.
Environment Variables:
MONGODB_USERNAME: MongoDB username (default: root)
MONGODB_PASSWORD: MongoDB password (default: wKZDbMNU1c)
MONGODB_URL: MongoDB connection URI (default: localhost:27018)
MONGODB_DATABASE_NAME: MongoDB database name (default: sientia)
MONGODB_TTL_INDEX_HOURS: TTL index duration in hours (default: 1)
Returns:
dict: MongoDB configuration dictionary with connection parameters
"""
username = getenv('MONGODB_USERNAME', 'root')
password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c')
uri = getenv('MONGODB_URL', 'localhost:27018')
connection_string = f'mongodb://{username}:{password}@{uri}'
return {
'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
}

View File

View File

@@ -0,0 +1,44 @@
from pandas import DataFrame
def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
"""
Filter to check if specific variables contain null values.
This function examines a DataFrame to determine if any of the specified variables
contain null (NaN) values. It returns True if null values are found for any of
the specified variables, False otherwise.
Args:
data (DataFrame): The pandas DataFrame to be examined. Must contain columns
named 'variable' and 'value'.
config (dict): Configuration dictionary containing the following key:
- variables (list): List of variable names to check for null values
Returns:
bool: True if any of the specified variables contain null values,
False if none of the specified variables contain null values.
"""
return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
def filter_empty_data(data: DataFrame, _config: dict) -> bool:
"""
Filter to check if the DataFrame is empty.
This function determines whether the provided DataFrame contains any data.
It's a simple utility function that can be used in conditional logic to
handle cases where no data is available.
Args:
data (DataFrame): The pandas DataFrame to be checked for emptiness.
_config (dict): Configuration dictionary (unused in this function).
The underscore prefix indicates this parameter is required for
interface consistency but not used in the implementation.
Returns:
bool: True if the DataFrame is empty (has no rows), False if it contains data.
"""
return data.empty

View File

@@ -0,0 +1,64 @@
import numpy as np
from pandas import DataFrame
def api_error_filter(response: dict, _config: dict) -> bool:
"""
Filter MLFlow API responses for error conditions.
This function analyzes MLFlow API responses to detect error conditions
and determine if the response should be filtered out due to quality
or reliability issues.
Args:
response: MLFlow API response data (dict)
_config: Filter configuration dictionary
Required keys:
- error_codes (list, optional): List of error codes to detect
- error_keywords (list, optional): List of error keywords to detect
- check_structure (bool, optional): Whether to validate response structure
Returns:
bool: True if data should be filtered (contains errors), False otherwise
"""
if not response:
return True
if not response['success']:
return True
return False
def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
"""
Filter data for NaN (Not a Number) values.
This function detects NaN values in MLFlow prediction results and
determines if the data quality is sufficient for further processing
or export operations.
Args:
predictions: DataFrame containing prediction data to check for NaN values
_config: Filter configuration dictionary
Required keys:
- max_nan_ratio (float, optional): Maximum allowed NaN value ratio (0.0 to 1.0)
- max_nan_count (int, optional): Maximum allowed NaN value count
- check_nested (bool, optional): Whether to check nested data structures
Returns:
bool: True if data should be filtered (too many NaN values), False otherwise
"""
data = (
predictions.replace({None: np.nan})
.infer_objects(copy=False)
.drop(columns=['timestamp'], errors='ignore')
)
if data.isna().all().all():
return True
return False

View File

@@ -0,0 +1,457 @@
"""
Model Monitoring Repository
This module contains the ModelMonitoringRepository class,
which is responsible for handling the communication with the Model Monitoring API.
It includes the methods that are used to answer ModelMonitoringService
requests using the Model Monitoring API functions.
By Monitoring we mean the evaluation of the performance of models, the generation of reports.
"""
import traceback
from datetime import datetime
from os import makedirs, path, remove
import mlflow
import pandas as pd
from sientia.ModelServing import ModelServing # type: ignore[import-untyped]
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
class MLFlowRepository:
def __init__(self, host, username, password, logger: Logger):
self.model_serving = ModelServing(
tracking_uri=host, username=username, password=password, logger=logger
)
self.logger = logger
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
"""
Detect and parse datetime index from data. index must be a timestamp like column.
This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
If another type or format, must raise an error.
"""
index = data.index
# Get type of first element of index
index_type = type(index[0])
self.logger.custom_info(f'Index type: {index_type}', metadata)
message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}'
# Check if all in index are of the same type
if not all(isinstance(i, index_type) for i in index):
raise ValueError(f'{message}')
# Check type and converts to DATETIME_FORMAT_WITH_TZ
if index_type is str:
# Validate format of string and return error if not valid
try:
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
except ValueError as e:
raise ValueError(f'{message}') from e
elif index_type == datetime or index_type == pd.Timestamp:
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined]
else:
raise ValueError(f'{message}')
return data
def transform(
self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
) -> dict:
"""
Transform data using a model.
Parameters:
- model_name (str): The name of the model to use for transformation.
- data (pandas.DataFrame): The data to transform.
- model_retention (int): The number of minutes to keep the model.
Returns:
- dict: A dictionary containing the transformed data.
"""
try:
self.logger.custom_debug(
f'Data received for model transformation: {data.to_csv()}', metadata
)
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('transform_flavor', 'sklearn')
compressed = model_config.get('is_compressed', False)
retention_target = model_config.get('retention_target', 'model')
transform_keyword = model_config.get('transform_function_keyword', 'predict')
transformed_data = self.model_serving.get_cached_transform(
model_name,
data,
model_retention,
flavor,
compressed,
retention_target,
transform_keyword,
)
self.logger.custom_debug(
f'Data received from model transformation: {transformed_data.to_csv()}', metadata
)
transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata)
return {'success': True, 'content': transformed_data.to_dict()}
except Exception as e: # noqa: BLE001
return {
'success': False,
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
def predict(
self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
) -> dict:
"""
Predict data using a model.
Parameters:
- model_name (str): The name of the model to use for prediction.
- data (pandas.DataFrame): The data to predict.
- model_retention (int): The number of minutes to keep the model.
Returns:
- dict: A dictionary containing the predicted data.
"""
try:
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('predict_flavor', 'pyfunc')
compressed = model_config.get('is_compressed', False)
retention_target = model_config.get('retention_target', 'model')
input_index = data.index
start_time = datetime.now()
self.logger.custom_debug(
f'Data received for model prediction: {data.to_csv()}', metadata
)
data = self.model_serving.get_cached_predict(
model_name, data, model_retention, flavor, compressed, retention_target
)
end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction'])
self.logger.custom_debug(
f'Data received from model prediction: {data.to_csv()}', metadata
)
data.index = input_index
data['response_time'] = (end_time - start_time).total_seconds()
return {'success': True, 'content': data.to_dict()}
except Exception as e: # noqa: BLE001
return {
'success': False,
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
def get_experiment_by_run_id(self, run_id: str) -> dict:
# Get the run information using the run_id
run = mlflow.get_run(run_id)
# Extract the experiment ID from the run
experiment_id = run.info.experiment_id
# Get the experiment details using the experiment ID
experiment = mlflow.get_experiment(experiment_id)
experiment_name = experiment.name
return experiment_name
def get_next_run_name(self, model_name: str) -> str:
"""
Generate the next run name for a specific MLFlow model.
This method calculates the next sequential run number for a model
by searching existing runs and incrementing the count. It ensures
unique run names for model training and retraining operations.
Args:
model_name (str): The name of the MLFlow model
Returns:
str: The next run name in format 'model_name-run_number'
"""
runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc'])
next_run_number = len(runs) + 1
return f'{model_name}-{next_run_number}'
def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
"""
Create a new MLFlow experiment for model retraining.
This method sets up the complete environment for model retraining by:
1. Loading the current production prediction model
2. Loading the current production transformation model
3. Fitting the transformation model with new data
4. Preparing data for prediction model retraining
5. Setting up the MLFlow experiment context
Args:
model_name (str): Name of the MLFlow model to retrain
data (pd.DataFrame): Training data for model retraining
Returns:
tuple: (prediction_model, data_model, experiment)
- prediction_model: Loaded prediction model for retraining
- data_model: Fitted transformation model
- experiment: MLFlow experiment name
"""
# load predictor model
predictor_uri = f'models:/{model_name}/production'
# load transform model
latest_production_id = self.model_serving.get_model_info(model_name) # type: ignore[no-any-return]
transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False)
# load
data_model = mlflow.sklearn.load_model(transform_uri)
prediction_model = mlflow.sklearn.load_model(predictor_uri)
data_model = data_model.fit(data)
treated_data = data_model.predict(data)
target_name = data_model.target_variable
y = data[target_name]
treated_data = pd.merge(treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment)
return prediction_model, data_model, experiment
def perform_model_retrain(
self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame
):
"""
Execute the complete model retraining process in MLFlow.
This method performs the actual model retraining by:
1. Starting a new MLFlow run with descriptive metadata
2. Logging model parameters and hyperparameters
3. Retraining both prediction and transformation models
4. Logging training data as artifacts
5. Saving retrained models to MLFlow registry
Args:
prediction_model: MLFlow prediction model to retrain
data_model: MLFlow transformation model to retrain
experiment (str): MLFlow experiment name for the retraining
model_name (str): Name of the model being retrained
data (pd.DataFrame): Training data used for retraining
Returns:
tuple: (status_message, experiment_name)
- status_message (str): Success confirmation message
- experiment_name (str): Name of the experiment
"""
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
experiment_description = f'Retrain model {model_name} with new data'
current_run_name = self.get_next_run_name(experiment)
with mlflow.start_run(
run_name=current_run_name, description=experiment_description
) as _run:
# update transfomation model
# fixed parameters
for name_atribute, val_atribute in pred_model_atributes.items():
if name_atribute != 'model':
mlflow.log_param(name_atribute, val_atribute)
# update prediction model
for name_atribute, val_atribute in data_model_atributes.items():
if name_atribute != 'model':
mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(data_model, 'data_model')
makedirs('temp', exist_ok=True)
file_path = f'temp/raw_data_{model_name}.csv'
data.to_csv(file_path, index=True)
# log the data raw
mlflow.log_artifact(file_path)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(prediction_model, 'prediction_model')
mlflow.log_param('retrain', True)
# clear temp file
if path.exists(file_path):
remove(file_path)
return 'Model retrained successfully', experiment
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
"""
Orchestrate the complete model retraining workflow.
This method coordinates the entire model retraining process by:
1. Creating the MLFlow experiment environment
2. Loading existing production models
3. Executing the retraining process
4. Returning comprehensive retraining results
Args:
data (pd.DataFrame): Training data for model retraining
model_name (str): Name of the MLFlow model to retrain
Returns:
tuple: (status_message, experiment_name)
- status_message (str): Retraining operation status
- experiment_name (str): MLFlow experiment identifier
"""
prediction_model, data_model, experiment = self.create_model_experiment(model_name, data)
retrain_result = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data
)
return retrain_result
def get_experiment(self, experiment_name: str) -> int:
"""
Retrieve MLFlow experiment ID by experiment name.
This method searches for an MLFlow experiment by name and
returns its unique identifier. It provides error handling
for non-existent experiments.
Args:
experiment_name (str): Name of the MLFlow experiment
Returns:
int: MLFlow experiment ID
Raises:
ValueError: If the experiment name is not found
"""
experiment = mlflow.get_experiment_by_name(experiment_name)
if experiment is None:
raise ValueError(f'Experiment {experiment_name} not found')
return experiment.experiment_id # type: ignore[no-any-return]
def get_experiment_last_run(self, experiment_id: int) -> str:
"""
Retrieve the most recent retraining run ID for an experiment.
This method searches for the latest run in an MLFlow experiment
that has been marked as a retraining run. It filters runs by
the 'retrain' parameter and orders them by completion time.
Args:
experiment_id (int): MLFlow experiment ID
Returns:
str: MLFlow run ID of the most recent retraining run
Raises:
ValueError: If runs data is not in expected DataFrame format
"""
runs = mlflow.search_runs(
experiment_ids=[experiment_id],
filter_string='', # Sem filtro no MLflow ainda
output_format='pandas',
)
if not isinstance(runs, pd.DataFrame):
raise ValueError('Runs is not a pandas DataFrame')
# Filtrar apenas as runs onde params.retrain == True
filtered_runs = runs[runs['params.retrain'] == 'True']
# Converter a coluna 'end_time' para datetime
filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
# Ordenar o DataFrame de forma descendente pela coluna 'end_time'
filtered_runs = filtered_runs.sort_values(by='end_time', ascending=False)
# Pegar a última run_id do DataFrame filtrado e ordenado
latest_run_id = filtered_runs.iloc[0]['run_id']
return latest_run_id
def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
"""
Update production model with a specific MLFlow run.
This method promotes a model from a specific MLFlow run to
production stage. It handles model registration, versioning,
and stage transitions with proper error handling.
Args:
run_id (str): MLFlow run ID containing the model to promote
model_name (str): Name of the MLFlow model
Returns:
dict: Model update metadata containing:
- model_name (str): Name of the updated model
- version (str): New model version number
- mlflow_run_id (str): Source run ID
Update Process:
1. Registers the model from the specified run
2. Retrieves the latest model version
3. Transitions the model to 'Production' stage
4. Archives existing production versions
"""
# Registrar o modelo
# Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
# Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)
# Colocar a versão do modelo em produção
# Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
client = mlflow.tracking.MlflowClient()
# Obter a versão mais recente registrada do modelo
model_versions = client.get_registered_model(model_name).latest_versions
if not isinstance(model_versions, list):
raise ValueError('Model versions is not a list')
max_version = max(model_versions, key=lambda x: int(x.version)).version
# Mover a versão mais recente do modelo para o estágio de 'Production'
client.transition_model_version_stage(
name=model_name, version=max_version, stage='Production', archive_existing_versions=True
)
return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}
def update_production_model(self, experiment: str, model_name: str) -> dict:
"""
Update production model using the latest retraining run.
This method orchestrates the complete production model update
process by identifying the most recent retraining run and
promoting it to production stage.
Args:
experiment (str): MLFlow experiment name
model_name (str): Name of the MLFlow model
Returns:
dict: Complete model update metadata containing:
- model_name (str): Name of the updated model
- version (str): New model version number
- mlflow_run_id (str): Source run ID
- mlflow_experiment_id (int): Experiment ID
"""
experiment_id = self.get_experiment(experiment)
run_id = self.get_experiment_last_run(experiment_id)
metadata = self.update_production_model_by_run_id(run_id, model_name)
metadata['mlflow_experiment_id'] = experiment_id
return metadata

View File

View File

@@ -0,0 +1,229 @@
"""
Model Manager Worker Module
This module provides the main worker implementation for the Sientia DataOps Model Manager system.
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
prediction and retraining workflows.
The worker supports two main task queues:
- predictions_batch-queue: Handles batch prediction workflows
- minimal_retrain-queue: Handles model retraining workflows
Key Features:
- Automatic scaling with PollerBehaviorAutoscaling
- Prometheus metrics integration
- Comprehensive error handling and logging
- Graceful shutdown with cleanup
- Multiple worker instances for different workflow types
Environment Variables:
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: model_manager)
- POD_ID: Kubernetes pod identifier for metrics
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
- PROJECT_NAME: Project name for notifications (default: model_manager)
"""
from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import PollerBehaviorAutoscaling, Worker
with workflow.unsafe.imports_passed_through():
import asyncio
import os
import sys
from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger
from model_manager import metrics
from model_manager.activities.activities import Activities
from model_manager.utils.connectors_config import (
build_mlflow_config,
build_mongodb_config,
build_postgres_config,
)
from model_manager.workflows.minimal_retrain import MinimalRetrain
from model_manager.workflows.predictions_batch import PredictionsBatch
from model_manager.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
POD_ID = os.getenv('POD_ID')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
async def main():
"""
Main entry point for the Model Manager worker application.
This function initializes and starts all components of the worker:
1. Sets up logging and metadata
2. Starts Prometheus metrics server
3. Initializes notification handler
4. Creates and configures activities
5. Starts Temporal client and workers
6. Manages worker lifecycle and graceful shutdown
The function runs indefinitely until interrupted or an error occurs.
On error, it performs cleanup and exits with a non-zero status code.
Raises:
Exception: Any unhandled exception during worker execution
SystemExit: On graceful shutdown or error conditions
"""
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
logger = get_logger(__name__)
metadata = {
'pod_id': POD_ID,
'model_name': '-',
'model_id': '-',
'workflow_name': '-',
'schedule_name': '-',
}
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
logger.custom_info('Starting prometheus client...', metadata)
start_prometheus_server()
logger.custom_info('Starting Notification Handler...', metadata)
mongo_config = build_mongodb_config()
notification_handler = NotificationHandler(
connection_string=mongo_config['connection_string'],
database=mongo_config['database_name'],
logger=logger,
project_name=os.getenv('PROJECT_NAME', 'model-manager'),
)
logger.custom_info('Starting Activities...', metadata)
activities = Activities(
postgres_config=build_postgres_config(),
mlflow_config=build_mlflow_config(),
logger=logger,
notification_handler=notification_handler,
)
logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
new_runtime = Runtime(
telemetry=TelemetryConfig(
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
)
)
logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
temporal_client = await client.Client.connect(
target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
runtime=new_runtime,
)
logger.custom_info('Starting Workers...', metadata)
workers = [
Worker(
temporal_client,
task_queue='minimal_retrain-queue',
workflows=[MinimalRetrain],
activities=[
activities.load_custom_query,
activities.retrain_model,
activities.update_production_model,
activities.export_data_to_postgres,
],
max_concurrent_workflow_tasks=50,
max_concurrent_activities=50,
max_concurrent_local_activities=50,
max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
Worker(
temporal_client,
task_queue='predictions_batch-queue',
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
activities=[
# MLFlow
activities.request_predict,
activities.request_transform,
# Gates
activities.input_gate,
activities.mlflow_response_gate,
activities.mlflow_content_gate,
activities.format_prediction,
activities.format_default_prediction,
activities.get_last_timestamp,
# Postgres
activities.load_custom_query,
activities.repeat_last_prediction,
activities.export_data_to_postgres,
activities.write_metrics,
],
max_concurrent_workflow_tasks=50,
max_concurrent_activities=50,
max_concurrent_local_activities=50,
max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
]
handlers = []
for w in workers:
handlers.append(w.run())
logger.custom_info('Workers started successfully', metadata)
try:
# This will run the workers and wait for them to complete.
# If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers)
except BaseException as e: # noqa: BLE001
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
finally:
if notification_handler:
notification_handler.shutdown()
if activities:
await activities.shutdown()
# Exit with a non-zero status code to indicate failure to Kubernetes
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
sys.exit(1)
def start_prometheus_server():
"""
Starts the Prometheus metrics server for monitoring and observability.
This function initializes the Prometheus HTTP server on the configured port
and sets the application health metric to indicate the service is running.
The server exposes metrics that can be scraped by Prometheus for monitoring
the health and performance of the Model Manager worker.
Environment Variables:
HTTP_METRICS_PORT: Port for the metrics server (default: 9090)
POD_ID: Pod identifier for metrics labeling
Raises:
SystemExit: If the metrics server fails to start
"""
try:
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port)
print(f'Prometheus server started on port {port}.')
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
except Exception as e: # noqa: BLE001
print(f'Failed to start Prometheus server: {e}')
os._exit(1)
if __name__ == '__main__':
asyncio.run(main())

View File

View File

@@ -0,0 +1,114 @@
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),
)

View File

@@ -0,0 +1,116 @@
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)

View File

@@ -0,0 +1,120 @@
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),
)

View File

@@ -0,0 +1,295 @@
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

158
pyproject.toml Normal file
View File

@@ -0,0 +1,158 @@
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "model-manager"
version = "0.0.0"
description = "Sientia DataOps Model Manager - ML Model Orchestration System"
readme = "README.md"
requires-python = ">=3.11"
authors = [
{name = "Aignosi", email = "dev@aignosi.com"}
]
[tool.ruff]
line-length = 100
target-version = "py311"
exclude = [
".git",
".venv",
"venv",
"__pycache__",
"*.pyc",
".pytest_cache",
"htmlcov",
]
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"I", # isort
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"UP", # pyupgrade
"N", # pep8-naming
"YTT", # flake8-2020
"S", # flake8-bandit
"BLE", # flake8-blind-except
"A", # flake8-builtins
"C90", # mccabe complexity
]
ignore = [
"E501", # line too long (handled by formatter)
"S101", # use of assert (needed for tests)
"S105", # possible hardcoded password (false positives)
"S106", # possible hardcoded password (false positives)
"N802", # function name should be lowercase (temporal decorators)
"N806", # variable in function should be lowercase
]
[tool.ruff.lint.per-file-ignores]
"tests/**/*.py" = [
"S101", # assert allowed in tests
"S105", # hardcoded passwords ok in tests
"S106", # hardcoded passwords ok in tests
]
[tool.ruff.lint.mccabe]
max-complexity = 15
[tool.ruff.format]
quote-style = "single"
indent-style = "space"
line-ending = "auto"
[tool.mypy]
python_version = "3.11"
warn_return_any = false
warn_unused_configs = true
disallow_untyped_defs = false
disallow_incomplete_defs = false
check_untyped_defs = true
no_implicit_optional = true
warn_redundant_casts = true
warn_unused_ignores = false
warn_no_return = true
strict_equality = true
ignore_missing_imports = false
# Ignore missing imports for external packages
[[tool.mypy.overrides]]
module = "temporalio.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia_do.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "mlflow.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "prometheus_client.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "pandas.*"
ignore_missing_imports = true
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = [
"-v",
"--strict-markers",
"--cov=model_manager",
"--cov-report=term-missing",
"--cov-report=html",
"--cov-report=xml",
]
markers = [
"asyncio: marks tests as async",
"integration: marks tests as integration tests",
"unit: marks tests as unit tests",
]
[tool.coverage.run]
source = ["model_manager"]
omit = [
"*/tests/*",
"*/venv/*",
"*/__pycache__/*",
"*/site-packages/*",
]
branch = true
[tool.coverage.report]
precision = 2
show_missing = true
skip_covered = false
exclude_lines = [
"pragma: no cover",
"def __repr__",
"def __str__",
"raise AssertionError",
"raise NotImplementedError",
"if __name__ == .__main__.:",
"if TYPE_CHECKING:",
"class .*\\bProtocol\\):",
"@(abc\\.)?abstractmethod",
]
[tool.coverage.html]
directory = "htmlcov"
[tool.bandit]
exclude_dirs = ["tests", "venv", ".venv"]
skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments

19
requirements-dev.txt Normal file
View File

@@ -0,0 +1,19 @@
# Development and Testing Dependencies
# These packages are only needed for development, testing, and code quality checks
# Install with: pip install -r requirements-dev.txt
# Code Quality & Linting
ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort)
mypy>=1.7.0 # Static type checker
bandit>=1.7.5 # Security vulnerability scanner
pandas-stubs>=2.0.0 # Type stubs for pandas
types-requests>=2.31.0 # Type stubs for requests
# Testing
pytest>=7.4.0 # Testing framework
pytest-cov>=4.1.0 # Coverage plugin for pytest
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
# Development Tools
ipython>=8.12.0 # Enhanced Python shell
ipdb>=0.13.13 # IPython debugger

6
requirements.txt Normal file
View File

@@ -0,0 +1,6 @@
temporalio
psycopg2-binary
sqlalchemy
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.6
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.39.0
prometheus-client

11
run_coverage.sh Executable file
View File

@@ -0,0 +1,11 @@
#!/bin/bash
# Exit on any error
set -e
echo "Activating virtual environment..."
source ./venv/bin/activate
pytest --cov=model_manager --cov-report=html
xdg-open htmlcov/index.html

21
run_local.sh Executable file
View File

@@ -0,0 +1,21 @@
#!/bin/bash
# Exit on any error
set -e
#echo "Activating virtual environment..."
#conda activate ./venv
echo "Loading environment variables from .env..."
if [ -f .env ]; then
export $(cat .env | grep -v '^#' | xargs)
echo "Environment variables loaded from .env"
else
echo "Warning: .env file not found. Continuing without environment variables."
fi
echo "Starting ingestor application..."
python -m model_manager.worker.worker

30
simulator/Dockerfile Normal file
View File

@@ -0,0 +1,30 @@
# syntax=docker/dockerfile:1.4
FROM python:3.11-slim
# Enable use of SSH agent/socket
# This line enables SSH during build
# (don't forget the syntax header above)
RUN apt-get update && apt-get install -y git openssh-client && rm -rf /var/lib/apt/lists/*
# Use build-time SSH mount for Git clone
# The SSH key will NOT remain in the image
# IMPORTANT: this block requires BuildKit
# and the --ssh flag during docker build
# SSH config to skip host key check (safe in CI/local dev)
RUN mkdir -p /root/.ssh && echo "StrictHostKeyChecking no" > /root/.ssh/config
WORKDIR /app
# Clone using SSH
ARG GIT_REPO
ARG GIT_BRANCH=main
# Mount SSH key just for this RUN
RUN --mount=type=ssh git clone --branch ${GIT_BRANCH} ${GIT_REPO} .
# Install requirements if exists
RUN if [ -f requirements.txt ]; then pip install --no-cache-dir -r requirements.txt; fi
CMD ["python", "server.py"]

9
sonar-project.properties Normal file
View File

@@ -0,0 +1,9 @@
sonar.projectKey=Aignosi_sientia-dataops-model-manager_dc6e051c-995a-4b57-9cde-f511018184c7
sonar.projectName=sientia-dataops-model-manager
sonar.sources=model_manager
sonar.tests=tests
sonar.qualitygate.wait=true
sonar.qualitygate.timeout=300
sonar.python.coverage.reportPaths=coverage.xml
sonar.python.xunit.reportPath=pytest.xml
sonar.python.version=3.11

0
tests/__init__.py Normal file
View File

View File

View File

View File

@@ -0,0 +1,97 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import mark
from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.activities import Activities
from model_manager.activities.gates import Gates
from model_manager.activities.mlflow import MLFlow
@patch('model_manager.activities.activities.Postgres.__init__')
@patch('model_manager.activities.activities.MLFlow.__init__')
@patch('model_manager.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10,
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
logger=logger,
notification_handler=notification_handler,
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, MLFlow)
assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler,
)
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler,
)
mock_gates_init.assert_called_once_with(
ANY, logger=logger, notification_handler=notification_handler
)
@mark.asyncio
@patch('model_manager.activities.activities.Postgres', return_value=MagicMock())
@patch('model_manager.activities.activities.MLFlow', return_value=MagicMock())
async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10,
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
logger=logger,
notification_handler=notification_handler,
)
await activities.shutdown()
mock_postgres_init.close.assert_called_once()

View File

@@ -0,0 +1,568 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from model_manager.activities.gates import Gates
@fixture
def gates_activity():
gates = Gates(
logger=MagicMock(),
notification_handler=MagicMock(),
)
gates.error = MagicMock()
gates.debug = MagicMock()
gates.info = MagicMock()
gates.warning = MagicMock()
gates.critical = MagicMock()
gates.send_notification = MagicMock()
return gates
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.error.assert_called_once_with(
'Filter INVALID_FILTER not found', metadata['metadata']
)
@mark.asyncio
@patch('model_manager.activities.gates.input_filter_functions')
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Input data with bad quality')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, '')
@mark.asyncio
@patch('model_manager.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filter_exception(
mock_mlflow_response_filter_functions, gates_activity
):
# Arrange
mock_mlflow_response_filter_functions.__contains__.return_value = True
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_mlflow_response_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': False,
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, 'API error occurred')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, '')
@mark.asyncio
@patch('model_manager.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filter_exception(
mock_mlflow_content_filter_functions, gates_activity
):
# Arrange
mock_mlflow_content_filter_functions.__contains__.return_value = True
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
'data': {
'success': False,
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_mlflow_content_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Transformed data not passed the content filter')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
def test_get_prediction_store_policy_invalid_policy(gates_activity):
# Arrange
prediction_store_policy = 'INVALID_POLICY'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
assert policy_value == 1
def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
# Arrange
prediction_store_policy = 'abc:INVALID_VALUE'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
assert policy_value == 1
def test_get_prediction_store_policy_valid_policy_type(gates_activity):
# Arrange
prediction_store_policy = 'abc:1'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
assert policy_value == 1
def test_get_prediction_store_policy_valid_policy(gates_activity):
# Arrange
prediction_store_policy = 'erl:1'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'erl'
assert policy_value == 1
@mark.asyncio
async def test_format_prediction_no_timestamp(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': {'2023-05-26 11:12:27': 1},
'response_time': {'2023-05-26 11:12:27': 0.1},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:1',
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 1}
assert result['response_time'] == {0: ANY}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ''}
@mark.asyncio
async def test_format_prediction_with_timestamp_erl(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1,
'2023-05-26 11:12:28': 2,
'2023-05-26 11:12:29': 3,
},
'response_time': {
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'erl:2',
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 2, 1: 1}
assert result['response_time'] == {0: 0.2, 1: 0.1}
assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
async def test_format_prediction_with_timestamp_lts(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1,
'2023-05-26 11:12:28': 2,
'2023-05-26 11:12:29': 3,
},
'response_time': {
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2',
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 3, 1: 2}
assert result['response_time'] == {0: 0.3, 1: 0.2}
assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [1, 2, 3],
'response_time': [0.1, 0.2, 0.3],
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2',
}
gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
try:
await gates_activity.format_prediction(input_data)
except ValueError as e:
assert str(e) == 'Invalid policy type: invalid'
else:
raise AssertionError('Expected ValueError')
@mark.asyncio
async def test_format_default_prediction(gates_activity):
# Arrange
input_data = {
**metadata,
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment',
}
# Act
result = await gates_activity.format_default_prediction(input_data)
# Assert
assert result['prediction'] == {0: 0}
assert result['response_time'] == {0: 0}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.1}
assert result['prediction_status'] == {0: 'Bad'}
assert result['comments'] == {0: 'Test comment'}
gates_activity.debug.assert_called()
@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert result == '2023-05-26 11:12:28'
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {'data': {}, **metadata}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert isinstance(result, str) # Should be a timestamp string
assert len(result) > 0
@mark.asyncio
@patch('model_manager.activities.gates.metrics')
async def test_write_metrics(mock_metrics, gates_activity):
"""Test write_metrics method."""
input_data = {
**metadata,
'prediction': {
'prediction': [1, 2, 3],
'prediction_confidence': [0.9, 0.8, 0.7],
'response_time': [0.1, 0.2, 0.3],
},
}
await gates_activity.write_metrics(input_data)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
0.1
)

View File

@@ -0,0 +1,301 @@
from unittest.mock import ANY, MagicMock, patch
import numpy as np
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.activities.mlflow import MLFlow
@patch('model_manager.activities.mlflow.MLFlowRepository')
def test___init__(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
logger=MagicMock(),
notification_handler=MagicMock(),
)
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
@fixture
@patch('model_manager.activities.mlflow.MLFlowRepository')
def mlflow(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
logger=MagicMock(),
notification_handler=MagicMock(),
)
mlflow.send_notification = MagicMock()
return mlflow
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch('model_manager.activities.mlflow.DataFrame')
@patch('model_manager.activities.mlflow.max')
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': [
{
'timestamp': '2024-01-01',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-01',
'variable': 'var2',
'value': 2.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 3.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 4.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
],
'model_name': 'test_model',
'model_config': {},
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
mlflow.model_monitoring_repository.transform.return_value = expected_response
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
# Call the method
response_data = await mlflow.request_transform(input_data)
# Verify the data was correctly transformed
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.pivot.assert_called_once_with(
index='timestamp', columns='variable', values='value'
)
mock_dataframe = mock_dataframe.return_value.pivot.return_value
mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
# mock_dataframe.reset_index.assert_called_once()
mock_dataframe.columns.name = None
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', mock_dataframe, {}, metadata['metadata']
)
@mark.asyncio
@patch('model_manager.activities.mlflow.DataFrame')
@patch('model_manager.activities.mlflow.to_datetime')
@patch('model_manager.activities.mlflow.max')
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': {
'variable': {
'2024-01-01': 'var1',
'2024-01-02': 'var2',
'2024-01-03': 'var1',
'2024-01-04': 'var2',
},
'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
},
'model_name': 'test_model',
'model_config': {},
}
# Mock the predict response
expected_response = {'prediction': [0.5, 0.6]}
mlflow.model_monitoring_repository.predict.return_value = expected_response
# Call the method
response_data = await mlflow.request_predict(input_data)
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', mock_dataframe.return_value, {}, metadata['metadata']
)
@mark.asyncio
async def test_retrain_model(mlflow):
data = {
'model_id': [4, 5, 6, 7],
'created_at': [1, 2, 3, 4],
'timestamp': [1, 1, 2, 2],
'variable': ['var1', 'var2', 'var1', 'var2'],
'value': [1, 2, 3, 4],
}
mlflow.model_monitoring_repository.retrain_model.return_value = (
'Model retrained successfully',
'test',
)
response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
mlflow.model_monitoring_repository.retrain_model.assert_called_once()
assert response == {
'status': 'Model retrained successfully',
'timestamp': 2,
'experiment': 'test',
}
@mark.asyncio
async def test_retrain_model_error(mlflow):
mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
'Error retraining model'
)
data = {
'model_id': [4, 5, 6, 7],
'created_at': [1, 2, 3, 4],
'timestamp': [1, 1, 2, 2],
'variable': ['var1', 'var2', 'var1', 'var2'],
'value': [1, 2, 3, 4],
}
try:
await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
except Exception as e: # noqa: BLE001
assert str(e) == 'Error retraining model'
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='RETRAIN_MODEL_ERROR',
message='Error retraining model test_model: Error retraining model',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('No exception raised')
@mark.asyncio
async def test_update_production_model(mlflow):
mlflow.model_monitoring_repository.update_production_model.return_value = {
'data1': 1,
'data2': 2,
}
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success',
}
response = await mlflow.update_production_model(input_data)
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
experiment='test', model_name='test_model'
)
assert response == {
'data1': {0: 1},
'data2': {0: 2},
'model_id': {0: 1},
'model_name': {0: 'test_model'},
'timestamp': {0: 2},
'status': {0: 'success'},
}
@mark.asyncio
async def test_update_production_model_error(mlflow):
mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
'Error updating production model'
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success',
}
try:
await mlflow.update_production_model(input_data)
except Exception as e: # noqa: BLE001
assert str(e) == 'Error updating production model'
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
message='Error updating production model test_model: Error updating production model',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('No exception raised')

View File

View File

@@ -0,0 +1,37 @@
from pandas import DataFrame
from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
)
def test_filter_specific_variables_null_values():
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']},
)
is False
)
def test_filter_specific_variables_null_values_with_null_values():
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']},
)
is True
)
def test_filter_empty_data():
assert filter_empty_data(DataFrame(), {}) is True
def test_filter_empty_data_with_data():
assert (
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
is False
)

View File

@@ -0,0 +1,23 @@
from pandas import DataFrame
from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response():
assert api_error_filter(None, {})
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {})
def test_api_error_filter_valid_response_success():
assert not api_error_filter({'success': True}, {})
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {})
def test_nan_values_filter_no_nan_values():
assert not nan_values_filter(DataFrame({'variable': [1, 2]}), {})

View File

@@ -0,0 +1,443 @@
from datetime import UTC, datetime
from unittest.mock import ANY, MagicMock, call, patch
import numpy as np
import pytest
from pandas import DataFrame, Timestamp
from model_manager.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch(
'model_manager.utils.repository.model_repository.ModelServing', autospec=True
) as mock_model_serving:
mock_instance = mock_model_serving.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
)
return repo
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
class Any:
pass
invalid_cases = [
({'value': {'2024-01-01 12:00:00': 1, 2024: 2}}),
({'value': {'2024-01-01': 1, '2024-01-02': 2}}),
({'value': {Any(): 1, Any(): 2}}),
]
@pytest.mark.parametrize('data', invalid_cases)
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
input_data = DataFrame(data)
with pytest.raises(ValueError) as e:
mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
assert (
str(e)
== 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S'
)
valid_cases = [
(
{'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}},
['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'],
),
(
{
'value': {
datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
}
},
['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
),
(
{
'value': {
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
}
},
['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'],
),
]
@pytest.mark.parametrize('data,expected', valid_cases)
def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
input_data = DataFrame(data)
response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
assert response.index.tolist() == expected
def test_transform_success(mlflow_repository):
data = MagicMock()
model_name = 'model'
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict'
)
mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata']
)
assert output == {
'success': True,
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value,
}
def test_transform_error(mlflow_repository):
data = MagicMock()
model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error')
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict'
)
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
def test_predict_success(mlflow_repository):
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3])
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model'
)
assert output['success'] is True
assert output['content'] == {
'prediction': {'index_1': 2, 'index_2': 3},
'response_time': {'index_1': ANY, 'index_2': ANY},
}
def test_predict_error(mlflow_repository):
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error'))
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model'
)
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
mlflow.get_run.return_value = MagicMock(
info=MagicMock(
experiment_id='0',
)
)
mlflow.get_experiment.return_value = MagicMock()
mlflow.get_experiment.return_value.name = 'test'
output = mlflow_repository.get_experiment_by_run_id('0')
assert output == 'test'
mlflow.get_run.assert_called_once_with('0')
mlflow.get_experiment.assert_called_once_with('0')
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_next_run_name(mlflow, mlflow_repository):
mlflow.search_runs.return_value = [1, 2, 3]
output = mlflow_repository.get_next_run_name('run')
assert output == 'run-4'
mlflow.search_runs.assert_called_once_with(
experiment_names=['run'],
order_by=['start_time desc'],
)
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0')
output = mlflow_repository.get_experiment('test')
assert output == '0'
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_error(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = None
try:
mlflow_repository.get_experiment('test')
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository):
mlflow.search_runs.return_value = DataFrame(
{
'params.retrain': ['True', 'False', 'True', 'False'],
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
'run_id': ['0', '1', '2', '3'],
}
)
output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with(
experiment_ids=[0],
filter_string='',
output_format='pandas',
)
assert output == '2'
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run_error(mlflow, mlflow_repository):
mlflow.search_runs.return_value = []
try:
mlflow_repository.get_experiment_last_run(0)
except ValueError as e:
assert str(e) == 'Runs is not a pandas DataFrame'
else:
raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
@patch('model_manager.utils.repository.model_repository.mlflow.set_experiment')
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
mlflow_repository.model_serving.get_model_info = MagicMock(return_value='0')
mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test')
mlflow_repository.get_experiment_by_run_id = MagicMock()
data_model_mock = MagicMock()
prediction_model_mock = MagicMock()
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
data_model_mock.fit.return_value = data_model_mock
data_model_mock.predict.return_value = DataFrame(
{
'x': [10, 20, 30],
}
)
data_model_mock.target_variable = 'y'
prediction_model_mock.fit.return_value = prediction_model_mock
data = DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
output = mlflow_repository.create_model_experiment('test', data)
mlflow_repository.model_serving.get_model_info.assert_called_once_with('test')
mlflow_repository.model_serving.get_model_uri.assert_called_once_with('0', prediction=False)
sklearn.load_model.assert_has_calls(
[
call(mlflow_repository.model_serving.get_model_uri.return_value),
call('models:/test/production'),
]
)
assert sklearn.load_model.call_count == 2
data_model_mock.fit.assert_called_once_with(data)
data_model_mock.predict.assert_called_once_with(data)
fit_args = prediction_model_mock.fit.call_args[0][0]
assert fit_args.equals(
DataFrame(
{
'x': [10, 20, 30],
'y': [4, 5, 6],
}
)
)
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
set_experiment.assert_called_once_with(mlflow_repository.get_experiment_by_run_id.return_value)
assert output == (
prediction_model_mock,
data_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value,
)
@patch('model_manager.utils.repository.model_repository.mlflow.start_run')
@patch('model_manager.utils.repository.model_repository.mlflow.log_param')
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model')
@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact')
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
prediction_model_mock = MagicMock()
data_model_mock = MagicMock()
experiment = 'test'
model_name = 'test'
data = MagicMock()
mlflow_repository.get_next_run_name = MagicMock(return_value='test-1')
run = MagicMock()
start_run.__enter__.return_value = run
output = mlflow_repository.perform_model_retrain(
prediction_model_mock, data_model_mock, experiment, model_name, data
)
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
start_run.assert_called_once_with(
run_name='test-1', description='Retrain model test with new data'
)
log_model.assert_has_calls(
[
call(data_model_mock, 'data_model'),
call(prediction_model_mock, 'prediction_model'),
]
)
data.to_csv.assert_called_once_with('temp/raw_data_test.csv', index=True)
log_artifact.assert_called_once_with('temp/raw_data_test.csv')
log_param.assert_has_calls(
[
call('retrain', True),
]
)
assert output == ('Model retrained successfully', experiment)
def test_retrain_model(mlflow_repository):
data = MagicMock()
model_name = 'test'
mlflow_repository.create_model_experiment = MagicMock(
return_value=('data_model', 'prediction_model', '0')
)
mlflow_repository.perform_model_retrain = MagicMock(return_value='Model retrained successfully')
output = mlflow_repository.retrain_model(data, model_name)
mlflow_repository.create_model_experiment.assert_called_once_with(model_name, data)
mlflow_repository.perform_model_retrain.assert_called_once_with(
'data_model', 'prediction_model', '0', model_name, data
)
assert output == 'Model retrained successfully'
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
client_mock = MagicMock()
mlflow.tracking.MlflowClient.return_value = client_mock
client_mock.get_registered_model.return_value = MagicMock(
latest_versions=[
MagicMock(version='1'),
MagicMock(version='2'),
MagicMock(version='3'),
]
)
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
'runs:/0/prediction_model',
'test',
)
mlflow.tracking.MlflowClient.assert_called_once()
client_mock.get_registered_model.assert_called_once_with('test')
client_mock.transition_model_version_stage.assert_called_once_with(
name='test',
version='3',
stage='Production',
archive_existing_versions=True,
)
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
}
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
mlflow.tracking.MlflowClient.return_value = MagicMock(
get_registered_model=MagicMock(return_value=MagicMock(latest_versions={}))
)
try:
mlflow_repository.update_production_model_by_run_id('0', 'test')
except Exception as e: # noqa: BLE001
assert str(e) == 'Model versions is not a list'
else:
raise AssertionError('Expected exception')
def test_update_production_model(mlflow_repository):
connector = mlflow_repository
with patch.object(connector, 'get_experiment', return_value='0') as get_experiment:
with patch.object(
connector, 'get_experiment_last_run', return_value='2'
) as get_experiment_last_run:
with patch.object(
connector,
'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3', 'mlflow_run_id': '0'},
) as update_production_model_by_run_id:
output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with('2', 'test')
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
'mlflow_experiment_id': '0',
}

View File

@@ -0,0 +1,115 @@
from os import environ
from model_manager.utils.connectors_config import (
build_mlflow_config,
build_mongodb_config,
build_postgres_config,
)
def test_build_mlflow_config_with_env_vars():
# Arrange
environ['MLFLOW_HOST'] = 'http://test-host'
environ['MLFLOW_PORT'] = '8080'
environ['MLFLOW_USERNAME'] = 'test-user'
environ['MLFLOW_PASSWORD'] = 'test-pass'
# Act
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://test-host'
assert config['port'] == 8080
assert config['username'] == 'test-user'
assert config['password'] == 'test-pass'
def test_build_mlflow_config_with_defaults():
# Arrange
# Clear any existing env vars
environ.pop('MLFLOW_HOST', None)
environ.pop('MLFLOW_PORT', None)
environ.pop('MLFLOW_USERNAME', None)
environ.pop('MLFLOW_PASSWORD', None)
# Act
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://localhost'
assert config['port'] == 5080
assert config['username'] == 'aignosi'
assert config['password'] == 'aignosi'
def test_build_postgres_config_with_env_vars():
# Arrange
environ['POSTGRES_HOST'] = 'test-host'
environ['POSTGRES_PORT'] = '5433'
environ['POSTGRES_USER'] = 'test-user'
environ['POSTGRES_PASSWORD'] = 'test-pass'
environ['POSTGRES_DBNAME'] = 'test-db'
environ['POSTGRES_MIN_CONNECTIONS'] = '10'
environ['POSTGRES_MAX_CONNECTIONS'] = '30'
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'test-host'
assert config['port'] == 5433
assert config['user'] == 'test-user'
assert config['password'] == 'test-pass'
assert config['dbname'] == 'test-db'
assert config['min_connections'] == 10
assert config['max_connections'] == 30
def test_build_postgres_config_with_defaults():
# Arrange
environ.pop('POSTGRES_HOST', None)
environ.pop('POSTGRES_PORT', None)
environ.pop('POSTGRES_USER', None)
environ.pop('POSTGRES_PASSWORD', None)
environ.pop('POSTGRES_DBNAME', None)
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'localhost'
assert config['port'] == 5432
assert config['user'] == 'sientia'
assert config['password'] == 'sientia'
assert config['dbname'] == 'sientia'
assert config['min_connections'] == 5
assert config['max_connections'] == 20
def test_build_mongo_db_config_with_env_vars():
environ['MONGODB_USERNAME'] = 'sientia1'
environ['MONGODB_PASSWORD'] = 'sientia1'
environ['MONGODB_URL'] = 'localhost:27018'
environ['MONGODB_DATABASE_NAME'] = 'test_db'
environ['MONGODB_TTL_INDEX_HOURS'] = '1'
assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db',
'ttl_index_seconds': 3600,
}
def test_build_mongo_db_config_with_defaults():
environ.pop('MONGODB_USERNAME', None)
environ.pop('MONGODB_PASSWORD', None)
environ.pop('MONGODB_DATABASE_NAME', None)
environ.pop('MONGODB_URL', None)
environ.pop('MONGODB_TTL_INDEX_HOURS', None)
assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia',
'ttl_index_seconds': 3600,
}

View File

@@ -0,0 +1,147 @@
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

View File

@@ -0,0 +1,780 @@
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()

View File

@@ -0,0 +1,106 @@
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,
)
]
)

View File

@@ -0,0 +1,77 @@
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)]
)

99
validate.sh Executable file
View File

@@ -0,0 +1,99 @@
#!/bin/bash
# Model Manager Code Validation Script
# This script runs all code quality checks before committing or deploying
set -e # Exit on any error
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color
echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Model Manager - Code Validation Suite ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""
# Check if virtual environment is activated
if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then
echo -e "${YELLOW}⚠️ Warning: No virtual environment detected${NC}"
echo -e "${YELLOW} Consider activating your venv/conda environment${NC}"
echo ""
fi
# Function to run a validation step
run_step() {
local step_name=$1
local step_command=$2
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${BLUE}${step_name}${NC}"
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
if eval "$step_command"; then
echo -e "${GREEN}${step_name} - PASSED${NC}"
echo ""
return 0
else
echo -e "${RED}${step_name} - FAILED${NC}"
echo ""
return 1
fi
}
# Track failures
FAILED_STEPS=()
# Step 1: Code Formatting Check (Ruff)
if ! run_step "1. Code Formatting (Ruff)" "ruff format --check model_manager/ tests/"; then
FAILED_STEPS+=("Code Formatting")
fi
# Step 2: Linting (Ruff)
if ! run_step "2. Code Linting (Ruff)" "ruff check model_manager/ tests/"; then
FAILED_STEPS+=("Linting")
fi
# Step 3: Type Checking (mypy)
if ! run_step "3. Type Checking (mypy)" "mypy model_manager/"; then
FAILED_STEPS+=("Type Checking")
fi
# Step 4: Security Analysis (Bandit)
if ! run_step "4. Security Analysis (Bandit)" "bandit -r model_manager/ -ll -q"; then
FAILED_STEPS+=("Security Analysis")
fi
# Step 5: Unit Tests (pytest)
if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-fail-under=80 -q"; then
FAILED_STEPS+=("Unit Tests")
fi
# Summary
echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Validation Summary ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""
if [ ${#FAILED_STEPS[@]} -eq 0 ]; then
echo -e "${GREEN}✅ All validation checks passed!${NC}"
echo -e "${GREEN} Your code is ready for commit/deployment.${NC}"
echo ""
exit 0
else
echo -e "${RED}❌ Validation failed for the following steps:${NC}"
for step in "${FAILED_STEPS[@]}"; do
echo -e "${RED}${step}${NC}"
done
echo ""
echo -e "${YELLOW}💡 Tips:${NC}"
echo -e "${YELLOW} • Run 'ruff format model_manager/ tests/' to auto-fix formatting${NC}"
echo -e "${YELLOW} • Run 'ruff check --fix model_manager/ tests/' to auto-fix linting issues${NC}"
echo -e "${YELLOW} • Review mypy errors and add type hints where needed${NC}"
echo -e "${YELLOW} • Check bandit warnings for security issues${NC}"
echo -e "${YELLOW} • Fix failing tests or improve test coverage${NC}"
echo ""
exit 1
fi

224
values.yaml Normal file
View File

@@ -0,0 +1,224 @@
# Default values for sientia-module.
# This is a YAML-formatted file.
# Declare variables to be passed into your templates.
# This will set the replicaset count more information can be found here: https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/
replicaCount: 1
# This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/
image:
repository: aignosi.azurecr.io/sientia-module-courier
# This sets the pull policy for images.
pullPolicy: Always
# Overrides the image tag whose default is the chart appVersion.
tag: "0.0.2"
0# This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
imagePullSecrets:
- name: docker-hub-secret
# This is to override the chart name.
nameOverride: "sientia-model-manager-worker"
fullnameOverride: "sientia-model-manager-worker"
namespace: sientia
# This section builds out the service account more information can be found here: https://kubernetes.io/docs/concepts/security/service-accounts/
serviceAccount:
# Specifies whether a service account should be created
create: true
# Automatically mount a ServiceAccount's API credentials?
automount: true
# Annotations to add to the service account
annotations: {}
# The name of the service account to use.
# If not set and create is true, a name is generated using the fullname template
name: "sientia-model-manager-worker"
# This is for setting Kubernetes Annotations to a Pod.
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
podAnnotations: {}
# This is for setting Kubernetes Labels to a Pod.
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/
podLabels: {}
podSecurityContext: {}
# fsGroup: 2000
securityContext: {}
# capabilities:
# drop:
# - ALL
# readOnlyRootFilesystem: true
# runAsNonRoot: true
# runAsUser: 1000
resources: {}
# We usually recommend not to specify default resources and to leave this as a conscious
# choice for the user. This also increases chances charts run on environments with little
# resources, such as Minikube. If you do want to specify resources, uncomment the following
# lines, adjust them as necessary, and remove the curly braces after 'resources:'.
# limits:
# cpu: 100m
# memory: 128Mi
# requests:
# cpu: 100m
# memory: 128Mi
# This is to setup the liveness and readiness probes more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/
livenessProbe:
exec:
command:
- sh
- -c
- pgrep -f "model_manager.worker.worker"
initialDelaySeconds: 20
periodSeconds: 30
readinessProbe:
exec:
command:
- sh
- -c
- pgrep -f "model_manager.worker.worker"
initialDelaySeconds: 10
periodSeconds: 15
# This section is for setting up autoscaling more information can be found here: https://kubernetes.io/docs/concepts/workloads/autoscaling/
autoscaling:
enabled: false
minReplicas: 1
maxReplicas: 100
targetCPUUtilizationPercentage: 80
# targetMemoryUtilizationPercentage: 80
# Additional volumes on the output Deployment definition.
volumes: []
# - name: foo
# secret:
# secretName: mysecret
# optional: false
# Additional volumeMounts on the output Deployment definition.
volumeMounts: []
# - name: foo
# mountPath: "/etc/foo"
# readOnly: true
nodeSelector: {}
tolerations: []
affinity: {}
services:
sdk-metrics:
enabled: true
type: ClusterIP
port: 9091
targetPort: 9091
name: sdk-metrics
metrics:
enabled: true
type: ClusterIP
port: 9090
targetPort: 9090
name: metrics
# Configuração do ServiceMonitor para o Prometheus Operator
# ref: https://github.com/prometheus-operator/prometheus-operator
serviceMonitor:
# Se true, um recurso ServiceMonitor será criado.
enabled: true
# O intervalo no qual as métricas devem ser coletadas (ex: 30s, 1m).
endpoints:
- port: metrics
path: /metrics
interval: 30s
relabelings: []
- port: sdk-metrics
path: /metrics
interval: 30s
relabelings: []
additionalLabels:
release: kube-prometheus-stack
env:
# Entrypoint variables
- name: GITHUB_REPO_URL
value: "git@github.com:Aignosi/sientia-dataops-model-manager.git"
- name: GITHUB_BRANCH
value: SIENTIAPDE-1222-ajustar-a-library-para-fazer-o-download-do-courier
- name: PYTHON_APP
value: "model_manager.worker.worker"
# Application variables
- name: POSTGRES_HOST
value: "paradedb-rw.paradedb.svc.cluster.local"
- name: POSTGRES_PORT
value: "5432"
- name: POSTGRES_USER
value: "sientia"
- name: POSTGRES_PASSWORD
value: "sientia"
- name: POSTGRES_DBNAME
value: "sientia"
- name: POSTGRES_MIN_CONNECTIONS
value: "10"
- name: POSTGRES_MAX_CONNECTIONS
value: "30"
- name: MLFLOW_HOST
value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
- name: MLFLOW_PORT
value: "80"
- name: MLFLOW_USERNAME
value: "aignosi"
- name: MLFLOW_PASSWORD
value: "1L0FP50j3ncp123"
- name: KAFKA_BOOTSTRAP_SERVERS
value: "kafka.kafka.svc.cluster.local:9092"
- name: LOG_LEVEL
value: "DEBUG"
- name: HTTP_METRICS_PORT
value: "9090"
- name: HTTP_SDK_METRICS_PORT
value: "9091"
- name: PROJECT_NAME
value: "sientia-model-manager"
- name: TEMPORAL_HOST
value: "temporal-frontend.temporal.svc.cluster.local:7233"
- name: TEMPORAL_NAMESPACE
value: "model-manager"
- name: MONGODB_USERNAME
value: "root"
- name: MONGODB_PASSWORD
value: "wKZDbMNU1c"
- name: MONGODB_URL
value: "my-release-mongodb.mongodb.svc.cluster.local:27017"
- name: MONGODB_DATABASE
value: "sientia"
- name: MONGODB_TTL_INDEX_HOURS
value: "1"
ssh:
enabled: true
secretName: git-ssh-key-sientia-model-manager-worker
sshPath: /mnt/.ssh
knownHostsPath: /mnt/known_hosts
# kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=5I5zpQ6sRaHqX1hD3dr+2mo647yO3FRc359/wu6gsP+ACRDRz5mp
# helm upgrade --install sientia-model-manager-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.5.0
# kubectl create secret generic git-ssh-key-sientia-model-manager-worker \
# --namespace sientia \
# --from-file=ssh-privatekey=git_key \
# --type=kubernetes.io/ssh-auth