diff --git a/.env.example b/.env.example index 865a73b..6e621e9 100644 --- a/.env.example +++ b/.env.example @@ -26,4 +26,10 @@ 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 +MONGODB_TTL_INDEX_HOURS="1" + +MINIO_ENDPOINT_URL="http://localhost:9000" +MINIO_ACCESS_KEY="sientia" +MINIO_SECRET_KEY="sientia" +MINIO_REGION_NAME="sa-east-1" +MINIO_DEFAULT_BUCKET="sientia" \ No newline at end of file diff --git a/.github/workflows/quality-gate.yml b/.github/workflows/quality-gate.yml index 9212f77..0d3ff14 100644 --- a/.github/workflows/quality-gate.yml +++ b/.github/workflows/quality-gate.yml @@ -30,12 +30,12 @@ jobs: owner: 'Aignosi' repositories: 'sientia-dataops-library,sientia-mlops-library' - - name: Prepare requirements.txt + - name: Prepare requirements-light.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 + requirements-light.txt > requirements_prepared.txt echo "PROCESSED_REQUIREMENTS_FILE=requirements_prepared.txt" >> $GITHUB_OUTPUT - name: Configure Git to use App Token diff --git a/.gitignore b/.gitignore index cfaf446..7cd11b3 100644 --- a/.gitignore +++ b/.gitignore @@ -37,6 +37,7 @@ __pycache__/ # Ignorar coverage htmlcov/ .coverage +coverage.xml # git keys git_key* @@ -45,4 +46,9 @@ git_log .env -tmp/ \ No newline at end of file +tmp/ +catboost_info/ + +.ruff_cache/ +.mypy_cache/ +mlruns/ \ No newline at end of file diff --git a/README.md b/README.md index 500153c..a941013 100644 --- a/README.md +++ b/README.md @@ -1,42 +1,108 @@ # Sientia DataOps Laborious -A high-performance, scalable machine learning prediction system built on Temporal.io for industrial data processing and ML model inference. The Laborious system provides enterprise-grade ML model management, batch prediction processing, and real-time data export capabilities with comprehensive data quality validation and monitoring. +A comprehensive, Temporal-based ML orchestration system for industrial data processing and model inference. Laborious delivers enterprise-grade batch prediction, model management, optional real-time export (OPC), and automated retraining with strong data quality validation and observability. + +## ๐Ÿ“‘ Table of Contents + +- [Features](#features) + - [Core Functionality](#core-functionality) + - [Advanced Capabilities](#advanced-capabilities) + - [Development & Quality Assurance](#development--quality-assurance) +- [Architecture](#architecture) + - [Architecture Principles](#architecture-principles) + - [Key Components](#key-components) + - [Data Flow Architecture](#data-flow-architecture) + - [Security Architecture](#security-architecture) +- [Workflows](#workflows) + - [Predictions Batch Workflow](#1-predictions-batch-workflow-predictions_batchpy) + - [Prediction Process Workflow](#2-prediction-process-workflow-prediction_processpy) + - [Format and Export Prediction Workflow](#3-format-and-export-prediction-workflow-format_and_export_predictionpy) + - [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy) +- [Installation & Setup](#installation--setup) + - [Prerequisites](#prerequisites) + - [Environment Setup](#environment-setup) + - [Temporal Namespace Setup](#temporal-namespace-setup) + - [Local Development Setup](#local-development-setup) +- [How to Run](#how-to-run) + - [Running the Laborious Application](#running-the-laborious-application) + - [Running Tests and Coverage](#running-tests-and-coverage) + - [Manual Test Execution](#manual-test-execution) + - [Manual Application Execution](#manual-application-execution) +- [Code Quality & Validation](#code-quality--validation) + - [Overview](#overview) + - [Validation Tools](#validation-tools) + - [Tools Installation](#tools-installation) + - [Complete Validation](#complete-validation) + - [Automatic Fixes](#automatic-fixes) + - [Configuration](#configuration) + - [CI/CD Integration](#cicd-integration) + - [Best Practices](#best-practices) +- [Testing](#testing) + - [Test Structure](#test-structure) + - [Test Execution](#test-execution) +- [Monitoring and Metrics](#monitoring-and-metrics) + - [Application Health Metrics](#application-health-metrics) + - [Prediction Operation Metrics](#prediction-operation-metrics) + - [OPC Export Metrics](#opc-export-metrics) + - [Data Quality Metrics](#data-quality-metrics) +- [Configuration](#configuration-1) + - [Environment Variables](#environment-variables) + - [OPC Configuration](#opc-configuration) + - [Workflow Configuration](#workflow-configuration) +- [Development](#development) + - [Project Structure](#project-structure) + - [Adding New Features](#adding-new-features) +- [Troubleshooting](#troubleshooting) + - [Common Issues](#common-issues) + - [Debug Mode](#debug-mode) +- [Performance Tuning](#performance-tuning) + - [Key Parameters](#key-parameters) + - [Scaling Considerations](#scaling-considerations) +- [Contributing](#contributing) + - [Code Quality Standards](#code-quality-standards) +- [License](#license) +- [Support](#support) ## 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 +- **Batch Prediction Processing**: High-throughput ML inference using MLFlow models +- **Temporal Workflow Orchestration**: Robust workflow management with retries and fault tolerance +- **Data Quality Gates**: Configurable filtering for input data and MLFlow API responses +- **Multi-Model Support**: Flexible model management with retention and versioning +- **Optional Real-time Export**: PostgreSQL persistence and OPC server integration for industrial systems +- **Comprehensive Monitoring**: Prometheus metrics and structured logging for observability ### Advanced Capabilities -- **Incremental Data Processing**: Timestamp-based data loading to avoid reprocessing +- **Incremental Data Processing**: Timestamp-based 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 +- **Notification System**: Integrated alerting via MongoDB +- **Scalable Architecture**: Kubernetes-ready with horizontal scaling +- **Model Retraining**: Automated retraining workflows with production model updates + +### Development & Quality Assurance +- **Code Quality Tools**: Ruff (lint/format), mypy (types), Bandit (security) +- **Automated Validation**: `validate.sh` and CI quality gates +- **Comprehensive Testing**: pytest with async support and high coverage +- **Type Safety**: Static type checking with mypy +- **Coverage Visualization**: Coverage Gutters integration ## 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. - +Laborious uses a Temporal-based architecture with strong separation of concerns and defensive error handling for production ML. ### 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 +- **Worker Layer**: Temporal workers, task queues, lifecycle +- **Workflow Layer**: Business orchestration and coordination +- **Activity Layer**: External system interactions and isolated operations +- **Data Layer**: Persistence, caching, connectors #### 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 +- **Automatic Retry Policies** for transient failures +- **Graceful Degradation** and circuit breaking for dependencies +- **Detailed Error Handling** with notifications #### 3. **Scalability & Performance** - **Horizontal Scaling**: Multiple worker instances for load distribution @@ -53,101 +119,57 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa ### 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` +- Temporal client setup, worker lifecycle, task queues +- Metrics server initialization, notification handler setup +- Graceful shutdown and autoscaling-friendly behavior #### **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 +- `predictions_batch.py`: Batch prediction entry point +- `sub_workflows/prediction_process.py`: Core prediction pipeline +- `sub_workflows/format_and_export_prediction.py`: Formatting and export +- `minimal_retrain.py`: Automated model retraining and production update #### **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 +- `gates.py`: Data quality validation and filtering +- `mlflow.py`: Transform and predict operations +- `opc.py`: OPC UA export to industrial systems (optional) +- `activities.py`: Aggregates activity interfaces #### **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 +- `connectors_config.py`: Env-driven configuration builders +- `repository/model_repository.py`: MLFlow operations and retraining +- `repository/opc_repository.py`: OPC communication and writes +- `filters/conditional_filters.py` and `filters/mlflow_filters.py` ### Data Flow Architecture #### **1. Batch Prediction Pipeline** ``` -Input Data (PostgreSQL) โ†’ Data Quality Gates โ†’ MLFlow Transform โ†’ -MLFlow Prediction โ†’ Response Validation โ†’ Export (PostgreSQL + OPC) +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 โ†’ +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 +- **MLFlow API Authentication**: Username/password +- **Database Security**: Encrypted connections and credential management +- **OPC Certificates** (if enabled): Client/server certs +- **Kubernetes Secrets**: Secure secret storage #### **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 +- TLS/SSL, network policies, service mesh, firewalls, VPN #### **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 +- At-rest/in-transit encryption, RBAC, audit logging, lifecycle management -## ๐Ÿ”„ Workflows +## Workflows ### 1. Predictions Batch Workflow (`predictions_batch.py`) @@ -351,8 +373,9 @@ flowchart LR - Temporal server/cluster - PostgreSQL database - MLFlow server -- OPC server(s) +- MinIO object storage (for MLFlow artifacts) - MongoDB server (for notifications) +- OPC server(s) if using OPC export **Note**: External dependencies must be available either through: - Kubernetes cluster deployment @@ -488,6 +511,68 @@ fi python -m laborious.worker.worker ``` +## Code Quality & Validation + +### Overview + +Since Python is not compiled, we validate quality, security, and correctness before execution. + +### Validation Tools + +- Ruff: Linting and formatting +- mypy: Static type checking +- Bandit: Security analysis +- pytest: Unit/integration testing with coverage + +### Tools Installation + +```bash +pip install -r requirements-dev.txt +``` + +### Complete Validation + +Option 1 (recommended): +```bash +./validate.sh +``` +The script runs, in order: +1. Format check (Ruff) +2. Linting (Ruff) +3. Type checking (mypy) +4. Security analysis (Bandit) +5. Tests with coverage (pytest) + +Option 2 (individual commands): +```bash +ruff format --check laborious/ tests/ +ruff check laborious/ tests/ +mypy laborious/ +bandit -r laborious/ -ll +pytest tests/ --cov=laborious --cov-report=term-missing +``` + +### Automatic Fixes + +```bash +ruff format laborious/ tests/ +ruff check --fix laborious/ tests/ +``` + +### Configuration + +All settings reside in `pyproject.toml` (Ruff, mypy, pytest, Bandit). + +### CI/CD Integration + +The workflow at `.github/workflows/quality-gate.yml` executes validations on each push/PR. + +### Best Practices + +- Run `./validate.sh` before committing +- Use `ruff check --watch` for continuous feedback +- Add type hints and tests for new code + ## ๐Ÿงช Testing ### Test Structure diff --git a/encode.sh b/encode.sh new file mode 100755 index 0000000..2beb0d1 --- /dev/null +++ b/encode.sh @@ -0,0 +1,13 @@ +source ./venv/bin/activate + +pip install pathspec +pip install pyyaml + +echo " +.git" >> .gitignore + +python encrypt.py ./ code --ignore .gitignore --chunk-size 100000 + +sed -i '/.git/d' .gitignore + +xdg-open . \ No newline at end of file diff --git a/encrypt.py b/encrypt.py new file mode 100644 index 0000000..6aff252 --- /dev/null +++ b/encrypt.py @@ -0,0 +1,112 @@ +import os +import argparse +from pathspec import PathSpec +import yaml + +''' +Usage: + python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000 +''' + + +def load_ignore_patterns(ignore_file, include_library): + # Ensure the .gitignore file exists + if not os.path.exists(ignore_file): + raise FileNotFoundError(f"Ignore file not found at {ignore_file}") + + # Load and parse the .gitignore patterns + with open(ignore_file, 'r') as file: + patterns = file.readlines() + if not include_library: + patterns.append('**/deploy/library/') + + spec = PathSpec.from_lines('gitwildmatch', patterns) + return spec + + +def is_ignored(file_path, spec): + """Check if a file should be ignored based on the ignore patterns.""" + return spec.match_file(file_path) if spec else False + + +def encode_file_tree_to_yaml(directory, ignore_file, include_library): + """Encode the file tree into a single YAML file.""" + ignore_patterns = load_ignore_patterns( + ignore_file, include_library) if ignore_file else None + file_tree = {} + + for root, dirs, files in os.walk(directory): + # Skip ignored directories + dirs[:] = [d for d in dirs if not is_ignored( + os.path.join(root, d), ignore_patterns)] + + for file in files: + file_path = os.path.join(root, file) + + # Skip ignored files + if is_ignored(file_path, ignore_patterns): + continue + + # Read file content + try: + with open(file_path, 'r', encoding='utf-8') as f: + content = f.read() + except Exception as e: + print(f"Error reading file {file_path}: {e}") + raise + + # Create nested dictionary structure + path_parts = os.path.relpath(file_path, directory).split(os.sep) + current_level = file_tree + + # all except the last part (the file name) + for part in path_parts[:-1]: + current_level = current_level.setdefault(part, {}) + + # Add the file and its content + current_level[path_parts[-1]] = content + return yaml.dump(file_tree, default_flow_style=False) + + +def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None): + """Chunk the YAML content and write it to the output file.""" + + chunks = [yaml_content] if chunk_size is None else [ + yaml_content[i:i + chunk_size] for i in range(0, len(yaml_content), chunk_size)] + + for i, chunk in enumerate(chunks): + chunk_file = f"{output_file}_{i}.yaml" + # Write the file tree to the output YAML file + with open(chunk_file, 'w', encoding='utf-8') as yaml_file: + yaml_file.write(chunk) + + +def main(): + parser = argparse.ArgumentParser( + description="Encrypts file tree to yaml file") + parser.add_argument("input_directory", help="Directory to encode") + parser.add_argument("output_yaml_file", help="Output YAML file") + parser.add_argument("--ignore", default=None, + help="Path to the ignore file") + parser.add_argument("--chunk-size", type=int, default=None, + help="Chunk size for the output YAML file") + parser.add_argument("--library", type=bool, default=False, + help="Incude the library in the output YAML file") + + # Parse arguments + args = parser.parse_args() + + # Example usage + directory_to_encode = args.input_directory + ignore_file_path = args.ignore + output_yaml_file = args.output_yaml_file + include_library = args.library + + content = encode_file_tree_to_yaml( + directory_to_encode, ignore_file_path, include_library) + chunk_and_write_file_tree_to_yaml( + content, output_yaml_file, args.chunk_size) + + +if __name__ == "__main__": + main() diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index ec5ae46..5487d2a 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -1,16 +1,18 @@ -from temporalio import activity, workflow +from temporalio import 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 + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.observability.logger import Logger -class Activities(Postgres, MLFlow, Gates, OPC): + from laborious.activities.gates import Gates + from laborious.activities.mlflow import MLFlow + from laborious.activities.opc import OPC + from laborious.activities.storage import Storage + + +class Activities(Storage, MLFlow, Gates, OPC): """ Main activities orchestrator for the Laborious system. @@ -32,12 +34,15 @@ class Activities(Postgres, MLFlow, Gates, OPC): 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): + def __init__( + self, + postgres_config: dict[str, Any], + mlflow_config: dict[str, Any], + minio_config: dict[str, Any], + opc_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize the Activities orchestrator with all required configurations. @@ -58,30 +63,36 @@ class Activities(Postgres, MLFlow, Gates, OPC): 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) + Storage.__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'], + minio_config=minio_config, + 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'], + minio_config=minio_config, + logger=logger, + notification_handler=notification_handler, + ) - Gates.__init__(self, 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) + OPC.__init__( + self, opc_servers=opc_config, logger=logger, notification_handler=notification_handler + ) async def shutdown(self): """ @@ -95,5 +106,5 @@ class Activities(Postgres, MLFlow, Gates, OPC): The method should be called before the application terminates to ensure proper resource cleanup and prevent resource leaks. """ - Postgres.close(self) + Storage.close(self) await OPC.shutdown(self) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 4a71455..8cf4afa 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -1,53 +1,64 @@ from temporalio import activity, workflow - with workflow.unsafe.imports_passed_through(): import traceback + from collections.abc import Callable, Mapping + 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 laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter - from typing import Any + + from laborious import metrics from laborious.utils.filters.conditional_filters import ( filter_empty_data, - filter_specific_variables_null_values + filter_specific_variables_null_values, ) - from pandas import DataFrame - from laborious import metrics + from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter + +# Strongly-typed filter function signatures +InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool] +ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool] +ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool] # Input filter function mappings -input_filter_functions = { +input_filter_functions: dict[str, InputFilterFunc] = { 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, 'EMPTY_DATA': filter_empty_data, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 2, - 'REPEAT': -1 - } +} + +# Confidence mappings kept separate from function maps to avoid Union types +input_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 2, + 'REPEAT': -1, } # MLFlow response filter function mappings -mlflow_response_filter_functions = { +mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = { 'API_ERROR': api_error_filter, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 10, - 'REPEAT': -1 - }, +} + +mlflow_response_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 10, + 'REPEAT': -1, } # MLFlow content filter function mappings -mlflow_content_filter_functions = { +mlflow_content_filter_functions: dict[str, ContentFilterFunc] = { 'NAN_VALUES': nan_values_filter, 'EMPTY_DATA': filter_empty_data, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 18, - 'REPEAT': -1 - } +} + +mlflow_content_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 18, + 'REPEAT': -1, } @@ -82,10 +93,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 +131,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 +139,38 @@ 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: 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_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 +205,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 +214,12 @@ 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=5)}', 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): @@ -221,34 +227,32 @@ 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: 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_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 +287,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 +296,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 {filters}', metadata) for fil, config in filters.items(): if fil not in mlflow_content_filter_functions: @@ -303,36 +307,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: 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_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 +364,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 +374,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 +414,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 +422,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 +488,56 @@ 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='format_retrain_report') + async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]: + """ + Format retrain report data according to configured storage policies. + """ + metadata = input_data['metadata'] + self.info('Formatting retrain report...', metadata) + + experiment_response = input_data['experiment_response'] + update_report = input_data['update_report'] + model_id = input_data['model_id'] + model_name = input_data['model_name'] + + report = DataFrame( + { + 'model_id': [model_id], + 'model_name': [model_name], + 'timestamp': [experiment_response['timestamp']], + 'status': [experiment_response['message']], + } + ) + + if experiment_response['success']: + # Retrain was successfull + report['version'] = update_report['version'] + report['mlflow_run_id'] = update_report['mlflow_run_id'] + report['mlflow_experiment_id'] = update_report['mlflow_experiment_id'] + + self.debug(f'Retrain report: {report.to_csv()}', metadata) + + return report.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. @@ -517,24 +562,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 +605,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/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index bd06283..a289c9f 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,20 +1,25 @@ 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_MS_WITH_TZ, + DATETIME_FORMAT_WITH_TZ, + now, + ) + + from laborious.utils.repository.minio_repository import MinioRepository 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): @@ -36,8 +41,16 @@ 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, + minio_config: dict[str, Any], + mlflow_password: str, + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize MLFlow activities with server configuration. @@ -52,18 +65,31 @@ 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") + if not hasattr(self, 'minio_repository'): + self.minio_repository: MinioRepository | None = None + + if self.minio_repository is None: + self.minio_repository = MinioRepository( + logger=logger, + notification_handler=notification_handler, + minio_endpoint_url=minio_config['endpoint_url'], + minio_access_key=minio_config['access_key'], + minio_secret_key=minio_config['secret_key'], + minio_region_name=minio_config['region_name'], + minio_default_bucket=minio_config['default_bucket'], + ) + + @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 +125,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 +134,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 +148,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 +197,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 +213,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. @@ -222,51 +253,88 @@ class MLFlow(BaseActivity): Raises: Exception: If retraining fails or encounters critical errors """ + + if self.minio_repository is None: + raise ValueError('Minio repository not initialized') + metadata = input_data['metadata'] - data = DataFrame(input_data['data']) + object_key = input_data['object_key'] + + self.info(f'Loading retrain data from Key: {object_key}', metadata) + + try: + data = self.minio_repository.get_parquet_as_dataframe( + object_key=object_key, metadata=metadata + ) + except Exception as e: + trace = traceback.format_exc() + self.send_notification( + metadata=metadata, + notification_id='ERROR_LOADING_RETRAIN_DATA', + message=f'Error loading retrain data: {e}', + block='retrain_model', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata) + return { + 'success': False, + 'message': f'Error loading retrain data: {e}', + 'traceback': trace, + 'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ), + } + + self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata) + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) self.info(f'Retraining model {model_name}...', metadata) timestamp = data['timestamp'].max() self.debug(f'Timestamp: {timestamp}', 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' + ) + 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() + # 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 - try: - retrain_output, experiment = self.model_monitoring_repository.retrain_model( - data=data, - model_name=model_name - ) + data['timestamp'] = data.index + data['timestamp'] = to_datetime( + data['timestamp'], format=DATETIME_FORMAT_WITH_TZ + ).dt.strftime(DATETIME_FORMAT) + data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT) - return { - 'status': retrain_output, - 'timestamp': timestamp, - 'experiment': experiment - } - except Exception as e: - trace = traceback.format_exc() + data.columns.name = None + + retrain_output = self.model_monitoring_repository.retrain_model( + data=data, model_name=model_name, model_config=model_config, metadata=metadata + ) + + if not retrain_output['success']: + trace = retrain_output['traceback'] self.send_notification( metadata=metadata, notification_id='RETRAIN_MODEL_ERROR', - message=f'Error retraining model {model_name}: {e}', + message=f'Error retraining model {model_name}: {retrain_output["message"]}', 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") + return {**retrain_output, 'timestamp': timestamp} + + @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. @@ -305,29 +373,18 @@ class MLFlow(BaseActivity): """ 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) + 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, metadata=metadata ) - 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() + self.info(f'Production model {model_name} updated successfully', metadata) + return response except Exception as e: trace = traceback.format_exc() @@ -337,7 +394,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/laborious/activities/opc.py b/laborious/activities/opc.py index 0b4ec9f..b4604d8 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -1,15 +1,16 @@ from temporalio import activity, workflow - with workflow.unsafe.imports_passed_through(): + import traceback + from typing import Any + + from pandas import DataFrame 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 laborious.utils.repository.opc_repository import OpcRepository - from typing import Any - import traceback - from pandas import DataFrame OPC_WRITTING_ERROR_CONFIDENCE = 12 @@ -33,15 +34,17 @@ class OPC(BaseActivity): notification_handler (NotificationHandler): Notification management instance """ - def __init__(self, opc_servers: dict[str, dict[str, Any]], - logger: Logger, notification_handler: NotificationHandler): - + 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) + BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True) self.opc_repository: dict[str, OpcRepository] = {} self.opc_servers = opc_servers @@ -70,10 +73,10 @@ class OPC(BaseActivity): 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'], + self.logger.info('Initializing OPC servers...') + for opc_id, server in self.opc_servers.items(): + self.opc_repository[opc_id] = OpcRepository( + opc_id=server['id'], url=server['url'], logger=self.logger, server_uri=server['server_uri'], @@ -82,30 +85,35 @@ class OPC(BaseActivity): server_cert_path=server['server_cert_path'], notification_handler=self.notification_handler, reconnection_interval=server['reconnection_interval'], - pod_id=self.pod_id + pod_id=self.pod_id, ) - is_connected, error_data = await self.opc_repository[id].connect() + is_connected, error_data = await self.opc_repository[opc_id].connect() if not is_connected: self.send_notification( metadata={ 'model_id': '-', 'model_name': '-', 'workflow_name': '-', - 'schedule_name': 'INITIALIZATION' + '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) + attachment_content=error_data.get('attachment_content', None), ) else: - self.logger.info( - f"OPC server {id} connected successfully.") + self.logger.info(f'OPC server {opc_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: + 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. @@ -127,7 +135,8 @@ class OPC(BaseActivity): try: is_success, error_data = await self.opc_repository[server_id].write_data( - tag, data, data_type, self.logger, metadata) + tag, data, data_type, self.logger, metadata + ) if not is_success: self.send_notification( metadata=metadata, @@ -135,8 +144,7 @@ class OPC(BaseActivity): message=error_data['message'], block=error_data['block'], level=error_data.get('level', NotificationLevel.ERROR), - attachment_content=error_data.get( - 'attachment_content', None) + attachment_content=error_data.get('attachment_content', None), ) return False return True @@ -144,11 +152,11 @@ class OPC(BaseActivity): 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", + 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 + attachment_content=trace, ) raise e @@ -174,21 +182,26 @@ class OPC(BaseActivity): 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." + message = f'OPC server {server_id} not found to perform write operation.' self.send_notification( metadata=metadata, - notification_id="OPC_SERVER_NOT_FOUND", + notification_id='OPC_SERVER_NOT_FOUND', message=message, - block="write_opc_data", + block='write_opc_data', level=NotificationLevel.ERROR, - attachment_content=f"OPC servers: {list(self.opc_repository.keys())}" + 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]: + 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. @@ -225,11 +238,13 @@ class OPC(BaseActivity): data=data.head(1)['prediction'].values[0], data_type=tag_config['data_type'], tag_type='prediction', - metadata=metadata + metadata=metadata, ) if local_success: self.info( - f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata) + f'Prediction data written to OPC server {server_id} for tag {tag}.', + metadata, + ) count += 1 success = success and local_success @@ -241,11 +256,13 @@ class OPC(BaseActivity): data=data.head(1)['prediction_confidence'].values[0], data_type=tag_config['data_type'], tag_type='confidence', - metadata=metadata + metadata=metadata, ) if local_success: self.info( - f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata) + f'Confidence data written to OPC server {server_id} for tag {tag}.', + metadata, + ) count += 1 success = success and local_success @@ -271,29 +288,33 @@ class OPC(BaseActivity): """ metadata = input_data['metadata'] - self.info("Writing data to OPC servers...", 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) + 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) + 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) + f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("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]: + def process_confidence( + self, data: DataFrame, success: bool, metadata: dict[str, Any] + ) -> dict[Any, Any]: """ Process prediction confidence based on OPC write operation success. @@ -323,12 +344,12 @@ class OPC(BaseActivity): 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 + 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) + self.debug('Data written to OPC servers successfully.', metadata) return data.to_dict() diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py new file mode 100644 index 0000000..8aec756 --- /dev/null +++ b/laborious/activities/storage.py @@ -0,0 +1,137 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + # Extend the Temporal Postgres activities for convenient query -> MinIO export + import traceback + from typing import Any + + import pandas as pd + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.temporal.activities.postgres import Postgres + from sientia_do.temporal.constants import now + + from laborious.utils.repository.minio_repository import MinioRepository + + +DATETIME_FILENAME_FORMAT = '%Y-%m-%d_%H-%M-%S' + + +class Storage(Postgres): + """ + Extensions for Postgres activities with a helper to export query results + directly to MinIO as Parquet and return the object name. + """ + + def __init__( + self, + host: str, + port: int, + user: str, + password: str, + dbname: str, + min_connections: int, + max_connections: int, + minio_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): + super().__init__( + host=host, + port=port, + user=user, + password=password, + dbname=dbname, + min_connections=min_connections, + max_connections=max_connections, + logger=logger, + notification_handler=notification_handler, + ) + + if not hasattr(self, 'minio_repository'): + self.minio_repository: MinioRepository | None = None + + if self.minio_repository is None: + self.minio_repository = MinioRepository( + logger=logger, + notification_handler=notification_handler, + minio_endpoint_url=minio_config['endpoint_url'], + minio_access_key=minio_config['access_key'], + minio_secret_key=minio_config['secret_key'], + minio_region_name=minio_config['region_name'], + minio_default_bucket=minio_config['default_bucket'], + ) + + @activity.defn(name='query_to_minio') + async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Execute SQL query, write result as Parquet to MinIO, and return object name. + + Args (input_data): + metadata (dict): Workflow metadata + query (str): SQL query + model_name (str): Model name for object naming + object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain) + + Returns: + dict: { success: bool, object_name: str, uri: str } + """ + + if self.minio_repository is None: + raise ValueError('Minio repository not initialized') + + metadata = input_data.get('metadata', {}) + object_prefix = input_data.get('object_prefix', 'datasets/retrain') + + timestamp = now().strftime(DATETIME_FILENAME_FORMAT) + object_name = f'{object_prefix}_{timestamp}.parquet' + uri = f's3://{self.minio_repository.minio_bucket}/{object_name}' + + try: + data = await self.load_custom_query(input_data) + if not data: + self.error('query_to_minio failed: No data returned from query', metadata) + return {'success': False, 'message': 'No data returned from query'} + + # Ensure we have a DataFrame + data = pd.DataFrame(data) + + # Write parquet to memory and upload via persistent client + self.minio_repository.store_dataframe_as_parquet( + dataframe=data, uri=uri, object_name=object_name, metadata=metadata + ) + + return {'success': True, 'object_key': object_name, 'uri': uri} + except Exception as e: + trace = traceback.format_exc() + self.send_notification( + metadata=metadata, + notification_id='ERROR_STORING_QUERY_TO_MINIO', + message=f'Error storing query to MinIO: {e}', + block='query_to_minio', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + self.error(trace, metadata) + + return {'success': False, 'message': str(e)} + + def close(self) -> None: + """Close Storage resources (MinIO client and Postgres engine).""" + try: + if hasattr(self, 'minio_repository') and self.minio_repository is not None: + try: + self.minio_repository.close() + finally: + self.minio_repository = None + finally: + # Ensure Postgres resources are disposed as well + try: + super().close() + except Exception: + self.logger.error('Error closing Postgres resources') + + def __del__(self): + self.close() diff --git a/laborious/metrics.py b/laborious/metrics.py index 97f7bb9..7826fac 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -23,50 +23,50 @@ Metric Labels: - opc_server_id: Identifier for OPC server operations """ -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( - "laborious_predictions_written_count", - "Number of predictions written to the database table predictions", + '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", + '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", + '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] + 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"], + '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] + '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/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index 145f248..7501a07 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -1,9 +1,9 @@ -from os import getenv import json -from typing import Dict, Any +from os import getenv +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 +30,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 +55,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_opc_config() -> Dict[str, Any]: +def build_opc_config() -> dict[str, Any]: """ Build OPC server configuration from environment variables. @@ -75,7 +75,7 @@ def build_opc_config() -> Dict[str, Any]: 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) + OPC_RECONNECTION_INTERVAL: Reconnection interval in seconds (fallback, default: 120) Returns: dict: OPC server configuration dictionary @@ -93,12 +93,12 @@ def build_opc_config() -> Dict[str, Any]: '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')) + 'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')), } } -def build_mongodb_config() -> Dict[str, Any]: +def build_mongodb_config() -> dict[str, Any]: """ Build MongoDB configuration from environment variables. @@ -125,5 +125,28 @@ 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, + } + + +def build_minio_config() -> dict[str, Any]: + """ + Build MinIO (S3-compatible) configuration from environment variables. + + Environment Variables: + MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000) + MINIO_ACCESS_KEY: Access key (default: minioadmin) + MINIO_SECRET_KEY: Secret key (default: minioadmin) + MINIO_REGION: Region name for S3 client (default: us-east-1) + MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious) + + Returns: + dict: MinIO configuration dictionary + """ + return { + 'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'), + 'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'), + 'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'), + 'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'), + 'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'), } diff --git a/laborious/utils/filters/conditional_filters.py b/laborious/utils/filters/conditional_filters.py index 2c51805..967e52c 100644 --- a/laborious/utils/filters/conditional_filters.py +++ b/laborious/utils/filters/conditional_filters.py @@ -20,8 +20,11 @@ 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 + + if data.empty: + return False + + return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty def filter_empty_data(data: DataFrame, _config: dict) -> bool: diff --git a/laborious/utils/filters/mlflow_filters.py b/laborious/utils/filters/mlflow_filters.py index f792a0e..82970e8 100644 --- a/laborious/utils/filters/mlflow_filters.py +++ b/laborious/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/laborious/utils/repository/minio_repository.py b/laborious/utils/repository/minio_repository.py new file mode 100644 index 0000000..03b0f15 --- /dev/null +++ b/laborious/utils/repository/minio_repository.py @@ -0,0 +1,146 @@ +""" +MinIO repository utilities. + +This module provides a lightweight repository around a MinIO/S3-compatible +object storage using boto3. It supports creating buckets on demand and +storing/loading pandas DataFrames in Parquet format. +""" + +from io import BytesIO +from typing import Any + +import boto3 +from botocore.config import Config +from botocore.exceptions import ClientError +from pandas import DataFrame, read_parquet +from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler +from sientia_do.observability.logger import Logger + + +class MinioRepository: + """ + Repository for interacting with a MinIO (S3-compatible) object storage. + + This class encapsulates a reusable `boto3` S3 client and convenience + helpers to persist and retrieve pandas DataFrames as Parquet files. + + Attributes: + storage_options (dict): Options compatible with pandas s3fs usage. + minio_bucket (str): Default bucket name used for operations. + minio_endpoint_url (str): MinIO endpoint URL. + minio_region_name (str): MinIO region name. + s3_client (Any): Reusable S3 client from `boto3`. + logger (Logger): Observability logger. + notification_handler (NotificationHandler): Notifications handler. + """ + def __init__( + self, + minio_endpoint_url: str, + minio_access_key: str, + minio_secret_key: str, + minio_region_name: str, + minio_default_bucket: str, + logger: Logger, + notification_handler: NotificationHandler, + ): + """Initialize the repository and S3 client. + + Args: + minio_endpoint_url (str): MinIO endpoint URL. + minio_access_key (str): Access key (AK). + minio_secret_key (str): Secret key (SK). + minio_region_name (str): Region name for the client. + minio_default_bucket (str): Default bucket name to operate on. + logger (Logger): Logger instance for structured logs. + notification_handler (NotificationHandler): Notification handler. + """ + # MinIO settings shared with pandas s3fs + self.storage_options = { + 'key': minio_access_key, + 'secret': minio_secret_key, + 'client_kwargs': {'endpoint_url': minio_endpoint_url}, + } + self.minio_bucket = minio_default_bucket + self.minio_endpoint_url = minio_endpoint_url + self.minio_region_name = minio_region_name + + logger.info( + f'Connecting to MinIO at {self.minio_endpoint_url}, default bucket: {self.minio_bucket}' + ) + + # Reusable MinIO client + self.s3_client: Any = boto3.client( + 's3', + endpoint_url=self.minio_endpoint_url, + aws_access_key_id=self.storage_options['key'], + aws_secret_access_key=self.storage_options['secret'], + region_name=self.minio_region_name, + config=Config( + signature_version='s3v4', + s3={'addressing_style': 'path'}, + retries={'max_attempts': 5, 'mode': 'standard'}, + connect_timeout=5, + read_timeout=120, + ), + ) + + self.logger = logger + self.notification_handler = notification_handler + + def close(self): + """Close the underlying S3 client.""" + self.s3_client.close() + + def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None: + """Ensure the default bucket exists; create it if missing. + + Args: + metadata (dict[str, Any]): Metadata used for structured logging. + """ + + try: + self.logger.custom_info(f"Checking if bucket '{self.minio_bucket}' exists", metadata) + self.s3_client.head_bucket(Bucket=self.minio_bucket) + except ClientError: + self.logger.custom_info(f"Creating bucket '{self.minio_bucket}'", metadata) + self.s3_client.create_bucket(Bucket=self.minio_bucket) + + def store_dataframe_as_parquet( + self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any] + ): + """Persist a DataFrame as a Parquet object in the default bucket. + + Args: + dataframe (DataFrame): DataFrame to persist. + uri (str): Human-friendly URI used for logging context. + object_name (str): Object key (path/key within the bucket). + metadata (dict[str, Any]): Metadata used for structured logging. + """ + self.ensure_bucket_exists(metadata) + + self.logger.custom_info(f'Storing dataframe as parquet in {uri}', metadata) + + buffer = BytesIO() + dataframe.to_parquet(buffer, engine='pyarrow', index=True) + buffer.seek(0) + self.s3_client.put_object(Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue()) + + self.logger.custom_info(f'Dataframe stored as parquet in {uri}', metadata) + + def get_parquet_as_dataframe(self, object_key: str, metadata: dict[str, Any]) -> DataFrame: + """Load a Parquet object from the default bucket into a DataFrame. + + Args: + object_key (str): Object key to retrieve from the bucket. + metadata (dict[str, Any]): Metadata used for structured logging. + + Returns: + DataFrame: Loaded DataFrame. + """ + self.logger.custom_info(f'Getting parquet as dataframe from {object_key}', metadata) + + response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key) + + # Read the content into a BytesIO buffer to support seek operations + buffer = BytesIO(response['Body'].read()) + return read_parquet(buffer) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 3107d1c..3361d9a 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1,203 +1,559 @@ """ -Model Monitoring Repository +MLflow repository utilities -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. +This module provides the `MLFlowRepository` class and helpers to interact with +an MLflow tracking server and model registry. It covers model discovery, +downloading/loading with multiple flavors, cached operations with retention +policies, transformation/prediction interfaces, retraining workflows, and +production model promotion. +Capabilities: +- Model loading and caching with retention policies +- Data transformation and prediction operations +- Model retraining workflows +- Production model updates and versioning """ -from datetime import datetime + +import ctypes +import gc import traceback -import pandas as pd +from datetime import datetime, timedelta +from os import environ, makedirs, path +from shutil import rmtree +from typing import Any, Literal, overload + import mlflow -from os import makedirs, path, remove -from sientia.ModelServing import ModelServing -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +import pandas as pd +from mlflow.entities import Experiment +from numpy import ndarray from sientia_do.observability.logger import Logger +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + +ARTIFACTS_PATH = './tmp/artifacts' +TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl' +PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl' + +INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." -class MLFlowRepository(): - def __init__(self, host, username, password, logger: Logger): +def force_memory_release(logger: Logger): + """Attempt to release memory from the Python process. - self.model_serving = ModelServing(tracking_uri=host, - username=username, password=password, - logger=logger) + Executes a garbage collection cycle and calls `malloc_trim(0)` on glibc + where available to return free memory to the OS. This may be a no-op on + non-glibc systems. + + Args: + logger (Logger): Logger for observability. + """ + gc.collect() + + try: + ctypes.CDLL('libc.so.6').malloc_trim(0) + logger.info('Memory released') + except Exception as e: + logger.info(f'Memory release failed: {e}') + + +class MLFlowRepository: + def __init__(self, host: str, username: str, password: str, logger: Logger): + """Initialize MLflow client and base state. + + Args: + host (str): MLflow tracking URI. + username (str): MLflow username. + password (str): MLflow password. + logger (Logger): Logger instance. + """ + # set tracking uri + mlflow.set_tracking_uri(host) + + environ['MLFLOW_TRACKING_USERNAME'] = username + environ['MLFLOW_TRACKING_PASSWORD'] = password + # Create an MLflow client + self.client = mlflow.tracking.MlflowClient() + + self.model_cache: dict[str, Any] = {} self.logger = logger - def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame: + """ + Functions related to get model registry parameters + """ + + def get_model_uri(self, run_id: str, prediction: bool = True): + """Build the artifact URI for a run's model. + + Args: + run_id (str): MLflow run identifier. + prediction (bool): If True, return `prediction_model` URI, + otherwise return `data_model` URI. + + Returns: + str: Artifact URI to the selected model within the run. """ - 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) + run_info = mlflow.get_run(run_id) + if prediction: + model_uri = run_info.info.artifact_uri + '/prediction_model' else: - raise ValueError( - f"{message}") + model_uri = run_info.info.artifact_uri + '/data_model' + return model_uri - return data + def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str: + """Resolve the run_id for a registered model at a given stage. - 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. + Args: + model_name (str): Registered model name. + stage (str): Desired stage (e.g., 'Production'). Returns: - - dict: A dictionary containing the transformed data. + str: Run ID for the latest version at the given stage. """ + # Use search_registered_models instead of deprecated get_latest_versions + registered_models = self.client.search_registered_models( + filter_string=f"name='{model_name}'" + ) - 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 + if not registered_models: + raise mlflow.exceptions.MlflowException( + f"Model '{model_name}' not found in the Model Registry." ) - self.logger.custom_debug( - f"Data received from model transformation: {transformed_data.to_csv()}", metadata) + # Get all versions of the model and filter by stage + model_versions = self.client.search_model_versions(filter_string=f"name='{model_name}'") - transformed_data = self.detect_and_parse_datetime_index( - transformed_data, metadata) + # Filter versions by the desired stage using current_stage attribute + stage_versions = [mv for mv in model_versions if mv.current_stage == stage] - 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 + if not stage_versions: + raise mlflow.exceptions.MlflowException( + f"Model '{model_name}' in stage '{stage}' not found in the Model Registry." ) - 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 + # Sort by version number to get the latest + latest_version = max(stage_versions, key=lambda v: int(v.version)) + run_id = latest_version.source.split('/') + return run_id[2] def get_next_run_name(self, model_name: str) -> str: """ - Generate the next run name for a specific MLFlow model. + 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 + 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"]) + 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: + def get_experiment( + self, experiment_name: str, create_if_not_exists: bool = False + ) -> Experiment: """ - Create a new MLFlow experiment for model retraining. + Retrieve MLflow experiment by name, optionally creating it. + + 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: + Experiment: MLflow experiment object + + Raises: + ValueError: If the experiment name is not found and creation is disabled + """ + experiment = mlflow.get_experiment_by_name(experiment_name) + + if experiment is None: + if create_if_not_exists: + experiment = mlflow.create_experiment(experiment_name) + else: + raise ValueError(f'Experiment {experiment_name} not found') + + return experiment + + def get_model_params(self, run_id: str): + """Fetch parameters associated with a given MLflow run. + + Args: + run_id (str): Run identifier to inspect. + + Returns: + dict: Mapping of parameter names to values. + """ + run_info = mlflow.get_run(run_id) + return run_info.data.params + + """ + Functions related to download and load models + """ + + def dowload_artifacts(self, model_name: str, artifact_path: str = 'data_model') -> str: + """ + Download artifacts from the latest production run of a model. + + Args: + model_name (str): Registered model name. + artifact_path (str): Relative path to artifacts within the run. + + Returns: + str: Local filesystem path where artifacts are saved. + """ + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + output_dir = f'{ARTIFACTS_PATH}/{model_name}' + + full_path = path.join(output_dir, artifact_path) + + if path.exists(full_path): + # Remove the directory and create a new one + rmtree(full_path) + makedirs(output_dir, exist_ok=True) + + self.logger.info(f'Downloading artifacts from {run_id} to {output_dir}') + + return self.client.download_artifacts(run_id, artifact_path, output_dir) + + def load_predict_model(self, model_name: str, flavor: str = 'sklearn') -> Any: + """ + Load a predictive model from the MLflow Model Registry. + + Args: + model_name (str): The name of the model to download from the registry. + flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch') + artifact_path (str | None): Path to compressed artifacts if model is compressed + + Returns: + mlflow.pyfunc.PyFuncModel: The loaded predictive model. + + Notes: + - The model is fetched from the "Production" stage of the MLflow Model Registry. + - Warnings during the model loading process are suppressed. + """ + model_uri = f'models:/{model_name}/production' + self.logger.info(f'Loading prediction model {model_name} from {model_uri}') + if flavor == 'pyfunc': + model = mlflow.pyfunc.load_model(model_uri) + elif flavor == 'sklearn': + model = mlflow.sklearn.load_model(model_uri) + elif flavor == 'pytorch': + model = mlflow.pytorch.load_model(model_uri) + else: + raise ValueError(INVALID_FLAVOR_MESSAGE) + + return model + + def load_transform_model(self, model_name: str, flavor: str) -> Any: + """ + Load the latest Production version of a transformation model. + + This method retrieves the latest production model run ID for the given + model name, constructs the model URI, and loads the model using MLflow. + + Args: + model_name (str): The name of the model to download. + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + artifact_path (str | None): Path to compressed artifacts if model is compressed + + Returns: + Any: The loaded model object, depending on the flavor used. + + Raises: + Exception: If the model run ID or URI cannot be retrieved, or if the + model cannot be loaded. + """ + + latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production') + model_uri = self.get_model_uri(latest_production_id, prediction=False) + + self.logger.info(f'Loading data model {model_name} from {model_uri}') + if flavor == 'sklearn': + model = mlflow.sklearn.load_model(model_uri) + elif flavor == 'pyfunc': + model = mlflow.pyfunc.load_model(model_uri) + elif flavor == 'pytorch': + model = mlflow.pytorch.load_model(model_uri) + else: + raise ValueError(INVALID_FLAVOR_MESSAGE) + return model + + def download_model( + self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False + ) -> tuple[Any, str | None]: + """ + Download model based on type ("predict" or "transform"). + + Args: + model_name (str): Name of the model to download + model_type (str): Type of model ('predict' or 'transform') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + load_wrapper (bool): Whether to load wrapper + + Returns: + tuple[Any, str | None]: Model object and optional artifact path. + """ + + self.logger.info( + f'Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}' + ) + + if model_type not in ['predict', 'transform']: + raise ValueError("Invalid model_type. Use 'predict' or 'transform'.") + + artifact_path = None + + if load_wrapper: + self.logger.info( + f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}' + ) + + target = 'prediction_model' if model_type == 'predict' else 'data_model' + + artifact_path = self.dowload_artifacts(model_name, target) + + self.logger.info( + f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' + ) + + raw_model = mlflow.pyfunc.load_model(artifact_path) + model = raw_model._model_impl.python_model + else: + if model_type == 'predict': + model = self.load_predict_model(model_name, flavor) + + else: + model = self.load_transform_model(model_name, flavor) + + return model, artifact_path + + """ + Functions related to data format + """ + + def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame: + """ + Normalize DataFrame index to the expected timestamp string format. + + The index must be timestamp-like. If the index is: + - string: it must match `DATETIME_FORMAT_WITH_TZ` + - datetime or pandas Timestamp: it will be converted to that format + Any other type raises a ValueError. + + Args: + data (pd.DataFrame): DataFrame with timestamp index. + metadata (dict): Metadata for structured logging. + + Returns: + pd.DataFrame: DataFrame with converted datetime index. + """ + 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): + types = map(str, map(type, index)) + raise ValueError(f'{message}. Elements are {",".join(types)}') + + # Check type and converts to DATETIME_FORMAT_WITH_TZ + if index_type is str: + # Validate format of string and return error if not valid + try: + pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ) + except ValueError as e: + raise ValueError(f'{message}. Unable to parse given date format: {e}') from e + + elif index_type == datetime or index_type == pd.Timestamp: + index = data.index + if hasattr(index, 'tz') and index.tz is None: + data.index = index.tz_localize('UTC') # type: ignore[attr-defined] + + data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined] + else: + raise ValueError(f'{message}. Got {index_type}.') + + return data + + """ + Functions related to cache management of models + """ + + def check_cache_retention(self, cache: dict, retention: int) -> bool: + """ + Check whether cached model data is still valid. + + Args: + cache (dict): Cached model data with a 'timestamp' key. + retention (int): Retention time in minutes. + + Returns: + bool: True if cache is still valid, False if expired. + """ + current_time = datetime.now() + cache_time = cache['timestamp'] + if current_time - cache_time >= timedelta(minutes=retention): + return False + return True + + def handle_valid_model(self, model_name: str, cache: dict) -> dict: + """ + Return the cached model configuration when retention is valid. + + Args: + model_name (str): Name of the model (for logging/consistency). + cache (dict): Cached model data structure. + + Returns: + dict: Model configuration. + """ + self.logger.debug(f'Model {model_name} is still valid, using cached version') + + return cache['target'] + + def handle_outdated_model(self, model_name: str, model_key: str) -> None: + """ + Clean up outdated cached model and its artifacts. + + Args: + model_name (str): Name of the model. + model_key (str): Cache key for the model. + + Returns: + None + """ + self.logger.debug(f'Model {model_name} is outdated, downloading a new one') + + del self.model_cache[model_key]['target']['model'] + del self.model_cache[model_key] + + def get_model(self, model_name: str, retention: int, model_type: str, flavor: str) -> Any: + """ + Retrieve a model with caching support based on retention policy. + + Args: + model_name (str): Name of the model to retrieve + retention (int): Cache retention time in minutes (0 = no cache). + model_type (str): Type of model ('predict' or 'transform') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + + Returns: + Any: Model object. + """ + # Retention is 0, download a new model + if retention <= 0: + model, _artifact_path = self.download_model( + model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False + ) + return model + + model_key = f'{model_name}_{model_type}' + + if model_key in self.model_cache: + cache = self.model_cache[model_key] + + # Check if config has changed or is outdated + if self.check_cache_retention(cache, retention): + return self.handle_valid_model(model_name=model_name, cache=cache) + else: + # Model is outdated, delete old model files + self.handle_outdated_model(model_name=model_name, model_key=model_key) + else: + self.logger.debug( + f'Model {model_name} is not in {model_type} cache, downloading a new one' + ) + + # Donwload new model + model, _artifact_path = self.download_model( + model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False + ) + + cache = {'target': model, 'timestamp': datetime.now()} + + self.model_cache[model_key] = cache + + return model + + @overload + def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: Literal['transform'], + retention: int, + flavor: str, + ) -> pd.DataFrame: ... + + @overload + def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: Literal['predict'], + retention: int, + flavor: str, + ) -> pd.DataFrame | ndarray: ... + + def get_cached_operation( + self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str + ) -> pd.DataFrame | ndarray: + """ + Execute a cached operation using the requested model. + + Args: + model_name (str): Registered model name. + data (pd.DataFrame): Input data. + retention (int): Cache retention in minutes. + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch'). + + Returns: + pd.DataFrame | ndarray: Operation result. + """ + if operation not in ['transform', 'predict']: + raise ValueError("Invalid operation. Use 'transform' or 'predict'.") + + model = self.get_model( + model_name=model_name, retention=retention, model_type=operation, flavor=flavor + ) + + prediction = model.predict(data) + + if retention == 0: + self.logger.info(f'Deleting model {model_name}:{operation} from memory') + del model + + force_memory_release(self.logger) + + return prediction + + """ + Functions related to model retraining + """ + + def fit_models( + self, + model_name: str, + data: pd.DataFrame, + latest_production_id: str, + metadata: dict, + transform_flavor: str = 'sklearn', + predict_flavor: str = 'sklearn', + target_name: str | None = None, + ) -> dict[str, dict[str, Any]]: + """ + Prepare models and data for a retraining run. This method sets up the complete environment for model retraining by: 1. Loading the current production prediction model @@ -207,48 +563,160 @@ class MLFlowRepository(): 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 + model_name (str): Name of the MLflow model to retrain. + data (pd.DataFrame): Training data for model retraining. + transform_flavor (str): Flavor for transformation model. + predict_flavor (str): Flavor for prediction model. + target_name (str | None): Optional target column; if None, use model target. + metadata (dict): Metadata for logging. Returns: - tuple: (prediction_model, data_model, experiment) - - prediction_model: Loaded prediction model for retraining - - data_model: Fitted transformation model - - experiment: MLFlow experiment name + dict[str, dict[str, Any]]: Mapping with prepared `prediction_model` and + `data_model`, including optional artifact paths. """ - # 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" + + self.logger.custom_info(f'Starting model experiment creation for {model_name}', metadata) + self.logger.custom_debug( + f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}', + metadata, ) - transform_uri = self.model_serving.get_model_uri( - latest_production_id, prediction=False + + # data.to_csv( + # f"tmp/retrain_data_{model_name}.csv", index=True) + + self.logger.custom_info( + f'Retrieved latest production run ID: {latest_production_id}', metadata ) - # 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) + self.logger.custom_info(f'Loading transformation model for {model_name}', metadata) - 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) + load_transform_wrapper = transform_flavor == 'pyfunc' - return prediction_model, data_model, experiment + data_model, data_artifact_path = self.download_model( + model_name=model_name, + model_type='transform', + flavor=transform_flavor, + load_wrapper=load_transform_wrapper, + ) - def perform_model_retrain(self, - prediction_model, - data_model, - experiment: str, - model_name: str, - data: pd.DataFrame): + self.logger.custom_info(f'Loading prediction model for {model_name}', metadata) + + load_predict_wrapper = predict_flavor == 'pyfunc' + + prediction_model, prediction_artifact_path = self.download_model( + model_name=model_name, + model_type='predict', + flavor=predict_flavor, + load_wrapper=load_predict_wrapper, + ) + + treated_data_candidate = data_model.fit(data) + + if not isinstance(treated_data_candidate, pd.DataFrame): + data_model = treated_data_candidate + treated_data = data_model.predict(data) + else: + treated_data = treated_data_candidate + + # Stores current index as timestamp, Courier model expects a timestamp column + # with specific format + treated_data['timestamp'] = treated_data.index + + # Parses timestamp column to datetime format to align with data + treated_data = self.detect_and_parse_datetime_index(treated_data, metadata) + + treated_data = treated_data.drop_duplicates(subset=['timestamp'], keep='first') + + self.logger.custom_debug(f'Treated data index: {treated_data.index}', metadata) + + # treated_data.to_csv( + # f"tmp/retrain_treated_data_{model_name}.csv", index=True) + + self.logger.custom_debug(f'Transformed data shape: {treated_data.shape}', metadata) + + if target_name is None: + target_name = data_model.target_variable + self.logger.custom_debug( + f'Using target variable from data model: {target_name}', metadata + ) + else: + self.logger.custom_debug(f'Using provided target variable: {target_name}', metadata) + + # Check if treated_data contains target variable + if target_name not in treated_data.columns: + self.logger.custom_debug( + f'Target variable {target_name} not found in treated data, aligning data with treated data indexes', + metadata, + ) + # Aligns data with treated data indexes to get target variable + aligned_data = data.loc[treated_data.index] + aligned_series = aligned_data[target_name] + retrain_dataset = pd.merge( # NOSONAR + treated_data, aligned_series, left_index=True, right_index=True + ) + else: + # Uses target variable from treated data + self.logger.custom_debug( + f'Target variable {target_name} found in treated data, using it', metadata + ) + retrain_dataset = treated_data + + # retrain_dataset.to_csv( + # f"tmp/retrain_retrain_dataset_{model_name}.csv", index=True) + + prediction_model.fit(retrain_dataset) + + self.logger.custom_info( + f'Model experiment creation completed successfully for {model_name}', metadata + ) + + retrain_data = { + 'prediction_model': { + 'model': prediction_model, + 'artifact_path': prediction_artifact_path, + }, + 'data_model': {'model': data_model, 'artifact_path': data_artifact_path}, + } + return retrain_data + + def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict): + """Log a model into the active MLflow run. + + Args: + model_data (dict): Model holder with keys 'model' and optional 'artifact_path'. + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch'). + model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'. + metadata (dict): Metadata for structured logging. """ - Execute the complete model retraining process in MLFlow. + model = model_data['model'] + + self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata) + if flavor == 'sklearn': + mlflow.sklearn.log_model(model, model_type) + elif flavor == 'pyfunc': + code_path = [path.join(model_data['artifact_path'], 'code', 'utils')] + + self.logger.custom_debug(f'Code path: {code_path}', metadata) + + model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False) + + self.logger.custom_debug('Model uploaded successfully', metadata) + elif flavor == 'pytorch': + mlflow.pytorch.log_model(model, model_type) + else: + raise ValueError(INVALID_FLAVOR_MESSAGE) + + def create_new_experiment( + self, + model_name: str, + data: pd.DataFrame, + retrain_data: dict, + latest_production_id: str, + metadata: dict, + transform_flavor: str = 'sklearn', + predict_flavor: str = 'sklearn', + ) -> dict: + """ + 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 @@ -260,153 +728,109 @@ class MLFlowRepository(): 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 + 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. + transform_flavor (str): Flavor for transformation model. + predict_flavor (str): Flavor for prediction model. + metadata (dict): Metadata for logging. Returns: - tuple: (status_message, experiment_name) - - status_message (str): Success confirmation message - - experiment_name (str): Name of the experiment + dict: Metadata about the created run and 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) + prediction_model = retrain_data['prediction_model'] + data_model = retrain_data['data_model'] - file_path = f"temp/raw_data_{model_name}.csv" - data.to_csv(file_path, index=True) + model_temp_path = path.join(ARTIFACTS_PATH, model_name) - # log the data raw - mlflow.log_artifact(file_path) + self.logger.custom_info(f'Starting model retraining process for {model_name}', metadata) - # dynamic parameters, including model itself - mlflow.sklearn.log_model(prediction_model, "prediction_model") - mlflow.log_param("retrain", True) + original_params = self.get_model_params(latest_production_id) + retrain_params = { + **original_params, + 'retrain': True, + 'retrain_date': datetime.now().isoformat(), + 'source_run_id': latest_production_id, + 'retrain_samples': str(data.shape), + } + experiment_description = f'Retrain model {model_name} with new data' - # clear temp file - if path.exists(file_path): - remove(file_path) + experiment = self.get_experiment(model_name, create_if_not_exists=True) + experiment_name = experiment.name - return "Model retrained successfully", experiment + current_run_name = self.get_next_run_name(experiment_name) - def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: - """ - Orchestrate the complete model retraining workflow. + self.logger.custom_debug(f'Attributes: {retrain_params}', metadata) - 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 + data_path = f'{model_temp_path}/retrain_data.csv' - Args: - data (pd.DataFrame): Training data for model retraining - model_name (str): Name of the MLFlow model to retrain + data.to_csv(data_path, index=True) - 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" + self.logger.custom_info( + f'Starting model upload for {experiment_name} with run name {current_run_name}', + metadata, ) - if not isinstance(runs, pd.DataFrame): - raise ValueError('Runs is not a pandas DataFrame') + with mlflow.start_run( + experiment_id=experiment.experiment_id, + run_name=current_run_name, + description=experiment_description, + ) as _run: + run_id = _run.info.run_id + self.logger.custom_info('Logging data model', metadata) + # dynamic parameters, including model itself + self.log_model(data_model, transform_flavor, 'data_model', metadata) - # Filtrar apenas as runs onde params.retrain == True - filtered_runs = runs[runs["params.retrain"] == 'True'] + # dynamic parameters, including model itself + self.logger.custom_info('Logging prediction model', metadata) + self.log_model(prediction_model, predict_flavor, 'prediction_model', metadata) - # Converter a coluna 'end_time' para datetime - filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time']) + self.logger.custom_info(f'Model logged successfully for {model_name}', metadata) - # Ordenar o DataFrame de forma descendente pela coluna 'end_time' - filtered_runs = filtered_runs.sort_values( - by='end_time', ascending=False) + self.logger.custom_info(f'Logging remaining parameters for {model_name}', metadata) - # Pegar a รบltima run_id do DataFrame filtrado e ordenado - latest_run_id = filtered_runs.iloc[0]['run_id'] + # update transfomation model + # fixed parameters + mlflow.log_params(retrain_params) - return latest_run_id + # log the data raw + mlflow.log_artifact(data_path) - def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict: + self.logger.custom_info('Deleting model from filesystem', metadata) + if path.exists(model_temp_path): + rmtree(model_temp_path) + + self.logger.custom_info('Deleting prediction model from memory', metadata) + del prediction_model['model'] + del prediction_model + + self.logger.custom_info('Deleting data model from memory', metadata) + del data_model['model'] + del data_model + + force_memory_release(self.logger) + + return { + 'run_id': run_id, + 'experiment_id': experiment.experiment_id, + 'experiment_name': experiment.name, + } + + def update_production_model_by_run_id( + self, run_id: str, model_name: str, metadata: dict + ) -> dict: """ - Update production model with a specific MLFlow run. + Promote a specific run's model to Production. 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 + run_id (str): MLflow run ID containing the model to promote. + model_name (str): Name of the MLflow model. + metadata (dict): Metadata for logging. Returns: dict: Model update metadata containing: @@ -420,19 +844,18 @@ class MLFlowRepository(): 3. Transitions the model to 'Production' stage 4. Archives existing production versions """ + + self.logger.custom_info( + f'Starting production model update for {model_name} with run ID: {run_id}', metadata + ) + # 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() + mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name) # Obter a versรฃo mais recente registrada do modelo - model_versions = client.get_registered_model( - model_name).latest_versions + model_versions = self.client.get_registered_model(model_name).latest_versions if not isinstance(model_versions, list): raise ValueError('Model versions is not a list') @@ -440,42 +863,323 @@ class MLFlowRepository(): 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 + self.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 - } + return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id} - def update_production_model(self, experiment: str, model_name: str) -> dict: + """ + Functions that provide the interface to model operations + """ + + def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict): + """ + Transform data using a cached transformation model. + + This method provides a high-level interface for data transformation operations + using MLFlow models. It handles model caching, error management, and data + format conversion automatically. + + Process Flow: + 1. Retrieves or downloads the transformation model using caching mechanism + 2. Applies the transformation model to the input data + 3. Converts the transformed data to dictionary format for API response + 4. Handles any exceptions and returns structured error information + 5. Manages model lifecycle based on retention policy (cleanup artifacts if needed) + + Parameters: + model_name (str): The name of the MLflow model to use for transformation. + data (pd.DataFrame): The input data to be transformed by the model. + model_config (dict): Model configuration parameters + metadata (dict): Metadata for logging + + Returns: + dict: Response dictionary containing: + - success (bool): Operation success status + - content (dict): Transformed data as dictionary, or error information + if operation failed. Error content includes: + - message (str): Error description + - traceback (str): Full exception traceback + + Raises: + Exception: Any exception during model loading or transformation is caught + and returned in the response structure rather than propagated. + """ + + self.logger.custom_debug( + f'Data received for model transformation: {data.head(5).to_csv()}', metadata + ) + + # data.to_csv( + # f"tmp/data_{model_name}.csv", index=True) + + model_retention = model_config.get('retention_minutes', 0) + flavor = model_config.get('transform_flavor', 'sklearn') + + try: + transformed_data: pd.DataFrame = self.get_cached_operation( + model_name, data, 'transform', model_retention, flavor + ) + + self.logger.custom_debug( + f'Data received from model transformation: {transformed_data.head(5).to_csv()}', + metadata, + ) + + # transformed_data.to_csv( + # f"tmp/transformed_data_{model_name}.csv", index=True) + + 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): + """ + Generate predictions using a cached prediction model. + + This method provides a high-level interface for model prediction operations + using MLFlow models. It handles model caching, performance monitoring, + response formatting, and error management automatically. + + Process Flow: + 1. Preserves input data index for result alignment + 2. Records prediction start time for performance measurement + 3. Retrieves or downloads the prediction model using caching mechanism + 4. Executes model prediction on the input data + 5. Formats predictions into DataFrame with proper column naming + 6. Restores original data index to maintain data alignment + 7. Calculates and adds response time measurement + 8. Converts results to dictionary format for API response + 9. Handles any exceptions and returns structured error information + + Parameters: + model_name (str): The name of the MLflow model to use for prediction. + data (pd.DataFrame): The input data to make predictions on. + model_retention (int): Cache retention time in minutes (0 = no caching). + model_config (dict): Model configuration parameters + metadata (dict): Metadata for logging + + Returns: + dict: Response dictionary containing: + - success (bool): Operation success status + - content (dict): Prediction results as dictionary with: + - prediction: Model predictions array + - response_time: Prediction execution time in seconds + Or error information if operation failed: + - message (str): Error description + - traceback (str): Full exception traceback + + Raises: + Exception: Any exception during model loading or prediction is caught + and returned in the response structure rather than propagated. + """ + + model_retention = model_config.get('retention_minutes', 0) + flavor = model_config.get('predict_flavor', 'sklearn') + + try: + input_index = data.index + start_time = datetime.now() + + self.logger.custom_debug( + f'Data received for model prediction: {data.head(5).to_csv()}', metadata + ) + + # data.to_csv( + # f"tmp/treated_data_{model_name}.csv", index=True) + + predict_data = self.get_cached_operation( + model_name, data, 'predict', model_retention, flavor + ) + + end_time = datetime.now() + + if isinstance(predict_data, pd.DataFrame): + self.logger.custom_debug( + f'Data received from model prediction: {data.head(5).to_csv()}', metadata + ) + + # predict_data.to_csv( + # f"tmp/predicted_data_{model_name}.csv", index=True) + predict_data.columns = pd.Index(['prediction']) + + else: + predict_data = pd.DataFrame(predict_data, columns=['prediction']) + # predict_data.to_csv( + # f"tmp/predicted_data_{model_name}.csv", index=True) + + predict_data.index = input_index + predict_data['response_time'] = (end_time - start_time).total_seconds() + + return {'success': True, 'content': predict_data.to_dict()} + + except Exception as e: + return { + 'success': False, + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, + } + + def retrain_model( + self, data: pd.DataFrame, model_name: str, model_config: dict, metadata: dict + ) -> dict[str, Any]: + """ + Orchestrate the complete model retraining workflow. + + This method coordinates the entire model retraining process by managing + the MLFlow experiment lifecycle, model loading, retraining execution, + and artifact management. It provides a comprehensive retraining solution + that maintains model versioning and experiment tracking. + + Process Flow: + 1. Creates MLFlow experiment environment: + - Loads current production prediction model + - Loads current production transformation model + - Fits transformation model with new training data + - Prepares transformed data for prediction model retraining + - Sets up MLFlow experiment context + 2. Executes model retraining: + - Starts new MLFlow run with descriptive metadata + - Logs model parameters and hyperparameters + - Retrains both prediction and transformation models + - Logs training data as artifacts + - Saves retrained models to MLFlow registry + - Cleans up temporary files + 3. Returns comprehensive retraining results + + Args: + data (pd.DataFrame): Training data for model retraining. Must contain + all features required by both transformation and + prediction models, including target variable. + model_name (str): Name of the MLflow model to retrain. Must exist + in the MLflow Model Registry in Production stage. + model_config (dict): Model configuration parameters + metadata (dict): Metadata for logging + + Returns: + dict: Retraining operation results and experiment details. + + Raises: + mlflow.exceptions.MlflowException: If model not found in registry + ValueError: If experiment cannot be created or models cannot be loaded + Exception: Any other exception during the retraining process + """ + + self.logger.custom_info(f'Starting model retraining workflow for {model_name}', metadata) + self.logger.custom_debug(f'Data received for model retraining: {data.to_csv()}', metadata) + + target_name = model_config.get('target', None) + + transform_flavor = model_config.get('transform_flavor', 'sklearn') + predict_flavor = model_config.get('predict_flavor', 'sklearn') + + self.logger.custom_debug( + f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}', + metadata, + ) + + try: + latest_production_id = self.get_model_run_id(model_name, stage='Production') + self.logger.custom_info('Creating model experiment environment', metadata) + retrain_data = self.fit_models( + model_name=model_name, + data=data, + transform_flavor=transform_flavor, + predict_flavor=predict_flavor, + target_name=target_name, + metadata=metadata, + latest_production_id=latest_production_id, + ) + self.logger.custom_info( + f'Model experiment created successfully: {retrain_data}', metadata + ) + + self.logger.custom_info('Saving model retrain', metadata) + experiment = self.create_new_experiment( + model_name=model_name, + data=data, + retrain_data=retrain_data, + transform_flavor=transform_flavor, + predict_flavor=predict_flavor, + metadata=metadata, + latest_production_id=latest_production_id, + ) + self.logger.custom_info( + f'Model retraining completed successfully for experiment: {experiment}', metadata + ) + + return { + 'success': True, + 'experiment': experiment, + 'message': 'Model retrained successfully.', + } + except Exception as e: + error_msg = f'Error retraining model {model_name}: {e}' + self.logger.custom_info(error_msg, metadata) + return { + 'success': False, + 'experiment': None, + 'message': error_msg, + 'traceback': traceback.format_exc(), + } + + def update_production_model( + self, experiment: dict[str, Any], model_name: str, metadata: dict + ) -> 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. + This method orchestrates the complete production model update process by + identifying the most recent retraining run and promoting it to production + stage. It handles model registration, versioning, and stage transitions + with comprehensive metadata tracking. + + Process Flow: + 1. Retrieves experiment information: + - Converts experiment name to MLFlow experiment ID + - Searches for the most recent retraining run in the experiment + - Filters runs by 'retrain' parameter and orders by completion time + 2. Promotes model to production: + - Registers the model from the specified run to MLFlow Model Registry + - Retrieves the latest model version number + - Transitions the model to 'Production' stage + - Archives existing production versions automatically + 3. Returns comprehensive update metadata Args: - experiment (str): MLFlow experiment name - model_name (str): Name of the MLFlow model + experiment (str): MLflow experiment name containing the retraining runs. + Must be a valid experiment that exists in MLFlow. + model_name (str): Name of the MLFlow model to update. Must exist + in the MLFlow Model Registry. + metadata (dict): Metadata for logging 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 + - version (str): New model version number (incremented automatically) + - mlflow_run_id (str): Source run ID of the promoted model + - mlflow_experiment_id (int): Experiment ID for tracking + + Raises: + ValueError: If experiment not found or model versions are invalid + mlflow.exceptions.MlflowException: If model registration or stage + transition fails + Exception: Any other exception during the update process + + Note: + This operation is irreversible. The previous production model will + be automatically archived when the new version is promoted. """ - 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) + run_id = experiment['run_id'] + experiment_id = experiment['experiment_id'] + metadata_result = self.update_production_model_by_run_id(run_id, model_name, metadata) - metadata['mlflow_experiment_id'] = experiment_id + metadata_result['mlflow_experiment_id'] = experiment_id - return metadata + return metadata_result diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 96ecfef..be4edb5 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/laborious/utils/repository/opc_repository.py @@ -1,16 +1,16 @@ -import asyncio -import traceback import time +import traceback 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 asyncua.ua import DataValue, DateTime, Variant, VariantType 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 = { @@ -33,17 +33,26 @@ data_type_map = { '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): +class OpcRepository: + def __init__( + self, + opc_id: str, + url: str, + logger: Logger, + notification_handler: NotificationHandler, + reconnection_interval: int = 60, + server_uri: str | None = None, + cert_path: str | None = None, + private_key_path: str | None = None, + server_cert_path: str | None = None, + pod_id: str | None = None, + ): self.url = url - self.id = id + self.id = opc_id self.server_uri = server_uri self.cert_path = cert_path self.private_key_path = private_key_path @@ -51,16 +60,16 @@ class OpcRepository(): self.logger = logger self.error_count = 0 self.reconnection_interval = reconnection_interval - self.last_reconnection_time = None + self.last_reconnection_time: None | datetime = None self.notification_handler = notification_handler - self.client = None + self.client: None | Client = None self.pod_id = pod_id self.metadata = { 'model_name': '-', 'model_id': '-', 'workflow_name': 'opc_repository', - 'schedule_name': '-' + 'schedule_name': '-', } async def set_security(self): @@ -83,13 +92,17 @@ class OpcRepository(): - Session Timeout: 10,000,000 ms """ - if not all([self.cert_path, self.private_key_path]): + if self.cert_path is None or self.private_key_path is None: raise ValueError( - "Certificate and private key paths must be provided for secure connection.") + '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 + server_cert = Path(self.server_cert_path) if self.server_cert_path else None + + if self.client is None: + raise ValueError('Client must be initialized before setting security') self.client.application_uri = self.server_uri self.logger.custom_info('Setting security...', self.metadata) @@ -97,7 +110,7 @@ class OpcRepository(): SecurityPolicyBasic256, certificate=str(cert), private_key=str(private_key), - server_certificate=str(server_cert) + server_certificate=str(server_cert) if server_cert else None, ) self.client.secure_channel_timeout = 10000000 self.client.session_timeout = 10000000 @@ -115,8 +128,7 @@ class OpcRepository(): 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) + 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]]: @@ -136,6 +148,13 @@ class OpcRepository(): try: self.last_reconnection_time = datetime.now() + if self.client is None: + return False, { + 'notification_id': f'OPC_CONNECTION_ERROR_{self.id}', + 'message': 'Client is not initialized', + 'block': 'opc_repository', + 'level': NotificationLevel.ERROR, + } await self.client.connect() return True, {} except Exception as e: @@ -143,11 +162,11 @@ class OpcRepository(): 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 + '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): @@ -162,11 +181,9 @@ class OpcRepository(): return try: await self.client.disconnect() - self.logger.custom_info( - 'Disconnected from OPC server', self.metadata) + 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.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]]: @@ -203,52 +220,62 @@ class OpcRepository(): if self.error_count > 5: self.logger.custom_warning( - f"OPC server {self.id} will be disconnected due to multiple errors", self.metadata) + 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) + 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) + 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": + 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: + 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) + 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 + '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}" + 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 + '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]]: + 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. @@ -286,42 +313,36 @@ class OpcRepository(): start_time = time.time() try: - node_obj = self.client.get_node(node) + # ignored because self.validate_connection is called before, so we know self.client is not None + node_obj = self.client.get_node(node) # type: ignore[union-attr] 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 + '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 + '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) + 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 - ) + now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond + ), ) try: @@ -331,7 +352,7 @@ class OpcRepository(): pod_id=self.pod_id, model_name=metadata['model_name'], pipeline_name=metadata['workflow_name'], - opc_server_id=self.id + opc_server_id=self.id, ).inc() end_time = time.time() @@ -340,7 +361,7 @@ class OpcRepository(): pod_id=self.pod_id, model_name=metadata['model_name'], pipeline_name=metadata['workflow_name'], - opc_server_id=self.id + opc_server_id=self.id, ).observe(response_time) except Exception as e: @@ -348,11 +369,11 @@ class OpcRepository(): 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 + '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 diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 567a257..781ad55 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -25,33 +25,37 @@ Environment Variables: - 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 +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 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 prometheus_client import start_http_server 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 + from laborious.activities.activities import Activities + from laborious.utils.connectors_config import ( + build_minio_config, + build_mlflow_config, + build_mongodb_config, + build_opc_config, + build_postgres_config, + ) + from laborious.workflows.minimal_retrain import MinimalRetrain + from laborious.workflows.predictions_batch import PredictionsBatch + from laborious.workflows.sub_workflows.format_and_export_prediction import ( + FormatAndExportPrediction, + ) + from laborious.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(): @@ -87,7 +91,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) @@ -97,7 +101,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', 'laborious'), ) logger.custom_info('Starting Activities...', metadata) @@ -105,21 +109,20 @@ async def main(): activities = Activities( postgres_config=build_postgres_config(), mlflow_config=build_mlflow_config(), + minio_config=build_minio_config(), opc_config=build_opc_config(), logger=logger, - notification_handler=notification_handler + 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) + 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}') ) ) @@ -128,7 +131,7 @@ async def main(): temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), - runtime=new_runtime + runtime=new_runtime, ) logger.custom_info('Starting Workers...', metadata) @@ -140,26 +143,28 @@ async def main(): workflows=[MinimalRetrain], activities=[ activities.load_custom_query, + activities.query_to_minio, activities.retrain_model, activities.update_production_model, - activities.export_data_to_postgres + activities.format_retrain_report, + 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, activities.request_transform, + activities.query_to_minio, # Gates activities.input_gate, activities.mlflow_response_gate, @@ -173,15 +178,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 = [] @@ -195,7 +200,7 @@ async def main(): # 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) + logger.custom_error(f'An unhandled exception occurred: {e}', metadata) finally: if notification_handler: notification_handler.shutdown() @@ -224,12 +229,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}") + print(f'Failed to start Prometheus server: {e}') os._exit(1) diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 1893e25..02878cf 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -1,14 +1,16 @@ 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 + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities -@workflow.defn(name="minimal_retrain") -class MinimalRetrain(): +@workflow.defn(name='minimal_retrain') +class MinimalRetrain: """ Automated model retraining workflow for the Laborious system. @@ -63,44 +65,62 @@ 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', } } model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) - data = await workflow.execute_local_activity_method( - Activities.load_custom_query, + storage_result = await workflow.execute_activity_method( + Activities.query_to_minio, { **metadata, 'query': input_data['query'], - 'datetime_columns': input_data.get('datetime_columns', []) + 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': model_name, + 'object_prefix': f'retrain_datasets/{model_name}/data', }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=600), ) + if not storage_result['success']: + return + experiment_response = await workflow.execute_activity_method( Activities.retrain_model, { **metadata, - 'data': data, - 'model_name': model_name + 'object_key': storage_result['object_key'], + 'model_name': model_name, + 'model_config': model_config, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(hours=1), ) - report = await workflow.execute_activity_method( - Activities.update_production_model, + if experiment_response['success']: + update_report = await workflow.execute_activity_method( + Activities.update_production_model, + {**metadata, 'model_name': model_name, **experiment_response}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + else: + update_report = {} + + report = await workflow.execute_local_activity_method( + Activities.format_retrain_report, { **metadata, + 'experiment_response': experiment_response, 'model_name': model_name, 'model_id': input_data['model_id'], - **experiment_response + 'update_report': update_report, }, 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 +129,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=600), ) diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index e522eaa..f79206f 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -1,14 +1,16 @@ 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 + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities -@workflow.defn(name="predictions_batch") -class PredictionsBatch(): +@workflow.defn(name='predictions_batch') +class PredictionsBatch: """ Main batch prediction workflow for the Laborious 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,18 @@ 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']), 'opc_output_config': input_data.get('opc_output_config', {}), - '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/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 8e7df07..72ff12d 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -1,15 +1,17 @@ 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 typing import Any + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.policies import retry_policy + + from laborious.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. @@ -79,10 +81,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: @@ -94,22 +96,18 @@ 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 opc prediction = await workflow.execute_activity_method( Activities.write_opc_data, - { - **metadata, - 'opc_output_config': input_data['opc_output_config'], - 'data': prediction - }, + {**metadata, 'opc_output_config': input_data['opc_output_config'], 'data': prediction}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) # write to postgres @@ -120,21 +118,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/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 777fa1c..7bd36e2 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -1,14 +1,16 @@ 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 + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities -@workflow.defn(name="prediction_process") -class PredictionProcess(): +@workflow.defn(name='prediction_process') +class PredictionProcess: """ Core prediction processing workflow for the Laborious 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), @@ -214,12 +208,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. @@ -265,7 +266,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), @@ -288,8 +289,8 @@ class PredictionProcess(): 'table_name': table_name, 'comment': comment, 'opc_output_config': input_data['opc_output_config'], - 'prediction_store_policy': input_data['prediction_store_policy'] - } + 'prediction_store_policy': input_data['prediction_store_policy'], + }, ) return True diff --git a/model_convert.ipynb b/model_convert.ipynb new file mode 100644 index 0000000..42f6c3d --- /dev/null +++ b/model_convert.ipynb @@ -0,0 +1,106 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 23, + "id": "e838ff21", + "metadata": {}, + "outputs": [], + "source": [ + "import csv\n", + "\n", + "def csv_to_tag_lists(csv_path: str) -> dict:\n", + " read_tags = []\n", + " write_tags = []\n", + "\n", + " def to_float(val):\n", + " try:\n", + " return float(str(val).strip())\n", + " except Exception:\n", + " return None\n", + "\n", + " with open(csv_path, newline=\"\", encoding=\"utf-8\") as f:\n", + " reader = csv.DictReader(f)\n", + " for row in reader:\n", + " # Basic normalization\n", + " op = (row.get(\"operation\") or \"\").strip()\n", + "\n", + " if op == \"READ\":\n", + " # Build common tag payload with required mappings\n", + " tag = {\n", + " \"server_id\": \"1\",\n", + " \"tag_address\": row.get(\"opc_tag\"),\n", + " \"tag_name\": row.get(\"name\"),\n", + " \"data_range\": [to_float(row.get(\"min_value\")), to_float(row.get(\"max_value\"))],\n", + " \"aggr_func\": row.get(\"aggregation_func\").lower(),\n", + " # keep other fields with their original names\n", + " \"frequency\": row.get(\"frequency\"),\n", + " \"local\": row.get(\"local\"),\n", + " \"area\": row.get(\"area\"),\n", + " \"description\": row.get(\"description\"),\n", + " }\n", + "\n", + " read_tags.append(tag)\n", + "\n", + " else:\n", + " tag = {\n", + " \"server_id\": \"1\",\n", + " \"addr\": row.get(\"opc_tag\"),\n", + " \"tag_name\": row.get(\"name\"),\n", + " \"local\": row.get(\"local\"),\n", + " \"area\": row.get(\"area\"),\n", + " \"description\": row.get(\"description\"),\n", + " }\n", + " \n", + " if op == \"WRITE_PREDICTION\":\n", + " tag[\"type\"] = \"prediction\"\n", + " write_tags.append(tag)\n", + " elif op == \"WRITE_CONFIDENCE\":\n", + " tag[\"type\"] = \"confidence\"\n", + " write_tags.append(tag)\n", + " # ignore any other operation values silently\n", + "\n", + " return {\"read_tags\": read_tags, \"write_tags\": write_tags}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4621cd43", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "file_names = [\"Courier - Pรกgina1.csv\"]\n", + "\n", + "for file_name in file_names:\n", + " write_file = file_name.replace(\".csv\", \".json\")\n", + "\n", + " with open(write_file, \"w\", encoding=\"utf-8\") as f:\n", + " json.dump(csv_to_tag_lists(file_name), f, indent=2, ensure_ascii=False)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5b153a3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,159 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "laborious" +version = "0.0.0" +description = "Sientia DataOps Laborious - 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 = [ + "BLE001", # ignore blind except, we need to send notifications with any error + "E501", # line too long (handled by formatter) + "S101", # use of assert (needed for tests) + "S105", # possible hardcoded password (false positives) + "S106", # possible hardcoded password (false positives) + "N802", # function name should be lowercase (temporal decorators) + "N806", # variable in function should be lowercase +] + +[tool.ruff.lint.per-file-ignores] +"tests/**/*.py" = [ + "S101", # assert allowed in tests + "S105", # hardcoded passwords ok in tests + "S106", # hardcoded passwords ok in tests +] + +[tool.ruff.lint.mccabe] +max-complexity = 15 + +[tool.ruff.format] +quote-style = "single" +indent-style = "space" +line-ending = "auto" + +[tool.mypy] +python_version = "3.11" +warn_return_any = false +warn_unused_configs = true +disallow_untyped_defs = false +disallow_incomplete_defs = false +check_untyped_defs = true +no_implicit_optional = true +warn_redundant_casts = true +warn_unused_ignores = false +warn_no_return = true +strict_equality = true +ignore_missing_imports = true + +# Ignore missing imports for external packages +[[tool.mypy.overrides]] +module = "temporalio.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "sientia_do.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "mlflow.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "prometheus_client.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "sientia.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "pandas.*" +ignore_missing_imports = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = [ + "-v", + "--strict-markers", + "--cov=model_manager", + "--cov-report=term-missing", + "--cov-report=html", + "--cov-report=xml", +] +markers = [ + "asyncio: marks tests as async", + "integration: marks tests as integration tests", + "unit: marks tests as unit tests", +] + +[tool.coverage.run] +source = ["model_manager"] +omit = [ + "*/tests/*", + "*/venv/*", + "*/__pycache__/*", + "*/site-packages/*", +] +branch = true + +[tool.coverage.report] +precision = 2 +show_missing = true +skip_covered = false +exclude_lines = [ + "pragma: no cover", + "def __repr__", + "def __str__", + "raise AssertionError", + "raise NotImplementedError", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", + "class .*\\bProtocol\\):", + "@(abc\\.)?abstractmethod", +] + +[tool.coverage.html] +directory = "htmlcov" + +[tool.bandit] +exclude_dirs = ["tests", "venv", ".venv"] +skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..56ab376 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,19 @@ +# Development and Testing Dependencies +# These packages are only needed for development, testing, and code quality checks +# Install with: pip install -r requirements-dev.txt + +# Code Quality & Linting +ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort) +mypy>=1.7.0 # Static type checker +bandit>=1.7.5 # Security vulnerability scanner +pandas-stubs>=2.0.0 # Type stubs for pandas +types-requests>=2.31.0 # Type stubs for requests + +# Testing +pytest>=7.4.0 # Testing framework +pytest-cov>=4.1.0 # Coverage plugin for pytest +pytest-asyncio>=0.21.0 # Async test support (already in main requirements) + +# Development Tools +ipython>=8.12.0 # Enhanced Python shell +ipdb>=0.13.13 # IPython debugger \ No newline at end of file diff --git a/requirements-light.txt b/requirements-light.txt new file mode 100644 index 0000000..c145097 --- /dev/null +++ b/requirements-light.txt @@ -0,0 +1,12 @@ +temporalio +psycopg2-binary +sqlalchemy +asyncua +redis +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.6 +prometheus-client +botocore +boto3 +s3fs +pyarrow +mlflow \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 1ce9f9c..3ba54e8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,7 @@ 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 +botocore +boto3 +s3fs +pyarrow diff --git a/simulator/Dockerfile b/simulator/Dockerfile deleted file mode 100644 index d467676..0000000 --- a/simulator/Dockerfile +++ /dev/null @@ -1,30 +0,0 @@ -# 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/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index de418be..1999e91 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -1,18 +1,19 @@ +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 laborious.activities.activities import Activities -from laborious.activities.mlflow import MLFlow from laborious.activities.gates import Gates +from laborious.activities.mlflow import MLFlow from laborious.activities.opc import OPC +from laborious.activities.storage import Storage -@patch('laborious.activities.activities.Postgres.__init__') +@patch('laborious.activities.activities.Storage.__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): - +def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_storage_init): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -20,20 +21,23 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' + minio_config = { + 'endpoint_url': 'localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', } + mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} + opc_config = { 'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, - 'group_id': 'test-group' + 'group_id': 'test-group', } logger = MagicMock() @@ -42,18 +46,19 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre activities = Activities( postgres_config=postgres_config, mlflow_config=mlflow_config, + minio_config=minio_config, opc_config=opc_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) assert isinstance(activities, Activities) - assert isinstance(activities, Postgres) + assert isinstance(activities, Storage) assert isinstance(activities, MLFlow) assert isinstance(activities, OPC) assert isinstance(activities, Gates) - mock_postgres_init.assert_called_once_with( + mock_storage_init.assert_called_once_with( ANY, host=postgres_config['host'], port=postgres_config['port'], @@ -62,8 +67,9 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre dbname=postgres_config['dbname'], min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], + minio_config=minio_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_mlflow_init.assert_called_once_with( @@ -72,30 +78,25 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre mlflow_port=mlflow_config['port'], mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], + minio_config=minio_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_opc_init.assert_called_once_with( - ANY, - opc_servers=opc_config, - logger=logger, - notification_handler=notification_handler + 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 + ANY, logger=logger, notification_handler=notification_handler ) @mark.asyncio -@patch('laborious.activities.activities.Postgres', return_value=MagicMock()) +@patch('laborious.activities.activities.Storage', 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): +async def test_shutdown(mock_opc_init, _mock_mlflow_init, mock_storage_init): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -103,20 +104,23 @@ async def test_shutdown(mock_opc_init, 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' + minio_config = { + 'endpoint_url': 'localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', } + mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} + opc_config = { 'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, - 'group_id': 'test-group' + 'group_id': 'test-group', } logger = MagicMock() @@ -125,11 +129,12 @@ async def test_shutdown(mock_opc_init, activities = Activities( postgres_config=postgres_config, mlflow_config=mlflow_config, + minio_config=minio_config, opc_config=opc_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) await activities.shutdown() mock_opc_init.shutdown.assert_called_once() - mock_postgres_init.close.assert_called_once() + mock_storage_init.close.assert_called_once() diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index a65c6d4..8d8822e 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 laborious.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,34 @@ 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() + + +@mark.asyncio +async def test_input_gate_with_filter_not_caught(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, + '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.debug.assert_called() @@ -125,51 +140,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('laborious.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 +194,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,86 +210,101 @@ 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() +@mark.asyncio +async def test_mlflow_response_gate_with_filter_not_caught(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'API_ERROR': {'policy': 'STOP'}}, + 'data': { + 'success': True, + '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.debug.assert_called() + + @mark.asyncio 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': { + 'success': True, + 'content': {'message': 'success'}, }, - '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('laborious.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 +316,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,31 +332,48 @@ 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() +@mark.asyncio +async def test_mlflow_content_gate_with_filter_not_caught(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, + 'data': {'content': {'message': 'success'}}, + '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() + + 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) + prediction_store_policy, metadata + ) # Assert assert policy_type == 'lts' @@ -341,7 +386,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 +400,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 +414,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 +428,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 +446,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 +464,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 +477,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 +499,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 +512,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 +524,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 +551,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 @@ -526,15 +569,42 @@ async def test_format_default_prediction(gates_activity): @mark.asyncio -async def test_get_last_timestamp_with_data(gates_activity): +async def test_format_retrain_report(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'] - } + 'experiment_response': { + 'success': True, + 'timestamp': '2023-05-26 11:12:27', + 'message': 'success', + }, + 'update_report': { + 'version': '1.0.0', + 'mlflow_run_id': 'test_mlflow_run_id', + 'mlflow_experiment_id': 'test_mlflow_experiment_id', + }, + 'model_id': 'test_model', + 'model_name': 'test_model', } + # Act + result = await gates_activity.format_retrain_report(input_data) + + # Assert + assert result['model_id'] == {0: 'test_model'} + assert result['model_name'] == {0: 'test_model'} + assert result['timestamp'] == {0: '2023-05-26 11:12:27'} + assert result['status'] == {0: 'success'} + assert result['version'] == {0: '1.0.0'} + assert result['mlflow_run_id'] == {0: 'test_mlflow_run_id'} + assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'} + + +@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) @@ -545,10 +615,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 +634,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 85cb7f1..662f0ba 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -1,45 +1,68 @@ -from datetime import datetime -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import ANY, MagicMock, call, 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 +from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + +from laborious.activities.mlflow import MLFlow -@patch("laborious.activities.mlflow.MLFlowRepository") -def test___init__(mock_mlflow_repository): +@patch('laborious.activities.mlflow.MLFlowRepository') +@patch('laborious.activities.mlflow.MinioRepository') +def test___init__(mock_minio_repository, 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', + minio_config={ + 'endpoint_url': 'http://localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', + }, 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) + + mock_minio_repository.assert_called_once_with( + logger=ANY, + notification_handler=ANY, + minio_endpoint_url='http://localhost:9000', + minio_access_key='minio', + minio_secret_key='minio123', + minio_region_name='us-east-1', + minio_default_bucket='test', ) @fixture -@patch("laborious.activities.mlflow.MLFlowRepository") -def mlflow(mock_mlflow_repository): +@patch('laborious.activities.mlflow.MLFlowRepository') +@patch('laborious.activities.mlflow.MinioRepository') +def mlflow(mock_minio_repository, 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', + minio_config={ + 'endpoint_url': 'http://localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', + }, logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) mlflow.send_notification = MagicMock() @@ -48,44 +71,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("laborious.activities.mlflow.DataFrame") -@patch("laborious.activities.mlflow.max") +@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'} + { + '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 +160,25 @@ 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('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" + '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,8 +189,9 @@ 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 ) mock_dataframe.return_value.__setitem__.assert_any_call( 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value @@ -158,9 +200,12 @@ 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) + + 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 @@ -172,98 +217,232 @@ 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] +@patch('laborious.activities.mlflow.to_datetime') +async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow): + mlflow.model_monitoring_repository.retrain_model.return_value = { + 'success': True, + 'experiment': 'test_experiment', + 'message': 'Model retrained successfully.', } - mlflow.model_monitoring_repository.retrain_model.return_value = ( - 'Model retrained successfully', 'test') + response = await mlflow.retrain_model( + { + **metadata, + 'object_key': 'test_object_key', + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } + ) - response = await mlflow.retrain_model({ - **metadata, - 'data': data, - 'model_name': 'test_model' - }) + raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value - mlflow.model_monitoring_repository.retrain_model.assert_called_once() + timestamp = raw_data.__getitem__.return_value.max.return_value + + raw_data.sort_values.assert_called_once_with('created_at', ascending=False) + raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with( + subset=['variable', 'timestamp'], keep='first' + ) + raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value + + raw_data.drop.assert_has_calls( + [ + call(columns=['model_id'], inplace=True, errors='ignore'), + call(columns=['created_at'], inplace=True, errors='ignore'), + ] + ) + raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value') + raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True) + + raw_data = raw_data.pivot.return_value + + raw_data.__setitem__.assert_has_calls( + [ + call('timestamp', raw_data.index), + call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value), + call('timestamp', mock_to_datetime.return_value), + ] + ) + mock_to_datetime.assert_has_calls( + [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)] + ) + mock_to_datetime.assert_has_calls( + [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)] + ) + + mlflow.model_monitoring_repository.retrain_model.assert_called_once_with( + data=raw_data, + model_name='test_model', + model_config={ + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + metadata=metadata['metadata'], + ) assert response == { - "status": 'Model retrained successfully', - "timestamp": 2, - "experiment": 'test' + 'success': True, + 'experiment': 'test_experiment', + 'message': 'Model retrained successfully.', + 'timestamp': timestamp, } @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] +@patch('laborious.activities.mlflow.to_datetime') +async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow): + mlflow.model_monitoring_repository.retrain_model.return_value = { + 'success': False, + 'traceback': 'test_traceback', + 'message': 'Model retrained failed.', } - try: - await mlflow.retrain_model({ + response = 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 + 'object_key': 'test_object_key', + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } + ) + + raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value + + timestamp = raw_data.__getitem__.return_value.max.return_value + + raw_data.sort_values.assert_called_once_with('created_at', ascending=False) + raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with( + subset=['variable', 'timestamp'], keep='first' + ) + raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value + + raw_data.drop.assert_has_calls( + [ + call(columns=['model_id'], inplace=True, errors='ignore'), + call(columns=['created_at'], inplace=True, errors='ignore'), + ] + ) + raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value') + raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True) + + raw_data = raw_data.pivot.return_value + + raw_data.__setitem__.assert_has_calls( + [ + call('timestamp', raw_data.index), + call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value), + call('timestamp', mock_to_datetime.return_value), + ] + ) + mock_to_datetime.assert_has_calls( + [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)] + ) + mock_to_datetime.assert_has_calls( + [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)] + ) + + mlflow.model_monitoring_repository.retrain_model.assert_called_once_with( + data=raw_data, + model_name='test_model', + model_config={ + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + metadata=metadata['metadata'], + ) + + mlflow.send_notification.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='RETRAIN_MODEL_ERROR', + message='Error retraining model test_model: Model retrained failed.', + block='retrain_model', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + + assert response == { + 'success': False, + 'traceback': 'test_traceback', + 'message': 'Model retrained failed.', + 'timestamp': timestamp, + } + + +@mark.asyncio +async def test_retrain_model_data_error(mlflow): + mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception( + 'Error loading retrain data' + ) + + response = await mlflow.retrain_model( + { + **metadata, + 'object_key': 'test_object_key', + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } + ) + + assert response == { + 'success': False, + 'message': 'Error loading retrain data: Error loading retrain data', + 'traceback': ANY, + 'timestamp': ANY, + } + + +@mark.asyncio +async def test_retrain_model_data_error_no_minio_repository(mlflow): + mlflow.minio_repository = None + + with raises(ValueError) as e: + await mlflow.retrain_model( + { + **metadata, + 'object_key': 'test_object_key', + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } ) - else: - assert False, "No exception raised" + + assert str(e.value) == 'Minio repository not initialized' @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' + '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', metadata=metadata['metadata'] + ) - assert response == { - 'data1': {0: 1}, - 'data2': {0: 2}, - 'model_id': {0: 1}, - 'model_name': {0: 'test_model'}, - 'timestamp': {0: 2}, - 'status': {0: 'success'} - } + assert response == mlflow.model_monitoring_repository.update_production_model.return_value @mark.asyncio @@ -278,7 +457,7 @@ async def test_update_production_model_error(mlflow): 'model_id': 1, 'experiment': 'test', 'timestamp': 2, - 'status': 'success' + 'status': 'success', } try: @@ -291,7 +470,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/activities/test_opc.py b/tests/laborious/activities/test_opc.py index 07f00b8..815df27 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -1,57 +1,55 @@ -from unittest.mock import patch, MagicMock, ANY, call, AsyncMock -from pandas import DataFrame -from pytest import fixture, mark +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch + import pytest_asyncio +from pandas import DataFrame +from pytest import mark 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", + '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() - ) + 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") +@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, {})) + connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {})) ) server2 = MagicMock( - connect=AsyncMock(return_value=(True, {})), - write_data=AsyncMock(return_value=(True, {})) + 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, {})) + 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() @@ -82,12 +80,10 @@ async def test_init_opc(mock_send_notification, mock_opc_repository): 'private_key_path': '', 'server_cert_path': '', 'reconnection_interval': 60, - } + }, } opc = OPC( - opc_servers=servers, - logger=mock_logger, - notification_handler=mock_notification_handler + opc_servers=servers, logger=mock_logger, notification_handler=mock_notification_handler ) await opc.init_opc() @@ -97,57 +93,63 @@ async def test_init_opc(mock_send_notification, mock_opc_repository): 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' - ) - ]) + mock_opc_repository.assert_has_calls( + [ + call( + opc_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( + opc_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 - ) - ]) + 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") +@patch('laborious.activities.opc.OpcRepository') async def opc(mock_opc_repository): servers = { 'server1': { @@ -161,17 +163,9 @@ async def opc(mock_opc_repository): } } - 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() - ) + 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 @@ -188,58 +182,79 @@ WRITE_DATA_CASES = [ @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) + 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) + 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' - }) + 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) + 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", + 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 + attachment_content=ANY, ) @mark.asyncio async def test_write_data_exception(opc): - opc.opc_repository['server1'].write_data.side_effect = Exception( - "Test error") + 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) + 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", + 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 + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio @@ -247,20 +262,13 @@ async def test_write_opc_data_success(opc): # Arrange input_data = { **metadata, - 'data': { - 'prediction': [0.75], - 'prediction_confidence': [0.95] - }, + 'data': {'prediction': [0.75], 'prediction_confidence': [0.95]}, 'opc_output_config': { 'server1': { - 'prediction_tags': { - 'tag1': {'data_type': 'float'} - }, - 'confidence_tags': { - 'tag2': {'data_type': 'float'} - } + 'prediction_tags': {'tag1': {'data_type': 'float'}}, + 'confidence_tags': {'tag2': {'data_type': 'float'}}, } - } + }, } # Act @@ -270,25 +278,30 @@ async def test_write_opc_data_success(opc): # 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'] - ) - ]) + 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 @@ -297,17 +310,9 @@ async def test_write_opc_data_empty_config(opc): # Arrange input_data = { **metadata, - 'data': { - 'prediction': [0.75], - 'prediction_confidence': [0.95] - }, + 'data': {'prediction': [0.75], 'prediction_confidence': [0.95]}, 'opc_servers': ['server1'], - 'opc_output_config': { - 'server1': { - 'prediction_tags': {}, - 'confidence_tags': {} - } - } + 'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}}, } # Act @@ -322,20 +327,13 @@ 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] - }, + 'data': {'prediction': [0.75], 'prediction_confidence': [0.95]}, 'opc_output_config': { 'server1': { - 'prediction_tags': { - 'tag1': {'data_type': 'float'} - }, - 'confidence_tags': { - 'tag2': {'data_type': 'float'} - } + 'prediction_tags': {'tag1': {'data_type': 'float'}}, + 'confidence_tags': {'tag2': {'data_type': 'float'}}, } - } + }, } # Act @@ -345,10 +343,13 @@ async def test_write_opc_data_no_validate_server(opc): 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), -]) +@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) diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py new file mode 100644 index 0000000..76b83e5 --- /dev/null +++ b/tests/laborious/activities/test_storage.py @@ -0,0 +1,213 @@ +import datetime +from unittest.mock import ANY, AsyncMock, MagicMock, patch + +from pytest import fixture, mark, raises +from sientia_do.notifications.models import NotificationLevel +from sientia_do.temporal.activities.postgres import Postgres + +from laborious.activities.storage import Storage + +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schedule_name': 'test_schedule', + } +} + + +@fixture +@patch('laborious.activities.storage.MinioRepository') +def storage(mock_minio_repository): + return Storage( + host='localhost', + port=5432, + user='postgres', + password='postgres', + dbname='postgres', + min_connections=1, + max_connections=10, + minio_config={ + 'endpoint_url': 'localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', + }, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + +@patch('laborious.activities.storage.MinioRepository') +def test___init___not_hasattr(mock_minio_repository): + logger = MagicMock() + notification_handler = MagicMock() + storage = Storage( + host='localhost', + port=5432, + user='postgres', + password='postgres', + dbname='postgres', + min_connections=1, + max_connections=10, + minio_config={ + 'endpoint_url': 'localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', + }, + logger=logger, + notification_handler=notification_handler, + ) + assert isinstance(storage, Postgres) + + mock_minio_repository.assert_called_once_with( + logger=logger, + notification_handler=notification_handler, + minio_endpoint_url='localhost:9000', + minio_access_key='minio', + minio_secret_key='minio123', + minio_region_name='us-east-1', + minio_default_bucket='test', + ) + + +@patch('laborious.activities.storage.MinioRepository') +def test___init___none_minio_repository(mock_minio_repository, storage): + storage.minio_repository = None + logger = MagicMock() + notification_handler = MagicMock() + + storage.__init__( + host='localhost', + port=5432, + user='postgres', + password='postgres', + dbname='postgres', + min_connections=1, + max_connections=10, + minio_config={ + 'endpoint_url': 'localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', + }, + logger=logger, + notification_handler=notification_handler, + ) + + mock_minio_repository.assert_called_once_with( + logger=logger, + notification_handler=notification_handler, + minio_endpoint_url='localhost:9000', + minio_access_key='minio', + minio_secret_key='minio123', + minio_region_name='us-east-1', + minio_default_bucket='test', + ) + + +@patch('laborious.activities.storage.MinioRepository') +def test___init___done_repository(mock_minio_repository, storage): + storage.__init__( + host='localhost', + port=5432, + user='postgres', + password='postgres', + dbname='postgres', + min_connections=1, + max_connections=10, + minio_config={ + 'endpoint_url': 'localhost:9000', + 'access_key': 'minio', + 'secret_key': 'minio123', + 'region_name': 'us-east-1', + 'default_bucket': 'test', + }, + logger=MagicMock(), + notification_handler=MagicMock(), + ) + mock_minio_repository.assert_not_called() + assert storage.minio_repository is not None + + +@mark.asyncio +async def test_query_to_minio_minio_repository_not_initialized(storage): + storage.minio_repository = None + + with raises(ValueError) as e: + await storage.query_to_minio({}) + + assert str(e.value) == 'Minio repository not initialized' + + +@mark.asyncio +async def test_query_to_minio_not_data(storage): + storage.load_custom_query = AsyncMock(return_value=None) + result = await storage.query_to_minio({}) + + storage.load_custom_query.assert_called_once_with({}) + assert result['success'] is False + assert result['message'] == 'No data returned from query' + + +@mark.asyncio +@patch('laborious.activities.storage.pd.DataFrame') +@patch('laborious.activities.storage.now') +async def test_query_to_minio_success(now, dataframe, storage): + data = [{'a': 1}, {'a': 2}, {'a': 3}] + storage.load_custom_query = AsyncMock(return_value=data) + now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0) + storage.minio_repository.minio_bucket = 'test' + + result = await storage.query_to_minio({'object_prefix': 'test', **metadata}) + + dataframe.assert_called_once_with(data) + + storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with( + dataframe=dataframe.return_value, + uri='s3://test/test_2024-01-01_00-00-00.parquet', + object_name='test_2024-01-01_00-00-00.parquet', + metadata=metadata['metadata'], + ) + + assert result['success'] is True + assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet' + assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet' + + +@mark.asyncio +async def test_query_to_minio_error(storage): + storage.send_notification = MagicMock() + storage.load_custom_query = AsyncMock(side_effect=Exception('test')) + result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'}) + assert result['success'] is False + assert result['message'] == 'test' + storage.send_notification.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='ERROR_STORING_QUERY_TO_MINIO', + message='Error storing query to MinIO: test', + block='query_to_minio', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + + +def test_close(storage): + storage.minio_repository = MagicMock() + + storage.close() + + assert storage.minio_repository is None + + +def test___del__(storage): + storage.close = MagicMock() + + storage.__del__() + + storage.close.assert_called_once() diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index 405bc9b..8a35801 100644 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -1,23 +1,36 @@ from pandas import DataFrame from laborious.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_empty_data(): + assert ( + filter_specific_variables_null_values(DataFrame(), 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 +38,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 f9c61e9..5353a6f 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 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 + assert api_error_filter(None, {}) is True # NOSONAR def test_api_error_filter_valid_response_fail(): - assert api_error_filter({'success': False}, {}) == True + assert api_error_filter({'success': False}, {}) is True def test_api_error_filter_valid_response_success(): - assert api_error_filter({'success': True}, {}) == False + assert api_error_filter({'success': True}, {}) is False 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]}), {}) is True def test_nan_values_filter_no_nan_values(): - assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False + assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False diff --git a/tests/laborious/utils/repository/test_minio_repository.py b/tests/laborious/utils/repository/test_minio_repository.py new file mode 100644 index 0000000..bebf743 --- /dev/null +++ b/tests/laborious/utils/repository/test_minio_repository.py @@ -0,0 +1,139 @@ +from unittest.mock import MagicMock, patch + +from botocore.utils import ClientError +from pytest import fixture, raises + +from laborious.utils.repository.minio_repository import MinioRepository + + +@patch('laborious.utils.repository.minio_repository.boto3') +@patch('laborious.utils.repository.minio_repository.Config') +def test___init___(mock_config, mock_boto3): + minio_repository = MinioRepository( + minio_endpoint_url='localhost:9000', + minio_access_key='minio', + minio_secret_key='minio123', + minio_region_name='us-east-1', + minio_default_bucket='test', + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + assert minio_repository.storage_options == { + 'key': 'minio', + 'secret': 'minio123', + 'client_kwargs': {'endpoint_url': 'localhost:9000'}, + } + assert minio_repository.minio_bucket == 'test' + assert minio_repository.minio_endpoint_url == 'localhost:9000' + assert minio_repository.minio_region_name == 'us-east-1' + + mock_config.assert_called_once_with( + signature_version='s3v4', + s3={'addressing_style': 'path'}, + retries={'max_attempts': 5, 'mode': 'standard'}, + connect_timeout=5, + read_timeout=120, + ) + + mock_boto3.client.assert_called_once_with( + 's3', + endpoint_url='localhost:9000', + aws_access_key_id='minio', + aws_secret_access_key='minio123', + region_name='us-east-1', + config=mock_config.return_value, + ) + + +@fixture +@patch('laborious.utils.repository.minio_repository.Config') +@patch('laborious.utils.repository.minio_repository.boto3') +def minio_repository(mock_boto3, mock_config): + return MinioRepository( + minio_endpoint_url='localhost:9000', + minio_access_key='minio', + minio_secret_key='minio123', + minio_region_name='us-east-1', + minio_default_bucket='test', + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + +def test_close(minio_repository): + minio_repository.close() + minio_repository.s3_client.close.assert_called_once() + + +def test_ensure_bucket_exists_bucket_exists(minio_repository): + assert minio_repository.ensure_bucket_exists({}) is None + + minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') + + +def test_ensure_bucket_exists_bucket_not_exists_create_success(minio_repository): + minio_repository.s3_client.head_bucket.side_effect = ClientError( + error_response={'Error': {'Code': '404'}}, operation_name='head_bucket' + ) + + assert minio_repository.ensure_bucket_exists({}) is None + + minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') + minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') + + +def test_ensure_bucket_exists_bucket_not_exists_create_error(minio_repository): + minio_repository.send_notification = MagicMock() + + minio_repository.s3_client.head_bucket.side_effect = ClientError( + error_response={'Error': {'Code': '404'}}, operation_name='head_bucket' + ) + minio_repository.s3_client.create_bucket.side_effect = ClientError( + error_response={'Error': {'Code': '404'}}, operation_name='create_bucket' + ) + + with raises(ClientError): + minio_repository.ensure_bucket_exists({}) + + minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') + minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') + + +@patch('laborious.utils.repository.minio_repository.BytesIO') +def test_store_dataframe_as_parquet(mock_bytesio, minio_repository): + input_data = MagicMock() + + minio_repository.ensure_bucket_exists = MagicMock(return_value=True) + + minio_repository.store_dataframe_as_parquet( + dataframe=input_data, uri='s3://test/test.parquet', object_name='test.parquet', metadata={} + ) + + minio_repository.ensure_bucket_exists.assert_called_once_with({}) + mock_bytesio.assert_called_once() + + input_data.to_parquet.assert_called_once_with( + mock_bytesio.return_value, engine='pyarrow', index=True + ) + mock_bytesio.return_value.seek.assert_called_once_with(0) + minio_repository.s3_client.put_object.assert_called_once_with( + Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value + ) + + +@patch('laborious.utils.repository.minio_repository.BytesIO') +@patch('laborious.utils.repository.minio_repository.read_parquet') +def test_get_parquet_as_dataframe(mock_read_parquet, mock_bytesio, minio_repository): + input_data = {'Body': MagicMock(read=MagicMock(return_value=b'test'))} + + minio_repository.s3_client.get_object.return_value = input_data + + output = minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={}) + + minio_repository.s3_client.get_object.assert_called_once_with(Bucket='test', Key='test.parquet') + + mock_bytesio.assert_called_once_with(input_data['Body'].read.return_value) + mock_read_parquet.assert_called_once_with(mock_bytesio.return_value) + + assert output == mock_read_parquet.return_value diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index cdc59b4..a437653 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,34 +1,60 @@ +from datetime import UTC, datetime from unittest.mock import ANY, MagicMock, call, patch + +import mlflow as mlflow_lib 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 +from pandas import DataFrame, Timestamp + +from laborious.utils.repository.model_repository import MLFlowRepository, force_memory_release + + +@patch('laborious.utils.repository.model_repository.ctypes') +@patch('laborious.utils.repository.model_repository.gc') +def test_force_memory_release_success(gc, ctypes): + logger = MagicMock() + + force_memory_release(logger) + + gc.collect.assert_called_once() + ctypes.CDLL.return_value.malloc_trim.assert_called_once_with(0) + logger.info.assert_called_once_with('Memory released') + + +@patch('laborious.utils.repository.model_repository.ctypes') +@patch('laborious.utils.repository.model_repository.gc') +def test_force_memory_release_error(gc, ctypes): + logger = MagicMock() + ctypes.CDLL.return_value.malloc_trim.side_effect = Exception('error') + force_memory_release(logger) + + gc.collect.assert_called_once() + ctypes.CDLL.return_value.malloc_trim.assert_called_once_with(0) + + logger.info.assert_called_once_with('Memory release failed: error') @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() - + with patch('laborious.utils.repository.model_repository.mlflow'): repo = MLFlowRepository( - host='http://localhost:5000', - username='admin', - password='admin', - logger=MagicMock() + host='http://localhost:5000', username='admin', password='admin', logger=MagicMock() ) return repo +@pytest.fixture +def mlflow(): + with patch('laborious.utils.repository.model_repository.mlflow') as mlflow: + yield mlflow + + 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', }, } @@ -37,203 +63,75 @@ 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' +def test_get_model_uri_prediction(mlflow, mlflow_repository): + mlflow.get_run.return_value = MagicMock(info=MagicMock(artifact_uri='test')) + output = mlflow_repository.get_model_uri('0', prediction=True) + assert output == 'test/prediction_model' 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_model_uri_transform(mlflow, mlflow_repository): + mlflow.get_run.return_value = MagicMock(info=MagicMock(artifact_uri='test')) + output = mlflow_repository.get_model_uri('0', prediction=False) + assert output == 'test/data_model' + mlflow.get_run.assert_called_once_with('0') + + +def test_get_model_run_id_not_registered_models(mlflow_repository): + mlflow_repository.client.search_registered_models.return_value = [] + with pytest.raises(mlflow_lib.exceptions.MlflowException) as e: + mlflow_repository.get_model_run_id('test') + + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + + assert str(e.value) == "Model 'test' not found in the Model Registry." + + +def test_get_model_run_id_not_stage_versions(mlflow_repository): + mlflow_repository.client.search_registered_models.return_value = [MagicMock(name='test')] + + mlflow_repository.client.search_model_versions.return_value = [ + MagicMock(current_stage='Staging'), + MagicMock(current_stage='Staging'), + MagicMock(current_stage='Archived'), + ] + + with pytest.raises(mlflow_lib.exceptions.MlflowException) as e: + mlflow_repository.get_model_run_id('test') + + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + mlflow_repository.client.search_model_versions.assert_called_once_with( + filter_string="name='test'" + ) + + assert str(e.value) == "Model 'test' in stage 'Production' not found in the Model Registry." + + +def test_get_model_run_id_success(mlflow_repository): + mlflow_repository.client.search_registered_models.return_value = [MagicMock(name='test')] + + mlflow_repository.client.search_model_versions.return_value = [ + MagicMock(current_stage='Production', version='1'), + MagicMock(current_stage='Production', version='2', source='runs/test/1'), + MagicMock(current_stage='Archived', version='3'), + ] + + output = mlflow_repository.get_model_run_id('test') + + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + mlflow_repository.client.search_model_versions.assert_called_once_with( + filter_string="name='test'" + ) + + assert output == '1' + + 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,17 +142,100 @@ def test_get_next_run_name(mlflow, mlflow_repository): ) -@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') +def test_get_experiment_experiment_exists(mlflow, mlflow_repository): + experiment = MagicMock(experiment_id='0') + + mlflow.get_experiment_by_name.return_value = experiment output = mlflow_repository.get_experiment('test') - assert output == 0 + assert output == experiment + + +def test_get_experiment_none_create(mlflow, mlflow_repository): + experiment = MagicMock(experiment_id='0') + + mlflow.get_experiment_by_name.return_value = None + + mlflow.create_experiment.return_value = experiment + + output = mlflow_repository.get_experiment('test', create_if_not_exists=True) + + assert output == experiment + + +def test_get_experiment_none_not_create(mlflow, mlflow_repository): + mlflow.get_experiment_by_name.return_value = None + + with pytest.raises(ValueError) as e: + mlflow_repository.get_experiment('test', create_if_not_exists=False) + + assert str(e) == 'Experiment test not found' + + +def test_get_model_params(mlflow, mlflow_repository): + output = mlflow_repository.get_model_params('test') + assert output == mlflow.get_run.return_value.data.params + + +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +@patch('laborious.utils.repository.model_repository.makedirs') +def test_download_artifacts_success(makedirs, rmtree, path, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='test') + + path.exists.return_value = True + + output = mlflow_repository.dowload_artifacts('test', 'path') + + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='test', stage='Production' + ) + + path.join.assert_called_once_with('./tmp/artifacts/test', 'path') + + path.exists.assert_called_once_with(path.join.return_value) + + rmtree.assert_called_once_with(path.join.return_value) + + makedirs.assert_called_once_with('./tmp/artifacts/test', exist_ok=True) + + mlflow_repository.client.download_artifacts.assert_called_once_with( + mlflow_repository.get_model_run_id.return_value, 'path', './tmp/artifacts/test' + ) + + assert output == mlflow_repository.client.download_artifacts.return_value + + +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +@patch('laborious.utils.repository.model_repository.makedirs') +def test_download_artifacts_success_path_false(makedirs, rmtree, path, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='test') + + path.exists.return_value = False + + output = mlflow_repository.dowload_artifacts('test', 'path') + + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='test', stage='Production' + ) + + path.join.assert_called_once_with('./tmp/artifacts/test', 'path') + + path.exists.assert_called_once_with(path.join.return_value) + + rmtree.assert_not_called() + + makedirs.assert_called_once_with('./tmp/artifacts/test', exist_ok=True) + + mlflow_repository.client.download_artifacts.assert_called_once_with( + mlflow_repository.get_model_run_id.return_value, 'path', './tmp/artifacts/test' + ) + + assert output == mlflow_repository.client.download_artifacts.return_value -@patch('laborious.utils.repository.model_repository.mlflow') def test_get_experiment_error(mlflow, mlflow_repository): mlflow.get_experiment_by_name.return_value = None @@ -263,188 +244,620 @@ 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 ValueError') -@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'], - }) +def test_load_predict_model_sklearn(mlflow, mlflow_repository): + result = mlflow_repository.load_predict_model('test_model', 'sklearn') - output = mlflow_repository.get_experiment_last_run(0) + assert result == mlflow.sklearn.load_model.return_value + mlflow.sklearn.load_model.assert_called_once_with('models:/test_model/production') - mlflow.search_runs.assert_called_once_with( - experiment_ids=[0], - filter_string="", - output_format="pandas", + +def test_load_predict_model_pyfunc(mlflow, mlflow_repository): + result = mlflow_repository.load_predict_model('test_model', 'pyfunc') + assert result == mlflow.pyfunc.load_model.return_value + mlflow.pyfunc.load_model.assert_called_once_with('models:/test_model/production') + + +def test_load_predict_model_pytorch(mlflow, mlflow_repository): + result = mlflow_repository.load_predict_model('test_model', 'pytorch') + assert result == mlflow.pytorch.load_model.return_value + mlflow.pytorch.load_model.assert_called_once_with('models:/test_model/production') + + +def test_load_predict_model_error(mlflow_repository): + with pytest.raises(ValueError) as e: + mlflow_repository.load_predict_model('test_model', 'invalid') + assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." + + +def validate_common_load_transform_model_mocks(mlflow_repository, model_name): + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name=model_name, stage='Production' + ) + mlflow_repository.get_model_uri.assert_called_once_with( + mlflow_repository.get_model_run_id.return_value, prediction=False ) - assert output == '2' + +def test_load_transform_model_sklearn(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() + + result = mlflow_repository.load_transform_model('test_model', 'sklearn') + + validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') + + assert result == mlflow.sklearn.load_model.return_value + mlflow.sklearn.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_last_run_error(mlflow, mlflow_repository): - mlflow.search_runs.return_value = [] +def test_load_transform_model_pyfunc(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() - try: - mlflow_repository.get_experiment_last_run(0) - except ValueError as e: - assert str(e) == 'Runs is not a pandas DataFrame' - else: - assert False + result = mlflow_repository.load_transform_model('test_model', 'pyfunc') + + validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') + + assert result == mlflow.pyfunc.load_model.return_value + mlflow.pyfunc.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) -@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): +def test_load_transform_model_pytorch(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() - 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() + result = mlflow_repository.load_transform_model('test_model', 'pytorch') + validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') - data_model_mock = MagicMock() - prediction_model_mock = MagicMock() + assert result == mlflow.pytorch.load_model.return_value + mlflow.pytorch.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) - 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' +def test_load_transform_model_error(mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() - prediction_model_mock.fit.return_value = prediction_model_mock + with pytest.raises(ValueError) as e: + mlflow_repository.load_transform_model('test_model', 'invalid') + assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." - data = DataFrame({ - 'x': [1, 2, 3], - 'y': [4, 5, 6] - }) - output = mlflow_repository.create_model_experiment('test', data) +def test_download_model_invalid_model_type(mlflow_repository): + with pytest.raises(ValueError) as e: + mlflow_repository.download_model('test_model', 'invalid', 'sklearn') + assert str(e) == "Invalid model_type. Use 'predict' or 'transform'." - 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 +@pytest.mark.parametrize( + 'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')] +) +def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type): + mlflow_repository.dowload_artifacts = MagicMock() - data_model_mock.fit.assert_called_once_with(data) - data_model_mock.predict.assert_called_once_with(data) + result = mlflow_repository.download_model('test_model', model_type[0], 'pyfunc', True) - 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.dowload_artifacts.assert_called_once_with('test_model', model_type[1]) + + mlflow.pyfunc.load_model.assert_called_once_with( + mlflow_repository.dowload_artifacts.return_value ) - 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 result == ( + mlflow.pyfunc.load_model.return_value._model_impl.python_model, + mlflow_repository.dowload_artifacts.return_value, ) - assert output == (prediction_model_mock, - data_model_mock, - mlflow_repository.get_experiment_by_run_id.return_value) + +def test_download_model_predict(mlflow_repository): + mlflow_repository.load_predict_model = MagicMock() + mlflow_repository.load_transform_model = MagicMock() + + result = mlflow_repository.download_model('test_model', 'predict', 'pyfunc', False) + + mlflow_repository.load_predict_model.assert_called_once_with('test_model', 'pyfunc') + mlflow_repository.load_transform_model.assert_not_called() + + assert result == (mlflow_repository.load_predict_model.return_value, None) -@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): +def test_download_model_transform(mlflow_repository): + mlflow_repository.load_predict_model = MagicMock() + mlflow_repository.load_transform_model = MagicMock() + + result = mlflow_repository.download_model('test_model', 'transform', 'pyfunc', False) + + mlflow_repository.load_predict_model.assert_not_called() + mlflow_repository.load_transform_model.assert_called_once_with('test_model', 'pyfunc') + + assert result == (mlflow_repository.load_transform_model.return_value, None) + + +invalid_cases = [ + ( + {'value': {'2024-01-01 12:00:00': 1, 2024: 2}}, + "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Elements are , .", + ), + ( + {'value': {'2024-01-01': 1, '2024-01-02': 2}}, + 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Unable to parse given date format', + ), + ( + {'value': {Any(): 1, Any(): 2}}, + 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Got .', + ), +] + + +@pytest.mark.parametrize('data', invalid_cases) +def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data): + input_data = DataFrame(data[0]) + + message = data[1] + + with pytest.raises(ValueError) as e: + mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) + + assert str(e) == message + + +valid_cases = [ + ( + {'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}}, + ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'], + ), + ( + { + 'value': { + datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1, + datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2, + } + }, + ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'], + ), + ( + { + 'value': { + datetime(2025, 1, 1, 12, 0, 0, tzinfo=None): 1, + datetime(2025, 1, 2, 12, 0, 0, tzinfo=None): 2, + } + }, + ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'], + ), + ( + { + 'value': { + Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1, + Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2, + } + }, + ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'], + ), +] + + +@pytest.mark.parametrize('data,expected', valid_cases) +def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected): + input_data = DataFrame(data) + + response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) + + assert response.index.tolist() == expected + + +@patch('laborious.utils.repository.model_repository.datetime') +def test_check_cache_retention_false(datetime_mock, mlflow_repository): + datetime_mock.now = MagicMock(return_value=datetime.strptime('2025-01-02', '%Y-%m-%d')) + cache = {'timestamp': datetime.strptime('2025-01-01', '%Y-%m-%d')} + + assert mlflow_repository.check_cache_retention(cache, 1) is False + + +@patch('laborious.utils.repository.model_repository.datetime') +def test_check_cache_retention_true(datetime_mock, mlflow_repository): + datetime_mock.now = MagicMock(return_value=datetime.strptime('2025-01-01', '%Y-%m-%d')) + cache = {'timestamp': datetime.strptime('2025-01-01', '%Y-%m-%d')} + assert mlflow_repository.check_cache_retention(cache, 1) is True + + +def test_handle_valid_model(mlflow_repository): + cache = {'target': {'model': 'model', 'artifact_path': 'test_artifact_path'}} + output = mlflow_repository.handle_valid_model('model_name', cache) + assert output == {'model': 'model', 'artifact_path': 'test_artifact_path'} + + +def test_handle_outdated_model(mlflow_repository): + mlflow_repository.model_cache = { + 'model_name_transform': { + 'target': {'model': 'model', 'artifact_path': 'test_artifact_path'} + } + } + mlflow_repository.handle_outdated_model('model_name', 'model_name_transform') + assert mlflow_repository.model_cache == {} + + +def test_get_model_retention_0(mlflow_repository): + model = MagicMock() + + mlflow_repository.download_model = MagicMock(return_value=(model, 'artifact_path')) + output = mlflow_repository.get_model('model_name', 0, 'predict', 'pyfunc') + + assert output == model + mlflow_repository.download_model.assert_called_once_with( + model_name='model_name', model_type='predict', flavor='pyfunc', load_wrapper=False + ) + + +def test_get_model_cached_valid(mlflow_repository): + mlflow_repository.check_cache_retention = MagicMock(return_value=True) + mlflow_repository.handle_valid_model = MagicMock() + mlflow_repository.handle_outdated_model = MagicMock() + mlflow_repository.model_cache = { + 'model_name_predict': { + 'target': 'cached_model', + } + } + + output = mlflow_repository.get_model('model_name', 1, 'predict', 'pyfunc') + + assert output == mlflow_repository.handle_valid_model.return_value + mlflow_repository.check_cache_retention.assert_called_once_with( + mlflow_repository.model_cache['model_name_predict'], 1 + ) + + mlflow_repository.handle_valid_model.assert_called_once_with( + model_name='model_name', cache=mlflow_repository.model_cache['model_name_predict'] + ) + + mlflow_repository.handle_outdated_model.assert_not_called() + + +def test_get_model_cached_outdated(mlflow_repository): + mlflow_repository.check_cache_retention = MagicMock(return_value=False) + mlflow_repository.handle_valid_model = MagicMock() + mlflow_repository.handle_outdated_model = MagicMock() + model = MagicMock() + mlflow_repository.download_model = MagicMock(return_value=(model, 'artifact_path')) + cache = { + 'model_name_predict': { + 'target': 'cached_model', + } + } + mlflow_repository.model_cache = cache + + output = mlflow_repository.get_model('model_name', 1, 'predict', 'pyfunc') + assert output == model + mlflow_repository.check_cache_retention.assert_called_once_with( + { + 'target': 'cached_model', + }, + 1, + ) + mlflow_repository.handle_valid_model.assert_not_called() + mlflow_repository.handle_outdated_model.assert_called_once_with( + model_name='model_name', model_key='model_name_predict' + ) + + +def test_get_model_cached_not_found(mlflow_repository): + mlflow_repository.check_cache_retention = MagicMock(return_value=False) + mlflow_repository.handle_valid_model = MagicMock() + mlflow_repository.handle_outdated_model = MagicMock() + mlflow_repository.model_cache = {} + model = MagicMock() + mlflow_repository.download_model = MagicMock(return_value=(model, 'artifact_path')) + output = mlflow_repository.get_model('model_name', 1, 'predict', 'pyfunc') + assert output == model + mlflow_repository.check_cache_retention.assert_not_called() + mlflow_repository.handle_valid_model.assert_not_called() + mlflow_repository.handle_outdated_model.assert_not_called() + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +def test_get_cached_operation_retention_0(force_memory_release, mlflow_repository): + model = MagicMock() + data = MagicMock() + mlflow_repository.get_model = MagicMock(return_value=model) + output = mlflow_repository.get_cached_operation('model_name', data, 'transform', 0, 'sklearn') + assert output == model.predict.return_value + force_memory_release.assert_called_once_with(mlflow_repository.logger) + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +def test_get_cached_predict_retention_not_0(force_memory_release, mlflow_repository): + model = MagicMock() + data = MagicMock() + mlflow_repository.get_model = MagicMock(return_value=model) + output = mlflow_repository.get_cached_operation('model_name', data, 'predict', 1, 'sklearn') + assert output == model.predict.return_value + force_memory_release.assert_not_called() + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +def test_get_cached_operation_invalid_operation(force_memory_release, mlflow_repository): + data = MagicMock() + with pytest.raises(ValueError) as e: + mlflow_repository.get_cached_operation('model_name', data, 'invalid', 0, 'sklearn') + assert str(e) == "Invalid operation. Use 'transform' or 'predict'." + + +@patch('laborious.utils.repository.model_repository.pd.merge') +@patch('laborious.utils.repository.model_repository.isinstance') +def test_fit_models_not_df_target_name_none_and_not_in_model( + isinstance_mock, pd_merge, mlflow_repository +): + isinstance_mock.return_value = False + + data_model = MagicMock() + prediction_model = MagicMock() + mlflow_repository.download_model = MagicMock( + side_effect=[(data_model, 'artifact_path'), (prediction_model, 'artifact_path')], + ) + mlflow_repository.detect_and_parse_datetime_index = MagicMock( + return_value=MagicMock(drop_duplicates=MagicMock(return_value=MagicMock(columns=[]))) + ) - 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.fit_models( + 'model_name', data, 'latest_production_id', metadata['metadata'], 'sklearn', 'pyfunc', None + ) - output = mlflow_repository.perform_model_retrain( - prediction_model_mock, data_model_mock, experiment, model_name, data) + mlflow_repository.download_model.assert_has_calls( + [ + call( + model_name='model_name', + model_type='transform', + flavor='sklearn', + load_wrapper=False, + ), + call(model_name='model_name', model_type='predict', flavor='pyfunc', load_wrapper=True), + ] + ) - 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') + data_model.fit.assert_called_once_with(data) - log_model.assert_has_calls([ - call(data_model_mock, "data_model"), - call(prediction_model_mock, "prediction_model"), - ]) + data_model.fit.return_value.predict.assert_called_once_with(data) - data.to_csv.assert_called_once_with( - "temp/raw_data_test.csv", index=True) + transformed_data = data_model.fit.return_value.predict.return_value - log_artifact.assert_called_once_with( - "temp/raw_data_test.csv") + transformed_data.__setitem__.assert_called_once_with('timestamp', transformed_data.index) - log_param.assert_has_calls([ - call("retrain", True), - ]) + mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( + transformed_data, metadata['metadata'] + ) - assert output == ("Model retrained successfully", experiment) + transformed_data = mlflow_repository.detect_and_parse_datetime_index.return_value + + transformed_data.drop_duplicates.assert_called_once_with(subset=['timestamp'], keep='first') + + transformed_data = transformed_data.drop_duplicates.return_value + + data.loc.__getitem__.assert_called_once_with(transformed_data.index) + + aligned_data = data.loc.__getitem__.return_value + + aligned_data.__getitem__.assert_called_once_with(data_model.fit.return_value.target_variable) + + pd_merge.assert_called_once_with( + transformed_data, aligned_data.__getitem__.return_value, left_index=True, right_index=True + ) + + prediction_model.fit.assert_called_once_with(pd_merge.return_value) + + assert output == { + 'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'}, + 'data_model': {'model': data_model.fit.return_value, 'artifact_path': 'artifact_path'}, + } -def test_retrain_model(mlflow_repository): +@patch('laborious.utils.repository.model_repository.pd.merge') +@patch('laborious.utils.repository.model_repository.isinstance') +def test_fit_models_df_target_name_not_none_and_in_model( + isinstance_mock, pd_merge, mlflow_repository +): + isinstance_mock.return_value = True + + data_model = MagicMock( + target_variable='feat_2', + ) + prediction_model = MagicMock() + mlflow_repository.download_model = MagicMock( + side_effect=[(data_model, 'artifact_path'), (prediction_model, 'artifact_path')], + ) + mlflow_repository.detect_and_parse_datetime_index = MagicMock( + return_value=MagicMock( + drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1'])) + ) + ) + data = MagicMock() - model_name = 'test' - mlflow_repository.create_model_experiment = MagicMock( - return_value=('data_model', 'prediction_model', '0')) + output = mlflow_repository.fit_models( + 'model_name', + data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + 'pyfunc', + 'feat_1', + ) - mlflow_repository.perform_model_retrain = MagicMock( - return_value='Model retrained successfully') + mlflow_repository.download_model.assert_has_calls( + [ + call( + model_name='model_name', + model_type='transform', + flavor='sklearn', + load_wrapper=False, + ), + call(model_name='model_name', model_type='predict', flavor='pyfunc', load_wrapper=True), + ] + ) - output = mlflow_repository.retrain_model(data, model_name) + data_model.fit.assert_called_once_with(data) - mlflow_repository.create_model_experiment.assert_called_once_with( - model_name, data) + transformed_data = data_model.fit.return_value - mlflow_repository.perform_model_retrain.assert_called_once_with( - 'data_model', 'prediction_model', '0', model_name, data) + transformed_data.__setitem__.assert_called_once_with('timestamp', transformed_data.index) - assert output == 'Model retrained successfully' + mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( + transformed_data, metadata['metadata'] + ) + + transformed_data = mlflow_repository.detect_and_parse_datetime_index.return_value + + transformed_data.drop_duplicates.assert_called_once_with(subset=['timestamp'], keep='first') + + transformed_data = transformed_data.drop_duplicates.return_value + + data.loc.__getitem__.assert_not_called() + + pd_merge.assert_not_called() + + prediction_model.fit.assert_called_once_with(transformed_data) + + assert output == { + 'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'}, + 'data_model': {'model': data_model, 'artifact_path': 'artifact_path'}, + } + + +def test_log_model_sklearn(mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + mlflow_repository.log_model(model_data, 'sklearn', 'prediction_model', metadata['metadata']) + mlflow.sklearn.log_model.assert_called_once_with(model_data['model'], 'prediction_model') + + +@patch('laborious.utils.repository.model_repository.path') +def test_log_model_pyfunc(path, mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + mlflow_repository.log_model(model_data, 'pyfunc', 'prediction_model', metadata['metadata']) + + mlflow.pyfunc.log_model.assert_not_called() + + path.join.assert_called_once_with('artifact_path', 'code', 'utils') + + model_data['model'].store_model.assert_called_once_with( + artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False + ) + + +def test_log_model_pytorch(mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + mlflow_repository.log_model(model_data, 'pytorch', 'prediction_model', metadata['metadata']) + mlflow.pytorch.log_model.assert_called_once_with(model_data['model'], 'prediction_model') + + +def test_log_model_error(mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + with pytest.raises(ValueError) as e: + mlflow_repository.log_model(model_data, 'invalid', 'prediction_model', metadata['metadata']) + assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +def test_create_new_experiment(_rmtree, path, force_memory_release, mlflow, mlflow_repository): + model_name = 'model_name' + data = MagicMock() + retrain_data = { + 'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + } + + mlflow_repository.get_model_params = MagicMock( + return_value={ + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + 'target_name': 'target_name', + } + ) + + mlflow_repository.get_experiment = MagicMock() + mlflow_repository.get_next_run_name = MagicMock() + mlflow_repository.log_model = MagicMock() + path.exists.return_value = True + path.join.return_value = './tmp/artifacts/model_name' + + report = mlflow_repository.create_new_experiment( + model_name, + data, + retrain_data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + 'pyfunc', + ) + + path.join.assert_called_once_with('./tmp/artifacts', 'model_name') + + mlflow_repository.get_model_params.assert_called_once_with('latest_production_id') + mlflow_repository.get_experiment.assert_called_once_with(model_name, create_if_not_exists=True) + mlflow_repository.get_next_run_name.assert_called_once_with( + mlflow_repository.get_experiment.return_value.name + ) + + data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=True) + + mlflow.start_run.assert_called_once_with( + experiment_id=mlflow_repository.get_experiment.return_value.experiment_id, + run_name=mlflow_repository.get_next_run_name.return_value, + description='Retrain model model_name with new data', + ) + + mlflow_repository.log_model.assert_has_calls( + [ + call(retrain_data['data_model'], 'sklearn', 'data_model', metadata['metadata']), + call( + retrain_data['prediction_model'], 'pyfunc', 'prediction_model', metadata['metadata'] + ), + ] + ) + + mlflow.log_params.assert_called_once_with( + { + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + 'target_name': 'target_name', + 'retrain': True, + 'retrain_date': ANY, + 'source_run_id': 'latest_production_id', + 'retrain_samples': data.shape.__str__.return_value, + } + ) + + mlflow.log_artifact.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv') + + force_memory_release.assert_called_once_with(mlflow_repository.logger) + + assert report == { + 'run_id': mlflow.start_run.return_value.__enter__.return_value.info.run_id, + 'experiment_id': mlflow_repository.get_experiment.return_value.experiment_id, + 'experiment_name': mlflow_repository.get_experiment.return_value.name, + } -@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( + mlflow_repository.client.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') + output = mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) mlflow.register_model.assert_called_once_with( - "runs:/0/prediction_model", + '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( + mlflow_repository.client.get_registered_model.assert_called_once_with('test') + mlflow_repository.client.transition_model_version_stage.assert_called_once_with( name='test', version='3', stage='Production', @@ -458,45 +871,198 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): } -@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={} - ) - ) + mlflow_repository.client.get_registered_model.return_value = MagicMock( + get_registered_model=MagicMock(return_value=MagicMock(latest_versions={})) ) try: - mlflow_repository.update_production_model_by_run_id('0', 'test') + mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) except Exception as e: assert str(e) == 'Model versions is not a list' else: - assert False + raise AssertionError('Expected Exception') + + +def test_transform_success(mlflow_repository): + data = MagicMock() + model_name = 'model' + model_config = { + 'retention_minutes': 60, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + } + + mlflow_repository.get_cached_operation = MagicMock() + + mlflow_repository.detect_and_parse_datetime_index = MagicMock() + + output = mlflow_repository.transform(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'transform', 60, 'sklearn' + ) + + mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( + mlflow_repository.get_cached_operation.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' + model_config = { + 'retention_minutes': 60, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + } + + mlflow_repository.get_cached_operation = MagicMock(side_effect=Exception('error')) + + output = mlflow_repository.transform(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'transform', 60, 'sklearn' + ) + + assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}} + + +def test_predict_success_array(mlflow_repository): + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + model_name = 'model' + mlflow_repository.get_cached_operation = MagicMock(return_value=np.array([2, 3])) + + output = mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'predict', 60, 'pyfunc' + ) + + 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_success_df(mlflow_repository): + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + model_name = 'model' + + mlflow_repository.get_cached_operation = MagicMock( + return_value=DataFrame({'feat_1': {'index_3': 2, 'index_4': 3}}) + ) + + output = mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'predict', 60, 'pyfunc' + ) + + 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' + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + + mlflow_repository.get_cached_operation = MagicMock(side_effect=Exception('error')) + + output = mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'predict', 60, 'pyfunc' + ) + + assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}} + + +def test_retrain_model(mlflow_repository): + data = MagicMock() + model_name = 'test' + model_config = {'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc'} + + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.fit_models = MagicMock() + mlflow_repository.create_new_experiment = MagicMock() + + output = mlflow_repository.retrain_model(data, model_name, model_config, metadata['metadata']) + + mlflow_repository.get_model_run_id.assert_called_once_with(model_name, stage='Production') + + mlflow_repository.fit_models.assert_called_once_with( + model_name=model_name, + data=data, + transform_flavor='sklearn', + predict_flavor='pyfunc', + target_name='target', + metadata=metadata['metadata'], + latest_production_id=mlflow_repository.get_model_run_id.return_value, + ) + + mlflow_repository.create_new_experiment.assert_called_once_with( + model_name=model_name, + data=data, + retrain_data=mlflow_repository.fit_models.return_value, + transform_flavor='sklearn', + predict_flavor='pyfunc', + metadata=metadata['metadata'], + latest_production_id=mlflow_repository.get_model_run_id.return_value, + ) + + assert output == { + 'success': True, + 'experiment': mlflow_repository.create_new_experiment.return_value, + 'message': 'Model retrained successfully.', + } + + +def test_retrain_model_error(mlflow_repository): + data = MagicMock() + model_name = 'test' + model_config = {'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc'} + mlflow_repository.get_model_run_id = MagicMock(side_effect=Exception('error')) + output = mlflow_repository.retrain_model(data, model_name, model_config, metadata['metadata']) + assert output == { + 'success': False, + 'experiment': None, + 'message': 'Error retraining model test: error', + 'traceback': ANY, + } def test_update_production_model(mlflow_repository): - connector = mlflow_repository + experiment = {'run_id': '0', 'experiment_id': '0'} + model_name = 'test' + mlflow_repository.update_production_model_by_run_id = MagicMock() + mlflow_repository.update_production_model_by_run_id.return_value = { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + } - 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 = mlflow_repository.update_production_model(experiment, model_name, metadata['metadata']) - output = connector.update_production_model('0', 'test') + mlflow_repository.update_production_model_by_run_id.assert_called_once_with( + '0', 'test', metadata['metadata'] + ) - 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', - } + 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 index bc84db8..6ec680d 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -1,9 +1,11 @@ -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 +from unittest.mock import ANY, AsyncMock, MagicMock, Mock, call, patch + +import pytest +from asyncua.crypto.security_policies import SecurityPolicyBasic256 +from sientia_do.notifications.models import NotificationLevel + +from laborious.utils.repository.opc_repository import OpcRepository @pytest.fixture @@ -14,15 +16,15 @@ def mock_logger(): @pytest.fixture def opc_repository(mock_logger): return OpcRepository( - id="test_repo", - url="opc.tcp://localhost:4840", + opc_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" + 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', ) @@ -35,22 +37,22 @@ def mock_client(): 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', }, } 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.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 @@ -62,12 +64,12 @@ 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.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" + 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 @@ -81,8 +83,16 @@ async def test_set_security_missing_certificates(opc_repository): try: await opc_repository.set_security() except ValueError as e: - assert str( - e) == "Certificate and private key paths must be provided for secure connection." + assert str(e) == 'Certificate and private key paths must be provided for secure connection.' + + +@pytest.mark.asyncio +async def test_set_security_missing_client(opc_repository): + opc_repository.client = None + try: + await opc_repository.set_security() + except ValueError as e: + assert str(e) == 'Client must be initialized before setting security' @pytest.mark.asyncio @@ -123,19 +133,34 @@ async def test_try_connect_success(opc_repository): 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") + 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['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_try_connect_no_client(opc_repository): + opc_repository.client = None + result = await opc_repository.try_connect() + assert result == ( + False, + { + 'notification_id': f'OPC_CONNECTION_ERROR_{opc_repository.id}', + 'message': 'Client is not initialized', + 'block': 'opc_repository', + 'level': NotificationLevel.ERROR, + }, + ) + + @pytest.mark.asyncio async def test_disconnect(opc_repository, mock_client): opc_repository.client = mock_client @@ -154,12 +179,11 @@ async def test_disconnect_no_client(opc_repository): @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") + 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 + 'Failed to disconnect from OPC server: Test error', ANY ) assert opc_repository.client is None @@ -177,9 +201,7 @@ async def test_validate_connection_none_client(opc_repository): 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.disconnect = AsyncMock(side_effect=Exception('Test error')) opc_repository.connect = AsyncMock(return_value=(True, {})) response = await opc_repository.validate_connection() @@ -188,34 +210,34 @@ async def test_validate_connection_error_count_disconnect_error(opc_repository): opc_repository.connect.assert_called_once() opc_repository.logger.custom_error.assert_has_calls( [ - call("Failed to disconnect from OPC server: Test error", ANY), + 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.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 - }) + 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)) + _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 @@ -224,19 +246,21 @@ async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, op 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 - }) + 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)) + 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 @@ -253,7 +277,7 @@ 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" + opc_repository.client.uaclient.protocol.state = 'open' output = await opc_repository.validate_connection() assert output == (True, {}) @@ -262,17 +286,16 @@ async def test_validate_connection_success(opc_repository): @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() - ) + 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']) + 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") + opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode') assert result == (True, {}) @@ -282,8 +305,9 @@ async def test_write_data_validate_connection_failed(opc_repository): 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']) + 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() @@ -295,18 +319,21 @@ 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")) + 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']) + 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") + 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['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 @@ -318,16 +345,20 @@ async def test_write_data_invalid_data_type(opc_repository, 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']) + 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") + 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['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 @@ -340,10 +371,11 @@ async def test_write_data(mock_metrics, opc_repository, 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']) + 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_client.get_node.assert_called_once_with('ns=2;s=TestNode') mock_node.write_value.assert_called_once() assert result == (True, {}) @@ -351,7 +383,7 @@ async def test_write_data(mock_metrics, opc_repository, mock_client): pod_id=opc_repository.pod_id, model_name=metadata['metadata']['model_name'], pipeline_name=metadata['metadata']['workflow_name'], - opc_server_id=opc_repository.id + opc_server_id=opc_repository.id, ) mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with() @@ -359,10 +391,11 @@ async def test_write_data(mock_metrics, opc_repository, mock_client): pod_id=opc_repository.pod_id, model_name=metadata['metadata']['model_name'], pipeline_name=metadata['metadata']['workflow_name'], - opc_server_id=opc_repository.id + opc_server_id=opc_repository.id, ) mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with( - ANY) + ANY + ) @pytest.mark.asyncio @@ -372,17 +405,21 @@ async def test_write_data_write_value_failed(opc_repository, 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") + 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']) + 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_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['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 910439c..80516ad 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,8 +1,12 @@ from os import environ -from laborious.utils.connectors_config import (build_mlflow_config, - build_opc_config, - build_postgres_config, - build_mongodb_config) + +from laborious.utils.connectors_config import ( + build_minio_config, + build_mlflow_config, + build_mongodb_config, + build_opc_config, + build_postgres_config, +) def test_build_mlflow_config_with_env_vars(): @@ -144,7 +148,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, } @@ -157,5 +161,35 @@ 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, + } + + +def test_build_minio_config_with_env_vars(): + environ['MINIO_ENDPOINT_URL'] = 'http://test-host' + environ['MINIO_ACCESS_KEY'] = 'test-key' + environ['MINIO_SECRET_KEY'] = 'test-secret' + environ['MINIO_REGION_NAME'] = 'test-region' + environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket' + assert build_minio_config() == { + 'endpoint_url': 'http://test-host', + 'access_key': 'test-key', + 'secret_key': 'test-secret', + 'region_name': 'test-region', + 'default_bucket': 'test-bucket', + } + + +def test_build_minio_config_with_defaults(): + environ.pop('MINIO_ENDPOINT_URL', None) + environ.pop('MINIO_ACCESS_KEY', None) + environ.pop('MINIO_SECRET_KEY', None) + environ.pop('MINIO_REGION_NAME', None) + environ.pop('MINIO_DEFAULT_BUCKET', None) + assert build_minio_config() == { + 'endpoint_url': 'http://localhost:9000', + 'access_key': 'minioadmin', + 'secret_key': 'minioadmin', + 'region_name': 'us-east-1', + 'default_bucket': 'laborious', } 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..006ef4e 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,10 @@ -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 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 @@ -12,149 +13,167 @@ 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("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@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" + '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_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.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 - )]) + 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) +@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" + '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_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.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 - ) - ]) + 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 df60ada..6baec30 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 laborious.activities.activities import Activities from laborious.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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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 @@ -34,26 +36,24 @@ 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'], 'opc_output_config': {'test': 'config'}, - '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 +62,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', @@ -129,13 +184,13 @@ async def test_run(workflow_mock, prediction_process): '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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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 @@ -149,17 +204,15 @@ 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'], - 'opc_output_config': {'test': 'config'} + '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 + ('stop', 0.95, 'Input data with bad quality'), # input_gate ] # Act @@ -167,23 +220,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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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 @@ -197,19 +262,17 @@ 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'], - 'opc_output_config': {'test': 'config'} + '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 + ('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 +280,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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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]) + prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True]) # Arrange input_data = { 'metadata': metadata, @@ -268,22 +357,20 @@ 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'], - 'opc_output_config': {'test': 'config'} + '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 + ('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 +379,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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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]) + prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True]) # Arrange input_data = { 'metadata': metadata, @@ -348,24 +472,22 @@ 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'], - 'opc_output_config': {'test': 'config'} + '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 + ('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 +495,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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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'} @@ -440,21 +616,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 +643,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('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'} @@ -475,21 +654,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 +686,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("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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'} @@ -521,14 +703,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, @@ -537,8 +719,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'model_name': model_name, 'model_config': model_config, 'opc_output_config': {'test': 'config'}, - 'prediction_store_policy': prediction_store_policy - }, confidence, last_timestamp, 'Prediction Process' + 'prediction_store_policy': prediction_store_policy, + }, + confidence, + last_timestamp, + 'Prediction Process', ) # Assert @@ -559,13 +744,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'table_name': table_name, 'comment': 'Prediction Process', 'opc_output_config': {'test': 'config'}, - 'prediction_store_policy': prediction_store_policy - } + 'prediction_store_policy': prediction_store_policy, + }, ) @mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@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'} @@ -576,13 +761,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, @@ -591,8 +776,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): 'model_name': model_name, 'model_config': model_config, 'opc_output_config': {'test': '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 b3b03b5..bcbad2d 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 laborious.activities.activities import Activities from laborious.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,76 +25,273 @@ metadata = { @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", + '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_config': { + 'target': 'test_target', + 'transform_flavor': 'test_transform_flavor', + 'predict_flavor': 'test_predict_flavor', + }, } workflow_mock.execute_activity_method = AsyncMock( - return_value={ - "data1": "1", - "data2": "2", - } + side_effect=[ + {'success': True, 'object_key': 'test_object_key'}, + {'success': True, 'experiment': 'test_experiment'}, + { + 'success': True, + 'version': 'test_version', + 'mlflow_run_id': 'test_mlflow_run_id', + 'mlflow_experiment_id': 'test_mlflow_experiment_id', + }, + {'report': 'test_report'}, + ] ) await minimal_retrain.run(input_data) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( - Activities.load_custom_query, + Activities.query_to_minio, { **metadata, - "query": input_data["query"], - 'datetime_columns': input_data.get('datetime_columns', []) + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], + 'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data', }, 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, + 'object_key': 'test_object_key', + 'model_name': input_data['model_name'], + 'model_config': input_data['model_config'], + }, + 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'], + 'success': True, + 'experiment': 'test_experiment', + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.format_retrain_report, + { + **metadata, + 'experiment_response': {'success': True, 'experiment': 'test_experiment'}, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'update_report': { + 'success': True, + 'version': 'test_version', + 'mlflow_run_id': 'test_mlflow_run_id', + 'mlflow_experiment_id': 'test_mlflow_experiment_id', + }, + }, + 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_local_activity_method.return_value, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + + +@mark.asyncio +@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock) +async def test_run_storage_fail(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_config': { + 'target': 'test_target', + 'transform_flavor': 'test_transform_flavor', + 'predict_flavor': 'test_predict_flavor', + }, + } + + workflow_mock.execute_activity_method = AsyncMock( + side_effect=[ + {'success': False, 'object_key': 'test_object_key'}, + {'success': True, 'experiment': 'test_experiment'}, { - **metadata, - 'data': workflow_mock.execute_activity_method.return_value, - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], + 'success': True, + 'version': 'test_version', + 'mlflow_run_id': 'test_mlflow_run_id', + 'mlflow_experiment_id': 'test_mlflow_experiment_id', }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + {'report': 'test_report'}, + ] + ) + + await minimal_retrain.run(input_data) + + workflow_mock.execute_activity_method.assert_called_once_with( + Activities.query_to_minio, + { + **metadata, + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], + 'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data', + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + + workflow_mock.execute_local_activity_method.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock) +async def test_run_fail_retrain(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_config': { + 'target': 'test_target', + 'transform_flavor': 'test_transform_flavor', + 'predict_flavor': 'test_predict_flavor', + }, + } + + workflow_mock.execute_activity_method = AsyncMock( + side_effect=[ + {'success': True, 'object_key': 'test_object_key'}, + {'success': False, 'experiment': 'test_experiment'}, + { + 'success': True, + 'version': 'test_version', + 'mlflow_run_id': 'test_mlflow_run_id', + 'mlflow_experiment_id': 'test_mlflow_experiment_id', + }, + {'report': 'test_report'}, + ] + ) + + await minimal_retrain.run(input_data) + + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.query_to_minio, + { + **metadata, + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], + 'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data', + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.retrain_model, + { + **metadata, + 'object_key': 'test_object_key', + '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.format_retrain_report, + { + **metadata, + 'experiment_response': {'success': False, 'experiment': 'test_experiment'}, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'update_report': {}, + }, + 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_local_activity_method.return_value, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + }, + 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/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 90d7d21..62cc43c 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 laborious.activities.activities import Activities from laborious.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('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' - } + workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'} input_data = { 'schedule_name': 'test_schedule', 'model_name': 'test_model', @@ -35,25 +35,25 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'opc_output_config': 'test_opc_output_config', '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 +61,19 @@ 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']), 'opc_output_config': input_data.get('opc_output_config', {}), - '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)] + ) diff --git a/validate.sh b/validate.sh new file mode 100755 index 0000000..aeed301 --- /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 laborious/ tests/ && ruff format --check laborious/ tests/"; then + FAILED_STEPS+=("Code Formatting") +fi + +# Step 2: Linting (Ruff) +if ! run_step "2. Code Linting (Ruff)" "ruff check --fix laborious/ tests/"; then + FAILED_STEPS+=("Linting") +fi + +# Step 3: Type Checking (mypy) +if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then + FAILED_STEPS+=("Type Checking") +fi + +# Step 4: Security Analysis (Bandit) +if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then + FAILED_STEPS+=("Security Analysis") +fi + +# Step 5: Unit Tests (pytest) +if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=laborious --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then + FAILED_STEPS+=("Unit Tests") +fi + +# Summary +echo -e "${BLUE}โ•”โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•—${NC}" +echo -e "${BLUE}โ•‘ Validation Summary โ•‘${NC}" +echo -e "${BLUE}โ•šโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•${NC}" +echo "" + +if [ ${#FAILED_STEPS[@]} -eq 0 ]; then + echo -e "${GREEN}โœ… All validation checks passed!${NC}" + echo -e "${GREEN} Your code is ready for commit/deployment.${NC}" + echo "" + exit 0 +else + echo -e "${RED}โŒ Validation failed for the following steps:${NC}" + for step in "${FAILED_STEPS[@]}"; do + echo -e "${RED} โ€ข ${step}${NC}" + done + echo "" + echo -e "${YELLOW}๐Ÿ’ก Tips:${NC}" + echo -e "${YELLOW} โ€ข Run 'ruff format laborious/ tests/' to auto-fix formatting${NC}" + echo -e "${YELLOW} โ€ข Run 'ruff check --fix laborious/ 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 \ No newline at end of file diff --git a/values.yaml b/values.yaml index 4eb83d7..1300546 100644 --- a/values.yaml +++ b/values.yaml @@ -3,7 +3,7 @@ # 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 +replicaCount: 2 # This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/ image: @@ -11,7 +11,7 @@ image: # This sets the pull policy for images. pullPolicy: Always # Overrides the image tag whose default is the chart appVersion. - tag: "0.0.2" + tag: "0.0.3" 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: @@ -151,7 +151,7 @@ env: - 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 + value: SIENTIAPDE-1231-ajustar-o-retreino-do-courier-no-laborious - name: PYTHON_APP value: "laborious.worker.worker" @@ -213,6 +213,17 @@ env: - name: MONGODB_TTL_INDEX_HOURS value: "1" + - name: MINIO_ENDPOINT_URL + value: "http://minio.minio.svc.cluster.local:9000" + - name: MINIO_ACCESS_KEY + value: "admin" + - name: MINIO_SECRET_KEY + value: "FvcxOPX55j" + - name: MINIO_REGION_NAME + value: "sa-east-1" + - name: MINIO_DEFAULT_BUCKET + value: "sientia" + ssh: enabled: true secretName: git-ssh-key-sientia-laborious-worker