Merge pull request #1 from Aignosi/release/SIENTIAPDE-1243
SIENTIAPDE-1243: Model Manager Project - Initial Implementation, Refactoring, and Enhancements
This commit is contained in:
26
.env.example
Normal file
26
.env.example
Normal 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
249
.github/workflows/quality-gate.yml
vendored
Normal 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
37
.gitignore
vendored
@@ -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
7
Makefile
Normal 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
972
README.md
@@ -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.
|
||||
|
||||
0
model_manager/__init__.py
Normal file
0
model_manager/__init__.py
Normal file
0
model_manager/activities/__init__.py
Normal file
0
model_manager/activities/__init__.py
Normal file
95
model_manager/activities/activities.py
Normal file
95
model_manager/activities/activities.py
Normal 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)
|
||||
583
model_manager/activities/gates.py
Normal file
583
model_manager/activities/gates.py
Normal 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)
|
||||
346
model_manager/activities/mlflow.py
Normal file
346
model_manager/activities/mlflow.py
Normal 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
57
model_manager/metrics.py
Normal 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],
|
||||
)
|
||||
0
model_manager/utils/__init__.py
Normal file
0
model_manager/utils/__init__.py
Normal file
89
model_manager/utils/connectors_config.py
Normal file
89
model_manager/utils/connectors_config.py
Normal 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,
|
||||
}
|
||||
0
model_manager/utils/filters/__init__.py
Normal file
0
model_manager/utils/filters/__init__.py
Normal file
44
model_manager/utils/filters/conditional_filters.py
Normal file
44
model_manager/utils/filters/conditional_filters.py
Normal 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
|
||||
64
model_manager/utils/filters/mlflow_filters.py
Normal file
64
model_manager/utils/filters/mlflow_filters.py
Normal 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
|
||||
457
model_manager/utils/repository/model_repository.py
Normal file
457
model_manager/utils/repository/model_repository.py
Normal 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
|
||||
0
model_manager/worker/__init__.py
Normal file
0
model_manager/worker/__init__.py
Normal file
229
model_manager/worker/worker.py
Normal file
229
model_manager/worker/worker.py
Normal 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())
|
||||
0
model_manager/workflows/__init__.py
Normal file
0
model_manager/workflows/__init__.py
Normal file
114
model_manager/workflows/minimal_retrain.py
Normal file
114
model_manager/workflows/minimal_retrain.py
Normal 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),
|
||||
)
|
||||
116
model_manager/workflows/predictions_batch.py
Normal file
116
model_manager/workflows/predictions_batch.py
Normal 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)
|
||||
0
model_manager/workflows/sub_workflows/__init__.py
Normal file
0
model_manager/workflows/sub_workflows/__init__.py
Normal 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),
|
||||
)
|
||||
295
model_manager/workflows/sub_workflows/prediction_process.py
Normal file
295
model_manager/workflows/sub_workflows/prediction_process.py
Normal 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
158
pyproject.toml
Normal 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
19
requirements-dev.txt
Normal 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
6
requirements.txt
Normal 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
11
run_coverage.sh
Executable 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
21
run_local.sh
Executable 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
30
simulator/Dockerfile
Normal 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
9
sonar-project.properties
Normal 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
0
tests/__init__.py
Normal file
0
tests/laborious/__init__.py
Normal file
0
tests/laborious/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
97
tests/laborious/activities/test_activities.py
Normal file
97
tests/laborious/activities/test_activities.py
Normal 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()
|
||||
568
tests/laborious/activities/test_gates.py
Normal file
568
tests/laborious/activities/test_gates.py
Normal 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
|
||||
)
|
||||
301
tests/laborious/activities/test_mlflow.py
Normal file
301
tests/laborious/activities/test_mlflow.py
Normal 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')
|
||||
0
tests/laborious/utils/__init__.py
Normal file
0
tests/laborious/utils/__init__.py
Normal file
0
tests/laborious/utils/filters/__init__.py
Normal file
0
tests/laborious/utils/filters/__init__.py
Normal file
37
tests/laborious/utils/filters/test_conditional_filters.py
Normal file
37
tests/laborious/utils/filters/test_conditional_filters.py
Normal 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
|
||||
)
|
||||
23
tests/laborious/utils/filters/test_mlflow_filters.py
Normal file
23
tests/laborious/utils/filters/test_mlflow_filters.py
Normal 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]}), {})
|
||||
443
tests/laborious/utils/repository/test_model_repository.py
Normal file
443
tests/laborious/utils/repository/test_model_repository.py
Normal 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',
|
||||
}
|
||||
115
tests/laborious/utils/test_connectors_config.py
Normal file
115
tests/laborious/utils/test_connectors_config.py
Normal 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,
|
||||
}
|
||||
@@ -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
|
||||
@@ -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()
|
||||
106
tests/laborious/workflows/test_minimal_retrain.py
Normal file
106
tests/laborious/workflows/test_minimal_retrain.py
Normal 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,
|
||||
)
|
||||
]
|
||||
)
|
||||
77
tests/laborious/workflows/test_predictions_batch.py
Normal file
77
tests/laborious/workflows/test_predictions_batch.py
Normal 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
99
validate.sh
Executable 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
224
values.yaml
Normal 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
|
||||
Reference in New Issue
Block a user