From 93d0849c80a88d425cf2f54bfedc67e2b33007bd Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Tue, 30 Sep 2025 14:55:38 -0300 Subject: [PATCH] SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations. This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment. --- .env.example | 29 + .github/workflows/quality-gate.yml | 74 ++ .gitignore | 33 +- Makefile | 7 + README.md | 830 +++++++++++++++++- model-manager/__init__.py | 0 model-manager/activities/__init__.py | 0 model-manager/activities/activities.py | 99 +++ model-manager/activities/gates.py | 587 +++++++++++++ model-manager/activities/mlflow.py | 343 ++++++++ model-manager/activities/opc.py | 356 ++++++++ model-manager/metrics.py | 72 ++ model-manager/utils/__init__.py | 0 model-manager/utils/connectors_config.py | 129 +++ model-manager/utils/filters/__init__.py | 0 .../utils/filters/conditional_filters.py | 45 + model-manager/utils/filters/mlflow_filters.py | 61 ++ .../utils/repository/model_repository.py | 481 ++++++++++ .../utils/repository/opc_repository.py | 359 ++++++++ model-manager/worker/__init__.py | 0 model-manager/worker/worker.py | 237 +++++ model-manager/workflows/__init__.py | 0 model-manager/workflows/minimal_retrain.py | 116 +++ model-manager/workflows/predictions_batch.py | 125 +++ .../workflows/sub_workflows/__init__.py | 0 .../format_and_export_prediction.py | 140 +++ .../sub_workflows/prediction_process.py | 296 +++++++ requirements.txt | 8 + run_coverage.sh | 11 + run_local.sh | 18 + simulator/Dockerfile | 30 + sonar-project.properties | 11 + tests/__init__.py | 0 tests/laborious/__init__.py | 0 tests/laborious/activities/__init__.py | 0 tests/laborious/activities/test_activities.py | 135 +++ tests/laborious/activities/test_gates.py | 597 +++++++++++++ tests/laborious/activities/test_mlflow.py | 297 +++++++ tests/laborious/activities/test_opc.py | 369 ++++++++ tests/laborious/utils/__init__.py | 0 tests/laborious/utils/filters/__init__.py | 0 .../utils/filters/test_conditional_filters.py | 30 + .../utils/filters/test_mlflow_filters.py | 22 + .../utils/repository/test_model_repository.py | 502 +++++++++++ .../utils/repository/test_opc_repository.py | 388 ++++++++ .../laborious/utils/test_connectors_config.py | 161 ++++ .../test_format_and_export_prediction.py | 160 ++++ .../subworkflows/test_prediction_process.py | 601 +++++++++++++ .../workflows/test_minimal_retrain.py | 98 +++ .../workflows/test_predictions_batch.py | 88 ++ values.yaml | 229 +++++ 51 files changed, 8171 insertions(+), 3 deletions(-) create mode 100644 .env.example create mode 100644 .github/workflows/quality-gate.yml create mode 100644 Makefile create mode 100644 model-manager/__init__.py create mode 100644 model-manager/activities/__init__.py create mode 100644 model-manager/activities/activities.py create mode 100644 model-manager/activities/gates.py create mode 100644 model-manager/activities/mlflow.py create mode 100644 model-manager/activities/opc.py create mode 100644 model-manager/metrics.py create mode 100644 model-manager/utils/__init__.py create mode 100644 model-manager/utils/connectors_config.py create mode 100644 model-manager/utils/filters/__init__.py create mode 100644 model-manager/utils/filters/conditional_filters.py create mode 100644 model-manager/utils/filters/mlflow_filters.py create mode 100644 model-manager/utils/repository/model_repository.py create mode 100644 model-manager/utils/repository/opc_repository.py create mode 100644 model-manager/worker/__init__.py create mode 100644 model-manager/worker/worker.py create mode 100644 model-manager/workflows/__init__.py create mode 100644 model-manager/workflows/minimal_retrain.py create mode 100644 model-manager/workflows/predictions_batch.py create mode 100644 model-manager/workflows/sub_workflows/__init__.py create mode 100644 model-manager/workflows/sub_workflows/format_and_export_prediction.py create mode 100644 model-manager/workflows/sub_workflows/prediction_process.py create mode 100644 requirements.txt create mode 100755 run_coverage.sh create mode 100755 run_local.sh create mode 100644 simulator/Dockerfile create mode 100644 sonar-project.properties create mode 100644 tests/__init__.py create mode 100644 tests/laborious/__init__.py create mode 100644 tests/laborious/activities/__init__.py create mode 100644 tests/laborious/activities/test_activities.py create mode 100644 tests/laborious/activities/test_gates.py create mode 100644 tests/laborious/activities/test_mlflow.py create mode 100644 tests/laborious/activities/test_opc.py create mode 100644 tests/laborious/utils/__init__.py create mode 100644 tests/laborious/utils/filters/__init__.py create mode 100644 tests/laborious/utils/filters/test_conditional_filters.py create mode 100644 tests/laborious/utils/filters/test_mlflow_filters.py create mode 100644 tests/laborious/utils/repository/test_model_repository.py create mode 100644 tests/laborious/utils/repository/test_opc_repository.py create mode 100644 tests/laborious/utils/test_connectors_config.py create mode 100644 tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py create mode 100644 tests/laborious/workflows/subworkflows/test_prediction_process.py create mode 100644 tests/laborious/workflows/test_minimal_retrain.py create mode 100644 tests/laborious/workflows/test_predictions_batch.py create mode 100644 values.yaml diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..865a73b --- /dev/null +++ b/.env.example @@ -0,0 +1,29 @@ +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" + +OPC_ID="1" +OPC_URL="opc.tcp://sientia-opc-simulator-opc.sientia.svc.cluster.local:4840" + +LOG_LEVEL="DEBUG" +HTTP_METRICS_PORT="9090" +HTTP_SDK_METRICS_PORT="9091" +PROJECT_NAME="sientia-laborious" + +TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233" +TEMPORAL_NAMESPACE="laborious" + +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" \ No newline at end of file diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml new file mode 100644 index 0000000..9212f77 --- /dev/null +++ b/.github/workflows/quality-gate.yml @@ -0,0 +1,74 @@ +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: 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: πŸ”§ 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 -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }} + pip install pytest pytest-cov pytest-asyncio + + - name: πŸ§ͺ Run Tests with Pytest + run: | + pytest tests --junitxml=pytest.xml --cov=laborious --cov-report=xml --cov-report=term + + - name: Run SonarQube Analysis + uses: SonarSource/sonarqube-scan-action@v5 + env: + SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }} + SONAR_HOST_URL: ${{ secrets.SONAR_HOST_URL }} diff --git a/.gitignore b/.gitignore index b7faf40..a450ce7 100644 --- a/.gitignore +++ b/.gitignore @@ -173,7 +173,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 +186,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 +205,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/ diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..66aae81 --- /dev/null +++ b/Makefile @@ -0,0 +1,7 @@ +VERSION = 1.0.8 +name = sientia-laborious +# ENVIRONMENT = production + +docker-hub: + @docker build --no-cache -t aignosi.azurecr.io/$(name):$(VERSION) . + @docker push aignosi.azurecr.io/$(name):$(VERSION) \ No newline at end of file diff --git a/README.md b/README.md index d3e0a11..3f86305 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,830 @@ -# 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 and OPC server integration for industrial systems +- **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 + +## Architecture + +The Laborious 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 (`laborious/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 + - OPC server connection management +- **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 (`laborious/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 (`laborious/activities/`)** +- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance +- **Gates**: Data quality validation and filtering mechanisms +- **MLFlow**: Model transformation and prediction operations +- **OPC**: Real-time data export to industrial OPC servers +- **Key Features**: + - Multiple inheritance pattern for unified activity interface + - Configurable filter policies and validation rules + - MLFlow model serving integration with configurable flavors + - OPC UA client with certificate-based authentication + - Comprehensive error handling and notification integration + - Support for multiple OPC servers with independent configurations + +#### **Data Services (`laborious/utils/`)** +- **Connectors Config**: Environment variable-based configuration management +- **Repository**: Data access layer for MLFlow and OPC operations + - `model_repository.py`: MLFlow model operations and retraining + - `opc_repository.py`: OPC server communication and data writing +- **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 multiple OPC servers and MLFlow model flavors + +### Data Flow Architecture + +#### **1. Batch Prediction Pipeline** +``` +Input Data (PostgreSQL) β†’ Data Quality Gates β†’ MLFlow Transform β†’ +MLFlow Prediction β†’ Response Validation β†’ Export (PostgreSQL + OPC) +``` + +#### **2. Model Retraining Pipeline** +``` +Training Data β†’ Model Retraining β†’ Quality Validation β†’ +Production Update β†’ Notification & Monitoring +``` + +#### **3. Real-time Export Pipeline** +``` +Prediction Results β†’ Data Formatting β†’ OPC Server Write β†’ +Success/Failure Metrics β†’ Notification System +``` + +### Security Architecture + +#### **Authentication & Authorization** +- **Certificate-based OPC Authentication**: Secure industrial communication +- **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"], + "opc_output_config": { + "server_id": "opc_server_1", + "tags": ["prediction_output"] + } +} +``` + +#### 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"], + "opc_output_config": {...} +} +``` + +#### 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πŸ”ƒ] + + A -.-> Redis[(Redis)] + 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 different output destinations +- **PostgreSQL Export**: Persists predictions to database with metrics +- **OPC Integration**: Writes predictions to OPC servers for real-time access +- **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. **OPC Export**: Writes predictions to OPC servers +5. **Metrics Recording**: Records export performance and success metrics + +#### Key Features +- **Flexible Formatting**: Configurable output formats for different destinations +- **Multi-Destination Export**: PostgreSQL and OPC server integration +- **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. write_opc_data] --> C[3. export_data_to_postgres] --> D[4. write_metrics] + + A -.-> Format[Data Formatting] + B -.-> OPC[OPC Servers] + C -.-> PostgreSQL[(PostgreSQL)] + D -.-> 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 +- OPC server(s) +- MongoDB server (for notifications) + +**Note**: External dependencies must be available either through: +- Kubernetes cluster deployment +- Docker Compose setup +- Cloud-managed services +- Local installations + +## πŸš€ Installation + +### Local Development Setup + +1. **Clone the repository** + ```bash + git clone + cd sientia-dataops-laborious + ``` + +2. **Create virtual environment** + ```bash + python3.11 -m venv venv + source ./venv/bin/activate + ``` + +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. **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 Laborious 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 laborious 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=laborious --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 laborious worker +python -m laborious.worker.worker +``` + +## πŸ§ͺ 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=laborious --cov-report=html + +# Run specific test modules +pytest tests/activities/test_gates.py +pytest tests/workflow/test_predictions_batch.py +``` + +## πŸ“Š Monitoring and Metrics + +The Laborious 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 +- `laborious_predictions_written_count`: Counter for successful prediction exports + - Labels: `pod_id`, `model_name`, `pipeline_name` +- `laborious_prediction_confidence_monitor`: Gauge for current prediction confidence levels + - Labels: `pod_id`, `model_name`, `pipeline_name` +- `laborious_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] + +### OPC Export Metrics +- `laborious_prediction_opc_writing_count`: Counter for OPC server write operations + - Labels: `pod_id`, `model_name`, `pipeline_name`, `opc_server_id` +- `laborious_prediction_opc_writing_response_time_monitor`: Histogram for OPC write response times + - Labels: `pod_id`, `model_name`, `pipeline_name`, `opc_server_id` + - 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 | `laborious` | 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 | +| `OPC_CONFIG` | OPC server configuration (JSON) | `{}` | No | +| `OPC_ID` | OPC server identifier | `1` | No | +| `OPC_URL` | OPC server URL | `opc.tcp://localhost:4840` | No | +| `OPC_SERVER_URI` | OPC server URI | `opc.tcp://localhost:4840` | No | +| `OPC_CERT_PATH` | OPC client certificate path | `None` | No | +| `OPC_PRIVATE_KEY_PATH` | OPC private key path | `None` | No | +| `OPC_SERVER_CERT_PATH` | OPC server certificate path | `None` | No | +| `OPC_RECONNECTION_INTERVAL` | OPC reconnection interval (ms) | `120` | No | +| `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 | `laborious` | 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 | + + + + +### OPC Configuration + +For multiple OPC servers, use the `OPC_CONFIG` environment variable: + +```json +{ + "opc_server_1": { + "url": "opc.tcp://server1:4840", + "name": "Server1", + "server_uri": "urn:server1:opcua", + "cert_path": "/path/to/cert.pem", + "private_key_path": "/path/to/key.pem", + "server_cert_path": "/path/to/server_cert.pem", + "reconnection_interval": 5000 + }, + "opc_server_2": { + "url": "opc.tcp://server2:4840", + "name": "Server2", + "server_uri": "urn:server2:opcua", + "cert_path": "/path/to/cert.pem", + "private_key_path": "/path/to/key.pem", + "server_cert_path": "/path/to/server_cert.pem", + "reconnection_interval": 5000 + } +} +``` + +For single OPC server, use individual environment variables: +- `OPC_URL` +- `OPC_NAME` +- `OPC_SERVER_URI` +- `OPC_CERT_PATH` +- `OPC_PRIVATE_KEY_PATH` +- `OPC_SERVER_CERT_PATH` +- `OPC_RECONNECTION_INTERVAL` + +### 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", + "opc_output_config":{}, + "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 +``` +laborious/ +β”œβ”€β”€ activities/ # Temporal activity implementations +β”‚ β”œβ”€β”€ activities.py # Main activities orchestrator +β”‚ β”œβ”€β”€ gates.py # Data quality gates and filtering +β”‚ β”œβ”€β”€ mlflow.py # MLFlow model operations +β”‚ └── opc.py # OPC server 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 +β”‚ └── opc_repository.py # OPC server 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. **OPC Connection Failures** + - Verify OPC server is accessible + - Check certificate and key file paths + - Review OPC server logs for connection issues + +5. **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 Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments. diff --git a/model-manager/__init__.py b/model-manager/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/activities/__init__.py b/model-manager/activities/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/activities/activities.py b/model-manager/activities/activities.py new file mode 100644 index 0000000..ec5ae46 --- /dev/null +++ b/model-manager/activities/activities.py @@ -0,0 +1,99 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + from sientia_do.temporal.activities.postgres import Postgres + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.observability.logger import Logger + from laborious.activities.mlflow import MLFlow + from laborious.activities.gates import Gates + from laborious.activities.opc import OPC + from typing import Any + + +class Activities(Postgres, MLFlow, Gates, OPC): + """ + Main activities orchestrator for the Laborious 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, data quality validation, and OPC server communications. + + 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 + - OPC: Real-time data export to OPC servers + + Attributes: + postgres_config (dict): PostgreSQL connection configuration + mlflow_config (dict): MLFlow server configuration + opc_config (dict): OPC 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], + opc_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 + opc_config: OPC server configuration dictionary + Can contain multiple server configurations + 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) + + OPC.__init__(self, + opc_servers=opc_config, + 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 + - OPC server connections + - 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) + await OPC.shutdown(self) diff --git a/model-manager/activities/gates.py b/model-manager/activities/gates.py new file mode 100644 index 0000000..4a71455 --- /dev/null +++ b/model-manager/activities/gates.py @@ -0,0 +1,587 @@ +from temporalio import activity, workflow + + +with workflow.unsafe.imports_passed_through(): + import traceback + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.temporal.activities.base import BaseActivity + from sientia_do.observability.logger import Logger + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now + from sientia_do.formatters import create_sample_dict + from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter + from typing import Any + from laborious.utils.filters.conditional_filters import ( + filter_empty_data, + filter_specific_variables_null_values + ) + from pandas import DataFrame + from laborious import metrics + +# 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 Laborious 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']): + self.debug( + f"Data not passed the input filter {fil}:{config}", metadata) + filter_output.append(config['policy']) + except Exception as e: + 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], \ + "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): + filter_output.append(config['policy']) + 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: + 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], \ + ", ".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): + filter_output.append(config['policy']) + 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: + 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], \ + "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) diff --git a/model-manager/activities/mlflow.py b/model-manager/activities/mlflow.py new file mode 100644 index 0000000..bd06283 --- /dev/null +++ b/model-manager/activities/mlflow.py @@ -0,0 +1,343 @@ +from temporalio import activity, workflow + + +with workflow.unsafe.imports_passed_through(): + from datetime import datetime + from pandas import Timestamp, to_datetime + from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.activities.base import BaseActivity + 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.formatters import create_sample_dict + from laborious.utils.repository.model_repository import MLFlowRepository + from typing import Any + import numpy as np + from pandas import DataFrame + import traceback + + +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() + + 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 diff --git a/model-manager/activities/opc.py b/model-manager/activities/opc.py new file mode 100644 index 0000000..0b4ec9f --- /dev/null +++ b/model-manager/activities/opc.py @@ -0,0 +1,356 @@ +from temporalio import activity, workflow + + +with workflow.unsafe.imports_passed_through(): + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.temporal.activities.base import BaseActivity + from sientia_do.observability.logger import Logger + from laborious.utils.repository.opc_repository import OpcRepository + from typing import Any + import traceback + from pandas import DataFrame + +OPC_WRITTING_ERROR_CONFIDENCE = 12 + + +class OPC(BaseActivity): + """ + OPC server integration activities for real-time data export. + + This class provides comprehensive OPC UA client functionality for connecting + to multiple OPC servers and writing prediction data in real-time. It implements + secure communication with certificate-based authentication and automatic + reconnection capabilities. + + The class supports multiple OPC servers with individual configurations and + provides robust error handling and monitoring for production environments. + + Attributes: + opc_servers (dict): Configuration for multiple OPC servers + opc_repository (dict): Active OPC repository connections + logger (Logger): Logging and observability instance + notification_handler (NotificationHandler): Notification management instance + """ + + def __init__(self, opc_servers: dict[str, dict[str, Any]], + logger: Logger, notification_handler: NotificationHandler): + + self.logger = logger + self.notification_handler = notification_handler + self.opc_servers = opc_servers + + BaseActivity.__init__( + self, logger, notification_handler, set_error_counter=True) + + self.opc_repository: dict[str, OpcRepository] = {} + self.opc_servers = opc_servers + + async def init_opc(self): + """ + Initialize OPC server connections and establish communication channels. + + This method iterates through all configured OPC servers and attempts to + establish secure connections using certificate-based authentication. + Each server connection is managed independently, and connection failures + are reported through the notification system. + + The method performs the following operations: + 1. Creates OpcRepository instances for each configured server + 2. Establishes secure connections with certificate validation + 3. Reports connection success/failure through notifications + 4. Logs connection status for operational visibility + + Raises: + Exception: If OPC repository initialization fails or connection + establishment encounters critical errors + + Note: + Connection failures are logged and reported but do not prevent + the initialization of other OPC servers. Each server is handled + independently to ensure maximum availability. + """ + self.logger.info("Initializing OPC servers...") + for id, server in self.opc_servers.items(): + self.opc_repository[id] = OpcRepository( + id=server['id'], + url=server['url'], + logger=self.logger, + server_uri=server['server_uri'], + cert_path=server['cert_path'], + private_key_path=server['private_key_path'], + server_cert_path=server['server_cert_path'], + notification_handler=self.notification_handler, + reconnection_interval=server['reconnection_interval'], + pod_id=self.pod_id + ) + is_connected, error_data = await self.opc_repository[id].connect() + if not is_connected: + self.send_notification( + metadata={ + 'model_id': '-', + 'model_name': '-', + 'workflow_name': '-', + 'schedule_name': 'INITIALIZATION' + }, + notification_id=error_data['notification_id'], + message=error_data['message'], + block=error_data['block'], + level=error_data.get('level', NotificationLevel.ERROR), + attachment_content=error_data.get( + 'attachment_content', None) + ) + else: + self.logger.info( + f"OPC server {id} connected successfully.") + + async def write_data(self, server_id: str, tag: str, data: Any, + data_type: str, tag_type: str, metadata: dict[str, Any]) -> bool: + """ + Write data to a specific OPC server tag with comprehensive error handling. + + This method provides a secure and reliable way to write data to OPC servers + with automatic error handling, notification integration, and detailed logging. + It validates server availability before attempting write operations and + provides comprehensive error reporting for operational monitoring. + + Args: + - server_id (str): The id of the OPC server. + - tag (str): The tag to write to. + - data (Any): The data to write. + - data_type (str): The data type. + - tag_type (str): The tag type. + + Returns: + - bool: True if the data was written successfully, False otherwise. + """ + + try: + is_success, error_data = await self.opc_repository[server_id].write_data( + tag, data, data_type, self.logger, metadata) + if not is_success: + self.send_notification( + metadata=metadata, + notification_id=error_data['notification_id'], + message=error_data['message'], + block=error_data['block'], + level=error_data.get('level', NotificationLevel.ERROR), + attachment_content=error_data.get( + 'attachment_content', None) + ) + return False + return True + except Exception as e: + trace = traceback.format_exc() + self.send_notification( + metadata=metadata, + notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR", + message=f"Error writing data to OPC server: {e}", + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=trace + ) + raise e + + def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool: + """ + Validate that an OPC server is available and configured for write operations. + + This method checks if the specified OPC server exists in the active + repository and is available for data writing operations. It provides + immediate feedback for server availability and logs validation failures + for operational monitoring. + + Args: + server_id (str): Unique identifier for the OPC server to validate + metadata (dict[str, Any]): Context metadata for logging and notifications + + Returns: + bool: True if server is available, False otherwise + + Note: + Server validation failures are automatically reported through the + notification system with detailed information about available servers. + This helps operators quickly identify configuration issues. + """ + if self.opc_repository.get(server_id) is None: + message = f"OPC server {server_id} not found to perform write operation." + self.send_notification( + metadata=metadata, + notification_id="OPC_SERVER_NOT_FOUND", + message=message, + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=f"OPC servers: {list(self.opc_repository.keys())}" + ) + return False + return True + + async def manage_output_tags( + self, server_id: str, config: dict[str, Any], data: DataFrame, + metadata: dict[str, Any], success: bool) -> tuple[bool, int]: + """ + Manage the writing of prediction and confidence data to OPC server tags. + + This method orchestrates the writing of multiple data types to OPC servers + based on configuration. It handles both prediction data and confidence + values independently, allowing for flexible tag configuration and + comprehensive error handling. + + The method supports two main tag types: + 1. Prediction tags: Write actual prediction values to configured OPC tags + 2. Confidence tags: Write confidence scores to separate OPC tags + + Args: + server_id (str): Unique identifier for the target OPC server + config (dict[str, Any]): OPC tag configuration containing: + - prediction_tags (dict, optional): Prediction tag configurations + - confidence_tags (dict, optional): Confidence tag configurations + data (DataFrame): DataFrame containing prediction and confidence data + metadata (dict[str, Any]): Context metadata for logging and notifications + success (bool): Current success status to maintain across operations + + Returns: + tuple[bool, int]: (overall_success, total_tags_written) + - overall_success: True if all configured tags were written successfully + - total_tags_written: Count of successfully written tags + """ + + count = 0 + if 'prediction_tags' in config: + for tag, tag_config in config['prediction_tags'].items(): + local_success = await self.write_data( + server_id=server_id, + tag=tag, + data=data.head(1)['prediction'].values[0], + data_type=tag_config['data_type'], + tag_type='prediction', + metadata=metadata + ) + if local_success: + self.info( + f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata) + count += 1 + success = success and local_success + + if 'confidence_tags' in config: + for tag, tag_config in config['confidence_tags'].items(): + local_success = await self.write_data( + server_id=server_id, + tag=tag, + data=data.head(1)['prediction_confidence'].values[0], + data_type=tag_config['data_type'], + tag_type='confidence', + metadata=metadata + ) + if local_success: + self.info( + f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata) + count += 1 + success = success and local_success + + return success, count + + @activity.defn(name='write_opc_data') + async def write_opc_data(self, input_data: dict[str, Any]) -> dict[Any, Any]: + """ + Write prediction and confidence data to OPC servers. The two writing + operations are optional and independent of each other. + + Args: + - input_data(dict[str, Any]): The input data. Contains the following keys: + - data(dict[str, Any]): The dataframe that contains the data to write + to the OPC servers. + - opc_output_config(dict[str, Any]): The OPC writing configuration. + The keys are the OPC server names and the values contain: + - prediction_tags(dict[str, Any]): The tags to write to the OPC servers. + - confidence_tags(dict[str, Any]): The tags to write to the OPC servers. + + Returns: + - dict[Any, Any]: The data that was written to the OPC servers. + + """ + metadata = input_data['metadata'] + self.info("Writing data to OPC servers...", metadata) + data = DataFrame(input_data['data']) + opc_output_config = input_data['opc_output_config'] + self.info(f"Data to write: {data.size} rows", metadata) + + success = True + + for server_id, config in opc_output_config.items(): + + if not self.validate_server(server_id, metadata): + success = False + continue + + local_success, local_count = await self.manage_output_tags( + server_id, config, data, metadata, success) + success = success and local_success + + self.info( + f"Process completed for OPC server {server_id}: {local_count} of {len(config['prediction_tags'])} prediction tags and {len(config['confidence_tags'])} confidence tags", metadata) + + return self.process_confidence(data, success, metadata) + + def process_confidence(self, data: DataFrame, success: bool, metadata: dict[str, Any]) -> dict[Any, Any]: + """ + Process prediction confidence based on OPC write operation success. + + This method updates the prediction confidence values in the DataFrame + based on the success status of OPC server write operations. If any + write operations failed, it sets the confidence to a predefined error + value to indicate data quality issues. + + The method implements a confidence degradation strategy: + - Success: Maintains original confidence values + - Failure: Sets confidence to error value for operational awareness + + Args: + data (DataFrame): DataFrame containing prediction and confidence data + success (bool): Overall success status of OPC write operations + metadata (dict[str, Any]): Context metadata for logging and notifications + + Returns: + dict[Any, Any]: Processed data as a dictionary with updated confidence values + + Note: + The error confidence value (OPC_WRITTING_ERROR_CONFIDENCE = 12) is + used to indicate that data was not successfully exported to OPC servers. + This allows downstream systems to handle data quality appropriately. + """ + + if not success: + data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE + self.debug( + f"Some data could not be written to OPC servers, setting confidence to {OPC_WRITTING_ERROR_CONFIDENCE}.", + metadata + ) + + else: + self.debug("Data written to OPC servers successfully.", metadata) + + return data.to_dict() + + async def shutdown(self): + """ + Gracefully shutdown all OPC server connections and cleanup resources. + + This method ensures proper cleanup of all active OPC server connections + by calling the disconnect method on each repository instance. It's + designed to be called during application shutdown to prevent resource + leaks and ensure clean termination. + + The method performs the following cleanup operations: + 1. Iterates through all active OPC repository connections + 2. Calls disconnect() on each repository instance + 3. Allows for graceful connection termination + 4. Prevents resource leaks and connection hanging + + Note: + This method should be called during application shutdown to ensure + proper cleanup. It handles all active connections regardless of + their current state and provides a clean shutdown experience. + """ + for opc in self.opc_repository.values(): + await opc.disconnect() diff --git a/model-manager/metrics.py b/model-manager/metrics.py new file mode 100644 index 0000000..97f7bb9 --- /dev/null +++ b/model-manager/metrics.py @@ -0,0 +1,72 @@ +""" +Laborious Metrics Module + +This module defines all Prometheus metrics used by the Sientia DataOps Laborious 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 and OPC 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 +- opc_server_id: Identifier for OPC server operations +""" + +from prometheus_client import Gauge, Counter, 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( + "laborious_predictions_written_count", + "Number of predictions written to the database table predictions", + CORE_LABELS, +) + +# Prediction quality metrics +PREDICTION_CONFIDENCE_MONITOR = Gauge( + "laborious_prediction_confidence_monitor", + "Current confidence of each prediction", + CORE_LABELS, +) + +# Performance monitoring metrics +PREDICTION_RESPONSE_TIME_MONITOR = Histogram( + "laborious_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] +) + +# OPC export metrics +PREDICTION_OPC_WRITING_COUNT = Counter( + "laborious_prediction_opc_writing_count", + "Number of predictions written to the OPC server", + [*CORE_LABELS, "opc_server_id"], +) + +PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram( + "laborious_prediction_opc_writing_response_time_monitor", + "Current response time of each prediction written to the OPC server", + [*CORE_LABELS, "opc_server_id"], + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0] +) diff --git a/model-manager/utils/__init__.py b/model-manager/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/utils/connectors_config.py b/model-manager/utils/connectors_config.py new file mode 100644 index 0000000..145f248 --- /dev/null +++ b/model-manager/utils/connectors_config.py @@ -0,0 +1,129 @@ +from os import getenv +import json +from typing import Dict, 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_opc_config() -> Dict[str, Any]: + """ + Build OPC server configuration from environment variables. + + This function constructs an OPC server configuration dictionary from + environment variables. It supports both single server and multi-server + configurations with flexible parameter handling. + + Environment Variables: + OPC_CONFIG: JSON string containing multiple OPC server configurations + OPC_ID: OPC server ID (fallback, default: 1) + OPC_URL: Single OPC server URL (fallback, default: opc.tcp://localhost:4840) + OPC_SERVER_URI: Single OPC server URI (fallback, default: opc.tcp://localhost:4840) + OPC_CERT_PATH: Client certificate path (fallback, default: None) + OPC_PRIVATE_KEY_PATH: Client private key path (fallback, default: None) + OPC_SERVER_CERT_PATH: Server certificate path (fallback, default: None) + OPC_RECONNECTION_INTERVAL: Reconnection interval in milliseconds (fallback, default: 120) + + Returns: + dict: OPC server configuration dictionary + """ + opc_raw = getenv('OPC_CONFIG', None) + + if opc_raw: + return json.loads(opc_raw) + + return { + getenv('OPC_ID', '1'): { + 'id': getenv('OPC_ID', '1'), + 'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'), + 'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'), + 'cert_path': getenv('OPC_CERT_PATH', None), + 'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None), + 'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None), + 'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')) + } + } + + +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 + } diff --git a/model-manager/utils/filters/__init__.py b/model-manager/utils/filters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/utils/filters/conditional_filters.py b/model-manager/utils/filters/conditional_filters.py new file mode 100644 index 0000000..2c51805 --- /dev/null +++ b/model-manager/utils/filters/conditional_filters.py @@ -0,0 +1,45 @@ +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 diff --git a/model-manager/utils/filters/mlflow_filters.py b/model-manager/utils/filters/mlflow_filters.py new file mode 100644 index 0000000..f792a0e --- /dev/null +++ b/model-manager/utils/filters/mlflow_filters.py @@ -0,0 +1,61 @@ +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}).drop( + columns=['timestamp'], errors='ignore').infer_objects() + + if data.isna().all().all(): + return True + + return False diff --git a/model-manager/utils/repository/model_repository.py b/model-manager/utils/repository/model_repository.py new file mode 100644 index 0000000..3107d1c --- /dev/null +++ b/model-manager/utils/repository/model_repository.py @@ -0,0 +1,481 @@ +""" +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. + +""" +from datetime import datetime +import traceback +import pandas as pd +import mlflow +from os import makedirs, path, remove +from sientia.ModelServing import ModelServing +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +from sientia_do.observability.logger import Logger + + +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 == str: + # Validate format of string and return error if not valid + try: + pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ) + except ValueError: + raise ValueError( + f"{message}") + + elif index_type == datetime or index_type == pd.Timestamp: + data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) + 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: + 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: + 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_run_id( + model_name, stage="Production" + ) + 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 int(experiment.experiment_id) + + 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 diff --git a/model-manager/utils/repository/opc_repository.py b/model-manager/utils/repository/opc_repository.py new file mode 100644 index 0000000..96ecfef --- /dev/null +++ b/model-manager/utils/repository/opc_repository.py @@ -0,0 +1,359 @@ +import asyncio +import traceback +import time +from datetime import datetime +from pathlib import Path +from typing import Any +from asyncua import Client +from asyncua.crypto.security_policies import SecurityPolicyBasic256 +from asyncua.ua import DataValue, Variant, VariantType, DateTime +from regex import F +from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler +from sientia_do.notifications.models import NotificationLevel +from sientia_do.observability.logger import Logger +from laborious import metrics + +data_type_map = { + 'float': { + 'converter': float, + 'opc_type': VariantType.Float, + }, + 'double': { + 'converter': float, + 'opc_type': VariantType.Double, + }, + 'int': { + 'converter': int, + 'opc_type': VariantType.Int32, + }, + 'bool': { + 'converter': bool, + 'opc_type': VariantType.Boolean, + }, + 'str': { + 'converter': str, + 'opc_type': VariantType.String, + } +} + + +class OpcRepository(): + def __init__(self, id: str, url: str, logger: Logger, + notification_handler: NotificationHandler, + reconnection_interval: int = 60, server_uri: str = None, cert_path: str = None, + private_key_path: str = None, server_cert_path: str = None, pod_id: str = None): + self.url = url + self.id = id + self.server_uri = server_uri + self.cert_path = cert_path + self.private_key_path = private_key_path + self.server_cert_path = server_cert_path + self.logger = logger + self.error_count = 0 + self.reconnection_interval = reconnection_interval + self.last_reconnection_time = None + self.notification_handler = notification_handler + self.client = None + self.pod_id = pod_id + + self.metadata = { + 'model_name': '-', + 'model_id': '-', + 'workflow_name': 'opc_repository', + 'schedule_name': '-' + } + + async def set_security(self): + """ + Configures the security settings for the OPC UA client. + This method sets up the security policy, certificates, and timeouts + required for establishing a secure connection with the OPC UA server. + Raises: + ValueError: If either the certificate path or private key path is not provided. + Attributes: + - cert_path (str): Path to the client's certificate file. + - private_key_path (str): Path to the client's private key file. + - server_cert_path (str, optional): Path to the server's certificate file. + - server_uri (str): The URI of the server to be used as the application URI. + - client (opcua.Client): The OPC UA client instance. + - logger (logging.Logger): Logger instance for logging information. + Security Settings: + - Security Policy: Basic256 + - Secure Channel Timeout: 10,000,000 ms + - Session Timeout: 10,000,000 ms + """ + + if not all([self.cert_path, self.private_key_path]): + raise ValueError( + "Certificate and private key paths must be provided for secure connection.") + cert = Path(self.cert_path) + private_key = Path(self.private_key_path) + server_cert = Path( + self.server_cert_path) if self.server_cert_path else None + + self.client.application_uri = self.server_uri + self.logger.custom_info('Setting security...', self.metadata) + await self.client.set_security( + SecurityPolicyBasic256, + certificate=str(cert), + private_key=str(private_key), + server_certificate=str(server_cert) + ) + self.client.secure_channel_timeout = 10000000 + self.client.session_timeout = 10000000 + + async def connect(self) -> tuple[bool, dict[str, Any]]: + """ + Establishes a connection to the OPC server. + This method initializes the OPC client using the provided URL and + sets up security if a certificate path is specified. It then + attempts to connect to the server and logs the connection status. + Raises: + Exception: If the connection to the OPC server fails. + """ + + self.client = Client(self.url) + if self.cert_path: + await self.set_security() + self.logger.custom_info( + f'Starting connection to OPC server {self.id}...', self.metadata) + return await self.try_connect() + + async def try_connect(self) -> tuple[bool, dict[str, Any]]: + """ + Attempt to establish connection to the OPC server. + + This method performs the actual connection attempt to the OPC server + and handles connection failures with comprehensive error reporting. + It updates reconnection timing and provides detailed error information + for operational monitoring and debugging. + + Returns: + tuple[bool, dict[str, Any]]: Connection result + - bool: True if connection successful, False otherwise + - dict: Error information if connection failed + """ + + try: + self.last_reconnection_time = datetime.now() + await self.client.connect() + return True, {} + except Exception as e: + trace = traceback.format_exc() + self.logger.custom_error(trace, self.metadata) + + return False, { + "notification_id": f"OPC_CONNECTION_ERROR_{self.id}", + "message": f"Failed to connect to OPC server: {e}", + "block": "opc_repository", + "level": NotificationLevel.ERROR, + "attachment_content": trace + } + + async def disconnect(self): + """ + Gracefully disconnect from the OPC server. + + This method safely terminates the connection to the OPC server + and cleans up client resources. It handles disconnection errors + gracefully and ensures proper resource cleanup. + """ + if self.client is None: + return + try: + await self.client.disconnect() + self.logger.custom_info( + 'Disconnected from OPC server', self.metadata) + except Exception as e: + self.logger.custom_error( + f"Failed to disconnect from OPC server: {e}", self.metadata) + self.client = None + + async def validate_connection(self) -> tuple[bool, dict[str, Any]]: + """ + Validate and maintain OPC server connection health. + + This method performs comprehensive connection validation and + implements automatic reconnection logic for production reliability. + It handles various connection states and implements intelligent + reconnection strategies with error counting and timing controls. + + Connection Validation: + 1. Checks client existence and connection state + 2. Implements error counting with automatic disconnection + 3. Enforces reconnection timing windows + 4. Provides detailed error reporting and notifications + + Reconnection Strategy: + - Error Count Threshold: Disconnects after 5 consecutive errors + - Reconnection Window: Enforces minimum intervals between attempts + - Automatic Recovery: Attempts reconnection when conditions allow + - State Monitoring: Continuously monitors connection health + + Args: + None + + Returns: + tuple[bool, dict[str, Any]]: Connection validation result + - bool: True if connection is healthy, False otherwise + - dict: Error information if validation fails + """ + if self.client is None: + return await self.connect() + + if self.error_count > 5: + self.logger.custom_warning( + f"OPC server {self.id} will be disconnected due to multiple errors", self.metadata) + try: + await self.disconnect() + except Exception as e: + trace = traceback.format_exc() + self.logger.custom_error( + f"Failed to disconnect from OPC server: {e}", self.metadata) + self.logger.custom_error(trace, self.metadata) + self.logger.custom_info( + f"Attempting to reconnect to OPC server {self.id}...", self.metadata) + return await self.connect() + + # Check if client is connected using asyncua's connection state + try: + if self.client.uaclient.protocol is None or self.client.uaclient.protocol.state == "closed": + # OPC server is not connected + self.logger.custom_error( + f"OPC server {self.id} is not connected", self.metadata) + if self.last_reconnection_time is None or (datetime.now() - self.last_reconnection_time).total_seconds( + ) > self.reconnection_interval: + await self.disconnect() + self.logger.custom_info( + f"Trying to reconnect to OPC server {self.id}...", self.metadata) + return await self.connect() + + return False, { + "notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{self.id}", + "message": f"OPC server {self.id} is not connected, waiting for next reconnection window...", + "block": "opc_repository", + "level": NotificationLevel.WARNING + } + return True, {} + except Exception as e: + trace = traceback.format_exc() + message = f"Failed to validate connection to OPC server: {e}" + self.logger.custom_error(message, self.metadata) + return False, { + "notification_id": f"OPC_CONNECTION_CHECK_ERROR_{self.id}", + "message": message, + "block": "opc_repository", + "level": NotificationLevel.ERROR, + "attachment_content": trace + } + + async def write_data(self, node: str, value: Any, data_type: str, + logger: Logger, metadata: dict[str, Any]) -> tuple[bool, dict[str, Any]]: + """ + Write data to OPC server with comprehensive validation and monitoring. + + This method provides secure and reliable data writing to OPC servers + with automatic connection validation, data type conversion, and + comprehensive error handling. It implements performance monitoring + and metrics collection for operational visibility. + + Data Writing Process: + 1. Connection validation and automatic reconnection + 2. Node validation and error handling + 3. Data type conversion and validation + 4. OPC data writing with timestamp + 5. Performance metrics collection + 6. Error handling and notification + + Args: + node (str): OPC node identifier to write data to + value (Any): Data value to write to the OPC node + data_type (str): Data type for OPC conversion + logger (Logger): Logger instance for operation logging + metadata (dict[str, Any]): Context metadata for logging and metrics + + Returns: + tuple[bool, dict[str, Any]]: Write operation result + - bool: True if write successful, False otherwise + - dict: Error information if write failed + """ + + is_connected, error = await self.validate_connection() + + if not is_connected: + return False, error + + start_time = time.time() + + try: + node_obj = self.client.get_node(node) + except Exception as e: + trace = traceback.format_exc() + logger.custom_error(trace, metadata.get('schedule_name', 'N/A')) + self.error_count += 1 + return False, { + "notification_id": f"OPC_WRITE_GET_NODE_ERROR_{self.id}", + "message": f"Failed to get node from OPC server: {e} | metadata: {metadata}", + "block": "opc_repository", + "level": NotificationLevel.ERROR, + "attachment_content": trace + } + + if data_type not in data_type_map: + return False, { + "notification_id": f"OPC_WRITE_DATA_TYPE_ERROR_{self.id}", + "message": f"Unsupported data type: {data_type} | metadata: {metadata}", + "block": "opc_repository", + "level": NotificationLevel.ERROR + } + + data = data_type_map[data_type]['converter'](value) + logger.custom_info( + f'Writing {data} - {type(data)} to {node}', metadata) + now = datetime.now() + ua_data = DataValue( + Variant(data, data_type_map[data_type]['opc_type']), + SourceTimestamp=DateTime( + now.year, + now.month, + now.day, + now.hour, + now.minute, + now.second, + now.microsecond + ) + ) + + try: + await node_obj.write_value(ua_data) + + metrics.PREDICTION_OPC_WRITING_COUNT.labels( + pod_id=self.pod_id, + model_name=metadata['model_name'], + pipeline_name=metadata['workflow_name'], + opc_server_id=self.id + ).inc() + + end_time = time.time() + response_time = end_time - start_time + metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels( + pod_id=self.pod_id, + model_name=metadata['model_name'], + pipeline_name=metadata['workflow_name'], + opc_server_id=self.id + ).observe(response_time) + + except Exception as e: + trace = traceback.format_exc() + logger.custom_error(trace, metadata) + self.error_count += 1 + return False, { + "notification_id": f"OPC_WRITE_DATA_ERROR_{self.id}", + "message": f"Failed to write data to OPC server: {e} | metadata: {metadata}", + "block": "opc_repository", + "level": NotificationLevel.ERROR, + "attachment_content": trace + } + self.error_count = 0 + + return True, {} diff --git a/model-manager/worker/__init__.py b/model-manager/worker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/worker/worker.py b/model-manager/worker/worker.py new file mode 100644 index 0000000..567a257 --- /dev/null +++ b/model-manager/worker/worker.py @@ -0,0 +1,237 @@ +""" +Laborious Worker Module + +This module provides the main worker implementation for the Sientia DataOps Laborious 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: laborious) +- 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: laborious) +""" + +from temporalio import workflow, client +from temporalio.worker import Worker, PollerBehaviorAutoscaling +from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig + +with workflow.unsafe.imports_passed_through(): + import os + import sys + import asyncio + from laborious.workflows.minimal_retrain import MinimalRetrain + from laborious.workflows.predictions_batch import PredictionsBatch + from laborious.workflows.sub_workflows.prediction_process import PredictionProcess + from laborious.workflows.sub_workflows.format_and_export_prediction import \ + FormatAndExportPrediction + from laborious.activities.activities import Activities + from laborious.utils.connectors_config import ( + build_postgres_config, + build_mlflow_config, + build_opc_config, + build_mongodb_config + ) + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.observability.logger import get_logger + from laborious import metrics + from prometheus_client import start_http_server + +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 Laborious 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. Initializes OPC connections + 6. Starts Temporal client and workers + 7. 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', 'laborious') + ) + + logger.custom_info('Starting Activities...', metadata) + + activities = Activities( + postgres_config=build_postgres_config(), + mlflow_config=build_mlflow_config(), + opc_config=build_opc_config(), + logger=logger, + notification_handler=notification_handler + ) + + logger.custom_info('Initializing OPC...', metadata) + await activities.init_opc() + + 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', 'laborious'), + 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, + # OPC + activities.write_opc_data, + # 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: # NOSONAR + 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 Laborious 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: + print(f"Failed to start Prometheus server: {e}") + os._exit(1) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/model-manager/workflows/__init__.py b/model-manager/workflows/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/workflows/minimal_retrain.py b/model-manager/workflows/minimal_retrain.py new file mode 100644 index 0000000..1893e25 --- /dev/null +++ b/model-manager/workflows/minimal_retrain.py @@ -0,0 +1,116 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from laborious.activities.activities import Activities + from typing import Any + from sientia_do.temporal.policies import retry_policy + from datetime import timedelta + + +@workflow.defn(name="minimal_retrain") +class MinimalRetrain(): + """ + Automated model retraining workflow for the Laborious 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) + ) diff --git a/model-manager/workflows/predictions_batch.py b/model-manager/workflows/predictions_batch.py new file mode 100644 index 0000000..e522eaa --- /dev/null +++ b/model-manager/workflows/predictions_batch.py @@ -0,0 +1,125 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from laborious.activities.activities import Activities + from typing import Any + from sientia_do.temporal.policies import retry_policy + from datetime import timedelta + + +@workflow.defn(name="predictions_batch") +class PredictionsBatch(): + """ + Main batch prediction workflow for the Laborious 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 + - opc_output_config (dict, optional): OPC server export 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']), + 'opc_output_config': input_data.get('opc_output_config', {}), + 'prediction_store_policy': input_data.get( + 'prediction_store_policy', 'lts:1') + } + + # Execute prediction process workflow + await workflow.execute_child_workflow( + 'prediction_process', prediction_input) diff --git a/model-manager/workflows/sub_workflows/__init__.py b/model-manager/workflows/sub_workflows/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model-manager/workflows/sub_workflows/format_and_export_prediction.py b/model-manager/workflows/sub_workflows/format_and_export_prediction.py new file mode 100644 index 0000000..8e7df07 --- /dev/null +++ b/model-manager/workflows/sub_workflows/format_and_export_prediction.py @@ -0,0 +1,140 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from laborious.activities.activities import Activities + from typing import Any + from datetime import timedelta + from sientia_do.temporal.policies import retry_policy + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + + +@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, OPC server export, and metrics recording. + It implements flexible formatting based on prediction quality and provides + comprehensive export capabilities to multiple destinations. + + 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 + - OPC Servers: Real-time industrial system integration + - 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. Exporting data to OPC servers for real-time industrial access + 4. Persisting data to PostgreSQL database with comprehensive metadata + 5. 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 + - Comprehensive export: Multi-destination data distribution + + 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 + - opc_output_config (dict[str, Any]): OPC server export configuration + - 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 opc + prediction = await workflow.execute_activity_method( + Activities.write_opc_data, + { + **metadata, + 'opc_output_config': input_data['opc_output_config'], + 'data': prediction + }, + 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) + ) diff --git a/model-manager/workflows/sub_workflows/prediction_process.py b/model-manager/workflows/sub_workflows/prediction_process.py new file mode 100644 index 0000000..777fa1c --- /dev/null +++ b/model-manager/workflows/sub_workflows/prediction_process.py @@ -0,0 +1,296 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from laborious.activities.activities import Activities + from typing import Any + from sientia_do.temporal.policies import retry_policy + from datetime import timedelta + + +@workflow.defn(name="prediction_process") +class PredictionProcess(): + """ + Core prediction processing workflow for the Laborious 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 + - opc_output_config (dict): OPC server export 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, + 'opc_output_config': input_data['opc_output_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, + 'opc_output_config': input_data['opc_output_config'], + 'prediction_store_policy': input_data['prediction_store_policy'] + } + ) + return True + + return False diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..1ce9f9c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +temporalio +psycopg2-binary +sqlalchemy +asyncua +redis +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 diff --git a/run_coverage.sh b/run_coverage.sh new file mode 100755 index 0000000..f9af4cb --- /dev/null +++ b/run_coverage.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Exit on any error +set -e + +echo "Activating virtual environment..." +source ./venv/bin/activate + +pytest --cov=laborious --cov-report=html + +xdg-open htmlcov/index.html \ No newline at end of file diff --git a/run_local.sh b/run_local.sh new file mode 100755 index 0000000..2bbd5c2 --- /dev/null +++ b/run_local.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +# Exit on any error +set -e + +echo "Activating virtual environment..." +source ./venv/bin/activate + +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 laborious.worker.worker diff --git a/simulator/Dockerfile b/simulator/Dockerfile new file mode 100644 index 0000000..d467676 --- /dev/null +++ b/simulator/Dockerfile @@ -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"] diff --git a/sonar-project.properties b/sonar-project.properties new file mode 100644 index 0000000..b332c70 --- /dev/null +++ b/sonar-project.properties @@ -0,0 +1,11 @@ +sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_beaec423-6c42-4f26-8134-b676287b499d +sonar.projectName=sientia-dataops-laborious_temporal +sonar.sources=laborious +sonar.tests=tests +sonar.projectVersion=1.0.0 +sonar.coverage.exclusions=laborious/worker/worker.py +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 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/__init__.py b/tests/laborious/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/activities/__init__.py b/tests/laborious/activities/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py new file mode 100644 index 0000000..de418be --- /dev/null +++ b/tests/laborious/activities/test_activities.py @@ -0,0 +1,135 @@ +from pytest import mark +from unittest.mock import patch, MagicMock, ANY +from sientia_do.temporal.activities.postgres import Postgres +from laborious.activities.activities import Activities +from laborious.activities.mlflow import MLFlow +from laborious.activities.gates import Gates +from laborious.activities.opc import OPC + + +@patch('laborious.activities.activities.Postgres.__init__') +@patch('laborious.activities.activities.MLFlow.__init__') +@patch('laborious.activities.activities.OPC.__init__') +@patch('laborious.activities.activities.Gates.__init__') +def test___init__(mock_gates_init, mock_opc_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' + } + + opc_config = { + 'bootstrap_servers': 'localhost:9092', + 'polling_time': 1000, + 'group_id': 'test-group' + } + + logger = MagicMock() + notification_handler = MagicMock() + + activities = Activities( + postgres_config=postgres_config, + mlflow_config=mlflow_config, + opc_config=opc_config, + logger=logger, + notification_handler=notification_handler + ) + + assert isinstance(activities, Activities) + assert isinstance(activities, Postgres) + assert isinstance(activities, MLFlow) + assert isinstance(activities, OPC) + 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_opc_init.assert_called_once_with( + ANY, + opc_servers=opc_config, + logger=logger, + notification_handler=notification_handler + ) + + mock_gates_init.assert_called_once_with( + ANY, + logger=logger, + notification_handler=notification_handler + ) + + +@mark.asyncio +@patch('laborious.activities.activities.Postgres', return_value=MagicMock()) +@patch('laborious.activities.activities.MLFlow', return_value=MagicMock()) +@patch('laborious.activities.activities.OPC', return_value=MagicMock()) +async def test_shutdown(mock_opc_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' + } + + opc_config = { + 'bootstrap_servers': 'localhost:9092', + 'polling_time': 1000, + 'group_id': 'test-group' + } + + logger = MagicMock() + notification_handler = MagicMock() + + activities = Activities( + postgres_config=postgres_config, + mlflow_config=mlflow_config, + opc_config=opc_config, + logger=logger, + notification_handler=notification_handler + ) + + await activities.shutdown() + mock_opc_init.shutdown.assert_called_once() + mock_postgres_init.close.assert_called_once() diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py new file mode 100644 index 0000000..a65c6d4 --- /dev/null +++ b/tests/laborious/activities/test_gates.py @@ -0,0 +1,597 @@ +from unittest.mock import MagicMock, ANY, patch +from pytest import fixture, mark +from sientia_do.notifications.models import NotificationLevel +from laborious.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('laborious.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('laborious.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('laborious.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: + result = await gates_activity.format_prediction(input_data) + except ValueError as e: + assert str(e) == "Invalid policy type: invalid" + else: + assert False, "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('laborious.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 + ) diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py new file mode 100644 index 0000000..85cb7f1 --- /dev/null +++ b/tests/laborious/activities/test_mlflow.py @@ -0,0 +1,297 @@ +from datetime import datetime +from unittest.mock import ANY, MagicMock, patch + +import numpy as np +from pandas import DataFrame, Timestamp +from pytest import fixture, mark, raises +from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ +from laborious.activities.mlflow import MLFlow +from sientia_do.notifications.models import NotificationLevel + + +@patch("laborious.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("laborious.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("laborious.activities.mlflow.DataFrame") +@patch("laborious.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("laborious.activities.mlflow.DataFrame") +@patch("laborious.activities.mlflow.to_datetime") +@patch("laborious.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: + 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: + assert False, "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: + 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: + assert False, "No exception raised" diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py new file mode 100644 index 0000000..07f00b8 --- /dev/null +++ b/tests/laborious/activities/test_opc.py @@ -0,0 +1,369 @@ +from unittest.mock import patch, MagicMock, ANY, call, AsyncMock +from pandas import DataFrame +from pytest import fixture, mark +import pytest_asyncio +from sientia_do.notifications.models import NotificationLevel + +from laborious.activities.opc import OPC + +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + + +def test__init__(): + servers = { + 'server1': 'config' + } + opc = OPC( + opc_servers=servers, + logger=MagicMock(), + notification_handler=MagicMock() + ) + + assert opc.opc_servers == servers + assert opc.opc_repository == {} + + +@mark.asyncio +@patch("laborious.activities.opc.OpcRepository") +@patch("laborious.activities.opc.OPC.send_notification") +async def test_init_opc(mock_send_notification, mock_opc_repository): + mock_logger = MagicMock() + server1 = MagicMock( + connect=AsyncMock(return_value=(True, {})), + write_data=AsyncMock(return_value=(True, {})) + ) + server2 = MagicMock( + connect=AsyncMock(return_value=(True, {})), + write_data=AsyncMock(return_value=(True, {})) + ) + server3 = MagicMock( + connect=AsyncMock(return_value=(False, { + 'notification_id': 'OPC_CONNECTION_ERROR_server3', + 'message': 'Failed to connect to OPC server: Test error', + 'block': 'opc_repository', + 'level': NotificationLevel.ERROR, + 'attachment_content': 'Test error' + })), + write_data=AsyncMock(return_value=(True, {})) + ) + mock_opc_repository.side_effect = [server1, server2, server3] + mock_notification_handler = MagicMock() + servers = { + 'server1': { + 'id': 'server1', + 'url': 'http://localhost:8080', + 'server_uri': 'opc.tcp://localhost:4840', + 'cert_path': '', + 'private_key_path': '', + 'server_cert_path': '', + 'reconnection_interval': 60, + }, + 'server2': { + 'id': 'server2', + 'url': 'http://localhost:8080', + 'server_uri': 'opc.tcp://localhost:4840', + 'cert_path': '', + 'private_key_path': '', + 'server_cert_path': '', + 'reconnection_interval': 60, + }, + 'server3': { + 'id': 'server3', + 'url': 'http://localhost:8080', + 'server_uri': 'opc.tcp://localhost:4840', + 'cert_path': '', + 'private_key_path': '', + 'server_cert_path': '', + 'reconnection_interval': 60, + } + } + opc = OPC( + opc_servers=servers, + logger=mock_logger, + notification_handler=mock_notification_handler + ) + await opc.init_opc() + + assert opc.opc_servers == servers + assert opc.logger == mock_logger + assert opc.notification_handler == mock_notification_handler + assert opc.opc_repository['server1'] == server1 + assert opc.opc_repository['server2'] == server2 + + mock_opc_repository.assert_has_calls([ + call( + id="server1", + url="http://localhost:8080", + logger=mock_logger, + server_uri="opc.tcp://localhost:4840", + cert_path="", + private_key_path="", + server_cert_path="", + notification_handler=mock_notification_handler, + reconnection_interval=60, + pod_id='localhost' + ), + ]) + mock_opc_repository.assert_has_calls([ + call( + id="server2", + url="http://localhost:8080", + logger=mock_logger, + server_uri="opc.tcp://localhost:4840", + cert_path="", + private_key_path="", + server_cert_path="", + notification_handler=mock_notification_handler, + reconnection_interval=60, + pod_id='localhost' + ) + ]) + + server1.connect.assert_called_once() + server2.connect.assert_called_once() + + mock_send_notification.assert_has_calls([ + call( + metadata={ + 'model_id': '-', + 'model_name': '-', + 'workflow_name': '-', + 'schedule_name': 'INITIALIZATION' + }, + notification_id="OPC_CONNECTION_ERROR_server3", + message="Failed to connect to OPC server: Test error", + block="opc_repository", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + ]) + + +@pytest_asyncio.fixture +@patch("laborious.activities.opc.OpcRepository") +async def opc(mock_opc_repository): + servers = { + 'server1': { + 'id': 'server1', + 'url': 'http://localhost:8080', + 'server_uri': 'opc.tcp://localhost:4840', + 'cert_path': '', + 'private_key_path': '', + 'server_cert_path': '', + 'reconnection_interval': 60, + } + } + + mock_opc_repository.return_value.write_data = AsyncMock( + return_value=(True, {}) + ) + mock_opc_repository.return_value.connect = AsyncMock( + return_value=(True, {}) + ) + opc = OPC( + opc_servers=servers, + logger=MagicMock(), + notification_handler=MagicMock() + ) + await opc.init_opc() + opc.send_notification = MagicMock() + return opc + + +WRITE_DATA_CASES = [ + ('tag1', 'int', 50), + ('tag2', 'float', 50.5), + ('tag3', 'bool', True), + ('tag4', 'string', 'test'), +] + + +@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES) +@mark.asyncio +async def test_write_data_success(opc, tag, data_type, data): + result = await opc.write_data(server_id='server1', tag=tag, data=data, + data_type=data_type, tag_type='prediction', metadata=metadata) + assert result is True + opc.opc_repository['server1'].write_data.assert_called_once_with( + tag, data, data_type, opc.logger, metadata) + + +@mark.asyncio +async def test_write_data_failed(opc): + opc.opc_repository['server1'].write_data.return_value = (False, { + 'notification_id': 'OPC_WRITE_DATA_ERROR_server1', + 'message': 'Failed to write data to OPC server: Test error', + 'block': 'opc_repository', + 'level': NotificationLevel.ERROR, + 'attachment_content': 'Test error' + }) + + result = await opc.write_data(server_id='server1', tag='tag1', data=50, + data_type='int', tag_type='prediction', metadata=metadata) + assert result is False + + opc.send_notification.assert_called_once_with( + metadata=metadata, + notification_id="OPC_WRITE_DATA_ERROR_server1", + message="Failed to write data to OPC server: Test error", + block="opc_repository", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + +@mark.asyncio +async def test_write_data_exception(opc): + opc.opc_repository['server1'].write_data.side_effect = Exception( + "Test error") + + try: + await opc.write_data(server_id='server1', tag='tag1', data=50, + data_type='int', tag_type='prediction', metadata=metadata) + + except Exception: + opc.send_notification.assert_called_once_with( + metadata=metadata, + notification_id="WRITE_OPC_PREDICTION_ERROR", + message="Error writing data to OPC server: Test error", + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + +@mark.asyncio +async def test_write_opc_data_success(opc): + # Arrange + input_data = { + **metadata, + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_output_config': { + 'server1': { + 'prediction_tags': { + 'tag1': {'data_type': 'float'} + }, + 'confidence_tags': { + 'tag2': {'data_type': 'float'} + } + } + } + } + + # Act + opc.write_data = AsyncMock(return_value=True) + opc.process_confidence = MagicMock(return_value={'data': 'data'}) + output = await opc.write_opc_data(input_data) + + # Assert + assert output == {'data': 'data'} + opc.write_data.assert_has_calls([ + call( + server_id='server1', + tag='tag1', + data=0.75, + data_type='float', + tag_type='prediction', + metadata=metadata['metadata'] + )]) + opc.write_data.assert_has_calls([ + call( + server_id='server1', + tag='tag2', + data=0.95, + data_type='float', + tag_type='confidence', + metadata=metadata['metadata'] + ) + ]) + assert opc.write_data.call_count == 2 + + +@mark.asyncio +async def test_write_opc_data_empty_config(opc): + # Arrange + input_data = { + **metadata, + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_servers': ['server1'], + 'opc_output_config': { + 'server1': { + 'prediction_tags': {}, + 'confidence_tags': {} + } + } + } + + # Act + await opc.write_opc_data(input_data) + + # Assert + opc.opc_repository['server1'].write_data.assert_not_called() + + +@mark.asyncio +async def test_write_opc_data_no_validate_server(opc): + opc.validate_server = MagicMock(return_value=False) + input_data = { + **metadata, + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_output_config': { + 'server1': { + 'prediction_tags': { + 'tag1': {'data_type': 'float'} + }, + 'confidence_tags': { + 'tag2': {'data_type': 'float'} + } + } + } + } + + # Act + await opc.write_opc_data(input_data) + + # Assert + opc.opc_repository['server1'].write_data.assert_not_called() + + +@mark.parametrize('data,success,expected', [ + (DataFrame({'prediction_confidence': [0]}), True, 0), + (DataFrame({'prediction_confidence': [0]}), False, 12), +]) +def test_process_confidence(opc, data, success, expected): + # Act + result = opc.process_confidence(data, success, metadata) + + # Assert + assert result['prediction_confidence'][0] == expected + + +def test_validate_server(opc): + assert opc.validate_server('server1', metadata) is True + assert opc.validate_server('server2', metadata) is False + + +@mark.asyncio +async def test_shutdown(opc): + opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True) + await opc.shutdown() + opc.opc_repository['server1'].disconnect.assert_called_once() diff --git a/tests/laborious/utils/__init__.py b/tests/laborious/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/utils/filters/__init__.py b/tests/laborious/utils/filters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py new file mode 100644 index 0000000..405bc9b --- /dev/null +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -0,0 +1,30 @@ +from pandas import DataFrame + +from laborious.utils.filters.conditional_filters import ( + filter_specific_variables_null_values, + filter_empty_data +) + + +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 diff --git a/tests/laborious/utils/filters/test_mlflow_filters.py b/tests/laborious/utils/filters/test_mlflow_filters.py new file mode 100644 index 0000000..f9c61e9 --- /dev/null +++ b/tests/laborious/utils/filters/test_mlflow_filters.py @@ -0,0 +1,22 @@ +from pandas import DataFrame +from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter + + +def test_api_error_filter_invalid_response(): + assert api_error_filter(None, {}) == True # NOSONAR + + +def test_api_error_filter_valid_response_fail(): + assert api_error_filter({'success': False}, {}) == True + + +def test_api_error_filter_valid_response_success(): + assert api_error_filter({'success': True}, {}) == False + + +def test_nan_values_filter_all_nan_values(): + assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True + + +def test_nan_values_filter_no_nan_values(): + assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py new file mode 100644 index 0000000..cdc59b4 --- /dev/null +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -0,0 +1,502 @@ +from unittest.mock import ANY, MagicMock, call, patch +import numpy as np +from pandas import DataFrame +import pytest +from datetime import datetime, timezone +from pandas import Timestamp +from laborious.utils.repository.model_repository import MLFlowRepository + + +@pytest.fixture +def mlflow_repository(): + with patch('laborious.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=timezone.utc): 1, + datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.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=timezone.utc): 1, + Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.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('laborious.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('laborious.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('laborious.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('laborious.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: + assert False + + +@patch('laborious.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('laborious.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: + assert False + + +@patch('laborious.utils.repository.model_repository.mlflow.sklearn') +@patch('laborious.utils.repository.model_repository.mlflow.set_experiment') +def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): + + mlflow_repository.model_serving.get_model_run_id = 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_run_id.assert_called_once_with( + 'test', stage='Production') + 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('laborious.utils.repository.model_repository.mlflow.start_run') +@patch('laborious.utils.repository.model_repository.mlflow.log_param') +@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model') +@patch('laborious.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('laborious.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('laborious.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: + assert str(e) == 'Model versions is not a list' + else: + assert False + + +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', + } diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py new file mode 100644 index 0000000..bc84db8 --- /dev/null +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -0,0 +1,388 @@ +import pytest +from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call +from asyncua.crypto.security_policies import SecurityPolicyBasic256 +from laborious.utils.repository.opc_repository import OpcRepository +from sientia_do.notifications.models import NotificationLevel +from datetime import datetime + + +@pytest.fixture +def mock_logger(): + return Mock() + + +@pytest.fixture +def opc_repository(mock_logger): + return OpcRepository( + id="test_repo", + url="opc.tcp://localhost:4840", + logger=mock_logger, + notification_handler=Mock(), + reconnection_interval=60, + server_uri="urn:test:server", + cert_path="/path/to/cert.pem", + private_key_path="/path/to/key.pem", + server_cert_path="/path/to/server_cert.pem" + ) + + +@pytest.fixture +def mock_client(): + with patch('laborious.utils.repository.opc_repository.Client') as mock: + client_instance = AsyncMock() + mock.return_value = client_instance + yield client_instance + + +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + + +def test_init(opc_repository): + assert opc_repository.id == "test_repo" + assert opc_repository.url == "opc.tcp://localhost:4840" + assert opc_repository.server_uri == "urn:test:server" + assert opc_repository.cert_path == "/path/to/cert.pem" + assert opc_repository.private_key_path == "/path/to/key.pem" + assert opc_repository.server_cert_path == "/path/to/server_cert.pem" + assert opc_repository.reconnection_interval == 60 + assert opc_repository.client is None + assert opc_repository.last_reconnection_time is None + assert opc_repository.error_count == 0 + + +@pytest.mark.asyncio +async def test_set_security(opc_repository, mock_client): + opc_repository.client = mock_client + await opc_repository.set_security() + + mock_client.application_uri = "urn:test:server" + mock_client.set_security.assert_called_once_with( + SecurityPolicyBasic256, + certificate="/path/to/cert.pem", + private_key="/path/to/key.pem", + server_certificate="/path/to/server_cert.pem" + ) + assert mock_client.secure_channel_timeout == 10000000 + assert mock_client.session_timeout == 10000000 + + +@pytest.mark.asyncio +async def test_set_security_missing_certificates(opc_repository): + opc_repository.cert_path = None + opc_repository.private_key_path = None + + try: + await opc_repository.set_security() + except ValueError as e: + assert str( + e) == "Certificate and private key paths must be provided for secure connection." + + +@pytest.mark.asyncio +async def test_connect_with_security(opc_repository, mock_client): + opc_repository.try_connect = AsyncMock(return_value=(True, {})) + result = await opc_repository.connect() + + opc_repository.try_connect.assert_called_once() + assert opc_repository.client == mock_client + assert result == (True, {}) + + +@pytest.mark.asyncio +async def test_connect_without_security(opc_repository, mock_client): + opc_repository.cert_path = None + opc_repository.try_connect = AsyncMock(return_value=(True, {})) + opc_repository.set_security = AsyncMock() + result = await opc_repository.connect() + + opc_repository.try_connect.assert_called_once() + opc_repository.set_security.assert_not_called() + assert opc_repository.client == mock_client + assert result == (True, {}) + + +@pytest.mark.asyncio +async def test_try_connect_success(opc_repository): + opc_repository.last_reconnection_time = None + opc_repository.client = AsyncMock() + result = await opc_repository.try_connect() + + opc_repository.client.connect.assert_called_once() + assert opc_repository.last_reconnection_time is not None + assert result == (True, {}) + + +@pytest.mark.asyncio +async def test_try_connect_fail(opc_repository): + opc_repository.last_reconnection_time = None + opc_repository.client = MagicMock() + opc_repository.client.connect.side_effect = Exception("Test error") + + is_connected, error_data = await opc_repository.try_connect() + + opc_repository.client.connect.assert_called_once() + assert is_connected is False + assert error_data['notification_id'] == f"OPC_CONNECTION_ERROR_{opc_repository.id}" + assert error_data['message'] == "Failed to connect to OPC server: Test error" + assert error_data['block'] == "opc_repository" + assert error_data['level'] == NotificationLevel.ERROR + assert error_data['attachment_content'] is not None + + +@pytest.mark.asyncio +async def test_disconnect(opc_repository, mock_client): + opc_repository.client = mock_client + await opc_repository.disconnect() + + mock_client.disconnect.assert_called_once() + assert opc_repository.client is None + + +@pytest.mark.asyncio +async def test_disconnect_no_client(opc_repository): + opc_repository.client = None + assert await opc_repository.disconnect() is None + + +@pytest.mark.asyncio +async def test_disconnect_error(opc_repository, mock_client): + opc_repository.client = mock_client + mock_client.disconnect.side_effect = Exception("Test error") + await opc_repository.disconnect() + + opc_repository.logger.custom_error.assert_called_once_with( + "Failed to disconnect from OPC server: Test error", + ANY + ) + assert opc_repository.client is None + + +@pytest.mark.asyncio +async def test_validate_connection_none_client(opc_repository): + opc_repository.client = None + opc_repository.connect = AsyncMock(return_value=(True, {})) + response = await opc_repository.validate_connection() + assert response == (True, {}) + opc_repository.connect.assert_called_once() + + +@pytest.mark.asyncio +async def test_validate_connection_error_count_disconnect_error(opc_repository): + opc_repository.error_count = 6 + opc_repository.client = AsyncMock() + opc_repository.disconnect = AsyncMock( + side_effect=Exception("Test error") + ) + opc_repository.connect = AsyncMock(return_value=(True, {})) + + response = await opc_repository.validate_connection() + assert response == opc_repository.connect.return_value + opc_repository.disconnect.assert_called_once() + opc_repository.connect.assert_called_once() + opc_repository.logger.custom_error.assert_has_calls( + [ + call("Failed to disconnect from OPC server: Test error", ANY), + ] + ) + + +@pytest.mark.asyncio +async def test_validate_connection_error_validate_connection_error(opc_repository): + opc_repository.client = MagicMock( + uaclient=Exception("Test error") + ) + opc_repository.error_count = 0 + + response = await opc_repository.validate_connection() + + assert response == (False, { + "notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}", + "message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'", + "block": "opc_repository", + "level": NotificationLevel.ERROR, + "attachment_content": ANY + }) + + +@pytest.mark.asyncio +@patch('laborious.utils.repository.opc_repository.datetime') +async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository): + _mock_datetime.now = MagicMock( + return_value=datetime(2025, 1, 1, 0, 0, 0)) + opc_repository.error_count = 0 + opc_repository.client = MagicMock() + opc_repository.client.uaclient.protocol = None + opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) + opc_repository.connect = MagicMock(return_value=(True, {})) + + response = await opc_repository.validate_connection() + opc_repository.connect.assert_not_called() + assert response == (False, { + "notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}", + "message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...", + "block": "opc_repository", + "level": NotificationLevel.WARNING + }) + + +@pytest.mark.asyncio +@patch('laborious.utils.repository.opc_repository.datetime') +async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository): + mock_datetime.now = MagicMock( + return_value=datetime(2025, 1, 1, 1, 0, 0)) + opc_repository.error_count = 0 + opc_repository.client = AsyncMock() + opc_repository.client.uaclient.protocol = None + opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) + opc_repository.connect = AsyncMock(return_value=(True, {})) + + response = await opc_repository.validate_connection() + opc_repository.connect.assert_called_once() + assert response == opc_repository.connect.return_value + + +@pytest.mark.asyncio +async def test_validate_connection_success(opc_repository): + opc_repository.client = MagicMock() + opc_repository.error_count = 0 + opc_repository.client.uaclient.protocol = MagicMock() + opc_repository.client.uaclient.protocol.state = "open" + + output = await opc_repository.validate_connection() + assert output == (True, {}) + + +@pytest.mark.asyncio +async def test_write_data_validate_connection_do_nothing(opc_repository): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = AsyncMock( + get_node=MagicMock() + ) + mock_node = AsyncMock() + opc_repository.client.get_node.return_value = mock_node + + result = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + + opc_repository.validate_connection.assert_called_once() + opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") + assert result == (True, {}) + + +@pytest.mark.asyncio +async def test_write_data_validate_connection_failed(opc_repository): + opc_repository.validate_connection = AsyncMock(return_value=(False, {})) + opc_repository.client = AsyncMock() + opc_repository.error_count = 0 + + result = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + + opc_repository.validate_connection.assert_called_once() + opc_repository.client.get_node.assert_not_called() + assert result == (False, {}) + + +@pytest.mark.asyncio +async def test_write_data_get_node_failed(opc_repository): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = AsyncMock() + opc_repository.error_count = 0 + opc_repository.client.get_node = MagicMock( + side_effect=Exception("Test error")) + + is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + + opc_repository.validate_connection.assert_called_once() + opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") + assert is_success is False + assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}" + assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}" + assert error_data['block'] == "opc_repository" + assert error_data['level'] == NotificationLevel.ERROR + assert error_data['attachment_content'] is not None + + +@pytest.mark.asyncio +async def test_write_data_invalid_data_type(opc_repository, mock_client): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = mock_client + mock_node = AsyncMock() + mock_client.get_node = MagicMock(return_value=mock_node) + + is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "invalid_type", opc_repository.logger, metadata['metadata']) + + opc_repository.validate_connection.assert_called_once() + mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") + + assert is_success is False + assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}" + assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}" + assert error_data['block'] == "opc_repository" + assert error_data['level'] == NotificationLevel.ERROR + assert error_data.get('attachment_content') is None + + +@pytest.mark.asyncio +@patch('laborious.utils.repository.opc_repository.metrics') +async def test_write_data(mock_metrics, opc_repository, mock_client): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = mock_client + mock_node = AsyncMock() + mock_client.get_node = MagicMock(return_value=mock_node) + + result = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + + mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") + mock_node.write_value.assert_called_once() + assert result == (True, {}) + + mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_called_once_with( + pod_id=opc_repository.pod_id, + model_name=metadata['metadata']['model_name'], + pipeline_name=metadata['metadata']['workflow_name'], + opc_server_id=opc_repository.id + ) + mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with() + + mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_called_once_with( + pod_id=opc_repository.pod_id, + model_name=metadata['metadata']['model_name'], + pipeline_name=metadata['metadata']['workflow_name'], + opc_server_id=opc_repository.id + ) + mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with( + ANY) + + +@pytest.mark.asyncio +async def test_write_data_write_value_failed(opc_repository, mock_client): + opc_repository.validate_connection = AsyncMock(return_value=(True, {})) + opc_repository.client = mock_client + mock_node = AsyncMock() + opc_repository.error_count = 0 + mock_client.get_node = MagicMock(return_value=mock_node) + mock_node.write_value.side_effect = Exception("Test error") + + is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata['metadata']) + + opc_repository.validate_connection.assert_called_once() + mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") + mock_node.write_value.assert_called_once() + assert is_success is False + assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}" + assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}" + assert error_data['block'] == "opc_repository" + assert error_data['level'] == NotificationLevel.ERROR + assert error_data['attachment_content'] is not None diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py new file mode 100644 index 0000000..910439c --- /dev/null +++ b/tests/laborious/utils/test_connectors_config.py @@ -0,0 +1,161 @@ +from os import environ +from laborious.utils.connectors_config import (build_mlflow_config, + build_opc_config, + build_postgres_config, + build_mongodb_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_opc_config_with_env_vars(): + # Arrange + environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}' + + # Act + config = build_opc_config() + + # Assert + assert config['opc']['name'] == 'test-opc' + assert config['opc']['url'] == 'opc.tcp://test:4840' + + +def test_build_opc_config_with_individual_env_vars(): + # Arrange + environ.pop('OPC_CONFIG', None) + environ['OPC_ID'] = '1' + environ['OPC_URL'] = 'opc.tcp://test:4840' + environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840' + environ['OPC_RECONNECTION_INTERVAL'] = '300' + + # Act + config = build_opc_config() + + # Assert + assert config['1']['id'] == '1' + assert config['1']['url'] == 'opc.tcp://test:4840' + assert config['1']['server_uri'] == 'opc.tcp://test:4840' + assert config['1']['reconnection_interval'] == 300 + + +def test_build_opc_config_with_defaults(): + # Arrange + environ.pop('OPC_CONFIG', None) + environ.pop('OPC_ID', None) + environ.pop('OPC_URL', None) + environ.pop('OPC_SERVER_URI', None) + environ.pop('OPC_RECONNECTION_INTERVAL', None) + + # Act + config = build_opc_config() + + # Assert + assert config['1']['id'] == '1' + assert config['1']['url'] == 'opc.tcp://localhost:4840' + assert config['1']['server_uri'] == 'opc.tcp://localhost:4840' + assert config['1']['reconnection_interval'] == 120 + + +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 + } diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py new file mode 100644 index 0000000..a8e6e20 --- /dev/null +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -0,0 +1,160 @@ +from unittest.mock import call, patch, AsyncMock, ANY +from pytest import mark, fixture + +from laborious.activities.activities import Activities +from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + + +@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("laborious.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", + "opc_servers": ["test_server"], + "opc_output_config": {"test": "config"}, + "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.write_opc_data, + { + 'opc_output_config': input_data['opc_output_config'], + 'data': workflow_mock.execute_local_activity_method.return_value, + **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_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 == 3 + assert workflow_mock.execute_local_activity_method.call_count == 1 + + +@mark.asyncio +@patch("laborious.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", + "opc_servers": ["test_server"], + "opc_output_config": {"test": "config"}, + "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.write_opc_data, + { + 'opc_output_config': input_data['opc_output_config'], + 'data': workflow_mock.execute_local_activity_method.return_value, + **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_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 == 3 + assert workflow_mock.execute_local_activity_method.call_count == 1 diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py new file mode 100644 index 0000000..df60ada --- /dev/null +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -0,0 +1,601 @@ +from unittest.mock import AsyncMock, patch, call, ANY +from pytest import fixture, mark +from laborious.activities.activities import Activities +from laborious.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("laborious.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'], + 'opc_output_config': {'test': 'config'}, + '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'], + 'opc_output_config': input_data['opc_output_config'], + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'comment': 'Error', + 'prediction_store_policy': input_data['prediction_store_policy'] + } + ) + + +@mark.asyncio +@patch("laborious.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'], + 'opc_output_config': {'test': 'config'} + } + + # 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("laborious.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'], + 'opc_output_config': {'test': 'config'} + } + + # 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("laborious.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'], + 'opc_output_config': {'test': 'config'} + } + + # 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("laborious.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'], + 'opc_output_config': {'test': 'config'} + } + + # 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("laborious.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("laborious.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("laborious.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, + 'opc_output_config': {'test': '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', + 'opc_output_config': {'test': 'config'}, + 'prediction_store_policy': prediction_store_policy + } + ) + + +@mark.asyncio +@patch("laborious.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, + 'opc_output_config': {'test': '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() diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py new file mode 100644 index 0000000..b3b03b5 --- /dev/null +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -0,0 +1,98 @@ +from unittest.mock import AsyncMock, MagicMock, call, patch, ANY +from pytest import fixture, mark +from laborious.activities.activities import Activities +from laborious.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('laborious.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 + ) + ]) diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py new file mode 100644 index 0000000..90d7d21 --- /dev/null +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -0,0 +1,88 @@ +from unittest.mock import AsyncMock, call, patch, ANY +from pytest import fixture, mark +from laborious.activities.activities import Activities +from laborious.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('laborious.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', + 'opc_output_config': 'test_opc_output_config', + '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']), + 'opc_output_config': input_data.get('opc_output_config', {}), + 'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1') + } + + workflow_mock.execute_child_workflow.assert_has_calls([ + call( + 'prediction_process', prediction_input) + ]) diff --git a/values.yaml b/values.yaml new file mode 100644 index 0000000..4eb83d7 --- /dev/null +++ b/values.yaml @@ -0,0 +1,229 @@ +# 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-laborious-worker" +fullnameOverride: "sientia-laborious-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-laborious-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 "laborious.worker.worker" + initialDelaySeconds: 20 + periodSeconds: 30 + +readinessProbe: + exec: + command: + - sh + - -c + - pgrep -f "laborious.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-laborious_temporal.git" + - name: GITHUB_BRANCH + value: SIENTIAPDE-1222-ajustar-a-library-para-fazer-o-download-do-courier + - name: PYTHON_APP + value: "laborious.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: OPC_ID + value: "1" + - name: OPC_URL + value: "opc.tcp://sientia-opc-simulator-opc.sientia.svc.cluster.local:4840" + + - 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-laborious" + + - name: TEMPORAL_HOST + value: "temporal-frontend.temporal.svc.cluster.local:7233" + - name: TEMPORAL_NAMESPACE + value: "laborious" + + - 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-laborious-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-laborious-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.5.0 + +# kubectl create secret generic git-ssh-key-sientia-laborious-worker \ +# --namespace sientia \ +# --from-file=ssh-privatekey=git_key \ +# --type=kubernetes.io/ssh-auth