From 93d0849c80a88d425cf2f54bfedc67e2b33007bd Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Tue, 30 Sep 2025 14:55:38 -0300 Subject: [PATCH 01/15] 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 From e83dd9d3c28044b545ffd91709de3bc888129768 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 11:41:24 -0300 Subject: [PATCH 02/15] SIENTIAPDE-1243: Implement automatic version calculation and branch validation in quality gate workflow, and update SonarQube project properties. (+190 -10 lines) --- .github/workflows/quality-gate.yml | 143 ++++++++++++++++++++++++++++- sonar-project.properties | 8 +- 2 files changed, 142 insertions(+), 9 deletions(-) diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml index 9212f77..257dcbf 100644 --- a/.github/workflows/quality-gate.yml +++ b/.github/workflows/quality-gate.yml @@ -1,9 +1,9 @@ name: Quality gate on: - push: - branches: - - main +# push: +# branches: +# - main pull_request: branches: - main @@ -21,6 +21,138 @@ jobs: fetch-depth: 0 persist-credentials: false + - name: Calculate Version + id: calculate-version + if: github.event_name == 'pull_request' + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + + // Função para calcular nova versΓ£o baseada no branch + function calculateVersion(lastVersion, branchName) { + const parseVersion = (v) => { + const match = v.match(/^(\d+)\.(\d+)\.(\d+)(?:-rc(\d+))?$/); + if (!match) throw new Error(`Invalid version format: ${v}`); + return { + major: parseInt(match[1]), + minor: parseInt(match[2]), + patch: parseInt(match[3]), + rc: match[4] ? parseInt(match[4]) : null + }; + }; + + const current = parseVersion(lastVersion); + + if (branchName.startsWith('release/')) { + return `${current.major + 1}.0.0`; + } else if (branchName.startsWith('feature/')) { + return `${current.major}.${current.minor + 1}.0`; + } else if (branchName.startsWith('fix/')) { + return `${current.major}.${current.minor}.${current.patch + 1}`; + } else if (branchName.startsWith('rc/')) { + if (current.rc !== null) { + return `${current.major}.${current.minor}.${current.patch}-rc${current.rc + 1}`; + } else { + return `${current.major}.${current.minor}.${current.patch}-rc1`; + } + } + + return null; // NΓ£o sugerir para outros tipos de branch + } + + try { + const branchName = context.payload.pull_request.head.ref; + console.log(`Branch name: ${branchName}`); + + // Validar se o branch segue os padrΓ΅es aceitos + const validPrefixes = ['release/', 'feature/', 'fix/', 'rc/']; + const isValidBranch = validPrefixes.some(prefix => branchName.startsWith(prefix)); + + if (!isValidBranch) { + const errorMessage = `## 🚨 Erro: Nome do Branch InvΓ‘lido\n\n` + + `O branch \`${branchName}\` nΓ£o segue os padrΓ΅es de nomenclatura aceitos.\n\n` + + `### πŸ“ PadrΓ΅es Aceitos:\n` + + `- \`release/*\`: Para releases de major version (ex: release/v2.0.0)\n` + + `- \`feature/*\`: Para novas funcionalidades (ex: feature/nova-funcionalidade)\n` + + `- \`fix/*\`: Para correçáes de bugs (ex: fix/correcao-bug)\n` + + `- \`rc/*\`: Para release candidates (ex: rc/v1.2.0-rc1)\n\n` + + `### πŸ”§ Como corrigir:\n` + + `1. Renomeie o branch para seguir um dos padrΓ΅es acima\n` + + `2. Ou crie um novo branch com o nome correto\n`; + + const prNumber = context.issue.number; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: prNumber, + body: errorMessage + }); + + core.setFailed(`Invalid branch name: ${branchName}. Must start with release/, feature/, fix/, or rc/`); + return; + } + + // Obter a ΓΊltima tag de release + console.log('Fetching latest release tag...'); + const { data: releases } = await github.rest.repos.listReleases({ + owner: context.repo.owner, + repo: context.repo.repo, + per_page: 1 + }); + + let lastReleaseVersion = '0.0.0'; + if (releases.length > 0) { + lastReleaseVersion = releases[0].tag_name.replace(/^v/, ''); + console.log(`Latest release version: ${lastReleaseVersion}`); + } else { + console.log('No releases found, using 0.0.0 as baseline'); + } + + // Calcular a nova versΓ£o baseada no branch + const newVersion = calculateVersion(lastReleaseVersion, branchName); + console.log(`Calculated version: ${newVersion}`); + + // Exportar a versΓ£o como output + core.setOutput('version', newVersion); + + // Adicionar comentΓ‘rio informativo no PR + const prNumber = context.issue.number; + const infoMessage = `## βœ… VersΓ£o Calculada Automaticamente\n\n` + + `**Branch:** \`${branchName}\`\n` + + `**Última release:** \`${lastReleaseVersion}\`\n` + + `**Nova versΓ£o:** \`${newVersion}\`\n\n` + + `### πŸ“ Regras Aplicadas:\n` + + `- \`release/*\`: Aumenta major, zera minor e patch (ex: 2.0.0)\n` + + `- \`feature/*\`: MantΓ©m major, aumenta minor, zera patch (ex: 1.2.0)\n` + + `- \`fix/*\`: MantΓ©m major e minor, aumenta patch (ex: 1.1.3)\n` + + `- \`rc/*\`: MantΓ©m versΓ£o base, aumenta RC (ex: 1.1.2-rc2)\n`; + + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: prNumber, + body: infoMessage + }); + + } catch (error) { + console.error('Error during version calculation:', error); + + const prNumber = context.issue.number; + const errorMessage = `## 🚨 Erro no CΓ‘lculo de VersΓ£o\n\n` + + `Ocorreu um erro durante o cΓ‘lculo da versΓ£o:\n\n\`\`\`\n${error.message}\n\`\`\`\n\n` + + `Por favor, verifique se o nome do branch estΓ‘ correto e tente novamente.`; + + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: prNumber, + body: errorMessage + }); + + core.setFailed(`Version calculation error: ${error.message}`); + } + - name: Generate App Token id: generate-app-token uses: actions/create-github-app-token@v1 @@ -65,10 +197,13 @@ jobs: - name: πŸ§ͺ Run Tests with Pytest run: | - pytest tests --junitxml=pytest.xml --cov=laborious --cov-report=xml --cov-report=term + pytest tests --junitxml=pytest.xml --cov=model-manager --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 }} + with: + args: > + -Dsonar.projectVersion=${{ steps.calculate-version.outputs.version || '0.0.0' }} diff --git a/sonar-project.properties b/sonar-project.properties index b332c70..0edb732 100644 --- a/sonar-project.properties +++ b/sonar-project.properties @@ -1,9 +1,7 @@ -sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_beaec423-6c42-4f26-8134-b676287b499d -sonar.projectName=sientia-dataops-laborious_temporal -sonar.sources=laborious +sonar.projectKey=Aignosi_sientia-dataops-model-manager_dc6e051c-995a-4b57-9cde-f511018184c7 +sonar.projectName=sientia-dataops-model-manager +sonar.sources=model-manager sonar.tests=tests -sonar.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 From e318b91c637db2461c4753a4b87bb45227ac9ccb Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 14:11:36 -0300 Subject: [PATCH 03/15] SIENTIAPDE-1243: Update README with correct repository name and installation instructions. --- README.md | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 3f86305..5f05ecc 100644 --- a/README.md +++ b/README.md @@ -367,13 +367,13 @@ flowchart LR 1. **Clone the repository** ```bash git clone - cd sientia-dataops-laborious + cd sientia-dataops-model-manager ``` 2. **Create virtual environment** ```bash - python3.11 -m venv venv - source ./venv/bin/activate + conda create -p ./venv python=3.11 + conda activate ./venv ``` 3. **Install dependencies** @@ -395,6 +395,17 @@ flowchart LR ./install_dependencies.sh ``` + 4. **Install Python dependencies** + ```bash + python -m pip install --upgrade pip + pip install -r requirements.txt + ``` + + 5. **Install test libraries** + ```bash + pip install pytest pytest-cov pytest-asyncio + ``` + 4. **Create environment configuration file** ```bash cp .env.example .env From aba4a3f1a5460a265d12f4fba6d31d06819c414d Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 14:29:24 -0300 Subject: [PATCH 04/15] SIENTIAPDE-1243: Refactor: Rename 'laborious' package to 'model_manager' This commit renames the 'laborious' package to 'model_manager' across the entire project. This includes renaming directories, modules, references in code, configuration files, and documentation to reflect the new package name. This change improves clarity and consistency within the project. --- .github/workflows/quality-gate.yml | 2 +- README.md | 16 +++++----- {model-manager => model_manager}/__init__.py | 0 .../activities/__init__.py | 0 .../activities/activities.py | 6 ++-- .../activities/gates.py | 8 ++--- .../activities/mlflow.py | 2 +- .../activities/opc.py | 2 +- {model-manager => model_manager}/metrics.py | 0 .../utils/__init__.py | 0 .../utils/connectors_config.py | 0 .../utils/filters/__init__.py | 0 .../utils/filters/conditional_filters.py | 0 .../utils/filters/mlflow_filters.py | 0 .../utils/repository/model_repository.py | 0 .../utils/repository/opc_repository.py | 2 +- .../worker/__init__.py | 0 .../worker/worker.py | 14 ++++---- .../workflows/__init__.py | 0 .../workflows/minimal_retrain.py | 2 +- .../workflows/predictions_batch.py | 2 +- .../workflows/sub_workflows/__init__.py | 0 .../format_and_export_prediction.py | 2 +- .../sub_workflows/prediction_process.py | 2 +- run_coverage.sh | 2 +- run_local.sh | 2 +- sonar-project.properties | 2 +- tests/laborious/activities/test_activities.py | 22 ++++++------- tests/laborious/activities/test_gates.py | 10 +++--- tests/laborious/activities/test_mlflow.py | 16 +++++----- tests/laborious/activities/test_opc.py | 8 ++--- .../utils/filters/test_conditional_filters.py | 2 +- .../utils/filters/test_mlflow_filters.py | 2 +- .../utils/repository/test_model_repository.py | 32 +++++++++---------- .../utils/repository/test_opc_repository.py | 10 +++--- .../laborious/utils/test_connectors_config.py | 2 +- .../test_format_and_export_prediction.py | 8 ++--- .../subworkflows/test_prediction_process.py | 22 ++++++------- .../workflows/test_minimal_retrain.py | 6 ++-- .../workflows/test_predictions_batch.py | 6 ++-- values.yaml | 6 ++-- 41 files changed, 109 insertions(+), 109 deletions(-) rename {model-manager => model_manager}/__init__.py (100%) rename {model-manager => model_manager}/activities/__init__.py (100%) rename {model-manager => model_manager}/activities/activities.py (96%) rename {model-manager => model_manager}/activities/gates.py (99%) rename {model-manager => model_manager}/activities/mlflow.py (99%) rename {model-manager => model_manager}/activities/opc.py (99%) rename {model-manager => model_manager}/metrics.py (100%) rename {model-manager => model_manager}/utils/__init__.py (100%) rename {model-manager => model_manager}/utils/connectors_config.py (100%) rename {model-manager => model_manager}/utils/filters/__init__.py (100%) rename {model-manager => model_manager}/utils/filters/conditional_filters.py (100%) rename {model-manager => model_manager}/utils/filters/mlflow_filters.py (100%) rename {model-manager => model_manager}/utils/repository/model_repository.py (100%) rename {model-manager => model_manager}/utils/repository/opc_repository.py (99%) rename {model-manager => model_manager}/worker/__init__.py (100%) rename {model-manager => model_manager}/worker/worker.py (94%) rename {model-manager => model_manager}/workflows/__init__.py (100%) rename {model-manager => model_manager}/workflows/minimal_retrain.py (98%) rename {model-manager => model_manager}/workflows/predictions_batch.py (98%) rename {model-manager => model_manager}/workflows/sub_workflows/__init__.py (100%) rename {model-manager => model_manager}/workflows/sub_workflows/format_and_export_prediction.py (98%) rename {model-manager => model_manager}/workflows/sub_workflows/prediction_process.py (99%) diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml index 257dcbf..3eaa27c 100644 --- a/.github/workflows/quality-gate.yml +++ b/.github/workflows/quality-gate.yml @@ -197,7 +197,7 @@ jobs: - name: πŸ§ͺ Run Tests with Pytest run: | - pytest tests --junitxml=pytest.xml --cov=model-manager --cov-report=xml --cov-report=term + pytest tests --junitxml=pytest.xml --cov=model_manager --cov-report=xml --cov-report=term - name: Run SonarQube Analysis uses: SonarSource/sonarqube-scan-action@v5 diff --git a/README.md b/README.md index 5f05ecc..42f49c2 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa ### Key Components -#### **Worker (`laborious/worker/worker.py`)** +#### **Worker (`model_manager/worker/worker.py`)** - **Purpose**: Main application orchestrator managing Temporal workers and task queues - **Responsibilities**: - Temporal client initialization and connection management @@ -68,7 +68,7 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa - Multi-instance deployment support - Two dedicated task queues: `predictions_batch-queue` and `minimal_retrain-queue` -#### **Workflows (`laborious/workflows/`)** +#### **Workflows (`model_manager/workflows/`)** - **PredictionsBatch**: Main entry point for batch prediction pipelines - **PredictionProcess**: Core prediction pipeline with MLFlow integration - **FormatAndExportPrediction**: Data formatting and export operations @@ -79,7 +79,7 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa - Comprehensive error handling and recovery - Configurable timeout and retry strategies -#### **Activities (`laborious/activities/`)** +#### **Activities (`model_manager/activities/`)** - **Activities**: Main activity orchestrator combining all functionality through multiple inheritance - **Gates**: Data quality validation and filtering mechanisms - **MLFlow**: Model transformation and prediction operations @@ -92,7 +92,7 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa - Comprehensive error handling and notification integration - Support for multiple OPC servers with independent configurations -#### **Data Services (`laborious/utils/`)** +#### **Data Services (`model_manager/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 @@ -475,7 +475,7 @@ source ./venv/bin/activate pytest # Run with coverage -pytest --cov=laborious --cov-report=html +pytest --cov=model_manager --cov-report=html # Run specific test categories pytest tests/activities/ @@ -496,7 +496,7 @@ if [ -f .env ]; then fi # Start the laborious worker -python -m laborious.worker.worker +python -m model_manager.worker.worker ``` ## πŸ§ͺ Testing @@ -516,7 +516,7 @@ tests/ pip install pytest pytest-cov pytest-asyncio # Run tests with coverage -pytest --cov=laborious --cov-report=html +pytest --cov=model_manager --cov-report=html # Run specific test modules pytest tests/activities/test_gates.py @@ -722,7 +722,7 @@ This is the configuration created by the Orchestrator in Temporal. ### Project Structure ``` -laborious/ +model_manager/ β”œβ”€β”€ activities/ # Temporal activity implementations β”‚ β”œβ”€β”€ activities.py # Main activities orchestrator β”‚ β”œβ”€β”€ gates.py # Data quality gates and filtering diff --git a/model-manager/__init__.py b/model_manager/__init__.py similarity index 100% rename from model-manager/__init__.py rename to model_manager/__init__.py diff --git a/model-manager/activities/__init__.py b/model_manager/activities/__init__.py similarity index 100% rename from model-manager/activities/__init__.py rename to model_manager/activities/__init__.py diff --git a/model-manager/activities/activities.py b/model_manager/activities/activities.py similarity index 96% rename from model-manager/activities/activities.py rename to model_manager/activities/activities.py index ec5ae46..1fe2d6f 100644 --- a/model-manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -4,9 +4,9 @@ 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 model_manager.activities.mlflow import MLFlow + from model_manager.activities.gates import Gates + from model_manager.activities.opc import OPC from typing import Any diff --git a/model-manager/activities/gates.py b/model_manager/activities/gates.py similarity index 99% rename from model-manager/activities/gates.py rename to model_manager/activities/gates.py index 4a71455..ec5e806 100644 --- a/model-manager/activities/gates.py +++ b/model_manager/activities/gates.py @@ -9,14 +9,14 @@ with workflow.unsafe.imports_passed_through(): 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 ( + from model_manager.utils.filters.mlflow_filters import nan_values_filter, api_error_filter + + from model_manager.utils.filters.conditional_filters import ( filter_empty_data, filter_specific_variables_null_values ) from pandas import DataFrame - from laborious import metrics + from model_manager import metrics # Input filter function mappings input_filter_functions = { diff --git a/model-manager/activities/mlflow.py b/model_manager/activities/mlflow.py similarity index 99% rename from model-manager/activities/mlflow.py rename to model_manager/activities/mlflow.py index bd06283..e55a55e 100644 --- a/model-manager/activities/mlflow.py +++ b/model_manager/activities/mlflow.py @@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through(): 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 model_manager.utils.repository.model_repository import MLFlowRepository from typing import Any import numpy as np from pandas import DataFrame diff --git a/model-manager/activities/opc.py b/model_manager/activities/opc.py similarity index 99% rename from model-manager/activities/opc.py rename to model_manager/activities/opc.py index 0b4ec9f..2222eed 100644 --- a/model-manager/activities/opc.py +++ b/model_manager/activities/opc.py @@ -6,7 +6,7 @@ with workflow.unsafe.imports_passed_through(): 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 model_manager.utils.repository.opc_repository import OpcRepository from typing import Any import traceback from pandas import DataFrame diff --git a/model-manager/metrics.py b/model_manager/metrics.py similarity index 100% rename from model-manager/metrics.py rename to model_manager/metrics.py diff --git a/model-manager/utils/__init__.py b/model_manager/utils/__init__.py similarity index 100% rename from model-manager/utils/__init__.py rename to model_manager/utils/__init__.py diff --git a/model-manager/utils/connectors_config.py b/model_manager/utils/connectors_config.py similarity index 100% rename from model-manager/utils/connectors_config.py rename to model_manager/utils/connectors_config.py diff --git a/model-manager/utils/filters/__init__.py b/model_manager/utils/filters/__init__.py similarity index 100% rename from model-manager/utils/filters/__init__.py rename to model_manager/utils/filters/__init__.py diff --git a/model-manager/utils/filters/conditional_filters.py b/model_manager/utils/filters/conditional_filters.py similarity index 100% rename from model-manager/utils/filters/conditional_filters.py rename to model_manager/utils/filters/conditional_filters.py diff --git a/model-manager/utils/filters/mlflow_filters.py b/model_manager/utils/filters/mlflow_filters.py similarity index 100% rename from model-manager/utils/filters/mlflow_filters.py rename to model_manager/utils/filters/mlflow_filters.py diff --git a/model-manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py similarity index 100% rename from model-manager/utils/repository/model_repository.py rename to model_manager/utils/repository/model_repository.py diff --git a/model-manager/utils/repository/opc_repository.py b/model_manager/utils/repository/opc_repository.py similarity index 99% rename from model-manager/utils/repository/opc_repository.py rename to model_manager/utils/repository/opc_repository.py index 96ecfef..94b6c46 100644 --- a/model-manager/utils/repository/opc_repository.py +++ b/model_manager/utils/repository/opc_repository.py @@ -11,7 +11,7 @@ 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 +from model_manager import metrics data_type_map = { 'float': { diff --git a/model-manager/worker/__init__.py b/model_manager/worker/__init__.py similarity index 100% rename from model-manager/worker/__init__.py rename to model_manager/worker/__init__.py diff --git a/model-manager/worker/worker.py b/model_manager/worker/worker.py similarity index 94% rename from model-manager/worker/worker.py rename to model_manager/worker/worker.py index 567a257..a10c4cb 100644 --- a/model-manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -33,13 +33,13 @@ 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 \ + from model_manager.workflows.minimal_retrain import MinimalRetrain + from model_manager.workflows.predictions_batch import PredictionsBatch + from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess + from model_manager.workflows.sub_workflows.format_and_export_prediction import \ FormatAndExportPrediction - from laborious.activities.activities import Activities - from laborious.utils.connectors_config import ( + from model_manager.activities.activities import Activities + from model_manager.utils.connectors_config import ( build_postgres_config, build_mlflow_config, build_opc_config, @@ -47,7 +47,7 @@ with workflow.unsafe.imports_passed_through(): ) from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import get_logger - from laborious import metrics + from model_manager import metrics from prometheus_client import start_http_server POD_ID = os.getenv('POD_ID') diff --git a/model-manager/workflows/__init__.py b/model_manager/workflows/__init__.py similarity index 100% rename from model-manager/workflows/__init__.py rename to model_manager/workflows/__init__.py diff --git a/model-manager/workflows/minimal_retrain.py b/model_manager/workflows/minimal_retrain.py similarity index 98% rename from model-manager/workflows/minimal_retrain.py rename to model_manager/workflows/minimal_retrain.py index 1893e25..4701c91 100644 --- a/model-manager/workflows/minimal_retrain.py +++ b/model_manager/workflows/minimal_retrain.py @@ -1,7 +1,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities + from model_manager.activities.activities import Activities from typing import Any from sientia_do.temporal.policies import retry_policy from datetime import timedelta diff --git a/model-manager/workflows/predictions_batch.py b/model_manager/workflows/predictions_batch.py similarity index 98% rename from model-manager/workflows/predictions_batch.py rename to model_manager/workflows/predictions_batch.py index e522eaa..330493d 100644 --- a/model-manager/workflows/predictions_batch.py +++ b/model_manager/workflows/predictions_batch.py @@ -1,7 +1,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities + from model_manager.activities.activities import Activities from typing import Any from sientia_do.temporal.policies import retry_policy from datetime import timedelta diff --git a/model-manager/workflows/sub_workflows/__init__.py b/model_manager/workflows/sub_workflows/__init__.py similarity index 100% rename from model-manager/workflows/sub_workflows/__init__.py rename to model_manager/workflows/sub_workflows/__init__.py diff --git a/model-manager/workflows/sub_workflows/format_and_export_prediction.py b/model_manager/workflows/sub_workflows/format_and_export_prediction.py similarity index 98% rename from model-manager/workflows/sub_workflows/format_and_export_prediction.py rename to model_manager/workflows/sub_workflows/format_and_export_prediction.py index 8e7df07..ef195c4 100644 --- a/model-manager/workflows/sub_workflows/format_and_export_prediction.py +++ b/model_manager/workflows/sub_workflows/format_and_export_prediction.py @@ -1,7 +1,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities + from model_manager.activities.activities import Activities from typing import Any from datetime import timedelta from sientia_do.temporal.policies import retry_policy diff --git a/model-manager/workflows/sub_workflows/prediction_process.py b/model_manager/workflows/sub_workflows/prediction_process.py similarity index 99% rename from model-manager/workflows/sub_workflows/prediction_process.py rename to model_manager/workflows/sub_workflows/prediction_process.py index 777fa1c..4addd75 100644 --- a/model-manager/workflows/sub_workflows/prediction_process.py +++ b/model_manager/workflows/sub_workflows/prediction_process.py @@ -1,7 +1,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities + from model_manager.activities.activities import Activities from typing import Any from sientia_do.temporal.policies import retry_policy from datetime import timedelta diff --git a/run_coverage.sh b/run_coverage.sh index f9af4cb..7d4f6f4 100755 --- a/run_coverage.sh +++ b/run_coverage.sh @@ -6,6 +6,6 @@ set -e echo "Activating virtual environment..." source ./venv/bin/activate -pytest --cov=laborious --cov-report=html +pytest --cov=model_manager --cov-report=html xdg-open htmlcov/index.html \ No newline at end of file diff --git a/run_local.sh b/run_local.sh index 2bbd5c2..680ff15 100755 --- a/run_local.sh +++ b/run_local.sh @@ -15,4 +15,4 @@ else fi echo "Starting ingestor application..." -python -m laborious.worker.worker +python -m model_manager.worker.worker diff --git a/sonar-project.properties b/sonar-project.properties index 0edb732..e8ae7ff 100644 --- a/sonar-project.properties +++ b/sonar-project.properties @@ -1,6 +1,6 @@ sonar.projectKey=Aignosi_sientia-dataops-model-manager_dc6e051c-995a-4b57-9cde-f511018184c7 sonar.projectName=sientia-dataops-model-manager -sonar.sources=model-manager +sonar.sources=model_manager sonar.tests=tests sonar.qualitygate.wait=true sonar.qualitygate.timeout=300 diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index de418be..106ccf7 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -1,16 +1,16 @@ 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 +from model_manager.activities.activities import Activities +from model_manager.activities.mlflow import MLFlow +from model_manager.activities.gates import Gates +from model_manager.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__') +@patch('model_manager.activities.activities.Postgres.__init__') +@patch('model_manager.activities.activities.MLFlow.__init__') +@patch('model_manager.activities.activities.OPC.__init__') +@patch('model_manager.activities.activities.Gates.__init__') def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init): postgres_config = { @@ -91,9 +91,9 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre @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()) +@patch('model_manager.activities.activities.Postgres', return_value=MagicMock()) +@patch('model_manager.activities.activities.MLFlow', return_value=MagicMock()) +@patch('model_manager.activities.activities.OPC', return_value=MagicMock()) async def test_shutdown(mock_opc_init, _mock_mlflow_init, mock_postgres_init): postgres_config = { diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index a65c6d4..b587b93 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -1,7 +1,7 @@ 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 +from model_manager.activities.gates import Gates @fixture @@ -52,7 +52,7 @@ async def test_input_gate_invalid_filter(gates_activity): @mark.asyncio -@patch('laborious.activities.gates.input_filter_functions') +@patch('model_manager.activities.gates.input_filter_functions') async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity): # Arrange mock_input_filter_functions.__contains__.return_value = True @@ -141,7 +141,7 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity): @mark.asyncio -@patch('laborious.activities.gates.mlflow_response_filter_functions') +@patch('model_manager.activities.gates.mlflow_response_filter_functions') async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions, gates_activity): # Arrange @@ -241,7 +241,7 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity): @mark.asyncio -@patch('laborious.activities.gates.mlflow_content_filter_functions') +@patch('model_manager.activities.gates.mlflow_content_filter_functions') async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions, gates_activity): # Arrange @@ -559,7 +559,7 @@ async def test_get_last_timestamp_no_data(gates_activity): @mark.asyncio -@patch('laborious.activities.gates.metrics') +@patch('model_manager.activities.gates.metrics') async def test_write_metrics(mock_metrics, gates_activity): """Test write_metrics method.""" input_data = { diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 85cb7f1..243982f 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -5,11 +5,11 @@ 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 model_manager.activities.mlflow import MLFlow from sientia_do.notifications.models import NotificationLevel -@patch("laborious.activities.mlflow.MLFlowRepository") +@patch("model_manager.activities.mlflow.MLFlowRepository") def test___init__(mock_mlflow_repository): mlflow = MLFlow( mlflow_host="http://localhost", @@ -31,7 +31,7 @@ def test___init__(mock_mlflow_repository): @fixture -@patch("laborious.activities.mlflow.MLFlowRepository") +@patch("model_manager.activities.mlflow.MLFlowRepository") def mlflow(mock_mlflow_repository): mlflow = MLFlow( mlflow_host="http://localhost:5000", @@ -58,8 +58,8 @@ metadata = { @mark.asyncio -@patch("laborious.activities.mlflow.DataFrame") -@patch("laborious.activities.mlflow.max") +@patch("model_manager.activities.mlflow.DataFrame") +@patch("model_manager.activities.mlflow.max") async def test_request_transform_success(mock_max, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data @@ -114,9 +114,9 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow): @mark.asyncio -@patch("laborious.activities.mlflow.DataFrame") -@patch("laborious.activities.mlflow.to_datetime") -@patch("laborious.activities.mlflow.max") +@patch("model_manager.activities.mlflow.DataFrame") +@patch("model_manager.activities.mlflow.to_datetime") +@patch("model_manager.activities.mlflow.max") async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index 07f00b8..dcff019 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -4,7 +4,7 @@ from pytest import fixture, mark import pytest_asyncio from sientia_do.notifications.models import NotificationLevel -from laborious.activities.opc import OPC +from model_manager.activities.opc import OPC metadata = { "metadata": { @@ -31,8 +31,8 @@ def test__init__(): @mark.asyncio -@patch("laborious.activities.opc.OpcRepository") -@patch("laborious.activities.opc.OPC.send_notification") +@patch("model_manager.activities.opc.OpcRepository") +@patch("model_manager.activities.opc.OPC.send_notification") async def test_init_opc(mock_send_notification, mock_opc_repository): mock_logger = MagicMock() server1 = MagicMock( @@ -147,7 +147,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository): @pytest_asyncio.fixture -@patch("laborious.activities.opc.OpcRepository") +@patch("model_manager.activities.opc.OpcRepository") async def opc(mock_opc_repository): servers = { 'server1': { diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index 405bc9b..a7f3efd 100644 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -1,6 +1,6 @@ from pandas import DataFrame -from laborious.utils.filters.conditional_filters import ( +from model_manager.utils.filters.conditional_filters import ( filter_specific_variables_null_values, filter_empty_data ) diff --git a/tests/laborious/utils/filters/test_mlflow_filters.py b/tests/laborious/utils/filters/test_mlflow_filters.py index f9c61e9..6fddf51 100644 --- a/tests/laborious/utils/filters/test_mlflow_filters.py +++ b/tests/laborious/utils/filters/test_mlflow_filters.py @@ -1,5 +1,5 @@ from pandas import DataFrame -from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter +from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter def test_api_error_filter_invalid_response(): diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index cdc59b4..ab65913 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -4,12 +4,12 @@ from pandas import DataFrame import pytest from datetime import datetime, timezone from pandas import Timestamp -from laborious.utils.repository.model_repository import MLFlowRepository +from model_manager.utils.repository.model_repository import MLFlowRepository @pytest.fixture def mlflow_repository(): - with patch('laborious.utils.repository.model_repository.ModelServing', + with patch('model_manager.utils.repository.model_repository.ModelServing', autospec=True) as mock_model_serving: mock_instance = mock_model_serving.return_value mock_instance.get_transformed_data = MagicMock() @@ -217,7 +217,7 @@ def test_predict_error(mlflow_repository): } -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_by_run_id(mlflow, mlflow_repository): mlflow.get_run.return_value = MagicMock( info=MagicMock( @@ -233,7 +233,7 @@ def test_get_experiment_by_run_id(mlflow, mlflow_repository): mlflow.get_experiment.assert_called_once_with('0') -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_get_next_run_name(mlflow, mlflow_repository): mlflow.search_runs.return_value = [1, 2, 3] output = mlflow_repository.get_next_run_name('run') @@ -244,7 +244,7 @@ def test_get_next_run_name(mlflow, mlflow_repository): ) -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_success(mlflow, mlflow_repository): mlflow.get_experiment_by_name.return_value = MagicMock( experiment_id='0') @@ -254,7 +254,7 @@ def test_get_experiment_success(mlflow, mlflow_repository): assert output == 0 -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_error(mlflow, mlflow_repository): mlflow.get_experiment_by_name.return_value = None @@ -266,7 +266,7 @@ def test_get_experiment_error(mlflow, mlflow_repository): assert False -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_last_run(mlflow, mlflow_repository): mlflow.search_runs.return_value = DataFrame({ 'params.retrain': ['True', 'False', 'True', 'False'], @@ -285,7 +285,7 @@ def test_get_experiment_last_run(mlflow, mlflow_repository): assert output == '2' -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_last_run_error(mlflow, mlflow_repository): mlflow.search_runs.return_value = [] @@ -297,8 +297,8 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository): assert False -@patch('laborious.utils.repository.model_repository.mlflow.sklearn') -@patch('laborious.utils.repository.model_repository.mlflow.set_experiment') +@patch('model_manager.utils.repository.model_repository.mlflow.sklearn') +@patch('model_manager.utils.repository.model_repository.mlflow.set_experiment') def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): mlflow_repository.model_serving.get_model_run_id = MagicMock( @@ -360,10 +360,10 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): 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') +@patch('model_manager.utils.repository.model_repository.mlflow.start_run') +@patch('model_manager.utils.repository.model_repository.mlflow.log_param') +@patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model') +@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact') def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository): prediction_model_mock = MagicMock() @@ -423,7 +423,7 @@ def test_retrain_model(mlflow_repository): assert output == 'Model retrained successfully' -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_update_production_model_by_run_id(mlflow, mlflow_repository): client_mock = MagicMock() mlflow.tracking.MlflowClient.return_value = client_mock @@ -458,7 +458,7 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): } -@patch('laborious.utils.repository.model_repository.mlflow') +@patch('model_manager.utils.repository.model_repository.mlflow') def test_update_production_model_by_run_id_error(mlflow, mlflow_repository): mlflow.tracking.MlflowClient.return_value = MagicMock( get_registered_model=MagicMock( diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index bc84db8..9795a10 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -1,7 +1,7 @@ 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 model_manager.utils.repository.opc_repository import OpcRepository from sientia_do.notifications.models import NotificationLevel from datetime import datetime @@ -28,7 +28,7 @@ def opc_repository(mock_logger): @pytest.fixture def mock_client(): - with patch('laborious.utils.repository.opc_repository.Client') as mock: + with patch('model_manager.utils.repository.opc_repository.Client') as mock: client_instance = AsyncMock() mock.return_value = client_instance yield client_instance @@ -212,7 +212,7 @@ async def test_validate_connection_error_validate_connection_error(opc_repositor @pytest.mark.asyncio -@patch('laborious.utils.repository.opc_repository.datetime') +@patch('model_manager.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)) @@ -233,7 +233,7 @@ async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, op @pytest.mark.asyncio -@patch('laborious.utils.repository.opc_repository.datetime') +@patch('model_manager.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)) @@ -333,7 +333,7 @@ async def test_write_data_invalid_data_type(opc_repository, mock_client): @pytest.mark.asyncio -@patch('laborious.utils.repository.opc_repository.metrics') +@patch('model_manager.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 diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 910439c..43a2ec9 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,5 +1,5 @@ from os import environ -from laborious.utils.connectors_config import (build_mlflow_config, +from model_manager.utils.connectors_config import (build_mlflow_config, build_opc_config, build_postgres_config, build_mongodb_config) diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index a8e6e20..356c32b 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -1,8 +1,8 @@ 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 model_manager.activities.activities import Activities +from model_manager.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ @@ -22,7 +22,7 @@ metadata = { @mark.asyncio -@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): input_data = { @@ -91,7 +91,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): @mark.asyncio -@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) async def test_run_default_path_flag(workflow_mock, format_and_export_prediction): input_data = { diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index df60ada..46b1565 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -1,7 +1,7 @@ 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 +from model_manager.activities.activities import Activities +from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess @fixture @@ -20,7 +20,7 @@ metadata = { @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_run(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=False) # Arrange @@ -135,7 +135,7 @@ async def test_run(workflow_mock, prediction_process): @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_run_stop_at_input_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=True) # Arrange @@ -183,7 +183,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True]) # Arrange @@ -253,7 +253,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock( side_effect=[False, False, True]) @@ -333,7 +333,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock( side_effect=[False, False, False, True]) @@ -429,7 +429,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_path_flag_handler_stop(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -464,7 +464,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process): @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_path_flag_handler_repeat(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -510,7 +510,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process): @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_path_flag_handler_continue(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -565,7 +565,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_path_flag_handler_unknown(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index b3b03b5..1d8e6d7 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -1,7 +1,7 @@ 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 +from model_manager.activities.activities import Activities +from model_manager.workflows.minimal_retrain import MinimalRetrain @fixture @@ -20,7 +20,7 @@ metadata = { @mark.asyncio -@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock) +@patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): input_data = { "model_id": "test_model_id", diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 90d7d21..0d7d477 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -1,7 +1,7 @@ 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 +from model_manager.activities.activities import Activities +from model_manager.workflows.predictions_batch import PredictionsBatch @fixture @@ -20,7 +20,7 @@ metadata = { @mark.asyncio -@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) +@patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): workflow_mock.execute_local_activity_method.return_value = { 'data': 'test_data' diff --git a/values.yaml b/values.yaml index 4eb83d7..3af8b24 100644 --- a/values.yaml +++ b/values.yaml @@ -70,7 +70,7 @@ livenessProbe: command: - sh - -c - - pgrep -f "laborious.worker.worker" + - pgrep -f "model_manager.worker.worker" initialDelaySeconds: 20 periodSeconds: 30 @@ -79,7 +79,7 @@ readinessProbe: command: - sh - -c - - pgrep -f "laborious.worker.worker" + - pgrep -f "model_manager.worker.worker" initialDelaySeconds: 10 periodSeconds: 15 @@ -153,7 +153,7 @@ env: - name: GITHUB_BRANCH value: SIENTIAPDE-1222-ajustar-a-library-para-fazer-o-download-do-courier - name: PYTHON_APP - value: "laborious.worker.worker" + value: "model_manager.worker.worker" # Application variables - name: POSTGRES_HOST From bc4d98f78d6b4f3251a13948e03bc746396a34fa Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 14:41:24 -0300 Subject: [PATCH 05/15] SIENTIAPDE-1243: Rename project from 'laborious' to 'model-manager' across codebase and configuration. This includes updating project names in environment variables, Makefiles, README, metrics, and Helm chart values. --- .env.example | 4 +-- Makefile | 2 +- README.md | 26 +++++++++---------- model_manager/activities/activities.py | 2 +- model_manager/activities/gates.py | 2 +- model_manager/metrics.py | 14 +++++----- model_manager/worker/worker.py | 16 ++++++------ model_manager/workflows/minimal_retrain.py | 2 +- model_manager/workflows/predictions_batch.py | 2 +- .../sub_workflows/prediction_process.py | 2 +- values.yaml | 18 ++++++------- 11 files changed, 45 insertions(+), 45 deletions(-) diff --git a/.env.example b/.env.example index 865a73b..8bd75e5 100644 --- a/.env.example +++ b/.env.example @@ -17,10 +17,10 @@ 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" +PROJECT_NAME="sientia-model-manager" TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233" -TEMPORAL_NAMESPACE="laborious" +TEMPORAL_NAMESPACE="model-manager" MONGODB_USERNAME="mongo_user" MONGODB_PASSWORD="mongo_db_password" diff --git a/Makefile b/Makefile index 66aae81..f6e4208 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ VERSION = 1.0.8 -name = sientia-laborious +name = sientia-model-manager # ENVIRONMENT = production docker-hub: diff --git a/README.md b/README.md index 42f49c2..61a50b9 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ A comprehensive AI model management platform for the complete machine learning l ## 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. +The Model Manager system uses a Temporal-based workflow architecture with clear separation of concerns and robust error handling. The architecture is designed for high availability, scalability, and operational excellence in production ML environments. ### Architecture Principles @@ -428,7 +428,7 @@ flowchart LR ## πŸ“¦ How to Run -### Running the Laborious Application +### Running the Model Manager Application Use the provided script to run the application locally: @@ -443,7 +443,7 @@ chmod +x run_local.sh The script will: - Activate the virtual environment - Load environment variables from `.env` -- Start the laborious worker application +- Start the model-manager worker application ### Running Tests and Coverage @@ -495,7 +495,7 @@ if [ -f .env ]; then export $(cat .env | grep -v '^#' | xargs) fi -# Start the laborious worker +# Start the model-manager worker python -m model_manager.worker.worker ``` @@ -525,25 +525,25 @@ pytest tests/workflow/test_predictions_batch.py ## πŸ“Š Monitoring and Metrics -The Laborious system exposes comprehensive Prometheus metrics for operational visibility and performance monitoring: +The Model Manager system exposes comprehensive Prometheus metrics for operational visibility and performance monitoring: ### Application Health Metrics - `app_up`: Application health status (1=healthy, 0=unhealthy) - Labels: `pod_id` ### Prediction Operation Metrics -- `laborious_predictions_written_count`: Counter for successful prediction exports +- `model_manager_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 +- `model_manager_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 +- `model_manager_prediction_response_time_monitor`: Histogram for prediction response times - Labels: `pod_id`, `model_name`, `pipeline_name` - Buckets: [0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0] ### OPC Export Metrics -- `laborious_prediction_opc_writing_count`: Counter for OPC server write operations +- `model_manager_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 +- `model_manager_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] @@ -559,7 +559,7 @@ The Laborious system exposes comprehensive Prometheus metrics for operational vi | Variable | Description | Default | Required | |----------|-------------|---------|----------| | `TEMPORAL_HOST` | Temporal server address | `localhost:7233` | Yes | -| `TEMPORAL_NAMESPACE` | Temporal namespace | `laborious` | No | +| `TEMPORAL_NAMESPACE` | Temporal namespace | `model-manager` | No | | `POSTGRES_HOST` | PostgreSQL hostname | `localhost` | Yes | | `POSTGRES_PORT` | PostgreSQL port | `5432` | Yes | | `POSTGRES_USER` | PostgreSQL username | `sientia` | Yes | @@ -585,7 +585,7 @@ The Laborious system exposes comprehensive Prometheus metrics for operational vi | `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 | +| `PROJECT_NAME` | Project name for metrics | `model-manager` | No | | `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No | | `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No | | `POD_ID` | Kubernetes pod identifier | `None` | No | @@ -838,4 +838,4 @@ For support and questions: --- -**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments. +**Note**: The Model Manager 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/activities/activities.py b/model_manager/activities/activities.py index 1fe2d6f..6a8656d 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -12,7 +12,7 @@ with workflow.unsafe.imports_passed_through(): class Activities(Postgres, MLFlow, Gates, OPC): """ - Main activities orchestrator for the Laborious system. + Main activities orchestrator for the Model Manager system. This class combines functionality from multiple activity classes to provide a unified interface for all workflow operations. It manages database connections, diff --git a/model_manager/activities/gates.py b/model_manager/activities/gates.py index ec5e806..338e212 100644 --- a/model_manager/activities/gates.py +++ b/model_manager/activities/gates.py @@ -53,7 +53,7 @@ mlflow_content_filter_functions = { class Gates(BaseActivity): """ - Data quality gates and filtering activities for the Laborious system. + Data quality gates and filtering activities for the Model Manager system. This class implements comprehensive data quality validation and filtering mechanisms that can be applied at different stages of the prediction pipeline. diff --git a/model_manager/metrics.py b/model_manager/metrics.py index 97f7bb9..96b806f 100644 --- a/model_manager/metrics.py +++ b/model_manager/metrics.py @@ -1,7 +1,7 @@ """ -Laborious Metrics Module +Model Manager Metrics Module -This module defines all Prometheus metrics used by the Sientia DataOps Laborious system +This module defines all Prometheus metrics used by the Sientia DataOps Model Manager system for monitoring and observability. The metrics provide insights into system performance, prediction quality, and operational health. @@ -37,21 +37,21 @@ CORE_LABELS = ["pod_id", "model_name", "pipeline_name"] # Prediction operation metrics PREDICTIONS_WRITTEN_COUNT = Counter( - "laborious_predictions_written_count", + "model_manager_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", + "model_manager_prediction_confidence_monitor", "Current confidence of each prediction", CORE_LABELS, ) # Performance monitoring metrics PREDICTION_RESPONSE_TIME_MONITOR = Histogram( - "laborious_prediction_response_time_monitor", + "model_manager_prediction_response_time_monitor", "Current response time of each prediction", CORE_LABELS, buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0] @@ -59,13 +59,13 @@ PREDICTION_RESPONSE_TIME_MONITOR = Histogram( # OPC export metrics PREDICTION_OPC_WRITING_COUNT = Counter( - "laborious_prediction_opc_writing_count", + "model_manager_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", + "model_manager_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/worker/worker.py b/model_manager/worker/worker.py index a10c4cb..5d1d0b1 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -1,7 +1,7 @@ """ -Laborious Worker Module +Model Manager Worker Module -This module provides the main worker implementation for the Sientia DataOps Laborious system. +This module provides the main worker implementation for the Sientia DataOps Model Manager system. It orchestrates Temporal workers, manages task queues, and handles the lifecycle of prediction and retraining workflows. @@ -18,11 +18,11 @@ Key Features: Environment Variables: - TEMPORAL_HOST: Temporal server address (default: localhost:7233) -- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious) +- TEMPORAL_NAMESPACE: Temporal namespace (default: model_manager) - POD_ID: Kubernetes pod identifier for metrics - HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090) - HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091) -- PROJECT_NAME: Project name for notifications (default: laborious) +- PROJECT_NAME: Project name for notifications (default: model_manager) """ from temporalio import workflow, client @@ -56,7 +56,7 @@ SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091")) async def main(): """ - Main entry point for the Laborious worker application. + Main entry point for the Model Manager worker application. This function initializes and starts all components of the worker: 1. Sets up logging and metadata @@ -97,7 +97,7 @@ async def main(): connection_string=mongo_config['connection_string'], database=mongo_config['database_name'], logger=logger, - project_name=os.getenv('PROJECT_NAME', 'laborious') + project_name=os.getenv('PROJECT_NAME', 'model-manager') ) logger.custom_info('Starting Activities...', metadata) @@ -127,7 +127,7 @@ async def main(): temporal_client = await client.Client.connect( target_host=host, - namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), + namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'), runtime=new_runtime ) @@ -214,7 +214,7 @@ def start_prometheus_server(): 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. + the health and performance of the Model Manager worker. Environment Variables: HTTP_METRICS_PORT: Port for the metrics server (default: 9090) diff --git a/model_manager/workflows/minimal_retrain.py b/model_manager/workflows/minimal_retrain.py index 4701c91..f6e2c80 100644 --- a/model_manager/workflows/minimal_retrain.py +++ b/model_manager/workflows/minimal_retrain.py @@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through(): @workflow.defn(name="minimal_retrain") class MinimalRetrain(): """ - Automated model retraining workflow for the Laborious system. + Automated model retraining workflow for the Model Manager system. This workflow implements a complete model retraining pipeline that loads training data, executes model retraining, updates production models, diff --git a/model_manager/workflows/predictions_batch.py b/model_manager/workflows/predictions_batch.py index 330493d..281b187 100644 --- a/model_manager/workflows/predictions_batch.py +++ b/model_manager/workflows/predictions_batch.py @@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through(): @workflow.defn(name="predictions_batch") class PredictionsBatch(): """ - Main batch prediction workflow for the Laborious system. + Main batch prediction workflow for the Model Manager system. This workflow orchestrates the complete batch prediction process, handling data loading, configuration management, and workflow delegation. It serves diff --git a/model_manager/workflows/sub_workflows/prediction_process.py b/model_manager/workflows/sub_workflows/prediction_process.py index 4addd75..e613ab5 100644 --- a/model_manager/workflows/sub_workflows/prediction_process.py +++ b/model_manager/workflows/sub_workflows/prediction_process.py @@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through(): @workflow.defn(name="prediction_process") class PredictionProcess(): """ - Core prediction processing workflow for the Laborious system. + Core prediction processing workflow for the Model Manager system. This workflow implements the complete ML model inference pipeline, handling data quality validation, MLFlow model interactions, and prediction processing. diff --git a/values.yaml b/values.yaml index 3af8b24..744c0c0 100644 --- a/values.yaml +++ b/values.yaml @@ -17,8 +17,8 @@ image: imagePullSecrets: - name: docker-hub-secret # This is to override the chart name. -nameOverride: "sientia-laborious-worker" -fullnameOverride: "sientia-laborious-worker" +nameOverride: "sientia-model-manager-worker" +fullnameOverride: "sientia-model-manager-worker" namespace: sientia # This section builds out the service account more information can be found here: https://kubernetes.io/docs/concepts/security/service-accounts/ @@ -31,7 +31,7 @@ serviceAccount: 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" + name: "sientia-model-manager-worker" # This is for setting Kubernetes Annotations to a Pod. # For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/ @@ -149,7 +149,7 @@ serviceMonitor: env: # Entrypoint variables - name: GITHUB_REPO_URL - value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git" + value: "git@github.com:Aignosi/sientia-dataops-model-manager.git" - name: GITHUB_BRANCH value: SIENTIAPDE-1222-ajustar-a-library-para-fazer-o-download-do-courier - name: PYTHON_APP @@ -195,12 +195,12 @@ env: - name: HTTP_SDK_METRICS_PORT value: "9091" - name: PROJECT_NAME - value: "sientia-laborious" + value: "sientia-model-manager" - name: TEMPORAL_HOST value: "temporal-frontend.temporal.svc.cluster.local:7233" - name: TEMPORAL_NAMESPACE - value: "laborious" + value: "model-manager" - name: MONGODB_USERNAME value: "root" @@ -215,15 +215,15 @@ env: ssh: enabled: true - secretName: git-ssh-key-sientia-laborious-worker + secretName: git-ssh-key-sientia-model-manager-worker sshPath: /mnt/.ssh knownHostsPath: /mnt/known_hosts # kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=5I5zpQ6sRaHqX1hD3dr+2mo647yO3FRc359/wu6gsP+ACRDRz5mp -# helm upgrade --install sientia-laborious-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.5.0 +# helm upgrade --install sientia-model-manager-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.5.0 -# kubectl create secret generic git-ssh-key-sientia-laborious-worker \ +# kubectl create secret generic git-ssh-key-sientia-model-manager-worker \ # --namespace sientia \ # --from-file=ssh-privatekey=git_key \ # --type=kubernetes.io/ssh-auth From b102f790878d4f300a230691745c9482fc13939f Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 16:25:08 -0300 Subject: [PATCH 06/15] SIENTIAPDE-1243: Remove OPC server integration and add code quality tools. This commit removes the OPC server integration from the Model Manager, including related activities, repositories, metrics, and configuration. It also adds code quality tools such as Ruff (linting/formatting), mypy (type checking), and Bandit (security analysis) along with a validation script and CI/CD integration for automated code validation. The README has been updated to reflect these changes. --- .env.example | 3 - .github/workflows/quality-gate.yml | 26 +- README.md | 235 ++++++----- model_manager/activities/activities.py | 17 +- model_manager/activities/opc.py | 356 ---------------- model_manager/metrics.py | 17 +- model_manager/utils/connectors_config.py | 39 -- .../utils/repository/opc_repository.py | 359 ---------------- model_manager/worker/worker.py | 12 +- model_manager/workflows/predictions_batch.py | 4 +- .../format_and_export_prediction.py | 24 +- .../sub_workflows/prediction_process.py | 4 +- pyproject.toml | 154 +++++++ requirements-dev.txt | 17 + requirements.txt | 1 - tests/laborious/activities/test_activities.py | 31 +- tests/laborious/activities/test_opc.py | 369 ----------------- .../utils/repository/test_opc_repository.py | 388 ------------------ .../laborious/utils/test_connectors_config.py | 49 --- .../test_format_and_export_prediction.py | 32 +- .../subworkflows/test_prediction_process.py | 18 +- .../workflows/test_predictions_batch.py | 4 +- validate.sh | 99 +++++ values.yaml | 5 - 24 files changed, 453 insertions(+), 1810 deletions(-) delete mode 100644 model_manager/activities/opc.py delete mode 100644 model_manager/utils/repository/opc_repository.py create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt delete mode 100644 tests/laborious/activities/test_opc.py delete mode 100644 tests/laborious/utils/repository/test_opc_repository.py create mode 100755 validate.sh diff --git a/.env.example b/.env.example index 8bd75e5..652ce20 100644 --- a/.env.example +++ b/.env.example @@ -11,9 +11,6 @@ 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" diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml index 3eaa27c..997202a 100644 --- a/.github/workflows/quality-gate.yml +++ b/.github/workflows/quality-gate.yml @@ -193,7 +193,31 @@ jobs: run: | python -m pip install --upgrade pip pip install -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }} - pip install pytest pytest-cov pytest-asyncio + pip install -r requirements-dev.txt + + - name: πŸ“ Code Formatting Check (Ruff) + run: | + echo "Checking code formatting..." + ruff format --check model_manager/ tests/ + continue-on-error: false + + - name: πŸ”Ž Code Linting (Ruff) + run: | + echo "Running linting checks..." + ruff check model_manager/ tests/ + continue-on-error: false + + - name: 🏷️ Type Checking (mypy) + run: | + echo "Running type checks..." + mypy model_manager/ + continue-on-error: true + + - name: πŸ”’ Security Analysis (Bandit) + run: | + echo "Running security analysis..." + bandit -r model_manager/ -ll -q + continue-on-error: true - name: πŸ§ͺ Run Tests with Pytest run: | diff --git a/README.md b/README.md index 61a50b9..28a38c3 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ A comprehensive AI model management platform for the complete machine learning l - **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 +- **Real-time Data Export**: PostgreSQL persistence for data storage - **Comprehensive Monitoring**: Prometheus metrics and detailed logging for operational visibility ### Advanced Capabilities @@ -19,6 +19,12 @@ A comprehensive AI model management platform for the complete machine learning l - **Scalable Architecture**: Kubernetes-ready deployment with horizontal scaling support - **Model Retraining**: Automated model retraining workflows with production model updates +### Development & Quality Assurance +- **Code Quality Tools**: Ruff (linting/formatting), mypy (type checking), Bandit (security analysis) +- **Automated Validation**: Pre-commit validation script and CI/CD integration +- **Comprehensive Testing**: pytest with async support and 70%+ code coverage +- **Type Safety**: Static type checking with mypy for improved code reliability + ## Architecture The Model Manager system uses a Temporal-based workflow architecture with clear separation of concerns and robust error handling. The architecture is designed for high availability, scalability, and operational excellence in production ML environments. @@ -60,7 +66,6 @@ The Model Manager system uses a Temporal-based workflow architecture with clear - 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 @@ -83,20 +88,16 @@ The Model Manager system uses a Temporal-based workflow architecture with clear - **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 (`model_manager/utils/`)** - **Connectors Config**: Environment variable-based configuration management -- **Repository**: Data access layer for MLFlow and OPC operations +- **Repository**: Data access layer for MLFlow 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 @@ -105,14 +106,14 @@ The Model Manager system uses a Temporal-based workflow architecture with clear - Connection pool management and optimization - Security credential management - Configuration validation and error handling - - Support for multiple OPC servers and MLFlow model flavors + - Support for MLFlow model flavors ### Data Flow Architecture #### **1. Batch Prediction Pipeline** ``` Input Data (PostgreSQL) β†’ Data Quality Gates β†’ MLFlow Transform β†’ -MLFlow Prediction β†’ Response Validation β†’ Export (PostgreSQL + OPC) +MLFlow Prediction β†’ Response Validation β†’ Export (PostgreSQL) ``` #### **2. Model Retraining Pipeline** @@ -121,16 +122,10 @@ 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 @@ -195,11 +190,7 @@ The **PredictionsBatch** workflow is the main entry point for batch prediction p "API_ERROR": {"POLICY": "STOP"} }, "model_retention": 60, - "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "opc_output_config": { - "server_id": "opc_server_1", - "tags": ["prediction_output"] - } + "path_priority": ["STOP", "CONTINUE", "REPEAT"] } ``` @@ -267,8 +258,7 @@ The **PredictionProcess** workflow implements the core prediction pipeline for M "NAN_VALUES": {"POLICY": "STOP"} }, "model_retention": 60, - "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "opc_output_config": {...} + "path_priority": ["STOP", "CONTINUE", "REPEAT"] } ``` @@ -288,33 +278,30 @@ flowchart LR The **FormatAndExportPrediction** workflow handles prediction data formatting and export operations to multiple destinations. #### Purpose -- **Data Formatting**: Formats prediction data for different output destinations +- **Data Formatting**: Formats prediction data for database storage - **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 +4. **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 +- **Database Export**: PostgreSQL integration for data persistence - **Performance Monitoring**: Comprehensive metrics for export operations - **Error Handling**: Robust error handling with notification integration #### Architecture Diagram ```mermaid flowchart LR - A[1. format_prediction/format_default_prediction] --> B[2. write_opc_data] --> C[3. export_data_to_postgres] --> D[4. write_metrics] + A[1. format_prediction/format_default_prediction] --> B[2. export_data_to_postgres] --> C[3. write_metrics] A -.-> Format[Data Formatting] - B -.-> OPC[OPC Servers] - C -.-> PostgreSQL[(PostgreSQL)] - D -.-> Prometheus[Prometheus] + B -.-> PostgreSQL[(PostgreSQL)] + C -.-> Prometheus[Prometheus] ``` ### 4. Minimal Retrain Workflow (`minimal_retrain.py`) @@ -351,7 +338,6 @@ flowchart LR - Temporal server/cluster - PostgreSQL database - MLFlow server -- OPC server(s) - MongoDB server (for notifications) **Note**: External dependencies must be available either through: @@ -398,12 +384,10 @@ flowchart LR 4. **Install Python dependencies** ```bash python -m pip install --upgrade pip + # Install production dependencies pip install -r requirements.txt - ``` - - 5. **Install test libraries** - ```bash - pip install pytest pytest-cov pytest-asyncio + # Install development and testing tools + pip install -r requirements-dev.txt ``` 4. **Create environment configuration file** @@ -499,6 +483,113 @@ fi python -m model_manager.worker.worker ``` +## πŸ” Code Quality & Validation + +### Overview + +Como Python nΓ£o Γ© uma linguagem compilada, utilizamos um conjunto robusto de ferramentas para validar a qualidade, seguranΓ§a e correção do cΓ³digo antes da execução. Estas ferramentas detectam erros, problemas de estilo, vulnerabilidades de seguranΓ§a e garantem a consistΓͺncia do cΓ³digo. + +### Ferramentas de Validação + +#### 1. **Ruff** - Linting e Formatação ⚑ +Ferramenta moderna e extremamente rΓ‘pida (escrita em Rust) que substitui mΓΊltiplas ferramentas: +- **Linting**: Detecta erros de cΓ³digo, problemas de estilo (PEP 8), bugs comuns +- **Formatação**: Formata cΓ³digo automaticamente de forma consistente +- **Velocidade**: 10-100x mais rΓ‘pido que Flake8/Black + +#### 2. **mypy** - Type Checking 🏷️ +Verificador de tipos estΓ‘ticos que analisa type hints: +- Detecta erros de tipo antes da execução +- Melhora a documentação do cΓ³digo +- Previne bugs relacionados a tipos incorretos + +#### 3. **Bandit** - AnΓ‘lise de SeguranΓ§a πŸ”’ +Scanner de vulnerabilidades de seguranΓ§a: +- Detecta padrΓ΅es inseguros de cΓ³digo +- Identifica hardcoded passwords, SQL injection, etc. +- Garante conformidade com prΓ‘ticas de seguranΓ§a + +#### 4. **pytest** - Testes Automatizados πŸ§ͺ +Framework de testes com cobertura de cΓ³digo: +- Executa testes unitΓ‘rios e de integração +- Mede cobertura de cΓ³digo +- Suporta testes assΓ­ncronos + +### Instalação das Ferramentas + +```bash +# Instalar dependΓͺncias de desenvolvimento +pip install -r requirements-dev.txt +``` + +### Validação Completa + +#### Opção 1: Script Automatizado (Recomendado) +```bash +# Executar todas as validaçáes de uma vez +./validate.sh +``` + +O script `validate.sh` executa automaticamente: +1. βœ… Verificação de formatação (Ruff) +2. βœ… Linting de cΓ³digo (Ruff) +3. βœ… Type checking (mypy) +4. βœ… AnΓ‘lise de seguranΓ§a (Bandit) +5. βœ… Testes unitΓ‘rios com cobertura (pytest) + +#### Opção 2: Comandos Individuais +```bash +# 1. Verificar formatação +ruff format --check model_manager/ tests/ + +# 2. Verificar linting +ruff check model_manager/ tests/ + +# 3. Verificar tipos +mypy model_manager/ + +# 4. AnΓ‘lise de seguranΓ§a +bandit -r model_manager/ -ll + +# 5. Executar testes +pytest tests/ --cov=model_manager --cov-report=term-missing +``` + +### Correção AutomΓ‘tica + +Algumas ferramentas podem corrigir problemas automaticamente: + +```bash +# Formatar cΓ³digo automaticamente +ruff format model_manager/ tests/ + +# Corrigir problemas de linting automaticamente +ruff check --fix model_manager/ tests/ +``` + +### Configuração + +Todas as ferramentas sΓ£o configuradas no arquivo `pyproject.toml`: +- **Ruff**: Regras de linting, formatação, complexidade +- **mypy**: Configuraçáes de type checking +- **pytest**: Opçáes de teste e cobertura +- **Bandit**: Regras de seguranΓ§a + +### Integração com CI/CD + +O workflow `.github/workflows/quality-gate.yml` executa automaticamente todas as validaçáes em cada push/PR: +- βœ… Formatação e linting bloqueiam merge se falharem +- ⚠️ Type checking e seguranΓ§a geram avisos mas nΓ£o bloqueiam +- βœ… Testes devem passar com cobertura mΓ­nima de 70% + +### Boas PrΓ‘ticas + +1. **Antes de Commit**: Execute `./validate.sh` para garantir qualidade +2. **Durante Desenvolvimento**: Use `ruff check --watch` para feedback em tempo real +3. **Type Hints**: Adicione type hints em funçáes novas para melhor validação +4. **Testes**: Mantenha cobertura acima de 70% +5. **SeguranΓ§a**: Revise e corrija todos os avisos do Bandit + ## πŸ§ͺ Testing ### Test Structure @@ -540,13 +631,6 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa - 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 -- `model_manager_prediction_opc_writing_count`: Counter for OPC server write operations - - Labels: `pod_id`, `model_name`, `pipeline_name`, `opc_server_id` -- `model_manager_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 @@ -571,14 +655,6 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa | `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 | @@ -590,45 +666,6 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa | `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No | | `POD_ID` | Kubernetes pod identifier | `None` | No | - - - -### 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: @@ -705,7 +742,6 @@ This is the configuration created by the Orchestrator in Temporal. }, "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;", @@ -726,8 +762,7 @@ model_manager/ β”œβ”€β”€ 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 +β”‚ └── mlflow.py # MLFlow model operations β”œβ”€β”€ workflows/ # Temporal workflow definitions β”‚ β”œβ”€β”€ predictions_batch.py # Main batch prediction workflow β”‚ β”œβ”€β”€ minimal_retrain.py # Model retraining workflow @@ -742,8 +777,7 @@ model_manager/ β”‚ β”‚ β”œβ”€β”€ 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 +β”‚ └── model_repository.py # MLFlow model operations β”œβ”€β”€ metrics.py # Prometheus metrics definitions └── __init__.py ``` @@ -775,12 +809,7 @@ model_manager/ - 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** +4. **Workflow Execution Failures** - Review activity error logs and notifications - Check data quality filter configurations - Verify input data format and required fields diff --git a/model_manager/activities/activities.py b/model_manager/activities/activities.py index 6a8656d..81407a7 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -6,28 +6,25 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.observability.logger import Logger from model_manager.activities.mlflow import MLFlow from model_manager.activities.gates import Gates - from model_manager.activities.opc import OPC from typing import Any -class Activities(Postgres, MLFlow, Gates, OPC): +class Activities(Postgres, MLFlow, Gates): """ Main activities orchestrator for the Model Manager system. This class combines functionality from multiple activity classes to provide a unified interface for all workflow operations. It manages database connections, - MLFlow model interactions, data quality validation, and OPC server communications. + MLFlow model interactions, and data quality validation. The class implements multiple inheritance to combine specialized functionality: - Postgres: Database operations and data persistence - MLFlow: Model inference and transformation operations - Gates: Data quality validation and filtering mechanisms - - 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 """ @@ -35,7 +32,6 @@ class Activities(Postgres, MLFlow, Gates, OPC): def __init__(self, postgres_config: dict[str, Any], mlflow_config: dict[str, Any], - opc_config: dict[str, Any], logger: Logger, notification_handler: NotificationHandler): """ @@ -49,8 +45,6 @@ class Activities(Postgres, MLFlow, Gates, OPC): 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 @@ -78,22 +72,15 @@ class Activities(Postgres, MLFlow, Gates, OPC): 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/opc.py b/model_manager/activities/opc.py deleted file mode 100644 index 2222eed..0000000 --- a/model_manager/activities/opc.py +++ /dev/null @@ -1,356 +0,0 @@ -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 model_manager.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 index 96b806f..a19b03c 100644 --- a/model_manager/metrics.py +++ b/model_manager/metrics.py @@ -13,14 +13,13 @@ 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 +- Export Operations: Database export performance - Response Times: Performance monitoring for various operations Metric Labels: - pod_id: Kubernetes pod identifier for multi-instance deployments - model_name: Name of the ML model being used - pipeline_name: Name of the prediction pipeline -- opc_server_id: Identifier for OPC server operations """ from prometheus_client import Gauge, Counter, Histogram @@ -56,17 +55,3 @@ PREDICTION_RESPONSE_TIME_MONITOR = Histogram( 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( - "model_manager_prediction_opc_writing_count", - "Number of predictions written to the OPC server", - [*CORE_LABELS, "opc_server_id"], -) - -PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram( - "model_manager_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/connectors_config.py b/model_manager/utils/connectors_config.py index 145f248..40e7c07 100644 --- a/model_manager/utils/connectors_config.py +++ b/model_manager/utils/connectors_config.py @@ -59,45 +59,6 @@ def build_mlflow_config() -> Dict[str, Any]: } -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. diff --git a/model_manager/utils/repository/opc_repository.py b/model_manager/utils/repository/opc_repository.py deleted file mode 100644 index 94b6c46..0000000 --- a/model_manager/utils/repository/opc_repository.py +++ /dev/null @@ -1,359 +0,0 @@ -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 model_manager 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/worker.py b/model_manager/worker/worker.py index 5d1d0b1..92631ab 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -42,7 +42,6 @@ with workflow.unsafe.imports_passed_through(): from model_manager.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 @@ -63,9 +62,8 @@ async def main(): 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 + 5. Starts Temporal client and workers + 6. Manages worker lifecycle and graceful shutdown The function runs indefinitely until interrupted or an error occurs. On error, it performs cleanup and exits with a non-zero status code. @@ -105,14 +103,10 @@ async def main(): 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) @@ -167,8 +161,6 @@ async def main(): 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, diff --git a/model_manager/workflows/predictions_batch.py b/model_manager/workflows/predictions_batch.py index 281b187..dec0fb7 100644 --- a/model_manager/workflows/predictions_batch.py +++ b/model_manager/workflows/predictions_batch.py @@ -58,7 +58,7 @@ class PredictionsBatch(): - 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: @@ -115,7 +115,7 @@ class PredictionsBatch(): }), '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') } diff --git a/model_manager/workflows/sub_workflows/format_and_export_prediction.py b/model_manager/workflows/sub_workflows/format_and_export_prediction.py index ef195c4..60fad76 100644 --- a/model_manager/workflows/sub_workflows/format_and_export_prediction.py +++ b/model_manager/workflows/sub_workflows/format_and_export_prediction.py @@ -14,9 +14,9 @@ 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. + data formatting, database persistence, and metrics recording. It implements flexible formatting based on prediction quality and provides - comprehensive export capabilities to multiple destinations. + comprehensive export capabilities. The workflow supports two main prediction paths: 1. Normal Prediction: Formats and exports successful prediction results @@ -24,7 +24,6 @@ class FormatAndExportPrediction(): Export Destinations: - PostgreSQL Database: Persistent storage with timestamp conversion - - OPC Servers: Real-time industrial system integration - Prometheus Metrics: Performance monitoring and operational visibility """ @@ -36,14 +35,12 @@ class FormatAndExportPrediction(): 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 + 3. Persisting data to PostgreSQL database with comprehensive metadata + 4. Recording performance metrics for operational monitoring The method implements flexible formatting strategies: - Normal predictions: Full data formatting with confidence scores - Error predictions: Default formatting with error indicators - - Comprehensive export: Multi-destination data distribution Args: input_data: Complete configuration for the export workflow @@ -58,7 +55,6 @@ class FormatAndExportPrediction(): - 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: @@ -100,18 +96,6 @@ class FormatAndExportPrediction(): 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, diff --git a/model_manager/workflows/sub_workflows/prediction_process.py b/model_manager/workflows/sub_workflows/prediction_process.py index e613ab5..c5407ec 100644 --- a/model_manager/workflows/sub_workflows/prediction_process.py +++ b/model_manager/workflows/sub_workflows/prediction_process.py @@ -64,7 +64,7 @@ class PredictionProcess(): - 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 @@ -210,7 +210,6 @@ class PredictionProcess(): '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, @@ -287,7 +286,6 @@ class PredictionProcess(): 'schema': schema, 'table_name': table_name, 'comment': comment, - 'opc_output_config': input_data['opc_output_config'], 'prediction_store_policy': input_data['prediction_store_policy'] } ) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..7b577c6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,154 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "model-manager" +version = "0.0.0" +description = "Sientia DataOps Model Manager - ML Model Orchestration System" +readme = "README.md" +requires-python = ">=3.11" +authors = [ + {name = "Aignosi", email = "dev@aignosi.com"} +] + +[tool.ruff] +line-length = 100 +target-version = "py311" +exclude = [ + ".git", + ".venv", + "venv", + "__pycache__", + "*.pyc", + ".pytest_cache", + "htmlcov", +] + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "B", # flake8-bugbear + "C4", # flake8-comprehensions + "UP", # pyupgrade + "N", # pep8-naming + "YTT", # flake8-2020 + "S", # flake8-bandit + "BLE", # flake8-blind-except + "A", # flake8-builtins + "C90", # mccabe complexity +] + +ignore = [ + "E501", # line too long (handled by formatter) + "S101", # use of assert (needed for tests) + "S105", # possible hardcoded password (false positives) + "S106", # possible hardcoded password (false positives) + "N802", # function name should be lowercase (temporal decorators) + "N806", # variable in function should be lowercase +] + +[tool.ruff.lint.per-file-ignores] +"tests/**/*.py" = [ + "S101", # assert allowed in tests + "S105", # hardcoded passwords ok in tests + "S106", # hardcoded passwords ok in tests +] + +[tool.ruff.lint.mccabe] +max-complexity = 15 + +[tool.ruff.format] +quote-style = "single" +indent-style = "space" +line-ending = "auto" + +[tool.mypy] +python_version = "3.11" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = false +disallow_incomplete_defs = false +check_untyped_defs = true +no_implicit_optional = true +warn_redundant_casts = true +warn_unused_ignores = true +warn_no_return = true +strict_equality = true +ignore_missing_imports = false + +# Ignore missing imports for external packages +[[tool.mypy.overrides]] +module = "temporalio.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "sientia_do.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "mlflow.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "prometheus_client.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "redis.*" +ignore_missing_imports = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = [ + "-v", + "--strict-markers", + "--cov=model_manager", + "--cov-report=term-missing", + "--cov-report=html", + "--cov-report=xml", +] +markers = [ + "asyncio: marks tests as async", + "integration: marks tests as integration tests", + "unit: marks tests as unit tests", +] + +[tool.coverage.run] +source = ["model_manager"] +omit = [ + "*/tests/*", + "*/venv/*", + "*/__pycache__/*", + "*/site-packages/*", +] +branch = true + +[tool.coverage.report] +precision = 2 +show_missing = true +skip_covered = false +exclude_lines = [ + "pragma: no cover", + "def __repr__", + "def __str__", + "raise AssertionError", + "raise NotImplementedError", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", + "class .*\\bProtocol\\):", + "@(abc\\.)?abstractmethod", +] + +[tool.coverage.html] +directory = "htmlcov" + +[tool.bandit] +exclude_dirs = ["tests", "venv", ".venv"] +skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..e65402b --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,17 @@ +# Development and Testing Dependencies +# These packages are only needed for development, testing, and code quality checks +# Install with: pip install -r requirements-dev.txt + +# Code Quality & Linting +ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort) +mypy>=1.7.0 # Static type checker +bandit>=1.7.5 # Security vulnerability scanner + +# Testing +pytest>=7.4.0 # Testing framework +pytest-cov>=4.1.0 # Coverage plugin for pytest +pytest-asyncio>=0.21.0 # Async test support (already in main requirements) + +# Development Tools +ipython>=8.12.0 # Enhanced Python shell +ipdb>=0.13.13 # IPython debugger diff --git a/requirements.txt b/requirements.txt index 1ce9f9c..8fcde0c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,6 @@ 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 diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 106ccf7..b6e53df 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -4,14 +4,12 @@ from sientia_do.temporal.activities.postgres import Postgres from model_manager.activities.activities import Activities from model_manager.activities.mlflow import MLFlow from model_manager.activities.gates import Gates -from model_manager.activities.opc import OPC @patch('model_manager.activities.activities.Postgres.__init__') @patch('model_manager.activities.activities.MLFlow.__init__') -@patch('model_manager.activities.activities.OPC.__init__') @patch('model_manager.activities.activities.Gates.__init__') -def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init): +def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): postgres_config = { 'host': 'localhost', @@ -30,19 +28,12 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre '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 ) @@ -50,7 +41,6 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre 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( @@ -76,13 +66,6 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre 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, @@ -93,9 +76,7 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre @mark.asyncio @patch('model_manager.activities.activities.Postgres', return_value=MagicMock()) @patch('model_manager.activities.activities.MLFlow', return_value=MagicMock()) -@patch('model_manager.activities.activities.OPC', return_value=MagicMock()) -async def test_shutdown(mock_opc_init, - _mock_mlflow_init, mock_postgres_init): +async def test_shutdown(_mock_mlflow_init, mock_postgres_init): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -113,23 +94,15 @@ async def test_shutdown(mock_opc_init, '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_opc.py b/tests/laborious/activities/test_opc.py deleted file mode 100644 index dcff019..0000000 --- a/tests/laborious/activities/test_opc.py +++ /dev/null @@ -1,369 +0,0 @@ -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 model_manager.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("model_manager.activities.opc.OpcRepository") -@patch("model_manager.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("model_manager.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/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py deleted file mode 100644 index 9795a10..0000000 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ /dev/null @@ -1,388 +0,0 @@ -import pytest -from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call -from asyncua.crypto.security_policies import SecurityPolicyBasic256 -from model_manager.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('model_manager.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('model_manager.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('model_manager.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('model_manager.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 index 43a2ec9..6179b33 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,6 +1,5 @@ from os import environ from model_manager.utils.connectors_config import (build_mlflow_config, - build_opc_config, build_postgres_config, build_mongodb_config) @@ -40,54 +39,6 @@ def test_build_mlflow_config_with_defaults(): 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' diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index 356c32b..3606d5d 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -34,8 +34,6 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): "prediction_confidence": 0, "schema": "test_schema", "table_name": "test_table", - "opc_servers": ["test_server"], - "opc_output_config": {"test": "config"}, "prediction_store_policy": "erl:1" } @@ -56,19 +54,6 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 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, @@ -86,7 +71,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): start_to_close_timeout=ANY )]) - assert workflow_mock.execute_activity_method.call_count == 3 + assert workflow_mock.execute_activity_method.call_count == 2 assert workflow_mock.execute_local_activity_method.call_count == 1 @@ -103,8 +88,6 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction "prediction_confidence": 0, "schema": "test_schema", "table_name": "test_table", - "opc_servers": ["test_server"], - "opc_output_config": {"test": "config"}, "comment": "test_comment" } @@ -125,19 +108,6 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction ) ]) - 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, diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 46b1565..6922ef1 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -38,7 +38,7 @@ async def test_run(workflow_mock, prediction_process): 'retention': '30' }, 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'}, + 'prediction_store_policy': 'lts:1' } @@ -125,7 +125,7 @@ async def test_run(workflow_mock, prediction_process): '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', @@ -153,7 +153,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): 'retention': '30' }, 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} + } # Mock the activity responses @@ -201,7 +201,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ 'retention': '30' }, 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} + } # Mock the activity responses @@ -272,7 +272,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process 'retention': '30' }, 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} + } # Mock the activity responses @@ -352,7 +352,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p 'retention': '30' }, 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} + } # Mock the activity responses @@ -536,7 +536,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): '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' ) @@ -558,7 +558,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'schema': schema, 'table_name': table_name, 'comment': 'Prediction Process', - 'opc_output_config': {'test': 'config'}, + 'prediction_store_policy': prediction_store_policy } ) @@ -590,7 +590,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): '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, "" ) diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 0d7d477..e71a462 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -32,7 +32,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch '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': { @@ -78,7 +78,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch }), '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') } diff --git a/validate.sh b/validate.sh new file mode 100755 index 0000000..6d5ab14 --- /dev/null +++ b/validate.sh @@ -0,0 +1,99 @@ +#!/bin/bash +# Model Manager Code Validation Script +# This script runs all code quality checks before committing or deploying + +set -e # Exit on any error + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' # No Color + +echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}" +echo -e "${BLUE}β•‘ Model Manager - Code Validation Suite β•‘${NC}" +echo -e "${BLUE}β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•${NC}" +echo "" + +# Check if virtual environment is activated +if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then + echo -e "${YELLOW}⚠️ Warning: No virtual environment detected${NC}" + echo -e "${YELLOW} Consider activating your venv/conda environment${NC}" + echo "" +fi + +# Function to run a validation step +run_step() { + local step_name=$1 + local step_command=$2 + + echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" + echo -e "${BLUE}β–Ά ${step_name}${NC}" + echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" + + if eval "$step_command"; then + echo -e "${GREEN}βœ… ${step_name} - PASSED${NC}" + echo "" + return 0 + else + echo -e "${RED}❌ ${step_name} - FAILED${NC}" + echo "" + return 1 + fi +} + +# Track failures +FAILED_STEPS=() + +# Step 1: Code Formatting Check (Ruff) +if ! run_step "1. Code Formatting (Ruff)" "ruff format --check model_manager/ tests/"; then + FAILED_STEPS+=("Code Formatting") +fi + +# Step 2: Linting (Ruff) +if ! run_step "2. Code Linting (Ruff)" "ruff check model_manager/ tests/"; then + FAILED_STEPS+=("Linting") +fi + +# Step 3: Type Checking (mypy) +if ! run_step "3. Type Checking (mypy)" "mypy model_manager/"; then + FAILED_STEPS+=("Type Checking") +fi + +# Step 4: Security Analysis (Bandit) +if ! run_step "4. Security Analysis (Bandit)" "bandit -r model_manager/ -ll -q"; then + FAILED_STEPS+=("Security Analysis") +fi + +# Step 5: Unit Tests (pytest) +if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-fail-under=70 -q"; then + FAILED_STEPS+=("Unit Tests") +fi + +# Summary +echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}" +echo -e "${BLUE}β•‘ Validation Summary β•‘${NC}" +echo -e "${BLUE}β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•${NC}" +echo "" + +if [ ${#FAILED_STEPS[@]} -eq 0 ]; then + echo -e "${GREEN}βœ… All validation checks passed!${NC}" + echo -e "${GREEN} Your code is ready for commit/deployment.${NC}" + echo "" + exit 0 +else + echo -e "${RED}❌ Validation failed for the following steps:${NC}" + for step in "${FAILED_STEPS[@]}"; do + echo -e "${RED} β€’ ${step}${NC}" + done + echo "" + echo -e "${YELLOW}πŸ’‘ Tips:${NC}" + echo -e "${YELLOW} β€’ Run 'ruff format model_manager/ tests/' to auto-fix formatting${NC}" + echo -e "${YELLOW} β€’ Run 'ruff check --fix model_manager/ tests/' to auto-fix linting issues${NC}" + echo -e "${YELLOW} β€’ Review mypy errors and add type hints where needed${NC}" + echo -e "${YELLOW} β€’ Check bandit warnings for security issues${NC}" + echo -e "${YELLOW} β€’ Fix failing tests or improve test coverage${NC}" + echo "" + exit 1 +fi diff --git a/values.yaml b/values.yaml index 744c0c0..fc71161 100644 --- a/values.yaml +++ b/values.yaml @@ -180,11 +180,6 @@ env: - 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" From dfc190c818608cee813bd0a48ab1aea97993ea5d Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 17:28:57 -0300 Subject: [PATCH 07/15] SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows This commit includes several changes: - Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability. - Improved error handling and logging in gates.py and mlflow.py. - Added input validation and filtering in gates.py to ensure data quality. - Enhanced prediction formatting and storage policy management in gates.py. - Updated metrics.py to use consistent naming conventions and labels. - Refactored connectors_config.py to use type hints and improve code clarity. - Updated conditional and MLFlow filters for better data quality checks. - Improved model repository logic for retraining and updating models. - Enhanced worker.py to include SDK metrics and improved error handling. - Refactored workflows for better modularity and error handling. - Updated tests to reflect the changes and improve test coverage. --- model_manager/activities/activities.py | 61 +- model_manager/activities/gates.py | 258 +++---- model_manager/activities/mlflow.py | 97 +-- model_manager/metrics.py | 24 +- model_manager/utils/connectors_config.py | 15 +- .../utils/filters/conditional_filters.py | 3 +- model_manager/utils/filters/mlflow_filters.py | 7 +- .../utils/repository/model_repository.py | 175 ++--- model_manager/worker/worker.py | 76 +- model_manager/workflows/minimal_retrain.py | 34 +- model_manager/workflows/predictions_batch.py | 47 +- .../format_and_export_prediction.py | 34 +- .../sub_workflows/prediction_process.py | 55 +- tests/laborious/activities/test_activities.py | 37 +- tests/laborious/activities/test_gates.py | 237 +++--- tests/laborious/activities/test_mlflow.py | 212 ++--- .../utils/filters/test_conditional_filters.py | 31 +- .../utils/filters/test_mlflow_filters.py | 11 +- .../utils/repository/test_model_repository.py | 335 ++++---- .../laborious/utils/test_connectors_config.py | 13 +- .../test_format_and_export_prediction.py | 197 ++--- .../subworkflows/test_prediction_process.py | 721 +++++++++++------- .../workflows/test_minimal_retrain.py | 120 +-- .../workflows/test_predictions_batch.py | 81 +- 24 files changed, 1482 insertions(+), 1399 deletions(-) diff --git a/model_manager/activities/activities.py b/model_manager/activities/activities.py index 81407a7..9fedd57 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -1,12 +1,14 @@ -from temporalio import activity, workflow +from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from sientia_do.temporal.activities.postgres import Postgres + from typing import Any + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger - from model_manager.activities.mlflow import MLFlow + from sientia_do.temporal.activities.postgres import Postgres + from model_manager.activities.gates import Gates - from typing import Any + from model_manager.activities.mlflow import MLFlow class Activities(Postgres, MLFlow, Gates): @@ -29,11 +31,13 @@ class Activities(Postgres, MLFlow, Gates): notification_handler (NotificationHandler): Notification management instance """ - def __init__(self, - postgres_config: dict[str, Any], - mlflow_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler): + def __init__( + self, + postgres_config: dict[str, Any], + mlflow_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize the Activities orchestrator with all required configurations. @@ -52,25 +56,30 @@ class Activities(Postgres, MLFlow, Gates): 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) + 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) + 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) + Gates.__init__(self, logger=logger, notification_handler=notification_handler) async def shutdown(self): """ diff --git a/model_manager/activities/gates.py b/model_manager/activities/gates.py index 338e212..091977c 100644 --- a/model_manager/activities/gates.py +++ b/model_manager/activities/gates.py @@ -1,53 +1,42 @@ from temporalio import activity, workflow - with workflow.unsafe.imports_passed_through(): import traceback + from typing import Any + + from pandas import DataFrame + from sientia_do.formatters import create_sample_dict from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel - from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.logger import Logger + from sientia_do.temporal.activities.base import BaseActivity from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now - from sientia_do.formatters import create_sample_dict - from model_manager.utils.filters.mlflow_filters import nan_values_filter, api_error_filter + from model_manager import metrics from model_manager.utils.filters.conditional_filters import ( filter_empty_data, - filter_specific_variables_null_values + filter_specific_variables_null_values, ) - from pandas import DataFrame - from model_manager import metrics + from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter # Input filter function mappings input_filter_functions = { 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, 'EMPTY_DATA': filter_empty_data, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 2, - 'REPEAT': -1 - } + '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 - }, + '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 - } + 'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1}, } @@ -82,10 +71,9 @@ class Gates(BaseActivity): Raises: Exception: If BaseActivity initialization fails """ - BaseActivity.__init__( - self, logger, notification_handler, set_error_counter=True) + BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True) - @activity.defn(name="input_gate") + @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. @@ -121,7 +109,7 @@ class Gates(BaseActivity): """ metadata = input_data['metadata'] - self.info("Performing input gate...", metadata) + self.info('Performing input gate...', metadata) filters = input_data['filters'] data = DataFrame(input_data['data']) @@ -129,40 +117,42 @@ class Gates(BaseActivity): filter_output = [] - self.debug(f"Input data: {data.head(5).to_string()}", metadata) - self.debug(f"Filters: {filters}", metadata) + 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) + 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) + self.debug(f'Data not passed the input filter {fil}:{config}', metadata) filter_output.append(config['policy']) - except Exception as e: + except Exception as e: # noqa: BLE001 trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id=f"INTPUT_GATE_ERROR__{fil}", - message=f"Error in filter {fil}:{config}: \n {e}", - block="input_gate", + 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 + 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(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, "" + self.info('Nothing was filtered by the input gate', metadata) + return None, 0, '' - @activity.defn(name="mlflow_response_gate") + @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. @@ -197,7 +187,7 @@ class Gates(BaseActivity): Exception: If response validation fails or configuration is invalid """ metadata = input_data['metadata'] - self.info("Performing mlflow response gate...", metadata) + self.info('Performing mlflow response gate...', metadata) filters = input_data['filters'] data = input_data['data'] @@ -206,14 +196,13 @@ class Gates(BaseActivity): 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) + 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) + self.error(f'Filter {fil} not found', metadata) continue try: if mlflow_response_filter_functions[fil](data, config): @@ -221,34 +210,36 @@ class Gates(BaseActivity): comments.append(data['content']['message']) self.send_notification( metadata=metadata, - notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}", + notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}', message=data['content']['message'], - block="mlflow_gate", + block='mlflow_gate', level=NotificationLevel.ERROR, - attachment_content=data['content']['traceback'] + attachment_content=data['content']['traceback'], ) - except Exception as e: + except Exception as e: # noqa: BLE001 trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}", - message=f"Error in filter {fil}:{config}: \n {e}", - block="mlflow_gate", + 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 + 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(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, "" + self.info('Nothing was filtered by the mlflow response gate', metadata) + return None, 0, '' - @activity.defn(name="mlflow_content_gate") + @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. @@ -283,7 +274,7 @@ class Gates(BaseActivity): Exception: If content validation fails or configuration is invalid """ metadata = input_data['metadata'] - self.info("Performing mlflow content gate...", metadata) + self.info('Performing mlflow content gate...', metadata) filters = input_data['filters'] data = DataFrame(input_data['data']) @@ -292,8 +283,8 @@ class Gates(BaseActivity): filter_output = [] - self.debug(f"Input data:\n {data.head(5).to_string()}", metadata) - self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata) + 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: @@ -303,36 +294,38 @@ class Gates(BaseActivity): 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", + 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() + attachment_content=data.to_string(), ) - except Exception as e: + except Exception as e: # noqa: BLE001 trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}", - message=f"Error in filter {fil}:{config}: \n {e}", - block="mlflow_gate", + 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 + 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(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, "" + 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]: + def get_prediction_store_policy( + self, prediction_store_policy: str, metadata: dict[str, Any] + ) -> tuple[str, int]: """ Parse and validate prediction store policy configuration. @@ -358,7 +351,9 @@ class Gates(BaseActivity): if len(policy_elements) < 2: self.error( - f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata) + f'Invalid prediction store policy: {prediction_store_policy}, using default policy', + metadata, + ) return 'lts', 1 policy_type = policy_elements[0] @@ -366,14 +361,20 @@ class Gates(BaseActivity): # 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: + 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) + 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") + @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. @@ -400,7 +401,7 @@ class Gates(BaseActivity): """ metadata = input_data['metadata'] prediction_store_policy = input_data['prediction_store_policy'] - self.info("Formatting prediction...", metadata) + self.info('Formatting prediction...', metadata) data = DataFrame(input_data['data']) @@ -408,48 +409,45 @@ class Gates(BaseActivity): 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) + 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) + 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) + 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) + 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) + 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}") + 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['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) + 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") + @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. @@ -477,22 +475,24 @@ class Gates(BaseActivity): """ metadata = input_data['metadata'] - self.debug("Formatting default prediction...", 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']] - }) + 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) + self.info(f'Default prediction formatted: {data.size} rows', metadata) return data.to_dict() - @activity.defn(name="get_last_timestamp") + @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. @@ -517,24 +517,22 @@ class Gates(BaseActivity): """ metadata = input_data['metadata'] - self.info("Getting last timestamp...", metadata) + self.info('Getting last timestamp...', metadata) data = DataFrame(input_data['data']) - self.debug(f"Input data: {data.head(5).to_string()}", metadata) + 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()) + max_timestamp = max(data['timestamp'].values.tolist()) - self.info( - f"Last timestamp: {max_timestamp}", metadata) + self.info(f'Last timestamp: {max_timestamp}', metadata) return max_timestamp - @activity.defn(name="write_metrics") + @activity.defn(name='write_metrics') async def write_metrics(self, input_data: dict[str, Any]): """ Write prediction performance metrics to Prometheus monitoring system. @@ -562,26 +560,24 @@ class Gates(BaseActivity): 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) + 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'] + 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'] + 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'] + pipeline_name=metadata['workflow_name'], ).observe(response_time) - self.info( - f"Metrics written for model {metadata['model_name']}", metadata) + 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 index e55a55e..4de863b 100644 --- a/model_manager/activities/mlflow.py +++ b/model_manager/activities/mlflow.py @@ -1,20 +1,19 @@ 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 + import traceback + from typing import Any + + import numpy as np + from pandas import DataFrame, to_datetime + from sientia_do.formatters import create_sample_dict from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel from sientia_do.observability.logger import Logger - from sientia_do.formatters import create_sample_dict + from sientia_do.temporal.activities.base import BaseActivity + from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + from model_manager.utils.repository.model_repository import MLFlowRepository - from typing import Any - import numpy as np - from pandas import DataFrame - import traceback class MLFlow(BaseActivity): @@ -36,8 +35,15 @@ class MLFlow(BaseActivity): 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): + 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. @@ -52,18 +58,17 @@ class MLFlow(BaseActivity): Raises: Exception: If MLFlowRepository initialization fails """ - BaseActivity.__init__( - self, logger, notification_handler, set_error_counter=True) + 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 + f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger ) - @activity.defn(name="request_transform") + @activity.defn(name='request_transform') async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Transform input data using MLFlow models. @@ -99,7 +104,7 @@ class MLFlow(BaseActivity): model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) - self.debug("Raw input data:", metadata) + 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 @@ -108,14 +113,12 @@ class MLFlow(BaseActivity): ) # Pivot data for model input format - data = data.pivot( - index='timestamp', columns='variable', - values='value') + 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('Processed input data:', metadata) self.debug(data.head(5).to_string(), metadata) # Request transformation from MLFlow model @@ -124,16 +127,20 @@ class MLFlow(BaseActivity): ) self.debug( - f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata) + 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) + f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}', + metadata, + ) - self.info("Data transformed successfully", metadata) + self.info('Data transformed successfully', metadata) return response_data - @activity.defn(name="request_predict") + @activity.defn(name='request_predict') async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Execute predictions using MLFlow models. @@ -169,14 +176,15 @@ class MLFlow(BaseActivity): 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) + 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) + data['timestamp'], format=DATETIME_FORMAT_WITH_TZ + ).dt.strftime(DATETIME_FORMAT) # Request prediction from MLFlow model response_data = self.model_monitoring_repository.predict( @@ -184,13 +192,15 @@ class MLFlow(BaseActivity): ) self.debug( - f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata) + f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}', + metadata, + ) - self.info("Data predicted successfully", metadata) + self.info('Data predicted successfully', metadata) return response_data - @activity.defn(name="retrain_model") + @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. @@ -234,8 +244,7 @@ class MLFlow(BaseActivity): 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 = data.pivot(index='timestamp', columns='variable', values='value') data.sort_index(inplace=True) data.reset_index(inplace=True) @@ -244,15 +253,10 @@ class MLFlow(BaseActivity): try: retrain_output, experiment = self.model_monitoring_repository.retrain_model( - data=data, - model_name=model_name + data=data, model_name=model_name ) - return { - 'status': retrain_output, - 'timestamp': timestamp, - 'experiment': experiment - } + return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment} except Exception as e: trace = traceback.format_exc() self.send_notification( @@ -261,12 +265,12 @@ class MLFlow(BaseActivity): message=f'Error retraining model {model_name}: {e}', block='retrain_model', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e - @activity.defn(name="update_production_model") + @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. @@ -311,12 +315,12 @@ class MLFlow(BaseActivity): status = input_data['status'] self.info( - f'Updating production model {model_name} from experiment {experiment}...', metadata) + 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 + experiment=experiment, model_name=model_name ) report = DataFrame([response]) @@ -325,8 +329,7 @@ class MLFlow(BaseActivity): report['timestamp'] = timestamp report['status'] = status - self.info( - f'Production model {model_name} updated successfully', metadata) + self.info(f'Production model {model_name} updated successfully', metadata) return report.to_dict() except Exception as e: @@ -337,7 +340,7 @@ class MLFlow(BaseActivity): message=f'Error updating production model {model_name}: {e}', block='update_production_model', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e diff --git a/model_manager/metrics.py b/model_manager/metrics.py index a19b03c..311f80f 100644 --- a/model_manager/metrics.py +++ b/model_manager/metrics.py @@ -22,36 +22,36 @@ Metric Labels: - pipeline_name: Name of the prediction pipeline """ -from prometheus_client import Gauge, Counter, Histogram +from prometheus_client import Counter, Gauge, Histogram # Application health metric APP_UP = Gauge( - "app_up", - "Indicates if the application is running (1) or shutting down (0)", - ["pod_id"], + '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"] +CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name'] # Prediction operation metrics PREDICTIONS_WRITTEN_COUNT = Counter( - "model_manager_predictions_written_count", - "Number of predictions written to the database table predictions", + 'model_manager_predictions_written_count', + 'Number of predictions written to the database table predictions', CORE_LABELS, ) # Prediction quality metrics PREDICTION_CONFIDENCE_MONITOR = Gauge( - "model_manager_prediction_confidence_monitor", - "Current confidence of each prediction", + 'model_manager_prediction_confidence_monitor', + 'Current confidence of each prediction', CORE_LABELS, ) # Performance monitoring metrics PREDICTION_RESPONSE_TIME_MONITOR = Histogram( - "model_manager_prediction_response_time_monitor", - "Current response time of each prediction", + 'model_manager_prediction_response_time_monitor', + 'Current response time of each prediction', CORE_LABELS, - buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0] + 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/connectors_config.py b/model_manager/utils/connectors_config.py index 40e7c07..aec7625 100644 --- a/model_manager/utils/connectors_config.py +++ b/model_manager/utils/connectors_config.py @@ -1,9 +1,8 @@ from os import getenv -import json -from typing import Dict, Any +from typing import Any -def build_postgres_config() -> Dict[str, Any]: +def build_postgres_config() -> dict[str, Any]: """ Build PostgreSQL database configuration from environment variables. @@ -30,11 +29,11 @@ def build_postgres_config() -> Dict[str, Any]: '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')) + 'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')), } -def build_mlflow_config() -> Dict[str, Any]: +def build_mlflow_config() -> dict[str, Any]: """ Build MLFlow server configuration from environment variables. @@ -55,11 +54,11 @@ def build_mlflow_config() -> Dict[str, Any]: 'host': getenv('MLFLOW_HOST', 'http://localhost'), 'port': int(getenv('MLFLOW_PORT', '5080')), 'username': getenv('MLFLOW_USERNAME', 'aignosi'), - 'password': getenv('MLFLOW_PASSWORD', 'aignosi') + 'password': getenv('MLFLOW_PASSWORD', 'aignosi'), } -def build_mongodb_config() -> Dict[str, Any]: +def build_mongodb_config() -> dict[str, Any]: """ Build MongoDB configuration from environment variables. @@ -86,5 +85,5 @@ def build_mongodb_config() -> Dict[str, Any]: return { 'connection_string': connection_string, 'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'), - 'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600 + 'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600, } diff --git a/model_manager/utils/filters/conditional_filters.py b/model_manager/utils/filters/conditional_filters.py index 2c51805..717a6f8 100644 --- a/model_manager/utils/filters/conditional_filters.py +++ b/model_manager/utils/filters/conditional_filters.py @@ -20,8 +20,7 @@ def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool False if none of the specified variables contain null values. """ - return not data[ - data['variable'].isin(config['variables']) & data['value'].isna()].empty + return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty def filter_empty_data(data: DataFrame, _config: dict) -> bool: diff --git a/model_manager/utils/filters/mlflow_filters.py b/model_manager/utils/filters/mlflow_filters.py index f792a0e..82970e8 100644 --- a/model_manager/utils/filters/mlflow_filters.py +++ b/model_manager/utils/filters/mlflow_filters.py @@ -52,8 +52,11 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool: 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() + data = ( + predictions.replace({None: np.nan}) + .drop(columns=['timestamp'], errors='ignore') + .infer_objects() + ) if data.isna().all().all(): return True diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py index 3107d1c..e2465a9 100644 --- a/model_manager/utils/repository/model_repository.py +++ b/model_manager/utils/repository/model_repository.py @@ -10,22 +10,23 @@ 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 datetime import datetime from os import makedirs, path, remove + +import mlflow +import pandas as pd from sientia.ModelServing import ModelServing -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from sientia_do.observability.logger import Logger +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ -class MLFlowRepository(): +class MLFlowRepository: def __init__(self, host, username, password, logger: Logger): - - self.model_serving = ModelServing(tracking_uri=host, - username=username, password=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: @@ -40,33 +41,32 @@ class MLFlowRepository(): # Get type of first element of index index_type = type(index[0]) - self.logger.custom_info(f"Index type: {index_type}", metadata) + 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}" + 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}") + raise ValueError(f'{message}') # Check type and converts to DATETIME_FORMAT_WITH_TZ - if index_type == str: + if index_type is str: # Validate format of string and return error if not valid try: pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ) - except ValueError: - raise ValueError( - f"{message}") + except ValueError as e: + raise ValueError(f'{message}') from e elif index_type == datetime or index_type == pd.Timestamp: data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) else: - raise ValueError( - f"{message}") + raise ValueError(f'{message}') return data - def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict: + def transform( + self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict + ) -> dict: """ Transform data using a model. @@ -81,41 +81,42 @@ class MLFlowRepository(): try: self.logger.custom_debug( - f"Data received for model transformation: {data.to_csv()}", metadata) + 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') + 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 + 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) + f'Data received from model transformation: {transformed_data.to_csv()}', metadata + ) - transformed_data = self.detect_and_parse_datetime_index( - transformed_data, metadata) + transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) - return { - 'success': True, - 'content': transformed_data.to_dict() - } + return {'success': True, 'content': transformed_data.to_dict()} - except Exception as e: + except Exception as e: # noqa: BLE001 return { 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, } - def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict: + def predict( + self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict + ) -> dict: """ Predict data using a model. @@ -137,31 +138,26 @@ class MLFlowRepository(): start_time = datetime.now() self.logger.custom_debug( - f"Data received for model prediction: {data.to_csv()}", metadata) + 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 + 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) + 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() - } + return {'success': True, 'content': data.to_dict()} - except Exception as e: + except Exception as e: # noqa: BLE001 return { 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, } def get_experiment_by_run_id(self, run_id: str) -> dict: @@ -190,10 +186,9 @@ class MLFlowRepository(): 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"]) + 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}" + return f'{model_name}-{next_run_number}' def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple: """ @@ -217,14 +212,10 @@ class MLFlowRepository(): - experiment: MLFlow experiment name """ # load predictor model - predictor_uri = f"models:/{model_name}/production" + 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 - ) + 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) @@ -233,20 +224,16 @@ class MLFlowRepository(): target_name = data_model.target_variable y = data[target_name] - treated_data = pd.merge( - treated_data, y, left_index=True, right_index=True) + 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): + 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. @@ -271,7 +258,7 @@ class MLFlowRepository(): """ 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" + 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 @@ -279,32 +266,32 @@ class MLFlowRepository(): # update transfomation model # fixed parameters for name_atribute, val_atribute in pred_model_atributes.items(): - if name_atribute != "model": + 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": + if name_atribute != 'model': mlflow.log_param(name_atribute, val_atribute) # dynamic parameters, including model itself - mlflow.sklearn.log_model(data_model, "data_model") + mlflow.sklearn.log_model(data_model, 'data_model') - makedirs("temp", exist_ok=True) + makedirs('temp', exist_ok=True) - file_path = f"temp/raw_data_{model_name}.csv" + 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) + 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 + return 'Model retrained successfully', experiment def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: """ @@ -325,10 +312,10 @@ class MLFlowRepository(): - status_message (str): Retraining operation status - experiment_name (str): MLFlow experiment identifier """ - prediction_model, data_model, experiment = self.create_model_experiment( - model_name, data) + 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) + prediction_model, data_model, experiment, model_name, data + ) return retrain_result def get_experiment(self, experiment_name: str) -> int: @@ -374,22 +361,21 @@ class MLFlowRepository(): """ runs = mlflow.search_runs( experiment_ids=[experiment_id], - filter_string="", # Sem filtro no MLflow ainda - output_format="pandas" + 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'] + 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) + 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'] @@ -423,16 +409,14 @@ class MLFlowRepository(): # 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) + 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 + model_versions = client.get_registered_model(model_name).latest_versions if not isinstance(model_versions, list): raise ValueError('Model versions is not a list') @@ -441,17 +425,10 @@ class MLFlowRepository(): # 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 + 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 - } + return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id} def update_production_model(self, experiment: str, model_name: str) -> dict: """ diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index 92631ab..76fc37e 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -25,32 +25,35 @@ Environment Variables: - PROJECT_NAME: Project name for notifications (default: model_manager) """ -from temporalio import workflow, client -from temporalio.worker import Worker, PollerBehaviorAutoscaling -from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig +from temporalio import client, workflow +from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig +from temporalio.worker import PollerBehaviorAutoscaling, Worker with workflow.unsafe.imports_passed_through(): + import asyncio import os import sys - import asyncio - from model_manager.workflows.minimal_retrain import MinimalRetrain - from model_manager.workflows.predictions_batch import PredictionsBatch - from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess - from model_manager.workflows.sub_workflows.format_and_export_prediction import \ - FormatAndExportPrediction - from model_manager.activities.activities import Activities - from model_manager.utils.connectors_config import ( - build_postgres_config, - build_mlflow_config, - build_mongodb_config - ) + + from prometheus_client import start_http_server from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import get_logger + from model_manager import metrics - from prometheus_client import start_http_server + from model_manager.activities.activities import Activities + from model_manager.utils.connectors_config import ( + build_mlflow_config, + build_mongodb_config, + build_postgres_config, + ) + from model_manager.workflows.minimal_retrain import MinimalRetrain + from model_manager.workflows.predictions_batch import PredictionsBatch + from model_manager.workflows.sub_workflows.format_and_export_prediction import ( + FormatAndExportPrediction, + ) + from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess POD_ID = os.getenv('POD_ID') -SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091")) +SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) async def main(): @@ -85,7 +88,7 @@ async def main(): logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata) - logger.custom_info("Starting prometheus client...", metadata) + logger.custom_info('Starting prometheus client...', metadata) start_prometheus_server() logger.custom_info('Starting Notification Handler...', metadata) @@ -95,7 +98,7 @@ async def main(): connection_string=mongo_config['connection_string'], database=mongo_config['database_name'], logger=logger, - project_name=os.getenv('PROJECT_NAME', 'model-manager') + project_name=os.getenv('PROJECT_NAME', 'model-manager'), ) logger.custom_info('Starting Activities...', metadata) @@ -104,16 +107,14 @@ async def main(): postgres_config=build_postgres_config(), mlflow_config=build_mlflow_config(), logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) - logger.custom_info( - f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata) + 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}") + metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}') ) ) @@ -122,7 +123,7 @@ async def main(): temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'), - runtime=new_runtime + runtime=new_runtime, ) logger.custom_info('Starting Workers...', metadata) @@ -136,20 +137,19 @@ async def main(): activities.load_custom_query, activities.retrain_model, activities.update_production_model, - activities.export_data_to_postgres + 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() + activity_task_poller_behavior=PollerBehaviorAutoscaling(), ), Worker( temporal_client, task_queue='predictions_batch-queue', - workflows=[PredictionsBatch, PredictionProcess, - FormatAndExportPrediction], + workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], activities=[ # MLFlow activities.request_predict, @@ -165,15 +165,15 @@ async def main(): activities.load_custom_query, activities.repeat_last_prediction, activities.export_data_to_postgres, - activities.write_metrics + 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() - ) + activity_task_poller_behavior=PollerBehaviorAutoscaling(), + ), ] handlers = [] @@ -186,8 +186,8 @@ async def main(): # 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) + except BaseException as e: # noqa: BLE001 + logger.custom_error(f'An unhandled exception occurred: {e}', metadata) finally: if notification_handler: notification_handler.shutdown() @@ -216,12 +216,12 @@ def start_prometheus_server(): SystemExit: If the metrics server fails to start """ try: - port = int(os.getenv("HTTP_METRICS_PORT", 9090)) + port = int(os.getenv('HTTP_METRICS_PORT', 9090)) start_http_server(port) - print(f"Prometheus server started on port {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}") + except Exception as e: # noqa: BLE001 + print(f'Failed to start Prometheus server: {e}') os._exit(1) diff --git a/model_manager/workflows/minimal_retrain.py b/model_manager/workflows/minimal_retrain.py index f6e2c80..03618bf 100644 --- a/model_manager/workflows/minimal_retrain.py +++ b/model_manager/workflows/minimal_retrain.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.activities.activities import Activities - from typing import Any - from sientia_do.temporal.policies import retry_policy from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from model_manager.activities.activities import Activities -@workflow.defn(name="minimal_retrain") -class MinimalRetrain(): +@workflow.defn(name='minimal_retrain') +class MinimalRetrain: """ Automated model retraining workflow for the Model Manager system. @@ -63,7 +65,7 @@ class MinimalRetrain(): 'schedule_name': input_data['schedule_name'], 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'workflow_name': 'minimal_retrain' + 'workflow_name': 'minimal_retrain', } } @@ -74,21 +76,17 @@ class MinimalRetrain(): { **metadata, 'query': input_data['query'], - 'datetime_columns': input_data.get('datetime_columns', []) + 'datetime_columns': input_data.get('datetime_columns', []), }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) experiment_response = await workflow.execute_activity_method( Activities.retrain_model, - { - **metadata, - 'data': data, - 'model_name': model_name - }, + {**metadata, 'data': data, 'model_name': model_name}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) report = await workflow.execute_activity_method( @@ -97,10 +95,10 @@ class MinimalRetrain(): **metadata, 'model_name': model_name, 'model_id': input_data['model_id'], - **experiment_response + **experiment_response, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) await workflow.execute_activity_method( @@ -109,8 +107,8 @@ class MinimalRetrain(): **metadata, 'data': report, 'schema': input_data['schema'], - 'table_name': input_data['table_name'] + 'table_name': input_data['table_name'], }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) diff --git a/model_manager/workflows/predictions_batch.py b/model_manager/workflows/predictions_batch.py index dec0fb7..b221e9f 100644 --- a/model_manager/workflows/predictions_batch.py +++ b/model_manager/workflows/predictions_batch.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.activities.activities import Activities - from typing import Any - from sientia_do.temporal.policies import retry_policy from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from model_manager.activities.activities import Activities -@workflow.defn(name="predictions_batch") -class PredictionsBatch(): +@workflow.defn(name='predictions_batch') +class PredictionsBatch: """ Main batch prediction workflow for the Model Manager system. @@ -74,7 +76,7 @@ class PredictionsBatch(): 'schedule_name': input_data['schedule_name'], 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'workflow_name': 'predictions_batch' + 'workflow_name': 'predictions_batch', } } @@ -84,10 +86,10 @@ class PredictionsBatch(): { **metadata, 'query': input_data['query'], - 'datetime_columns': input_data.get('datetime_columns', []) + 'datetime_columns': input_data.get('datetime_columns', []), }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=300) + start_to_close_timeout=timedelta(seconds=300), ) # Prepare input for prediction_process workflow @@ -98,28 +100,17 @@ class PredictionsBatch(): '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' - } - }), + 'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}), + 'mlflow_transform_filters': input_data.get( + 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}} + ), + 'mlflow_predict_filters': input_data.get( + 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}} + ), 'model_config': input_data.get('model_config', {}), 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), - - 'prediction_store_policy': input_data.get( - 'prediction_store_policy', 'lts:1') + 'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'), } # Execute prediction process workflow - await workflow.execute_child_workflow( - 'prediction_process', prediction_input) + await workflow.execute_child_workflow('prediction_process', prediction_input) diff --git a/model_manager/workflows/sub_workflows/format_and_export_prediction.py b/model_manager/workflows/sub_workflows/format_and_export_prediction.py index 60fad76..2a4266a 100644 --- a/model_manager/workflows/sub_workflows/format_and_export_prediction.py +++ b/model_manager/workflows/sub_workflows/format_and_export_prediction.py @@ -1,15 +1,17 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.activities.activities import Activities - from typing import Any from datetime import timedelta - from sientia_do.temporal.policies import retry_policy + from typing import Any + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.policies import retry_policy + + from model_manager.activities.activities import Activities -@workflow.defn(name="format_and_export_prediction") -class FormatAndExportPrediction(): +@workflow.defn(name='format_and_export_prediction') +class FormatAndExportPrediction: """ Data formatting and export workflow for prediction results. @@ -75,10 +77,10 @@ class FormatAndExportPrediction(): 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': prediction_confidence, - 'prediction_store_policy': input_data['prediction_store_policy'] + 'prediction_store_policy': input_data['prediction_store_policy'], }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) else: @@ -90,10 +92,10 @@ class FormatAndExportPrediction(): 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': prediction_confidence, - 'comment': input_data['comment'] + 'comment': input_data['comment'], }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) # write to postgres @@ -104,21 +106,15 @@ class FormatAndExportPrediction(): 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction, - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } + 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) await workflow.execute_activity_method( Activities.write_metrics, - { - **metadata, - 'prediction': prediction - }, + {**metadata, 'prediction': prediction}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + 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 index c5407ec..bac0490 100644 --- a/model_manager/workflows/sub_workflows/prediction_process.py +++ b/model_manager/workflows/sub_workflows/prediction_process.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.activities.activities import Activities - from typing import Any - from sientia_do.temporal.policies import retry_policy from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from model_manager.activities.activities import Activities -@workflow.defn(name="prediction_process") -class PredictionProcess(): +@workflow.defn(name='prediction_process') +class PredictionProcess: """ Core prediction processing workflow for the Model Manager system. @@ -84,10 +86,7 @@ class PredictionProcess(): # Get last timestamp for incremental processing last_timestamp = await workflow.execute_local_activity_method( Activities.get_last_timestamp, - { - **metadata, - 'data': data - }, + {**metadata, 'data': data}, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=1), ) @@ -97,7 +96,7 @@ class PredictionProcess(): **metadata, 'filters': input_data['input_filters'], 'data': data, - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], } path_flag, confidence, comment = await workflow.execute_local_activity_method( @@ -116,12 +115,7 @@ class PredictionProcess(): # 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 - }, + {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config}, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=5), ) @@ -134,7 +128,7 @@ class PredictionProcess(): 'filters': input_data['mlflow_transform_filters'], 'data': response_data, 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=1), @@ -155,7 +149,7 @@ class PredictionProcess(): 'filters': input_data['mlflow_transform_filters'], 'data': transformed_data, 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=1), @@ -172,7 +166,7 @@ class PredictionProcess(): **metadata, 'data': transformed_data, 'model_name': model_name, - 'model_config': model_config + 'model_config': model_config, }, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=5), @@ -186,7 +180,7 @@ class PredictionProcess(): 'filters': input_data['mlflow_predict_filters'], 'data': response_data, 'type': 'predict', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=1), @@ -213,12 +207,19 @@ class PredictionProcess(): 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'comment': comment, - 'prediction_store_policy': input_data['prediction_store_policy'] - } + '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: + 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. @@ -264,7 +265,7 @@ class PredictionProcess(): 'schema': schema, 'table_name': table_name, 'model': model_id, - 'last_timestamp': last_timestamp + 'last_timestamp': last_timestamp, }, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=1), @@ -286,8 +287,8 @@ class PredictionProcess(): 'schema': schema, 'table_name': table_name, 'comment': comment, - 'prediction_store_policy': input_data['prediction_store_policy'] - } + 'prediction_store_policy': input_data['prediction_store_policy'], + }, ) return True diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index b6e53df..b0cd4e6 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -1,16 +1,17 @@ +from unittest.mock import ANY, MagicMock, patch + from pytest import mark -from unittest.mock import patch, MagicMock, ANY from sientia_do.temporal.activities.postgres import Postgres + from model_manager.activities.activities import Activities -from model_manager.activities.mlflow import MLFlow from model_manager.activities.gates import Gates +from model_manager.activities.mlflow import MLFlow @patch('model_manager.activities.activities.Postgres.__init__') @patch('model_manager.activities.activities.MLFlow.__init__') @patch('model_manager.activities.activities.Gates.__init__') def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): - postgres_config = { 'host': 'localhost', 'port': 5432, @@ -18,15 +19,10 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } + mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} logger = MagicMock() notification_handler = MagicMock() @@ -35,7 +31,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): postgres_config=postgres_config, mlflow_config=mlflow_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) assert isinstance(activities, Activities) @@ -53,7 +49,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_mlflow_init.assert_called_once_with( @@ -63,13 +59,11 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_gates_init.assert_called_once_with( - ANY, - logger=logger, - notification_handler=notification_handler + ANY, logger=logger, notification_handler=notification_handler ) @@ -84,15 +78,10 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init): 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } + mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} logger = MagicMock() notification_handler = MagicMock() @@ -101,7 +90,7 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init): postgres_config=postgres_config, mlflow_config=mlflow_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) await activities.shutdown() diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index b587b93..ee68007 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -1,6 +1,8 @@ -from unittest.mock import MagicMock, ANY, patch +from unittest.mock import ANY, MagicMock, patch + from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel + from model_manager.activities.gates import Gates @@ -20,11 +22,11 @@ def gates_activity(): metadata = { - "metadata": { - "model_id": "test_model", - "model_name": "test_model", - "workflow_name": "test_workflow", - "schema_name": "test_schedule", + 'metadata': { + 'model_id': 'test_model', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schema_name': 'test_schedule', }, } @@ -34,20 +36,18 @@ async def test_input_gate_invalid_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, + 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, 'data': {'value': [1, 2, 3]}, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.input_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') gates_activity.error.assert_called_once_with( - "Filter INVALID_FILTER not found", metadata['metadata'] + 'Filter INVALID_FILTER not found', metadata['metadata'] ) @@ -57,28 +57,27 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac # Arrange mock_input_filter_functions.__contains__.return_value = True mock_input_filter_functions.__getitem__.return_value = MagicMock( - side_effect=Exception("Test error")) + side_effect=Exception('Test error') + ) input_data = { **metadata, - 'filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}} - }, + 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, 'data': {'value': []}, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.input_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') gates_activity.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="INTPUT_GATE_ERROR__EMPTY_DATA", + notification_id='INTPUT_GATE_ERROR__EMPTY_DATA', message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error", - block="input_gate", + block='input_gate', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) @@ -89,14 +88,14 @@ async def test_input_gate_no_filters(gates_activity): **metadata, 'filters': {}, 'data': {'value': [1, 2, 3]}, - 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] + 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } # Act result = await gates_activity.input_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') gates_activity.debug.assert_called() @@ -105,18 +104,16 @@ async def test_input_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}} - }, + 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, 'data': {'value': []}, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.input_gate(input_data) # Assert - assert result == ('STOP', -1, "Input data with bad quality") + assert result == ('STOP', -1, 'Input data with bad quality') gates_activity.debug.assert_called() @@ -125,51 +122,49 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, + 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, 'data': {'content': {'message': 'success'}}, 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.mlflow_response_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') @mark.asyncio @patch('model_manager.activities.gates.mlflow_response_filter_functions') -async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions, - gates_activity): +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")) + side_effect=Exception('Test error') + ) input_data = { **metadata, - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, + 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, 'data': {'content': {'message': 'success'}}, 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.mlflow_response_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') gates_activity.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER", + notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER', message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error", - block="mlflow_gate", + block='mlflow_gate', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) @@ -181,14 +176,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity): 'filters': {}, 'data': {'content': {'message': 'success'}}, 'type': 'test', - 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] + 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } # Act result = await gates_activity.mlflow_response_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') gates_activity.debug.assert_called() @@ -197,25 +192,20 @@ async def test_mlflow_response_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': { - 'API_ERROR': {'policy': 'STOP'} - }, + 'filters': {'API_ERROR': {'policy': 'STOP'}}, 'data': { 'success': False, - 'content': { - 'message': 'API error occurred', - 'traceback': 'error trace' - } + 'content': {'message': 'API error occurred', 'traceback': 'error trace'}, }, 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.mlflow_response_gate(input_data) # Assert - assert result == ('STOP', -1, "API error occurred") + assert result == ('STOP', -1, 'API error occurred') gates_activity.debug.assert_called() gates_activity.send_notification.assert_called() @@ -225,58 +215,53 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, + 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, 'data': {'value': [1, 2, 3]}, 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.mlflow_content_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') @mark.asyncio @patch('model_manager.activities.gates.mlflow_content_filter_functions') -async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions, - gates_activity): +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")) + side_effect=Exception('Test error') + ) input_data = { **metadata, - 'filters': { - 'API_ERROR': {'POLICY': 'STOP'} - }, + 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, 'data': { 'success': False, - 'content': { - 'message': 'API error occurred', - 'traceback': 'error trace' - } + 'content': {'message': 'API error occurred', 'traceback': 'error trace'}, }, 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } # Act result = await gates_activity.mlflow_content_gate(input_data) # Assert - assert result == (None, 0, "") + 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", + notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR', message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error", - block="mlflow_gate", + block='mlflow_gate', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) @@ -288,14 +273,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity): 'filters': {}, 'data': {'value': [1, 2, 3]}, 'type': 'test', - 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] + 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } # Act result = await gates_activity.mlflow_content_gate(input_data) # Assert - assert result == (None, 0, "") + assert result == (None, 0, '') gates_activity.debug.assert_called() @@ -304,20 +289,17 @@ async def test_mlflow_content_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': { - 'NAN_VALUES': {'policy': 'STOP', 'config': {}} - }, + 'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}}, 'data': {'value': [None, None, None]}, 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] + '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") + assert result == ('STOP', -1, 'Transformed data not passed the content filter') gates_activity.debug.assert_called() gates_activity.send_notification.assert_called() @@ -328,7 +310,8 @@ def test_get_prediction_store_policy_invalid_policy(gates_activity): # Act policy_type, policy_value = gates_activity.get_prediction_store_policy( - prediction_store_policy, metadata) + prediction_store_policy, metadata + ) # Assert assert policy_type == 'lts' @@ -341,7 +324,8 @@ def test_get_prediction_store_policy_invalid_policy_value(gates_activity): # Act policy_type, policy_value = gates_activity.get_prediction_store_policy( - prediction_store_policy, metadata) + prediction_store_policy, metadata + ) # Assert assert policy_type == 'lts' @@ -354,7 +338,8 @@ def test_get_prediction_store_policy_valid_policy_type(gates_activity): # Act policy_type, policy_value = gates_activity.get_prediction_store_policy( - prediction_store_policy, metadata) + prediction_store_policy, metadata + ) # Assert assert policy_type == 'lts' @@ -367,7 +352,8 @@ def test_get_prediction_store_policy_valid_policy(gates_activity): # Act policy_type, policy_value = gates_activity.get_prediction_store_policy( - prediction_store_policy, metadata) + prediction_store_policy, metadata + ) # Assert assert policy_type == 'erl' @@ -380,16 +366,12 @@ async def test_format_prediction_no_timestamp(gates_activity): input_data = { **metadata, 'data': { - 'prediction': { - '2023-05-26 11:12:27': 1 - }, - 'response_time': { - '2023-05-26 11:12:27': 0.1 - } + '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' + 'prediction_store_policy': 'lts:1', } # Act @@ -402,7 +384,7 @@ async def test_format_prediction_no_timestamp(gates_activity): assert result['model_id'] == {0: 'test_model'} assert result['prediction_confidence'] == {0: 0.9} assert result['prediction_status'] == {0: 'Good'} - assert result['comments'] == {0: ""} + assert result['comments'] == {0: ''} @mark.asyncio @@ -420,11 +402,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity): '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' + 'prediction_store_policy': 'erl:2', } # Act @@ -433,12 +415,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity): # 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['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: ""} + assert result['comments'] == {0: '', 1: ''} @mark.asyncio @@ -456,11 +437,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity): '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' + 'prediction_store_policy': 'lts:2', } # Act @@ -469,12 +450,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity): # 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['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: ""} + assert result['comments'] == {0: '', 1: ''} @mark.asyncio @@ -482,22 +462,23 @@ 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']}, + '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' + 'prediction_store_policy': 'lts:2', } - gates_activity.get_prediction_store_policy = MagicMock( - return_value=('invalid', 1)) + gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1)) try: - result = await gates_activity.format_prediction(input_data) + await gates_activity.format_prediction(input_data) except ValueError as e: - assert str(e) == "Invalid policy type: invalid" + assert str(e) == 'Invalid policy type: invalid' else: - assert False, "Expected ValueError" + raise AssertionError('Expected ValueError') @mark.asyncio @@ -508,7 +489,7 @@ async def test_format_default_prediction(gates_activity): 'timestamp': '2023-05-26 11:12:27', 'model_id': 'test_model', 'prediction_confidence': 0.1, - 'comment': 'Test comment' + 'comment': 'Test comment', } # Act @@ -528,12 +509,7 @@ async def test_format_default_prediction(gates_activity): @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'] - } - } + 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) @@ -545,10 +521,7 @@ async def test_get_last_timestamp_with_data(gates_activity): @mark.asyncio async def test_get_last_timestamp_no_data(gates_activity): # Arrange - input_data = { - 'data': {}, - **metadata - } + input_data = {'data': {}, **metadata} # Act result = await gates_activity.get_last_timestamp(input_data) @@ -567,30 +540,28 @@ async def test_write_metrics(mock_metrics, gates_activity): 'prediction': { 'prediction': [1, 2, 3], 'prediction_confidence': [0.9, 0.8, 0.7], - 'response_time': [0.1, 0.2, 0.3] - } + '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'] + 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 + 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'] + 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 index 243982f..80b7e85 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -1,45 +1,42 @@ -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 model_manager.activities.mlflow import MLFlow +from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel +from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + +from model_manager.activities.mlflow import MLFlow -@patch("model_manager.activities.mlflow.MLFlowRepository") +@patch('model_manager.activities.mlflow.MLFlowRepository') def test___init__(mock_mlflow_repository): mlflow = MLFlow( - mlflow_host="http://localhost", + mlflow_host='http://localhost', mlflow_port=5000, - mlflow_username="admin", - mlflow_password="admin", + mlflow_username='admin', + mlflow_password='admin', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) - assert mlflow.mlflow_host == "http://localhost" + assert mlflow.mlflow_host == 'http://localhost' assert mlflow.mlflow_port == 5000 - assert mlflow.mlflow_username == "admin" - assert mlflow.mlflow_password == "admin" + assert mlflow.mlflow_username == 'admin' + assert mlflow.mlflow_password == 'admin' - mock_mlflow_repository.assert_called_once_with( - "http://localhost:5000", "admin", "admin", ANY - ) + mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY) @fixture -@patch("model_manager.activities.mlflow.MLFlowRepository") +@patch('model_manager.activities.mlflow.MLFlowRepository') def mlflow(mock_mlflow_repository): mlflow = MLFlow( - mlflow_host="http://localhost:5000", + mlflow_host='http://localhost:5000', mlflow_port=5000, - mlflow_username="admin", - mlflow_password="admin", + mlflow_username='admin', + mlflow_password='admin', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) mlflow.send_notification = MagicMock() @@ -48,44 +45,67 @@ def mlflow(mock_mlflow_repository): metadata = { - "metadata": { - "model_id": "test_model", - "model_name": "test_model", - "workflow_name": "test_workflow", - "schema_name": "test_schedule", + 'metadata': { + 'model_id': 'test_model', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schema_name': 'test_schedule', }, } @mark.asyncio -@patch("model_manager.activities.mlflow.DataFrame") -@patch("model_manager.activities.mlflow.max") +@patch('model_manager.activities.mlflow.DataFrame') +@patch('model_manager.activities.mlflow.max') async def test_request_transform_success(mock_max, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data input_data = { **metadata, 'data': [ - {'timestamp': '2024-01-01', 'variable': 'var1', - 'value': 1.0, 'created_at': '2024-01-01 12:00:00'}, - {'timestamp': '2024-01-01', 'variable': 'var2', - 'value': 2.0, 'created_at': '2024-01-01 12:00:00'}, - {'timestamp': '2024-01-02', 'variable': 'var1', - 'value': 3.0, 'created_at': '2024-01-02 12:00:00'}, - {'timestamp': '2024-01-02', 'variable': 'var2', - 'value': 4.0, 'created_at': '2024-01-02 12:00:00'}, - {'timestamp': '2024-01-02', 'variable': 'var1', - 'value': 1.0, 'created_at': '2024-01-01 12:00:00'}, - {'timestamp': '2024-01-02', 'variable': 'var2', - 'value': 1.0, 'created_at': '2024-01-01 12:00:00'} + { + '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': {} + 'model_config': {}, } # Mock the transform response - expected_response = {'prediction': [0.5, 0.6], 'timestamp': [ - '2024-01-01', '2024-01-02']} + 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 @@ -114,30 +134,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow): @mark.asyncio -@patch("model_manager.activities.mlflow.DataFrame") -@patch("model_manager.activities.mlflow.to_datetime") -@patch("model_manager.activities.mlflow.max") +@patch('model_manager.activities.mlflow.DataFrame') +@patch('model_manager.activities.mlflow.to_datetime') +@patch('model_manager.activities.mlflow.max') async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data input_data = { **metadata, 'data': { - "variable": { - "2024-01-01": "var1", - "2024-01-02": "var2", - "2024-01-03": "var1", - "2024-01-04": "var2" + '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 - } + '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': {} + 'model_config': {}, } # Mock the predict response @@ -148,9 +163,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo 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.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 ) @@ -158,9 +171,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo 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 - ) + mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) # Verify the response assert response_data == expected_response @@ -174,28 +185,26 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo @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] + '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') + 'Model retrained successfully', + 'test', + ) - response = await mlflow.retrain_model({ - **metadata, - 'data': data, - 'model_name': 'test_model' - }) + 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' + 'status': 'Model retrained successfully', + 'timestamp': 2, + 'experiment': 'test', } @@ -206,20 +215,16 @@ async def test_retrain_model_error(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] + '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: + await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'}) + except Exception as e: # noqa: BLE001 assert str(e) == 'Error retraining model' mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], @@ -227,20 +232,18 @@ async def test_retrain_model_error(mlflow): message='Error retraining model test_model: Error retraining model', block='retrain_model', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "No exception raised" + raise AssertionError('No exception raised') @mark.asyncio async def test_update_production_model(mlflow): - mlflow.model_monitoring_repository.update_production_model.return_value = ( - { - "data1": 1, - "data2": 2 - } - ) + mlflow.model_monitoring_repository.update_production_model.return_value = { + 'data1': 1, + 'data2': 2, + } input_data = { **metadata, @@ -248,13 +251,14 @@ async def test_update_production_model(mlflow): 'model_id': 1, 'experiment': 'test', 'timestamp': 2, - 'status': 'success' + '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') + experiment='test', model_name='test_model' + ) assert response == { 'data1': {0: 1}, @@ -262,7 +266,7 @@ async def test_update_production_model(mlflow): 'model_id': {0: 1}, 'model_name': {0: 'test_model'}, 'timestamp': {0: 2}, - 'status': {0: 'success'} + 'status': {0: 'success'}, } @@ -278,12 +282,12 @@ async def test_update_production_model_error(mlflow): 'model_id': 1, 'experiment': 'test', 'timestamp': 2, - 'status': 'success' + 'status': 'success', } try: await mlflow.update_production_model(input_data) - except Exception as e: + except Exception as e: # noqa: BLE001 assert str(e) == 'Error updating production model' mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], @@ -291,7 +295,7 @@ async def test_update_production_model_error(mlflow): message='Error updating production model test_model: Error updating production model', block='update_production_model', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "No exception raised" + raise AssertionError('No exception raised') diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index a7f3efd..e23105e 100644 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -1,23 +1,29 @@ from pandas import DataFrame from model_manager.utils.filters.conditional_filters import ( + filter_empty_data, 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 + 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 + assert ( + filter_specific_variables_null_values( + DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}), + config={'variables': ['variable2']}, + ) + is True + ) def test_filter_empty_data(): @@ -25,6 +31,7 @@ def test_filter_empty_data(): def test_filter_empty_data_with_data(): - assert filter_empty_data( - DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), - {}) is False + 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 index 6fddf51..43e96a5 100644 --- a/tests/laborious/utils/filters/test_mlflow_filters.py +++ b/tests/laborious/utils/filters/test_mlflow_filters.py @@ -1,22 +1,23 @@ from pandas import DataFrame + from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter def test_api_error_filter_invalid_response(): - assert api_error_filter(None, {}) == True # NOSONAR + assert api_error_filter(None, {}) def test_api_error_filter_valid_response_fail(): - assert api_error_filter({'success': False}, {}) == True + assert api_error_filter({'success': False}, {}) def test_api_error_filter_valid_response_success(): - assert api_error_filter({'success': True}, {}) == False + assert not api_error_filter({'success': True}, {}) def test_nan_values_filter_all_nan_values(): - assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True + assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) def test_nan_values_filter_no_nan_values(): - assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False + assert not nan_values_filter(DataFrame({'variable': [1, 2]}), {}) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index ab65913..58fcf7e 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,34 +1,33 @@ +from datetime import UTC, datetime 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 pandas import DataFrame, Timestamp + from model_manager.utils.repository.model_repository import MLFlowRepository @pytest.fixture def mlflow_repository(): - with patch('model_manager.utils.repository.model_repository.ModelServing', - autospec=True) as mock_model_serving: + with patch( + 'model_manager.utils.repository.model_repository.ModelServing', autospec=True + ) as mock_model_serving: mock_instance = mock_model_serving.return_value mock_instance.get_transformed_data = MagicMock() repo = MLFlowRepository( - host='http://localhost:5000', - username='admin', - password='admin', - logger=MagicMock() + 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", + 'metadata': { + 'model_id': 'test_model', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schema_name': 'test_schedule', }, } @@ -38,80 +37,56 @@ class Any: 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 - } - } - ) + ({'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) +@pytest.mark.parametrize('data', invalid_cases) def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data): - input_data = DataFrame( - data - ) + input_data = DataFrame(data) with pytest.raises(ValueError) as e: - mlflow_repository.detect_and_parse_datetime_index( - input_data, metadata['metadata']) + 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" + 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': {'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 + datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1, + datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2, } - }, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'] + }, + ['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 + Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1, + Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2, } - }, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'] + }, + ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'], ), ] -@pytest.mark.parametrize("data,expected", valid_cases) +@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']) + response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) assert response.index.tolist() == expected @@ -122,18 +97,19 @@ def test_transform_success(mlflow_repository): mlflow_repository.detect_and_parse_datetime_index = MagicMock() - output = mlflow_repository.transform( - model_name, data, {}, metadata['metadata']) + 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') + 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']) + 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 + 'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value, } @@ -141,80 +117,48 @@ def test_transform_error(mlflow_repository): data = MagicMock() model_name = 'model' - mlflow_repository.model_serving.get_cached_transform.side_effect = Exception( - 'error') + mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error') - output = mlflow_repository.transform( - model_name, data, {}, metadata['metadata']) + 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') + model_name, data, 0, 'sklearn', False, 'model', 'predict' + ) - assert output == { - 'success': False, - 'content': { - 'message': 'error', - 'traceback': ANY - } - } + 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 - } - }) + 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] - ) + mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3]) - output = mlflow_repository.predict( - model_name, data, {}, metadata['metadata']) + 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') + 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 - } + '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 - } - }) + 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') - ) + mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error')) - output = mlflow_repository.predict( - model_name, data, {}, metadata['metadata']) + 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') + model_name, data, 0, 'pyfunc', False, 'model' + ) - assert output == { - 'success': False, - 'content': { - 'message': 'error', - 'traceback': ANY - } - } + assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}} @patch('model_manager.utils.repository.model_repository.mlflow') @@ -246,8 +190,7 @@ def test_get_next_run_name(mlflow, mlflow_repository): @patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_success(mlflow, mlflow_repository): - mlflow.get_experiment_by_name.return_value = MagicMock( - experiment_id='0') + mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0') output = mlflow_repository.get_experiment('test') @@ -263,23 +206,25 @@ def test_get_experiment_error(mlflow, mlflow_repository): except ValueError as e: assert str(e) == 'Experiment test not found' else: - assert False + raise AssertionError('Expected exception') @patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_last_run(mlflow, mlflow_repository): - mlflow.search_runs.return_value = DataFrame({ - 'params.retrain': ['True', 'False', 'True', 'False'], - 'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'], - 'run_id': ['0', '1', '2', '3'], - }) + 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", + filter_string='', + output_format='pandas', ) assert output == '2' @@ -294,17 +239,14 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository): except ValueError as e: assert str(e) == 'Runs is not a pandas DataFrame' else: - assert False + raise AssertionError('Expected exception') @patch('model_manager.utils.repository.model_repository.mlflow.sklearn') @patch('model_manager.utils.repository.model_repository.mlflow.set_experiment') def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): - - mlflow_repository.model_serving.get_model_run_id = MagicMock( - return_value='0') - mlflow_repository.model_serving.get_model_uri = MagicMock( - return_value='test') + 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() @@ -313,29 +255,30 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): 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.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] - }) + 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) + '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"), - ]) + 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) @@ -343,21 +286,23 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): fit_args = prediction_model_mock.fit.call_args[0][0] assert fit_args.equals( - DataFrame({ - 'x': [10, 20, 30], - 'y': [4, 5, 6], - }) + 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 - ) + 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) + assert output == ( + prediction_model_mock, + data_model_mock, + mlflow_repository.get_experiment_by_run_id.return_value, + ) @patch('model_manager.utils.repository.model_repository.mlflow.start_run') @@ -365,41 +310,43 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): @patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model') @patch('model_manager.utils.repository.model_repository.mlflow.log_artifact') def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository): - prediction_model_mock = MagicMock() data_model_mock = MagicMock() experiment = 'test' model_name = 'test' data = MagicMock() - mlflow_repository.get_next_run_name = MagicMock( - return_value='test-1') + 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) + 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') + 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"), - ]) + 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) + 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_artifact.assert_called_once_with('temp/raw_data_test.csv') - log_param.assert_has_calls([ - call("retrain", True), - ]) + log_param.assert_has_calls( + [ + call('retrain', True), + ] + ) - assert output == ("Model retrained successfully", experiment) + assert output == ('Model retrained successfully', experiment) def test_retrain_model(mlflow_repository): @@ -407,18 +354,18 @@ def test_retrain_model(mlflow_repository): model_name = 'test' mlflow_repository.create_model_experiment = MagicMock( - return_value=('data_model', 'prediction_model', '0')) + return_value=('data_model', 'prediction_model', '0') + ) - mlflow_repository.perform_model_retrain = MagicMock( - return_value='Model retrained successfully') + 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.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) + 'data_model', 'prediction_model', '0', model_name, data + ) assert output == 'Model retrained successfully' @@ -438,7 +385,7 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): output = mlflow_repository.update_production_model_by_run_id('0', 'test') mlflow.register_model.assert_called_once_with( - "runs:/0/prediction_model", + 'runs:/0/prediction_model', 'test', ) @@ -461,38 +408,34 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): @patch('model_manager.utils.repository.model_repository.mlflow') def test_update_production_model_by_run_id_error(mlflow, mlflow_repository): mlflow.tracking.MlflowClient.return_value = MagicMock( - get_registered_model=MagicMock( - return_value=MagicMock( - latest_versions={} - ) - ) + get_registered_model=MagicMock(return_value=MagicMock(latest_versions={})) ) try: mlflow_repository.update_production_model_by_run_id('0', 'test') - except Exception as e: + except Exception as e: # noqa: BLE001 assert str(e) == 'Model versions is not a list' else: - assert False + raise AssertionError('Expected exception') def test_update_production_model(mlflow_repository): connector = mlflow_repository - with patch.object(connector, 'get_experiment', - return_value='0') as get_experiment: - with patch.object(connector, 'get_experiment_last_run', - return_value='2') as get_experiment_last_run: - with patch.object(connector, 'update_production_model_by_run_id', - return_value={'model_name': 'test', 'version': '3', - 'mlflow_run_id': '0'}) as update_production_model_by_run_id: - + 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') + update_production_model_by_run_id.assert_called_once_with('2', 'test') assert output == { 'model_name': 'test', diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 6179b33..0f5e005 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,7 +1,10 @@ from os import environ -from model_manager.utils.connectors_config import (build_mlflow_config, - build_postgres_config, - build_mongodb_config) + +from model_manager.utils.connectors_config import ( + build_mlflow_config, + build_mongodb_config, + build_postgres_config, +) def test_build_mlflow_config_with_env_vars(): @@ -95,7 +98,7 @@ def test_build_mongo_db_config_with_env_vars(): assert build_mongodb_config() == { 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', 'database_name': 'test_db', - 'ttl_index_seconds': 3600 + 'ttl_index_seconds': 3600, } @@ -108,5 +111,5 @@ def test_build_mongo_db_config_with_defaults(): assert build_mongodb_config() == { 'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018', 'database_name': 'sientia', - 'ttl_index_seconds': 3600 + '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 index 3606d5d..fb83fac 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -1,9 +1,12 @@ -from unittest.mock import call, patch, AsyncMock, ANY -from pytest import mark, fixture +from unittest.mock import ANY, AsyncMock, call, patch + +from pytest import fixture, mark +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from model_manager.activities.activities import Activities -from model_manager.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +from model_manager.workflows.sub_workflows.format_and_export_prediction import ( + FormatAndExportPrediction, +) @fixture @@ -12,119 +15,133 @@ def format_and_export_prediction(): metadata = { - "metadata": { - "model_id": "test_model", - "model_name": "test_model", - "workflow_name": "test_workflow", - "schema_name": "test_schedule", + 'metadata': { + 'model_id': 'test_model', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schema_name': 'test_schedule', }, } @mark.asyncio -@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@patch( + 'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow', + new_callable=AsyncMock, +) async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): - input_data = { 'metadata': metadata, - "path_flag": None, - "data": {"test": "data"}, - "timestamp": "2021-01-01", - "model_id": 1, - "prediction_confidence": 0, - "schema": "test_schema", - "table_name": "test_table", - "prediction_store_policy": "erl:1" + 'path_flag': None, + 'data': {'test': 'data'}, + 'timestamp': '2021-01-01', + 'model_id': 1, + 'prediction_confidence': 0, + 'schema': 'test_schema', + 'table_name': 'test_table', + 'prediction_store_policy': 'erl:1', } await format_and_export_prediction.run(input_data) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.format_prediction, - { - 'data': input_data['data'], - 'timestamp': input_data['timestamp'], - 'model_id': input_data['model_id'], - 'prediction_confidence': input_data['prediction_confidence'], - 'prediction_store_policy': input_data['prediction_store_policy'], - **metadata - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.format_prediction, + { + 'data': input_data['data'], + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': input_data['prediction_confidence'], + 'prediction_store_policy': input_data['prediction_store_policy'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value, - **metadata, - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } - }, - 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 == 2 assert workflow_mock.execute_local_activity_method.call_count == 1 @mark.asyncio -@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@patch( + 'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow', + new_callable=AsyncMock, +) async def test_run_default_path_flag(workflow_mock, format_and_export_prediction): - input_data = { 'metadata': metadata, - "path_flag": "default", - "data": {"test": "data"}, - "timestamp": "2021-01-01", - "model_id": 1, - "prediction_confidence": 0, - "schema": "test_schema", - "table_name": "test_table", - "comment": "test_comment" + 'path_flag': 'default', + 'data': {'test': 'data'}, + 'timestamp': '2021-01-01', + 'model_id': 1, + 'prediction_confidence': 0, + 'schema': 'test_schema', + 'table_name': 'test_table', + 'comment': 'test_comment', } await format_and_export_prediction.run(input_data) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.format_default_prediction, - { - 'timestamp': input_data['timestamp'], - 'model_id': input_data['model_id'], - 'prediction_confidence': input_data['prediction_confidence'], - 'comment': input_data['comment'], - **metadata - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.format_default_prediction, + { + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': input_data['prediction_confidence'], + 'comment': input_data['comment'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value, - **metadata, - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } - }, - 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 index 6922ef1..d45585a 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -1,5 +1,7 @@ -from unittest.mock import AsyncMock, patch, call, ANY +from unittest.mock import ANY, AsyncMock, call, patch + from pytest import fixture, mark + from model_manager.activities.activities import Activities from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess @@ -10,17 +12,17 @@ def prediction_process(): metadata = { - "metadata": { - "model_id": "test_model", - "model_name": "test_model", - "workflow_name": "test_workflow", - "schema_name": "test_schedule", + 'metadata': { + 'model_id': 'test_model', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schema_name': 'test_schedule', }, } @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_run(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=False) # Arrange @@ -34,26 +36,23 @@ async def test_run(workflow_mock, prediction_process): 'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'}, 'model_name': 'test_model_name', - 'model_config': { - 'retention': '30' - }, + 'model_config': {'retention': '30'}, 'path_priority': ['continue', 'repeat', 'stop'], - - 'prediction_store_policy': 'lts:1' + '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 + ('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"), + ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) - ('continue', 0.95, "Transformed data not passed the content filter"), + ('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"), + ('continue', 0.95, 'Error'), ] # Act @@ -62,57 +61,112 @@ async def test_run(workflow_mock, prediction_process): # 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_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', @@ -125,17 +179,16 @@ async def test_run(workflow_mock, prediction_process): 'model_id': 1, 'model_name': 'test_model_name', 'model_config': input_data['model_config'], - 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'comment': 'Error', - 'prediction_store_policy': input_data['prediction_store_policy'] - } + 'prediction_store_policy': input_data['prediction_store_policy'], + }, ) @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_run_stop_at_input_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=True) # Arrange @@ -149,17 +202,14 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): 'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'}, 'model_name': 'test_model_name', - 'model_config': { - 'retention': '30' - }, + 'model_config': {'retention': '30'}, 'path_priority': ['continue', 'repeat', 'stop'], - } # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ '2024-01-01', # get_last_timestamp - ('stop', 0.95, "Input data with bad quality"), # input_gate + ('stop', 0.95, 'Input data with bad quality'), # input_gate ] # Act @@ -167,23 +217,35 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): # 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_local_activity_method.assert_has_calls( + [ + call( + Activities.get_last_timestamp, + { + 'data': input_data['data'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + call( + Activities.input_gate, + { + 'filters': input_data['input_filters'], + 'data': input_data['data'], + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + ] + ) workflow_mock.execute_child_workflow.assert_not_called() @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True]) # Arrange @@ -197,19 +259,16 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ 'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'}, 'model_name': 'test_model_name', - 'model_config': { - 'retention': '30' - }, + 'model_config': {'retention': '30'}, 'path_priority': ['continue', 'repeat', 'stop'], - } # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ '2024-01-01', # get_last_timestamp - ('repeat', 0.95, "Input data with bad quality"), # input_gate + ('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) + ('continue', 0.95, 'Error'), # mlflow_response_gate (transform) ] # Act @@ -217,46 +276,72 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ # 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_local_activity_method.assert_has_calls( + [ + call( + Activities.get_last_timestamp, + { + 'data': input_data['data'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.input_gate, + { + 'filters': input_data['input_filters'], + 'data': input_data['data'], + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.request_transform, + { + 'data': input_data['data'], + 'model_name': input_data['model_name'], + 'model_config': input_data['model_config'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.mlflow_response_gate, + { + 'filters': input_data['mlflow_transform_filters'], + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'type': 'transform', + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) workflow_mock.execute_child_workflow.assert_not_called() @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock( - side_effect=[False, False, True]) + prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True]) # Arrange input_data = { 'metadata': metadata, @@ -268,22 +353,19 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process 'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'}, 'model_name': 'test_model_name', - 'model_config': { - 'retention': '30' - }, + 'model_config': {'retention': '30'}, 'path_priority': ['continue', 'repeat', 'stop'], - } # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ '2024-01-01', # get_last_timestamp - ('continue', 0.95, "Input data with bad quality"), # input_gate + ('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"), + ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) - ('continue', 0.95, "Transformed data not passed the content filter"), + ('continue', 0.95, 'Transformed data not passed the content filter'), ] # Act @@ -292,51 +374,88 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process # 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_local_activity_method.assert_has_calls( + [ + call( + Activities.get_last_timestamp, + { + 'data': input_data['data'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.input_gate, + { + 'filters': input_data['input_filters'], + 'data': input_data['data'], + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.request_transform, + { + 'data': input_data['data'], + 'model_name': input_data['model_name'], + 'model_config': input_data['model_config'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.mlflow_response_gate, + { + 'filters': input_data['mlflow_transform_filters'], + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'type': 'transform', + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.mlflow_content_gate, + { + 'filters': input_data['mlflow_transform_filters'], + 'data': 'transformed_data', + 'type': 'transform', + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) workflow_mock.execute_child_workflow.assert_not_called() @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock( - side_effect=[False, False, False, True]) + prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True]) # Arrange input_data = { 'metadata': metadata, @@ -348,24 +467,21 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p 'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'}, 'model_name': 'test_model_name', - 'model_config': { - 'retention': '30' - }, + 'model_config': {'retention': '30'}, 'path_priority': ['continue', 'repeat', 'stop'], - } # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ '2024-01-01', # get_last_timestamp - ('continue', 0.95, "Input data with bad quality"), # input_gate + ('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"), + ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) - ('continue', 0.95, "Transformed data not passed the content filter"), + ('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) + ('continue', 0.95, 'Error'), # mlflow_response_gate (predict) ] # Act @@ -373,63 +489,117 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p # 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_local_activity_method.assert_has_calls( + [ + call( + Activities.get_last_timestamp, + { + 'data': input_data['data'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.input_gate, + { + 'filters': input_data['input_filters'], + 'data': input_data['data'], + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.request_transform, + { + 'data': input_data['data'], + 'model_name': input_data['model_name'], + 'model_config': input_data['model_config'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.mlflow_response_gate, + { + 'filters': input_data['mlflow_transform_filters'], + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'type': 'transform', + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.mlflow_content_gate, + { + 'filters': input_data['mlflow_transform_filters'], + 'data': 'transformed_data', + 'type': 'transform', + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.request_predict, + { + 'data': 'transformed_data', + 'model_name': input_data['model_name'], + 'model_config': input_data['model_config'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.mlflow_response_gate, + { + 'filters': input_data['mlflow_predict_filters'], + 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + 'type': 'predict', + 'path_priority': input_data['path_priority'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) workflow_mock.execute_child_workflow.assert_not_called() @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_path_flag_handler_stop(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -440,21 +610,24 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process): model = 'test_model' last_timestamp = '2024-01-01' model_name = 'test_model_name' - model_config = { - 'retention': '30' - } + model_config = {'retention': '30'} # Act result = await prediction_process.path_flag_handler( - data, path_flag, { + 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, "" + 'model_config': model_config, + }, + confidence, + last_timestamp, + '', ) # Assert @@ -464,7 +637,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process): @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_path_flag_handler_repeat(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -475,21 +648,24 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process): model = 'test_model' last_timestamp = '2024-01-01' model_name = 'test_model_name' - model_config = { - 'retention': '30' - } + model_config = {'retention': '30'} # Act result = await prediction_process.path_flag_handler( - data, path_flag, { + 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, "" + 'model_config': model_config, + }, + confidence, + last_timestamp, + '', ) # Assert @@ -504,13 +680,13 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process): 'last_timestamp': last_timestamp, }, retry_policy=ANY, - start_to_close_timeout=ANY + start_to_close_timeout=ANY, ) workflow_mock.execute_child_workflow.assert_not_called() @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_path_flag_handler_continue(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -521,14 +697,14 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): model = 'test_model' last_timestamp = '2024-01-01' model_name = 'test_model_name' - model_config = { - 'retention': '30' - } + model_config = {'retention': '30'} prediction_store_policy = 'erl:1' # Act result = await prediction_process.path_flag_handler( - data, path_flag, { + data, + path_flag, + { 'metadata': metadata, 'schema': schema, 'table_name': table_name, @@ -536,9 +712,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'last_timestamp': last_timestamp, 'model_name': model_name, 'model_config': model_config, - - 'prediction_store_policy': prediction_store_policy - }, confidence, last_timestamp, 'Prediction Process' + 'prediction_store_policy': prediction_store_policy, + }, + confidence, + last_timestamp, + 'Prediction Process', ) # Assert @@ -558,14 +736,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'schema': schema, 'table_name': table_name, 'comment': 'Prediction Process', - - 'prediction_store_policy': prediction_store_policy - } + 'prediction_store_policy': prediction_store_policy, + }, ) @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) async def test_path_flag_handler_unknown(workflow_mock, prediction_process): # Arrange data = {'test': 'data'} @@ -576,13 +753,13 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): model = 'test_model' last_timestamp = '2024-01-01' model_name = 'test_model_name' - model_config = { - 'retention': '30' - } + model_config = {'retention': '30'} prediction_store_policy = 'erl:1' # Act result = await prediction_process.path_flag_handler( - data, path_flag, { + data, + path_flag, + { **metadata, 'schema': schema, 'table_name': table_name, @@ -590,9 +767,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): 'last_timestamp': last_timestamp, 'model_name': model_name, 'model_config': model_config, - - 'prediction_store_policy': prediction_store_policy - }, confidence, last_timestamp, "" + 'prediction_store_policy': prediction_store_policy, + }, + confidence, + last_timestamp, + '', ) # Assert diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index 1d8e6d7..5296e80 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -1,5 +1,7 @@ -from unittest.mock import AsyncMock, MagicMock, call, patch, ANY +from unittest.mock import ANY, AsyncMock, call, patch + from pytest import fixture, mark + from model_manager.activities.activities import Activities from model_manager.workflows.minimal_retrain import MinimalRetrain @@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain: metadata = { - "metadata": { - "model_id": "test_model_id", - "model_name": "test_model", - "workflow_name": "minimal_retrain", - "schedule_name": "test_schedule", + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'minimal_retrain', + 'schedule_name': 'test_schedule', }, } @@ -23,19 +25,19 @@ metadata = { @patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): input_data = { - "model_id": "test_model_id", - "model_name": "test_model", - "workflow_name": "minimal_retrain", - "schedule_name": "test_schedule", - "query": "test_query", - "schema": "test_schema", - "table_name": "test_table", + '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", + 'data1': '1', + 'data2': '2', } ) @@ -47,52 +49,58 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): Activities.load_custom_query, { **metadata, - "query": input_data["query"], - 'datetime_columns': input_data.get('datetime_columns', []) + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), }, retry_policy=ANY, - start_to_close_timeout=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.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.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 - ) - ]) + 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 index e71a462..c49ffea 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -1,5 +1,7 @@ -from unittest.mock import AsyncMock, call, patch, ANY +from unittest.mock import ANY, AsyncMock, call, patch + from pytest import fixture, mark + from model_manager.activities.activities import Activities from model_manager.workflows.predictions_batch import PredictionsBatch @@ -10,11 +12,11 @@ def predictions_batch() -> PredictionsBatch: metadata = { - "metadata": { - "model_id": "test_model_id", - "model_name": "test_model", - "workflow_name": "predictions_batch", - "schedule_name": "test_schedule", + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'predictions_batch', + 'schedule_name': 'test_schedule', }, } @@ -22,9 +24,7 @@ metadata = { @mark.asyncio @patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): - workflow_mock.execute_local_activity_method.return_value = { - 'data': 'test_data' - } + workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'} input_data = { 'schedule_name': 'test_schedule', 'model_name': 'test_model', @@ -32,28 +32,27 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'query': 'SELECT * FROM test', 'schema': 'test_schema', 'table_name': 'test_table', - 'datetime_columns': ['timestamp', 'created_at'], 'prediction_store_policy': 'erl:1', - 'model_config': { - 'retention': '30' - } + '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 - ) - ]) + 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'}, @@ -61,28 +60,18 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch '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' - } - }), + 'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}), + 'mlflow_transform_filters': input_data.get( + 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}} + ), + 'mlflow_predict_filters': input_data.get( + 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}} + ), 'model_config': input_data.get('model_config', {}), 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), - - 'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1') + 'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'), } - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'prediction_process', prediction_input) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [call('prediction_process', prediction_input)] + ) From 8bcee4c6a072bb76f2e958e9e2c45dc6e61cd376 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 17:36:47 -0300 Subject: [PATCH 08/15] SIENTIAPDE-1243: Fix: Resolved mypy errors and added pandas stubs. Addressed type hinting issues and suppressed mypy warnings to improve code quality and maintainability. --- model_manager/activities/mlflow.py | 2 +- model_manager/utils/repository/model_repository.py | 7 +++---- pyproject.toml | 12 ++++++++++-- requirements-dev.txt | 2 ++ 4 files changed, 16 insertions(+), 7 deletions(-) diff --git a/model_manager/activities/mlflow.py b/model_manager/activities/mlflow.py index 4de863b..9e5a900 100644 --- a/model_manager/activities/mlflow.py +++ b/model_manager/activities/mlflow.py @@ -330,7 +330,7 @@ class MLFlow(BaseActivity): report['status'] = status self.info(f'Production model {model_name} updated successfully', metadata) - return report.to_dict() + return report.to_dict() # type: ignore[no-any-return] except Exception as e: trace = traceback.format_exc() diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py index e2465a9..c52b3f3 100644 --- a/model_manager/utils/repository/model_repository.py +++ b/model_manager/utils/repository/model_repository.py @@ -17,7 +17,7 @@ from os import makedirs, path, remove import mlflow import pandas as pd -from sientia.ModelServing import ModelServing +from sientia.ModelServing import ModelServing # type: ignore[import-untyped] from sientia_do.observability.logger import Logger from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ @@ -214,7 +214,7 @@ class MLFlowRepository: # 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') + latest_production_id = self.model_serving.get_model_info(model_name) # type: ignore[no-any-return] transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False) # load data_model = mlflow.sklearn.load_model(transform_uri) @@ -336,11 +336,10 @@ class MLFlowRepository: 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) + return experiment.experiment_id # type: ignore[no-any-return] def get_experiment_last_run(self, experiment_id: int) -> str: """ diff --git a/pyproject.toml b/pyproject.toml index 7b577c6..85823aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,14 +68,14 @@ line-ending = "auto" [tool.mypy] python_version = "3.11" -warn_return_any = true +warn_return_any = false warn_unused_configs = true disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = true no_implicit_optional = true warn_redundant_casts = true -warn_unused_ignores = true +warn_unused_ignores = false warn_no_return = true strict_equality = true ignore_missing_imports = false @@ -101,6 +101,14 @@ ignore_missing_imports = true module = "redis.*" ignore_missing_imports = true +[[tool.mypy.overrides]] +module = "sientia.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "pandas.*" +ignore_missing_imports = true + [tool.pytest.ini_options] testpaths = ["tests"] python_files = ["test_*.py"] diff --git a/requirements-dev.txt b/requirements-dev.txt index e65402b..50bf389 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -6,6 +6,8 @@ ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort) mypy>=1.7.0 # Static type checker bandit>=1.7.5 # Security vulnerability scanner +pandas-stubs>=2.0.0 # Type stubs for pandas +types-requests>=2.31.0 # Type stubs for requests # Testing pytest>=7.4.0 # Testing framework From ef3c2a0c8a56569ff9a7cbaa983e1b2642b69203 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 17:44:45 -0300 Subject: [PATCH 09/15] SIENTIAPDE-1243: Add type ignores and fix datetime index formatting in gates and model repository. --- model_manager/activities/gates.py | 18 +++++++++--------- .../utils/repository/model_repository.py | 2 +- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/model_manager/activities/gates.py b/model_manager/activities/gates.py index 091977c..897d57d 100644 --- a/model_manager/activities/gates.py +++ b/model_manager/activities/gates.py @@ -126,9 +126,9 @@ class Gates(BaseActivity): self.error(f'Filter {fil} not found', metadata) continue try: - if input_filter_functions[fil](data, config['config']): + if input_filter_functions[fil](data, config['config']): # type: ignore[operator] self.debug(f'Data not passed the input filter {fil}:{config}', metadata) - filter_output.append(config['policy']) + filter_output.append(config['policy']) # type: ignore[index] except Exception as e: # noqa: BLE001 trace = traceback.format_exc() self.send_notification( @@ -145,7 +145,7 @@ class Gates(BaseActivity): self.info(f'Input gate result: {path_flag}', metadata) return ( path_flag, - input_filter_functions['path_confidence'][path_flag], + input_filter_functions['path_confidence'][path_flag], # type: ignore[index] 'Input data with bad quality', ) @@ -205,8 +205,8 @@ class Gates(BaseActivity): self.error(f'Filter {fil} not found', metadata) continue try: - if mlflow_response_filter_functions[fil](data, config): - filter_output.append(config['policy']) + if mlflow_response_filter_functions[fil](data, config): # type: ignore[operator] + filter_output.append(config['policy']) # type: ignore[index] comments.append(data['content']['message']) self.send_notification( metadata=metadata, @@ -232,7 +232,7 @@ class Gates(BaseActivity): self.info(f'Mlflow response gate result: {path_flag}', metadata) return ( path_flag, - mlflow_response_filter_functions['path_confidence'][path_flag], + mlflow_response_filter_functions['path_confidence'][path_flag], # type: ignore[index] ', '.join(comments), ) @@ -290,8 +290,8 @@ class Gates(BaseActivity): if fil not in mlflow_content_filter_functions: continue try: - if mlflow_content_filter_functions[fil](data, config): - filter_output.append(config['policy']) + if mlflow_content_filter_functions[fil](data, config): # type: ignore[operator] + filter_output.append(config['policy']) # type: ignore[index] self.send_notification( metadata=metadata, notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}', @@ -316,7 +316,7 @@ class Gates(BaseActivity): self.info(f'Mlflow content gate result: {path_flag}', metadata) return ( path_flag, - mlflow_content_filter_functions['path_confidence'][path_flag], + mlflow_content_filter_functions['path_confidence'][path_flag], # type: ignore[index] 'Transformed data not passed the content filter', ) diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py index c52b3f3..5dafc76 100644 --- a/model_manager/utils/repository/model_repository.py +++ b/model_manager/utils/repository/model_repository.py @@ -58,7 +58,7 @@ class MLFlowRepository: raise ValueError(f'{message}') from e elif index_type == datetime or index_type == pd.Timestamp: - data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) + data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined] else: raise ValueError(f'{message}') From 07e699aae2669526115f6c56a67620414c30532f Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 17:54:02 -0300 Subject: [PATCH 10/15] SIENTIAPDE-1243: Fix: Corrected dataframe inference and mlflow repository tests, and updated activity method calls in format and export prediction tests. --- model_manager/utils/filters/mlflow_filters.py | 2 +- tests/laborious/utils/repository/test_model_repository.py | 8 +++----- .../subworkflows/test_format_and_export_prediction.py | 6 +++--- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/model_manager/utils/filters/mlflow_filters.py b/model_manager/utils/filters/mlflow_filters.py index 82970e8..493ed9c 100644 --- a/model_manager/utils/filters/mlflow_filters.py +++ b/model_manager/utils/filters/mlflow_filters.py @@ -54,8 +54,8 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool: """ data = ( predictions.replace({None: np.nan}) + .infer_objects(copy=False) .drop(columns=['timestamp'], errors='ignore') - .infer_objects() ) if data.isna().all().all(): diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 58fcf7e..116f1e4 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -194,7 +194,7 @@ def test_get_experiment_success(mlflow, mlflow_repository): output = mlflow_repository.get_experiment('test') - assert output == 0 + assert output == '0' @patch('model_manager.utils.repository.model_repository.mlflow') @@ -245,7 +245,7 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository): @patch('model_manager.utils.repository.model_repository.mlflow.sklearn') @patch('model_manager.utils.repository.model_repository.mlflow.set_experiment') def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): - mlflow_repository.model_serving.get_model_run_id = MagicMock(return_value='0') + mlflow_repository.model_serving.get_model_info = MagicMock(return_value='0') mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test') mlflow_repository.get_experiment_by_run_id = MagicMock() @@ -268,9 +268,7 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): 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_info.assert_called_once_with('test') mlflow_repository.model_serving.get_model_uri.assert_called_once_with('0', prediction=False) sklearn.load_model.assert_has_calls( diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index fb83fac..e1fd744 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -69,7 +69,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): { 'schema': input_data['schema'], 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value, + 'data': workflow_mock.execute_local_activity_method.return_value, **metadata, 'timestamp_conversion': { 'column': 'timestamp', @@ -130,7 +130,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction { 'schema': input_data['schema'], 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value, + 'data': workflow_mock.execute_local_activity_method.return_value, **metadata, 'timestamp_conversion': { 'column': 'timestamp', @@ -143,5 +143,5 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction ] ) - assert workflow_mock.execute_activity_method.call_count == 3 + assert workflow_mock.execute_activity_method.call_count == 2 assert workflow_mock.execute_local_activity_method.call_count == 1 From c12973410c7655426afca7c196f0710966fd6592 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 18:02:39 -0300 Subject: [PATCH 11/15] SIENTIAPDE-1243: Add .bandit/ and validate.txt to .gitignore to exclude code quality tool cache files. --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index a450ce7..312e581 100644 --- a/.gitignore +++ b/.gitignore @@ -51,6 +51,10 @@ coverage.xml .pytest_cache/ cover/ +# Code quality tools cache +.bandit/ +validate.txt + # Translations *.mo *.pot From 8f4ac3284ef1d037bded9b60eb93051c1490f5c2 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Wed, 1 Oct 2025 18:16:21 -0300 Subject: [PATCH 12/15] SIENTIAPDE-1243: Remove Redis dependency and related configurations. --- README.md | 1 - pyproject.toml | 4 ---- requirements.txt | 1 - 3 files changed, 6 deletions(-) diff --git a/README.md b/README.md index 28a38c3..ce79521 100644 --- a/README.md +++ b/README.md @@ -267,7 +267,6 @@ The **PredictionProcess** workflow implements the core prediction pipeline for M 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] diff --git a/pyproject.toml b/pyproject.toml index 85823aa..a7d7439 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,10 +97,6 @@ ignore_missing_imports = true module = "prometheus_client.*" ignore_missing_imports = true -[[tool.mypy.overrides]] -module = "redis.*" -ignore_missing_imports = true - [[tool.mypy.overrides]] module = "sientia.*" ignore_missing_imports = true diff --git a/requirements.txt b/requirements.txt index 8fcde0c..685fa78 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,6 @@ temporalio psycopg2-binary sqlalchemy -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 From 9a1355fe9b2dc294e254468f35258611d8041fca Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Thu, 2 Oct 2025 10:08:02 -0300 Subject: [PATCH 13/15] SIENTIAPDE-1243: Add Temporal namespace setup guide to README and improve validation script. --- README.md | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++ run_local.sh | 7 +++- validate.sh | 2 +- 3 files changed, 109 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index ce79521..5a007cc 100644 --- a/README.md +++ b/README.md @@ -345,6 +345,109 @@ flowchart LR - Cloud-managed services - Local installations +### Temporal Namespace Setup + +The Model Manager requires a dedicated Temporal namespace to isolate workflows and maintain proper execution history. The namespace must be created **before** starting the application. + +#### Why Create a Namespace? + +- **Isolation**: Separates Model Manager workflows from other applications +- **Retention Control**: Configures workflow history retention (default: 7 days) +- **Multi-tenancy**: Enables multiple environments (dev, staging, prod) on same cluster +- **Security**: Allows namespace-level access control and permissions + +#### When to Create? + +- βœ… **Before first deployment** in any environment +- βœ… **Once per environment** (dev, staging, production) +- βœ… **After Temporal cluster setup** or upgrade + +#### How to Create the Namespace + +**Option 1: Using Temporal Admin Tools Pod (Recommended for Kubernetes)** + +```bash +# 1. List Temporal pods +kubectl get pods -n temporal + +# 2. Connect to admin tools pod +kubectl exec -it -n temporal -- bash + +# 3. Create namespace +tctl --namespace model-manager namespace register \ + --retention 7 \ + --description "Model Manager - ML Model Orchestration Namespace" + +# 4. Verify creation +tctl --namespace model-manager namespace describe + +# 5. Exit pod +exit +``` + +**Option 2: Using Port Forward (Local Development)** + +```bash +# 1. Port forward Temporal frontend +kubectl port-forward -n temporal svc/temporal-frontend 7233:7233 + +# 2. In another terminal, create namespace +tctl --address localhost:7233 \ + --namespace model-manager \ + namespace register \ + --retention 7 \ + --description "Model Manager - ML Model Orchestration Namespace" + +# 3. Verify +tctl --address localhost:7233 --namespace model-manager namespace describe +``` + +**Option 3: Direct kubectl exec (One-liner)** + +```bash +kubectl exec -n temporal -- \ + tctl --namespace model-manager namespace register \ + --retention 7 \ + --description "Model Manager - ML Model Orchestration Namespace" +``` + +#### Namespace Configuration + +| Parameter | Value | Description | +|-----------|-------|-------------| +| **Name** | `model-manager` | Namespace identifier (configurable via `TEMPORAL_NAMESPACE` env var) | +| **Retention** | `7 days` | Workflow history retention period | +| **Description** | `Model Manager - ML Model Orchestration Namespace` | Human-readable description | + +#### Verification + +To verify the namespace was created successfully: + +```bash +# List all namespaces +kubectl exec -n temporal -- tctl namespace list + +# Describe specific namespace +kubectl exec -n temporal -- \ + tctl --namespace model-manager namespace describe +``` + +#### Troubleshooting + +**Error: "namespace already exists"** +- βœ… This is fine! The namespace is already configured +- No action needed, proceed with application deployment + +**Error: "connection refused"** +- ❌ Temporal server is not accessible +- Verify Temporal cluster is running: `kubectl get pods -n temporal` +- Check network connectivity and port forwarding + +**Error: "permission denied"** +- ❌ Insufficient permissions to create namespace +- Contact cluster administrator for namespace creation +- Or request elevated permissions for your service account + ## πŸš€ Installation ### Local Development Setup diff --git a/run_local.sh b/run_local.sh index 680ff15..d4c8c10 100755 --- a/run_local.sh +++ b/run_local.sh @@ -3,10 +3,12 @@ # Exit on any error set -e -echo "Activating virtual environment..." -source ./venv/bin/activate +#echo "Activating virtual environment..." + +#conda activate ./venv echo "Loading environment variables from .env..." + if [ -f .env ]; then export $(cat .env | grep -v '^#' | xargs) echo "Environment variables loaded from .env" @@ -15,4 +17,5 @@ else fi echo "Starting ingestor application..." + python -m model_manager.worker.worker diff --git a/validate.sh b/validate.sh index 6d5ab14..4bab810 100755 --- a/validate.sh +++ b/validate.sh @@ -67,7 +67,7 @@ if ! run_step "4. Security Analysis (Bandit)" "bandit -r model_manager/ -ll -q"; fi # Step 5: Unit Tests (pytest) -if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-fail-under=70 -q"; then +if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-fail-under=80 -q"; then FAILED_STEPS+=("Unit Tests") fi From 4285ddb59dcc61a0bac0750391ff7c459b51ec3c Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Thu, 2 Oct 2025 12:18:52 -0300 Subject: [PATCH 14/15] SIENTIAPDE-1243: Improve versioning logic in quality gate and optimize workflow This commit refactors the version calculation logic in the quality gate workflow to handle initial releases and improve accuracy. It also adds disk space cleanup to the workflow and uses --no-cache-dir when installing dependencies to prevent caching issues. --- .github/workflows/quality-gate.yml | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml index 997202a..4cb59fc 100644 --- a/.github/workflows/quality-gate.yml +++ b/.github/workflows/quality-gate.yml @@ -101,16 +101,20 @@ jobs: per_page: 1 }); - let lastReleaseVersion = '0.0.0'; + let lastReleaseVersion = null; + let newVersion = null; + if (releases.length > 0) { lastReleaseVersion = releases[0].tag_name.replace(/^v/, ''); console.log(`Latest release version: ${lastReleaseVersion}`); + // Calcular a nova versΓ£o baseada na ΓΊltima release + newVersion = calculateVersion(lastReleaseVersion, branchName); } else { - console.log('No releases found, using 0.0.0 as baseline'); + console.log('No releases found, starting from 0.0.0'); + lastReleaseVersion = 'N/A'; + // Primeira release: comeΓ§ar com 0.0.0 independente do tipo de branch + newVersion = '0.0.0'; } - - // Calcular a nova versΓ£o baseada no branch - const newVersion = calculateVersion(lastReleaseVersion, branchName); console.log(`Calculated version: ${newVersion}`); // Exportar a versΓ£o como output @@ -176,6 +180,18 @@ jobs: run: | git config --global url."https://oauth2:${GH_APP_TOKEN}@github.com/".insteadOf "https://github.com/" + - name: 🧹 Free Disk Space + run: | + echo "Disk space before cleanup:" + df -h + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf /opt/hostedtoolcache/CodeQL + sudo docker image prune --all --force + echo "Disk space after cleanup:" + df -h + - name: πŸ”§ Setup Python uses: actions/setup-python@v4 with: @@ -192,8 +208,8 @@ jobs: - name: πŸ“¦ Install Dependencies run: | python -m pip install --upgrade pip - pip install -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }} - pip install -r requirements-dev.txt + pip install --no-cache-dir -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }} + pip install --no-cache-dir -r requirements-dev.txt - name: πŸ“ Code Formatting Check (Ruff) run: | From 2f84d7ac2096ea1ad1015c2de2f17a2f43e85075 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Thu, 2 Oct 2025 12:37:55 -0300 Subject: [PATCH 15/15] SIENTIAPDE-1243: Update SonarQube Scan Action to v6 in quality-gate workflow. --- .github/workflows/quality-gate.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml index 4cb59fc..7d78999 100644 --- a/.github/workflows/quality-gate.yml +++ b/.github/workflows/quality-gate.yml @@ -240,7 +240,7 @@ jobs: pytest tests --junitxml=pytest.xml --cov=model_manager --cov-report=xml --cov-report=term - name: Run SonarQube Analysis - uses: SonarSource/sonarqube-scan-action@v5 + uses: SonarSource/sonarqube-scan-action@v6 env: SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }} SONAR_HOST_URL: ${{ secrets.SONAR_HOST_URL }}