diff --git a/README.md b/README.md index 3041bfc..b63be57 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,8 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc - [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) + - [Drift Workflow](#5-drift-workflow-driftpy) + - [Simple Metrics Workflow](#6-simple-metrics-workflow-simple_metricspy) - [Installation & Setup](#installation--setup) - [Prerequisites](#prerequisites) - [Environment Setup](#environment-setup) @@ -77,13 +79,16 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc ### Advanced Capabilities - **Incremental Data Processing**: Timestamp-based loading to avoid reprocessing - **Configurable Data Retention**: Model retention policies with automatic cleanup +- **MinIO Payload Offload**: Automatic offload of large DataFrames to MinIO with retention cleanup +- **Data Drift Detection**: Univariate and multivariate drift monitoring against reference data +- **Regression Metrics**: Automated RMSE, MSE, MAE, R² calculation and export - **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 +- **Automated Validation**: CI quality gates and individual tool commands - **Comprehensive Testing**: pytest with async support and high coverage - **Type Safety**: Static type checking with mypy - **Coverage Visualization**: Coverage Gutters integration @@ -129,6 +134,8 @@ Laborious uses a Temporal-based architecture with strong separation of concerns - `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 +- `drift.py`: Data drift detection and monitoring +- `simple_metrics.py`: Regression metrics calculation (RMSE, MSE, MAE, R²) #### **Activities (`laborious/activities/`)** - `gates.py`: Data quality validation, filtering, and data formatting operations @@ -139,16 +146,26 @@ Laborious uses a Temporal-based architecture with strong separation of concerns - MLFlow model transformation and prediction - Model retraining and production updates - Reference data retrieval from MLflow Model Registry +- `storage.py`: PostgreSQL queries and MinIO-aware data loading + - `load_query_with_minio_offload`: SQL load with automatic MinIO offload + - `export_payload_to_postgres`: Resolve MinIO payloads and export to Postgres + - `cleanup_minio_objects_expired`: Retention-based MinIO object cleanup + - `query_to_minio`: Legacy parquet upload for retraining data +- `model_metrics.py`: Drift detection and regression metrics + - Univariate and multivariate drift calculation + - Simple metrics (RMSE, MSE, MAE, R²) - `opc.py`: OPC UA export to industrial systems (optional) - `api.py`: PI Web API export operations (optional) - Prediction and confidence data writing to PI Web API - Error handling and notification integration -- `activities.py`: Aggregates activity interfaces +- `activities.py`: Aggregates all activity interfaces (Storage, MLFlow, Gates, OPC, ModelMetrics, API) #### **Data Services (`laborious/utils/`)** - `connectors_config.py`: Env-driven configuration builders +- `models/minio_dataframe_payload.py`: MinIO-offloaded DataFrame payload model - `repository/model_repository.py`: MLFlow operations and retraining - `repository/opc_repository.py`: OPC communication and writes +- `repository/minio_manager.py`: MinIO object storage operations - `filters/conditional_filters.py` and `filters/mlflow_filters.py` ### Data Flow Architecture @@ -167,6 +184,18 @@ Training Data → Model Retraining → Quality Validation → Production Update → Notification & Monitoring ``` +#### **3. Drift Detection Pipeline** +``` +Target Data (PostgreSQL) + Reference Data (MLFlow) → +Drift Calculation (univariate + multivariate) → PostgreSQL Export +``` + +#### **4. Simple Metrics Pipeline** +``` +Predictions + Targets (PostgreSQL JOIN) → +Metrics Calculation (RMSE, MSE, MAE, R²) → PostgreSQL Export +``` + ### Security Architecture #### **Authentication & Authorization** @@ -262,20 +291,21 @@ The **PredictionProcess** workflow implements the core prediction pipeline for M - **Prediction Export**: Delegates prediction formatting and export operations #### Execution Flow -1. **Timestamp Retrieval**: Gets the last processed timestamp for incremental processing -2. **Input Data Gate**: Applies configured filters for data quality validation -3. **Path Decision**: Determines processing path based on filter results -4. **MLFlow Transform**: Requests data transformation using MLFlow models -5. **Response Validation**: Filters transform responses for quality assurance +1. **Input Data Gate**: Applies configured filters for data quality validation +2. **Path Decision**: Determines processing path based on filter results +3. **MLFlow Transform**: Requests data transformation using MLFlow models +4. **Response Validation**: Filters transform responses for quality assurance +5. **Content Validation**: Filters transformed data content for quality check 6. **MLFlow Prediction**: Executes prediction using transformed data -7. **Content Validation**: Filters prediction responses for final quality check +7. **Prediction Response Validation**: Filters prediction responses for final quality check 8. **Export Delegation**: Delegates to FormatAndExportPrediction workflow +9. **MinIO Cleanup**: Cleans up expired offloaded payloads (if any, in `finally` block) #### Key Features - **Configurable Quality Gates**: Multiple filter types with policy-based configuration - **Flexible Path Handling**: Configurable decision paths (STOP, CONTINUE, REPEAT) - **MLFlow Integration**: Comprehensive model management and inference -- **Incremental Processing**: Timestamp-based data processing optimization +- **MinIO Cleanup**: Automatic retention-based cleanup of offloaded payloads - **Comprehensive Monitoring**: Detailed metrics and error reporting #### Input Parameters @@ -320,12 +350,12 @@ The **PredictionProcess** workflow implements the core prediction pipeline for M #### Architecture Diagram ```mermaid flowchart LR - A[1. get_last_timestamp] --> B[2. input_gate] --> C[3. request_transform] --> D[4. mlflow_response_gate] --> E[5. mlflow_content_gate] --> F[6. request_predict] --> G[7. mlflow_response_gate] --> H[8. format_and_export_prediction🔃] + A[1. input_gate] --> B[2. request_transform] --> C[3. mlflow_response_gate] --> D[4. mlflow_content_gate] --> E[5. request_predict] --> F[6. mlflow_response_gate] --> G[7. format_and_export_prediction🔃] + G --> H[8. cleanup_minio_objects_expired] - A -.-> Redis[(Redis)] - C -.-> MLFlow[MLFlow] - F -.-> MLFlow[MLFlow] - G -.-> Filters[MLFlow Filters] + B -.-> MLFlow[MLFlow] + E -.-> MLFlow[MLFlow] + H -.-> MinIO[(MinIO)] ``` ### 3. Format and Export Prediction Workflow (`format_and_export_prediction.py`) @@ -403,6 +433,76 @@ flowchart LR D -.-> PostgreSQL[(PostgreSQL)] ``` +### 5. Drift Workflow (`drift.py`) + +The **Drift** workflow detects data drift by comparing current data against a reference dataset from the MLflow Model Registry. + +#### Execution Flow +1. **Data Loading**: Loads target data and reference data in parallel +2. **Drift Calculation**: Calculates univariate and multivariate drift metrics +3. **Data Export**: Exports drift metrics to PostgreSQL + +#### Architecture Diagram +```mermaid +flowchart LR + A[1. load_custom_query] --> C[3. calculate_drift] --> D[4. export_data_to_postgres] + B[2. get_reference_data] --> C + + A -.-> Database[(Database)] + B -.-> MLFlow[MLFlow] + D -.-> PostgreSQL[(PostgreSQL)] +``` + +#### Input Parameters +```json +{ + "schedule_name": "hourly_drift", + "model_name": "temperature_model", + "model_id": "temp_001", + "schema": "sientia_data", + "source_table_name": "laborious_data", + "target_table_name": "drift_metrics", + "interval": 60, + "model_config": { "target": "temperature" }, + "drift_metrics": ["kolmogorov_smirnov", "jensen_shannon", "wasserstein"], + "chunk_period": "min" +} +``` + +### 6. Simple Metrics Workflow (`simple_metrics.py`) + +The **SimpleMetrics** workflow calculates regression metrics (RMSE, MSE, MAE, R²) by comparing predictions against actual target values. + +#### Execution Flow +1. **Data Loading**: Loads prediction vs target data via a JOIN query +2. **Metrics Calculation**: Calculates configured regression metrics +3. **Data Export**: Exports metrics to PostgreSQL + +#### Architecture Diagram +```mermaid +flowchart LR + A[1. load_custom_query] --> B[2. calculate_simple_metrics] --> C[3. export_data_to_postgres] + + A -.-> Database[(Database)] + C -.-> PostgreSQL[(PostgreSQL)] +``` + +#### Input Parameters +```json +{ + "schedule_name": "hourly_metrics", + "model_name": "temperature_model", + "model_id": "temp_001", + "schema": "sientia_data", + "predictions_table_name": "predictions", + "data_table_name": "laborious_data", + "target_table_name": "simple_metrics", + "interval_minutes": 60, + "model_config": { "target": "temperature" }, + "metrics": ["rmse", "mse", "mae", "r2"] +} +``` + ## 📋 Prerequisites - Python 3.11+ @@ -527,8 +627,8 @@ pytest pytest --cov=laborious --cov-report=html # Run specific test categories -pytest tests/activities/ -pytest tests/workflow/ +pytest tests/laborious/activities/ +pytest tests/laborious/workflows/ ``` ### Manual Application Execution @@ -569,18 +669,7 @@ 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): +Run each validation step individually: ```bash ruff format --check laborious/ tests/ ruff check laborious/ tests/ @@ -606,7 +695,7 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac ### Best Practices -- Run `./validate.sh` before committing +- Run all validation steps before committing - Use `ruff check --watch` for continuous feedback - Add type hints and tests for new code @@ -615,15 +704,34 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac ### Test Structure ``` tests/ -├── activities/ # Activity implementation tests -│ ├── test_gates.py # Data quality gates and formatting tests -│ ├── test_mlflow.py # MLFlow operations and reference data tests -│ └── ... # Other activity tests -├── workflows/ # Workflow orchestration tests -│ └── subworkflows/ # Sub-workflow tests -│ └── test_format_and_export_prediction.py # Export workflow tests -├── utils/ # Utility function tests -└── integration/ # End-to-end workflow tests +├── conftest.py # Global fixtures and env setup +├── laborious/ +│ ├── activities/ # Activity implementation tests +│ │ ├── test_activities.py # Activities aggregator tests +│ │ ├── test_gates.py # Data quality gates and formatting tests +│ │ ├── test_mlflow.py # MLFlow operations and reference data tests +│ │ ├── test_storage.py # Storage and MinIO offload tests +│ │ ├── test_model_metrics.py # Drift and simple metrics tests +│ │ ├── test_opc.py # OPC operations tests +│ │ └── test_api.py # PI Web API operations tests +│ ├── workflows/ # Workflow orchestration tests +│ │ ├── test_predictions_batch.py +│ │ ├── test_minimal_retrain.py +│ │ ├── test_drift.py +│ │ ├── test_simple_metrics.py +│ │ └── subworkflows/ +│ │ ├── test_prediction_process.py +│ │ └── test_format_and_export_prediction.py +│ └── utils/ # Utility function tests +│ ├── test_connectors_config.py +│ ├── models/ +│ │ └── test_minio_dataframe_payload.py +│ ├── filters/ +│ │ ├── test_conditional_filters.py +│ │ └── test_mlflow_filters.py +│ └── repository/ +│ ├── test_model_repository.py +│ └── test_opc_repository.py ``` ### Test Coverage @@ -643,8 +751,8 @@ pip install pytest pytest-cov pytest-asyncio pytest --cov=laborious --cov-report=html # Run specific test modules -pytest tests/activities/test_gates.py -pytest tests/workflow/test_predictions_batch.py +pytest tests/laborious/activities/test_gates.py +pytest tests/laborious/workflows/test_predictions_batch.py ``` ## 📊 Monitoring and Metrics @@ -716,10 +824,85 @@ The Laborious system exposes comprehensive Prometheus metrics for operational vi | `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No | | `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No | | `POD_ID` | Kubernetes pod identifier | `None` | No | +| `MINIO_ENDPOINT_URL` | MinIO endpoint URL | `http://localhost:9000` | Yes | +| `MINIO_ACCESS_KEY` | MinIO access key | `minioadmin` | Yes | +| `MINIO_SECRET_KEY` | MinIO secret key | `minioadmin` | Yes | +| `MINIO_REGION_NAME` | MinIO region name | `us-east-1` | No | +| `MINIO_DEFAULT_BUCKET` | Default MinIO bucket | `laborious` | No | +| `MINIO_RETENTION_HOURS` | Retention window (hours) for offloaded MinIO objects | `24` | No | +| `SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES` | Offload threshold for DataFrame-derived payloads | `int(1.5 * 1024 * 1024)` | No | +### MinIO Payload Offload & Retention + +Laborious uses MinIO to prevent Temporal workflow history from carrying very large in-memory payloads (pandas `DataFrame`-derived dicts). +Whenever a payload exceeds a configurable size threshold, it is stored as a parquet file in MinIO and the workflow history only keeps a lightweight reference. + +Notes: +- `SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES` supports: + - Integer bytes (e.g. `"1572864"`) + - Float MiB (e.g. `"1.5"`), converted to bytes as `MiB * 1024 * 1024` +- Fallback behavior uses `1.5 MiB` when the env var is missing or invalid. + +#### Wire Contract: `MinioDataFramePayload` + +The payload is implemented in `laborious/utils/models/minio_dataframe_payload.py`. +The dataclass does **not** store a pandas `DataFrame` field. +Instead, the `DataFrame` is only used at build time by: +- `MinioDataFramePayload.from_dataframe(...)` +- `MinioDataFramePayload.from_dataframe_to_dict(...)` + +After evaluation, the payload is serialized for Temporal as a flat dict: +- **Inline path**: `data` contains `df.to_dict()`, and MinIO keys (`object_key`, `bucket`, ...) are absent / `None`. +- **MinIO path**: the dict contains: + - `bucket` + - `object_key` (full MinIO object name returned by `MinioRepository.upload_file`) + - `object_prefix` (directory prefix used for cleanup listing; relative to the repository namespace) + - `uri` (best-effort `s3:///<...>` string) + - `data` is omitted / set to `None`. + +When an activity needs pandas operations, it resolves references using: +- `MinioDataFramePayload.retrieve(minio_repo)` — downloads from MinIO or returns inline data as a DataFrame + +#### MinIO Object Naming (Retention Parsing) + +MinIO object basename (required convention): +`{model_name}-{operation}-{timestamp}.parquet` + +Where: +- `model_name`: model identifier used by the pipeline +- `operation`: `initial` (SQL/query load before transform) or `transform` (after MLFlow transform) +- `timestamp`: `DATETIME_FORMAT_FILENAME` from `sientia_do.temporal.constants` + +The relative object key (under the repository namespace) is always shaped as: +`training_datasets/{model_name}/{basename}` + +Retention cleanup parses timestamps from the basename using the `-initial-` / `-transform-` anchors. +`model_name` may contain hyphens; parsing is resilient to it. + +#### Workflows / Activities Integration + +Predictions batch uses MinIO offload as follows: +1. `predictions_batch` calls `Activities.load_query_with_minio_offload` + - On success, it puts the serialized `MinioDataFramePayload` dict into `prediction_input["data"]`. +2. `sub_workflows/prediction_process` + - Tracks which MinIO prefixes were referenced for offloaded payloads. + - Runs `Activities.cleanup_minio_objects_expired` in a `finally` block (only when MinIO offload happened). +3. `laborious/activities/gates.py` and `laborious/activities/mlflow.py` + - Resolve offloaded payloads transparently before constructing pandas `DataFrame` objects. + +#### Legacy: `query_to_minio` (Minimal Retrain) + +`Storage.query_to_minio` is intentionally kept with its legacy behavior for `minimal_retrain`. +It always uploads parquet and returns `{success, object_key, uri}`. +It is not used by predictions batch MinIO offload, and its objects are not part of the retention parser described above. + +Legacy MinIO object layout (relative key): +`training_datasets/{model_name}/{object_prefix}_{timestamp}.parquet` where `object_prefix` is sanitized +(slashes replaced by underscores) to keep a stable model-level directory. + ### OPC Configuration For multiple OPC servers, use the `OPC_CONFIG` environment variable: @@ -882,27 +1065,35 @@ This is the configuration created by the Orchestrator in Temporal. ``` laborious/ ├── activities/ # Temporal activity implementations -│ ├── activities.py # Main activities orchestrator +│ ├── activities.py # Main activities aggregator │ ├── gates.py # Data quality gates and filtering │ ├── mlflow.py # MLFlow model operations +│ ├── storage.py # PostgreSQL queries and MinIO offload +│ ├── model_metrics.py # Drift and regression metrics │ ├── opc.py # OPC server operations │ └── api.py # PI Web API operations ├── workflows/ # Temporal workflow definitions │ ├── predictions_batch.py # Main batch prediction workflow │ ├── minimal_retrain.py # Model retraining workflow +│ ├── drift.py # Data drift detection workflow +│ ├── simple_metrics.py # Regression metrics workflow │ └── sub_workflows/ # Sub-workflow implementations │ ├── prediction_process.py # Core prediction workflow │ └── format_and_export_prediction.py # Export workflow ├── worker/ # Worker implementation -│ └── worker.py # Main worker orchestrator +│ ├── worker.py # Main worker orchestrator +│ └── prepare_worker.py # Worker factory with autoscaling config ├── utils/ # Utility functions -│ ├── connectors_config.py # Database configuration +│ ├── connectors_config.py # Environment-driven config builders +│ ├── models/ # Data models +│ │ └── minio_dataframe_payload.py # MinIO-offloaded DataFrame payload │ ├── filters/ # Data quality filters │ │ ├── conditional_filters.py # Conditional data filters │ │ └── mlflow_filters.py # MLFlow response filters │ └── repository/ # Data access layer │ ├── model_repository.py # MLFlow model operations -│ └── opc_repository.py # OPC server operations +│ ├── opc_repository.py # OPC server operations +│ └── minio_manager.py # MinIO object storage operations ├── metrics.py # Prometheus metrics definitions └── __init__.py ``` @@ -1004,4 +1195,4 @@ For support and questions: --- -**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments. \ No newline at end of file +**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments. diff --git a/e2e/conftest.py b/e2e/conftest.py index 5dd208e..7825335 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -3,11 +3,13 @@ Pytest configuration and fixtures for E2E tests. """ from unittest.mock import AsyncMock, MagicMock, patch +from io import BytesIO import pandas as pd import pytest import pytest_asyncio from sqlalchemy import create_engine, text +from testcontainers.minio import MinioContainer from testcontainers.postgres import PostgresContainer from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker @@ -25,6 +27,17 @@ TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017' TEST_DATABASE_NAME = 'test_db' +@pytest_asyncio.fixture(scope='session') +def minio_container(): + """ + MinIO S3-compatible storage for E2E tests that exercise real offload uploads. + """ + minio = MinioContainer() + minio.start() + yield minio + minio.stop() + + @pytest_asyncio.fixture(scope='session') def postgres_container(): """ @@ -186,6 +199,20 @@ def mock_mongo_client(): return mock_client +@pytest.fixture +def notification_inserts(mock_mongo_client): + """ + Mongo insert_one mock used by CoreNotificationHandler for notification persistence. + + Yields: + MagicMock for insert_one, reset before each test. + """ + mock_db = mock_mongo_client.__getitem__.return_value + mock_collection = mock_db.__getitem__.return_value + mock_collection.insert_one.reset_mock() + yield mock_collection.insert_one + + @pytest_asyncio.fixture def notification_handler(mock_logger, mock_mongo_client): """ @@ -216,11 +243,28 @@ def metrics_controller(mock_logger): def mock_minio_repository(): """Mock MinIO repository for object storage operations.""" mock_repo = MagicMock() + + # Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve() + # can decode the payload if offloading is exercised in an integration scenario. + parquet_df = pd.DataFrame({'a': [1]}) + parquet_buffer = BytesIO() + parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True) + parquet_bytes = parquet_buffer.getvalue() - # Mock repository methods - mock_repo.put_parquet_from_dataframe = AsyncMock(return_value='test-object-key') - mock_repo.get_parquet_as_dataframe = AsyncMock(return_value=pd.DataFrame()) - mock_repo.minio_bucket = 'test-bucket' + # sientia_do MinioRepository API + mock_repo.bucket = 'test-bucket' + mock_repo.upload_file = AsyncMock( + side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: { + 'minio_object_name': f'sientia/streamlit-connectors/{relative_key}', + 'original_filename': relative_key.rsplit('/', 1)[-1], + 'uploaded_at': '2024-01-01T00:00:00Z', + 'sha256_hash': 'deadbeef', + } + ) + mock_repo.download_file = AsyncMock(return_value=parquet_bytes) + mock_repo.list_objects = AsyncMock(return_value=[]) + mock_repo.delete_file = AsyncMock() + mock_repo.close = MagicMock() return mock_repo @@ -229,15 +273,17 @@ def mock_minio_repository(): def mock_pi_web_api_repository(): """Mock PI Web API repository for PI Web API operations.""" mock_repo = MagicMock() - mock_repo.write_value = AsyncMock( - return_value={ - 'Items': [ - { - 'WebId': 'web_id_1' - } - ] - } - ) + + async def _write_value(web_ids, value, metadata=None, **kwargs): + """ + Mirror successful PI writes: one response item per requested web_id. + + write_pi_web_api_data passes the list into process_pi_web_api_response (not a + wrapped {'Items': ...} envelope). + """ + return [{'WebId': wid, 'Errors': []} for wid in web_ids] + + mock_repo.write_value = AsyncMock(side_effect=_write_value) mock_repo.close = MagicMock() return mock_repo @@ -248,7 +294,7 @@ def mock_opc_repository(): mock_repo.write_data = AsyncMock( return_value=(True, {'response_time': 0.1}) ) - mock_repo.disconnect = MagicMock() + mock_repo.disconnect = AsyncMock() return mock_repo @pytest_asyncio.fixture @@ -261,7 +307,8 @@ def patch_create_engine(postgres_engine): @pytest_asyncio.fixture def patch_minio_repository(mock_minio_repository): """Patch MinioRepository to return mock.""" - with patch('laborious.utils.repository.minio_repository.MinioRepository', return_value=mock_minio_repository): + # Patch where Activities resolves the symbol (import binds the original class). + with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository): yield @pytest_asyncio.fixture @@ -389,11 +436,13 @@ async def test_activities( 'password': 'test', }, minio_config={ - 'endpoint_url': 'http://localhost:9000', + # Host:port only; Minio() prepends http(s):// from the secure flag. + 'endpoint_url': 'localhost:9000', 'access_key': 'test', 'secret_key': 'test', - 'region_name': 'us-east-1', 'default_bucket': 'test-bucket', + 'retention_hours': 24, + 'secure': False, }, opc_config={}, pi_web_api_config={ @@ -416,6 +465,88 @@ async def test_activities( await activities.shutdown() +@pytest_asyncio.fixture(scope='function') +async def test_activities_real_minio( + postgres_engine, + postgres_container, + minio_container, + mock_logger, + notification_handler, + metrics_controller, + patch_create_engine, + patch_mlflow, + patch_pi_web_api_repository, + mock_opc_repository, +): + """ + Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests. + """ + minio_client = minio_container.get_client() + if not minio_client.bucket_exists('test-bucket'): + minio_client.make_bucket('test-bucket') + minio_port = minio_container.get_exposed_port(9000) + activities = Activities( + postgres_config={ + 'host': 'localhost', + 'port': postgres_container.get_exposed_port(5432), + 'user': 'test', + 'password': 'test', + 'dbname': 'test', + 'min_connections': 1, + 'max_connections': 5, + }, + mlflow_config={ + 'host': 'http://localhost', + 'port': '5000', + 'username': 'test', + 'password': 'test', + }, + minio_config={ + 'endpoint_url': f'localhost:{minio_port}', + 'access_key': 'minioadmin', + 'secret_key': 'minioadmin', + 'default_bucket': 'test-bucket', + 'retention_hours': 24, + 'secure': False, + }, + opc_config={}, + pi_web_api_config={ + 'base_url': 'http://localhost:8080', + 'auth_type': 'bearer', + 'auth_token': 'test_token', + }, + logger=mock_logger, + notification_handler=notification_handler, + ) + activities.opc_repository = {'1': mock_opc_repository} + try: + yield activities + finally: + await activities.shutdown() + + +def _worker_activity_list(test_activities: Activities): + return [ + test_activities.load_custom_query, + test_activities.load_query_with_minio_offload, + test_activities.cleanup_minio_objects_expired, + test_activities.input_gate, + test_activities.request_transform, + test_activities.mlflow_response_gate, + test_activities.mlflow_content_gate, + test_activities.request_predict, + test_activities.repeat_last_prediction, + test_activities.format_prediction, + test_activities.format_transformed_data, + test_activities.format_default_prediction, + test_activities.write_pi_web_api_data, + test_activities.write_opc_data, + test_activities.export_data_to_postgres, + test_activities.export_payload_to_postgres, + test_activities.write_metrics, + ] + + @pytest_asyncio.fixture(scope='function') async def temporal_test_env(): """Create Temporal test environment.""" @@ -431,22 +562,18 @@ async def temporal_worker(temporal_test_env, test_activities): temporal_test_env.client, task_queue='test-queue', workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], - activities=[ - test_activities.load_custom_query, - test_activities.get_last_timestamp, - test_activities.input_gate, - test_activities.request_transform, - test_activities.mlflow_response_gate, - test_activities.mlflow_content_gate, - test_activities.request_predict, - test_activities.repeat_last_prediction, - test_activities.format_prediction, - test_activities.format_transformed_data, - test_activities.format_default_prediction, - test_activities.write_pi_web_api_data, - test_activities.write_opc_data, - test_activities.export_data_to_postgres, - test_activities.write_metrics, - ], + activities=_worker_activity_list(test_activities), + ) as worker: + yield worker + + +@pytest_asyncio.fixture(scope='function') +async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio): + """Temporal worker backed by Activities using real MinIO testcontainer.""" + async with Worker( + temporal_test_env.client, + task_queue='test-queue', + workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], + activities=_worker_activity_list(test_activities_real_minio), ) as worker: yield worker diff --git a/e2e/helpers.py b/e2e/helpers.py new file mode 100644 index 0000000..1aaa4f2 --- /dev/null +++ b/e2e/helpers.py @@ -0,0 +1,174 @@ +""" +Shared helpers for E2E tests (Temporal workflows + PostgreSQL). +""" + +import asyncio +from datetime import datetime +from decimal import Decimal +from typing import Any + +from sqlalchemy import text +from sqlalchemy.engine import Engine + + +async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0): + """ + Start a workflow and wait for its result. + + Args: + client: Temporal client from WorkflowEnvironment. + workflow_run: Workflow run method (e.g. PredictionsBatch.run). + input_data: Workflow input payload. + workflow_id: Unique workflow id. + timeout: Max seconds to wait for completion. + + Return: + Workflow result value. + """ + handle = await client.start_workflow( + workflow_run, + input_data, + id=workflow_id, + task_queue='test-queue', + ) + return await asyncio.wait_for(handle.result(), timeout=timeout) + + +def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None: + """ + Replace laborious_data rows for a model_id with one row per value (sensor_1..n). + + Args: + postgres_engine: SQLAlchemy engine. + model_id: Model id column value. + values: Per-sensor values; use string 'NULL' for SQL NULL. + """ + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + values_sql = [] + for i, value in enumerate(values): + values_sql.append(f""" + ({model_id}, 'sensor_{i + 1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') + """) + insert_sql = f""" + INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) + VALUES + {', '.join(values_sql)} + """ + conn.execute(text(insert_sql)) + + +def assert_prediction( + postgres_engine: Engine, + model_id: int, + prediction: float = 0.5, + prediction_confidence: int | Decimal = 0, + prediction_status: str = 'Good', + comments: str = '', +) -> None: + """ + Assert exactly one prediction row exists for model_id with expected columns. + + Args: + postgres_engine: SQLAlchemy engine. + model_id: Expected model_id. + prediction: Expected prediction value. + prediction_confidence: Expected confidence (int or Decimal for numeric column). + prediction_status: Expected status string. + comments: Expected comments string. + """ + import pytest + + with postgres_engine.connect() as conn: + result_query = conn.execute( + text( + f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id} ' + f'ORDER BY created_at ASC' + ) + ) + prediction_rows = result_query.fetchall() + assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}' + row = prediction_rows[0] + assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}' + assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), ( + f'Expected prediction={prediction}, got {row[1]}' + ) + assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal( + str(prediction_confidence) + ), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}' + assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}" + assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" + + +def assert_continue( + postgres_engine: Engine, + model_id: int, + prediction_confidence: Decimal = Decimal(2), + comments: str = 'Input data with bad quality', +) -> None: + """Assert one default-style prediction row after CONTINUE gate path.""" + with postgres_engine.connect() as conn: + result_query = conn.execute( + text( + f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id}' + ) + ) + prediction_rows = result_query.fetchall() + assert len(prediction_rows) == 1, 'Expected one prediction record despite warnings' + row = prediction_rows[0] + assert row[1] == 0, f'Expected prediction=0, got {row[1]}' + assert row[2] == prediction_confidence, ( + f'Expected prediction_confidence={prediction_confidence}, got {row[2]}' + ) + assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}" + assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" + + +def assert_stop(postgres_engine: Engine, model_id: int) -> None: + """Assert no prediction rows for model_id.""" + import pytest + + with postgres_engine.connect() as conn: + result_query = conn.execute( + text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) + count = result_query.scalar() + assert count == 0, f'Expected no predictions, but found {count} records' + + +def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None: + """ + Assert two prediction rows for model_id both match last_prediction. + + Rows are compared in created_at order for stability. + + Args: + postgres_engine: SQLAlchemy engine. + model_id: Model id. + last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows. + """ + import pytest + + with postgres_engine.connect() as conn: + result_query = conn.execute( + text( + f'SELECT model_id, prediction, prediction_confidence, prediction_status ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id} ' + f'ORDER BY created_at ASC' + ) + ) + prediction_rows = result_query.fetchall() + assert len(prediction_rows) == 2, 'Expected two prediction records' + assert prediction_rows[0] == last_prediction, ( + f'Expected first row {last_prediction}, got {prediction_rows[0]}' + ) + assert prediction_rows[1] == last_prediction, ( + f'Expected second row {last_prediction}, got {prediction_rows[1]}' + ) + + +def make_workflow_id(prefix: str) -> str: + """Build a unique workflow id using a prefix and current timestamp.""" + return f'{prefix}-{datetime.now().timestamp()}' diff --git a/e2e/scenarios.md b/e2e/scenarios.md index eb9c606..5732cd3 100644 --- a/e2e/scenarios.md +++ b/e2e/scenarios.md @@ -2,6 +2,13 @@ This document describes all possible test scenarios for the `predictions_batch` workflow and its child workflows `prediction_process` and `format_and_export_prediction`. +## Running automated E2E tests (`e2e/`) + +- **Runtime**: Docker (or a Docker-compatible daemon) must be available so [testcontainers](https://testcontainers.com/) can start **PostgreSQL** and **MinIO** containers. +- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio]`). +- **Invocation**: run only integration-marked tests, for example: `pytest e2e/ -m integration`. +- **MinIO tests**: `e2e/test_minio_offload.py` exercises real S3 uploads; other E2E modules continue to mock MinIO on the worker used by most scenarios. + ## Workflow Overview The `predictions_batch` workflow: @@ -445,6 +452,8 @@ The `predictions_batch` workflow: ### 3.2 Error Scenarios +These paths do **not** rely on Temporal activity retries for export failures: the write activities run once, errors are handled inside the activity, and the **workflow completes successfully** with degraded metadata on the persisted prediction (`prediction_confidence` and `comments`). + #### Scenario 3.2.1: PI Web API Write Error **Description**: PI Web API export fails @@ -453,15 +462,16 @@ The `predictions_batch` workflow: - PI Web API service unavailable or invalid config **Expected Behavior**: -- `write_pi_web_api_data` raises exception -- Notification sent -- Workflow fails after retries -- PostgreSQL export may not execute (depends on execution order) +- `write_pi_web_api_data` surfaces the failure (exception handled in the activity layer) +- Notification may be sent +- Workflow **completes** (does not fail) +- Prediction row is still written to PostgreSQL with error confidence **13** and a comment describing the PI error +- Subsequent steps (e.g. OPC, Postgres) still run per workflow order with the updated prediction payload **Assertions**: -- PI Web API error notification sent -- Workflow fails -- May impact subsequent exports +- PI Web API error notification sent (when applicable) +- Workflow completes +- PostgreSQL contains the prediction with `prediction_confidence` 13 and expected `comments` --- @@ -473,14 +483,15 @@ The `predictions_batch` workflow: - OPC server unavailable or invalid configuration **Expected Behavior**: -- `write_opc_data` raises exception -- Notification sent -- Workflow fails after retries +- `write_opc_data` reports failure without aborting the workflow +- Notification may be sent +- Workflow **completes** (does not fail) +- Prediction row is written to PostgreSQL with OPC error confidence **12** and a comment indicating OPC write issues **Assertions**: -- OPC error notification sent -- Workflow fails -- PostgreSQL export may not execute +- OPC error notification sent (when applicable) +- Workflow completes +- PostgreSQL contains the prediction with `prediction_confidence` 12 and expected `comments` --- @@ -497,7 +508,7 @@ The `predictions_batch` workflow: - `process_pi_web_api_response` detects partial failure - Error confidence set (13) - Notification sent for failed tag -- Workflow completes with error confidence +- Workflow completes with error confidence (single activity attempt; no retry loop) **Assertions**: - One tag written successfully diff --git a/e2e/test_child_workflows_e2e.py b/e2e/test_child_workflows_e2e.py new file mode 100644 index 0000000..52a4b3f --- /dev/null +++ b/e2e/test_child_workflows_e2e.py @@ -0,0 +1,74 @@ +""" +Direct E2E execution of child workflows (smaller surface than PredictionsBatch). +""" + +from decimal import Decimal + +import pytest +from sqlalchemy import text +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from e2e.helpers import make_workflow_id, start_and_await_workflow +from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_format_and_export_prediction_default_path_e2e( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + postgres_engine, +): + """ + Run FormatAndExportPrediction with path_flag set (format_default_prediction path). + """ + client = temporal_test_env.client + model_id = 401 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + + metadata = { + 'metadata': { + 'model_id': model_id, + 'model_name': 'test_model', + 'schedule_name': 'test-schedule', + 'workflow_name': 'subworkflow.format_and_export_prediction', + } + } + input_data = { + 'metadata': metadata, + 'path_flag': 'CONTINUE', + 'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'}, + 'prediction_confidence': 2, + 'timestamp': '2024-01-01 12:00:00+00:00', + 'model_id': model_id, + 'model_name': 'test_model', + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'comment': 'e2e child workflow default path', + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'prediction_store_policy': 'lts:1', + } + + await start_and_await_workflow( + client, + FormatAndExportPrediction.run, + input_data, + make_workflow_id('e2e-format-export-child'), + ) + + with postgres_engine.connect() as conn: + row = conn.execute( + text( + f'SELECT prediction, prediction_confidence, prediction_status, comments ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id}' + ) + ).fetchone() + assert row is not None + assert row[0] == 0 + assert row[1] == Decimal(2) + assert row[2] == 'Bad' + assert row[3] == 'e2e child workflow default path' diff --git a/e2e/test_minio_offload.py b/e2e/test_minio_offload.py new file mode 100644 index 0000000..f8c0122 --- /dev/null +++ b/e2e/test_minio_offload.py @@ -0,0 +1,124 @@ +""" +E2E-style tests for MinIO offload using a real MinIO testcontainer. +""" + +from unittest.mock import patch + +import pytest +from sqlalchemy import text +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow +from laborious.activities.activities import Activities +from laborious.utils.models import minio_dataframe_payload as mdp +from laborious.workflows.predictions_batch import PredictionsBatch + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_load_query_with_minio_offload_writes_object_to_bucket( + postgres_engine, + minio_container, + test_activities_real_minio: Activities, +): + """ + With a tiny offload threshold, query results are uploaded as Parquet to MinIO. + + Uses real MinioRepository against testcontainers MinIO (no MinIO mock). + """ + model_id = 501 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + insert_sample_data(postgres_engine, model_id, [1.0, 2.0]) + + metadata = { + 'metadata': { + 'schedule_name': 'test-schedule', + 'model_name': 'test_model', + 'model_id': model_id, + 'workflow_name': 'predictions_batch', + } + } + with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1): + payload = await test_activities_real_minio.load_query_with_minio_offload( + { + **metadata, + 'query': ( + 'SELECT timestamp, variable, value, created_at ' + f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}' + ), + 'model_name': 'test_model', + 'datetime_columns': ['timestamp', 'created_at'], + } + ) + assert payload.object_key, 'offloaded payload must reference a MinIO object' + assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict' + + df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata']) + assert len(df) >= 1 + + client = minio_container.get_client() + listed = list(client.list_objects('test-bucket', recursive=True)) + names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed] + assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}' + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_predictions_batch_with_minio_offload_path( + temporal_test_env: WorkflowEnvironment, + temporal_worker_real_minio: Worker, + postgres_engine, + test_activities_real_minio: Activities, +): + """ + Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes. + """ + model_id = 502 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')) + insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0]) + + input_data = { + 'schedule_name': 'test-schedule', + 'model_name': 'test_model', + 'model_id': model_id, + 'query': ( + 'SELECT timestamp, variable, value, created_at ' + f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}' + ), + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'input_filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'mlflow_transform_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'mlflow_predict_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'save_transform': True, + 'prediction_store_policy': 'lts:1', + 'model_config': { + 'retention_minutes': 0, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'sklearn', + }, + 'datetime_columns': ['timestamp', 'created_at'], + } + + with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1): + await start_and_await_workflow( + temporal_test_env.client, + PredictionsBatch.run, + input_data, + make_workflow_id('test-batch-minio-offload'), + ) + + with postgres_engine.connect() as conn: + count = conn.execute( + text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}') + ).scalar() + assert count == 1 diff --git a/e2e/test_predictions_batch_format_export.py b/e2e/test_predictions_batch_format_export.py index e7bdd91..35543f8 100644 --- a/e2e/test_predictions_batch_format_export.py +++ b/e2e/test_predictions_batch_format_export.py @@ -2,17 +2,17 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios. """ -import asyncio -from datetime import datetime -from unittest.mock import ANY, AsyncMock, patch, call +from decimal import Decimal +from typing import Any, cast +from unittest.mock import ANY, AsyncMock, call -import pandas as pd import pytest from sientia_do.notifications.models import NotificationLevel from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch @@ -25,13 +25,13 @@ base_input_data = { 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -55,70 +55,6 @@ def get_base_input_data(model_id): 'query': base_query.format(model_id=model_id), } -def insert_sample_data(postgres_engine, model_id, values: list): - with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}")) - - values_sql = [] - for i, value in enumerate(values): - values_sql.append(f""" - ({model_id}, 'sensor_{i+1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') - """) - - insert_sql = f""" - INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES - {', '.join(values_sql)} - """ - conn.execute(text(insert_sql)) - -async def start_and_await_workflow(client, input_data, workflow_id): - handle = await client.start_workflow( - PredictionsBatch.run, - input_data, - id=workflow_id, - task_queue='test-queue', - ) - print("[TEST] ✓ Workflow started") - - print("\n[TEST] 3. Waiting for workflow completion...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - print("[TEST] ✓ Workflow completed successfully") - except asyncio.TimeoutError: - pytest.fail("Workflow execution timed out after 60 seconds") - -def assert_prediction( - postgres_engine, model_id, prediction: float = 0.5, - prediction_confidence: int = 0, prediction_status: str = 'Good', - comments: str = '', -): - """ - Verify prediction was created with correct values in database - - Args: - postgres_engine: Database engine - model_id: Model ID to check - prediction: Expected prediction value (default 0.5 from mock) - prediction_confidence: Expected confidence value (default 0 for normal predictions) - prediction_status: Expected status (default 'Good') - comments: Expected comments (default empty string) - """ - print("\n[TEST] 4. Verifying prediction was created with correct values...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}") - ) - prediction_rows = result_query.fetchall() - assert len(prediction_rows) == 1, f"Expected one prediction record, got {len(prediction_rows)}" - - row = prediction_rows[0] - assert row[0] == model_id, f"Expected model_id={model_id}, got {row[0]}" - assert row[1] == prediction, f"Expected prediction={prediction}, got {row[1]}" - assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}" - assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}" - assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" - @pytest.mark.asyncio @pytest.mark.integration @@ -129,35 +65,31 @@ async def test_scenario_3_1_1_default_prediction_export( postgres_engine, ): """ - Scenario 3.1.1: Default Prediction Export - - Description: - Error prediction path creates default prediction. - - Expected Behavior: - - format_default_prediction called instead of format_prediction - - Default prediction created with error metadata - - Exported to PostgreSQL only - - Transformed data NOT processed - - Metrics written - - Assertions: - - format_default_prediction called - - format_prediction NOT called - - format_transformed_data NOT called - - One PostgreSQL export only - - Default values in prediction data - - Comment included + Scenario 3.1.1: Default prediction export (non-None path_flag). + + Triggers input_gate CONTINUE via SPECIFIC_VARIABLES_NULL_VALUES so + PredictionProcess calls FormatAndExportPrediction with path_flag set. + That workflow uses format_default_prediction (not format_prediction) and + skips format_transformed_data / transform Postgres export. + + Optional PI Web API and OPC outputs still run when configured. """ client = temporal_test_env.client model_id = 311 - print("\n[TEST] 1. Inserting test data...") - insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") + with postgres_engine.begin() as conn: + conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}")) + conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) + insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) input_data = get_base_input_data(model_id) + input_data['input_filters'] = { + 'SPECIFIC_VARIABLES_NULL_VALUES': { + 'POLICY': 'CONTINUE', + 'CONFIG': {'variables': ['sensor_1']}, + }, + } input_data['pi_web_api_output_config'] = { 'endpoint': 'test_endpoint', 'prediction_tags': {'tag_1': 'web_id_1'}, @@ -178,10 +110,9 @@ async def test_scenario_3_1_1_default_prediction_export( } } - print("\n[TEST] 2. Starting workflow that should create default prediction...") - workflow_id = f'test-default-prediction-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + wid = make_workflow_id('test-default-prediction') + + await start_and_await_workflow(client, PredictionsBatch.run, input_data, wid) test_activities.pi_web_api_client.write_value.assert_has_calls( [ @@ -189,9 +120,8 @@ async def test_scenario_3_1_1_default_prediction_export( web_ids=['web_id_1'], value={ 'Timestamp': '2024-01-01 12:00:00+0000', - 'Value': 0.5, + 'Value': 0, }, - endpoint='test_endpoint', metadata={ 'model_id': 311, 'model_name': 'test_model', @@ -203,9 +133,8 @@ async def test_scenario_3_1_1_default_prediction_export( web_ids=['web_id_2'], value={ 'Timestamp': '2024-01-01 12:00:00+0000', - 'Value': 0, + 'Value': 2, }, - endpoint='test_endpoint', metadata={ 'model_id': 311, 'model_name': 'test_model', @@ -217,28 +146,51 @@ async def test_scenario_3_1_1_default_prediction_export( any_order=True, ) - test_activities.opc_repository['1'].write_data.assert_has_calls( + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_has_calls( [ - call('addr_1', 0.5, 'float', ANY, - { - 'model_id': 311, - 'model_name': 'test_model', - 'schedule_name': 'test-schedule', - 'workflow_name': 'predictions_batch', - }), - call('addr_2', 0, 'float', ANY, - { - 'model_id': 311, - 'model_name': 'test_model', - 'schedule_name': 'test-schedule', - 'workflow_name': 'predictions_batch', - }), + call( + 'addr_1', + 0, + 'float', + ANY, + { + 'model_id': 311, + 'model_name': 'test_model', + 'schedule_name': 'test-schedule', + 'workflow_name': 'predictions_batch', + }, + ), + call( + 'addr_2', + 2, + 'float', + ANY, + { + 'model_id': 311, + 'model_name': 'test_model', + 'schedule_name': 'test-schedule', + 'workflow_name': 'predictions_batch', + }, + ), ] ) - assert_prediction(postgres_engine, model_id) + with postgres_engine.connect() as conn: + tf_count = conn.execute( + text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}") + ).scalar() + assert tf_count == 0, 'transform export must be skipped when path_flag is set' + + assert_prediction( + postgres_engine, + model_id, + prediction=0, + prediction_confidence=Decimal(2), + prediction_status='Bad', + comments='Input data with bad quality', + ) - print("\n[TEST] ✓ All assertions passed!") @@ -274,9 +226,7 @@ async def test_scenario_3_1_2_export_with_opc_only( model_id = 312 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['opc_output_config'] = { @@ -295,12 +245,12 @@ async def test_scenario_3_1_2_export_with_opc_only( } input_data['pi_web_api_output_config'] = None # No PI Web API config - print("\n[TEST] 2. Starting workflow with OPC only...") - workflow_id = f'test-opc-only-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-only') + ) - test_activities.opc_repository['1'].write_data.assert_has_calls( + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_has_calls( [ call('addr_1', 0.5, 'float', ANY, { @@ -323,7 +273,6 @@ async def test_scenario_3_1_2_export_with_opc_only( assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -357,9 +306,7 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( model_id = 313 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['pi_web_api_output_config'] = { @@ -369,10 +316,9 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( } input_data['opc_output_config'] = None # No OPC config - print("\n[TEST] 2. Starting workflow with PI Web API only...") - workflow_id = f'test-pi-api-only-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-only') + ) test_activities.pi_web_api_client.write_value.assert_has_calls( [ @@ -382,7 +328,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0.5, }, - endpoint='test_endpoint', metadata={ 'model_id': 313, 'model_name': 'test_model', @@ -396,7 +341,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0, }, - endpoint='test_endpoint', metadata={ 'model_id': 313, 'model_name': 'test_model', @@ -408,11 +352,11 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( any_order=True, ) - test_activities.opc_repository['1'].write_data.assert_not_called() + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_not_called() assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -445,25 +389,22 @@ async def test_scenario_3_1_4_export_without_optional_outputs( model_id = 314 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['opc_output_config'] = None # No OPC config input_data['pi_web_api_output_config'] = None # No PI Web API config - print("\n[TEST] 2. Starting workflow without optional outputs...") - workflow_id = f'test-no-optional-outputs-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-no-optional-outputs') + ) test_activities.pi_web_api_client.write_value.assert_not_called() - test_activities.opc_repository['1'].write_data.assert_not_called() + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_not_called() assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -494,11 +435,9 @@ async def test_scenario_3_1_5_export_without_transformed_data( model_id = 315 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) with postgres_engine.begin() as conn: conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['save_transform'] = False # Don't save transformed data @@ -522,10 +461,9 @@ async def test_scenario_3_1_5_export_without_transformed_data( } } - print("\n[TEST] 2. Starting workflow without transformed data export...") - workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-no-transform-export') + ) test_activities.pi_web_api_client.write_value.assert_has_calls( [ @@ -535,7 +473,6 @@ async def test_scenario_3_1_5_export_without_transformed_data( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0.5, }, - endpoint='test_endpoint', metadata={ 'model_id': 315, 'model_name': 'test_model', @@ -549,7 +486,6 @@ async def test_scenario_3_1_5_export_without_transformed_data( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0, }, - endpoint='test_endpoint', metadata={ 'model_id': 315, 'model_name': 'test_model', @@ -561,7 +497,8 @@ async def test_scenario_3_1_5_export_without_transformed_data( any_order=True, ) - test_activities.opc_repository['1'].write_data.assert_has_calls( + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_has_calls( [ call('addr_1', 0.5, 'float', ANY, { @@ -589,7 +526,6 @@ async def test_scenario_3_1_5_export_without_transformed_data( assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -599,31 +535,20 @@ async def test_scenario_3_2_1_pi_web_api_write_error( temporal_worker: Worker, test_activities: Activities, postgres_engine, + notification_inserts, ): """ Scenario 3.2.1: PI Web API Write Error - - Description: - PI Web API export fails. - - Expected Behavior: - - write_pi_web_api_data raises exception - - Notification sent - - Workflow fails after retries - - PostgreSQL export may not execute (depends on execution order) - - Assertions: - - PI Web API error notification sent - - Workflow fails - - May impact subsequent exports + + Export failure is handled inside the activity; there is no retry loop. The + workflow completes and PostgreSQL stores prediction_confidence 13 and the + error message in comments. """ client = temporal_test_env.client model_id = 321 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") test_activities.pi_web_api_client.write_value.side_effect = Exception( "PI Web API service unavailable") @@ -649,18 +574,17 @@ async def test_scenario_3_2_1_pi_web_api_write_error( } } - print("\n[TEST] 2. Starting workflow that should fail on PI Web API write...") - workflow_id = f'test-pi-api-error-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-error') + ) assert_prediction( postgres_engine, model_id, prediction_confidence=13, comments='PI Web API service unavailable', ) - - print("\n[TEST] ✓ All assertions passed!") + assert notification_inserts.call_count >= 1 + @pytest.mark.asyncio @@ -673,29 +597,18 @@ async def test_scenario_3_2_2_opc_write_error( ): """ Scenario 3.2.2: OPC Write Error - - Description: - OPC server write fails. - - Expected Behavior: - - write_opc_data raises exception - - Notification sent - - Workflow fails after retries - - Assertions: - - OPC error notification sent - - Workflow fails - - PostgreSQL export may not execute + + OPC failure is reported without failing the workflow; there is no retry + loop. PostgreSQL stores prediction_confidence 12 and OPC error comments. """ client = temporal_test_env.client model_id = 322 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") - test_activities.opc_repository['1'].write_data.return_value = (False, { + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.return_value = (False, { 'notification_id': 'OPC_WRITE_DATA_ERROR_1', 'message': 'OPC server unavailable', 'block': 'opc_repository', @@ -724,10 +637,9 @@ async def test_scenario_3_2_2_opc_write_error( 'confidence_tags': {'tag_2': 'web_id_2'}, } - print("\n[TEST] 2. Starting workflow that should fail on OPC write...") - workflow_id = f'test-opc-error-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error') + ) assert_prediction( postgres_engine, model_id, @@ -735,7 +647,6 @@ async def test_scenario_3_2_2_opc_write_error( comments='Some data could not be written to OPC servers', ) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -748,51 +659,24 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( ): """ Scenario 3.2.3: PI Web API Partial Write Error - - Description: - Two prediction tags attempt to be written to PI Web API, but only one succeeds. - - Expected Behavior: - - write_pi_web_api_data processes response - - process_pi_web_api_response detects partial failure - - Error confidence set (13) - - Notification sent for failed tag - - Workflow completes with error confidence - - Assertions: - - One tag written successfully - - One tag failed - - Error confidence set in prediction - - Error notification sent - - Workflow completes + + Partial PI write: confidence 13, descriptive comments, workflow completes + without an activity retry loop. """ client = temporal_test_env.client model_id = 323 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") - test_activities.pi_web_api_client.write_value = AsyncMock(side_effect=[ - { - 'Items': [ - { - 'WebId': 'web_id_1', - 'Errors': [], - }, - ] - }, - Exception('Tag write failed'), - { - 'Items': [ - { - 'WebId': 'web_id_2', - 'Errors': [], - }, - ] - }, - ]) + test_activities.pi_web_api_client.write_value = AsyncMock( + side_effect=[ + # Prediction batch: two web_ids requested, only one acknowledged. + [{'WebId': 'web_id_1', 'Errors': []}], + # Confidence write succeeds. + [{'WebId': 'web_id_2', 'Errors': []}], + ] + ) input_data = get_base_input_data(model_id) input_data['pi_web_api_output_config'] = { @@ -815,10 +699,9 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( } } - print("\n[TEST] 2. Starting workflow with partial PI Web API write error...") - workflow_id = f'test-pi-api-partial-error-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error') + ) assert_prediction( postgres_engine, model_id, @@ -826,4 +709,3 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.", ) - print("\n[TEST] ✓ All assertions passed!") diff --git a/e2e/test_predictions_batch_main_workflow.py b/e2e/test_predictions_batch_main_workflow.py index f8469ef..3f8d37f 100644 --- a/e2e/test_predictions_batch_main_workflow.py +++ b/e2e/test_predictions_batch_main_workflow.py @@ -3,14 +3,14 @@ End-to-end tests for PredictionsBatch workflow - Main workflow scenarios. """ import asyncio -from datetime import datetime -import pandas as pd -import pytest from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +import pytest + +from e2e.helpers import make_workflow_id, start_and_await_workflow from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch @@ -23,66 +23,20 @@ async def test_scenario_1_1_1_happy_path_complete_success( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.1.1: Happy Path - Complete Success - - Description: - Workflow completes successfully with valid SQL query and all activities succeed. - - Process Flow: - 1. load_custom_query returns DataFrame with sensor data - 2. Workflow prepares prediction input with all configurations - 3. prediction_process child workflow executes: - - get_last_timestamp retrieves last processing timestamp - - input_gate validates data quality (passes) - - request_transform calls MLFlow transform (mocked, returns features) - - mlflow_response_gate validates transform response (passes) - - mlflow_content_gate validates transform content (passes) - - request_predict calls MLFlow predict (mocked, returns predictions) - - mlflow_response_gate validates predict response (passes) - - mlflow_content_gate validates predict content (passes) - 4. format_and_export_prediction child workflow executes: - - format_prediction formats the prediction data - - format_transformed_data formats transformed data (if save_transform=True) - - export_data_to_postgres saves to database - - write_metrics records execution metrics - - Expected Behavior: - - All activities execute successfully without errors - - All gates pass with no quality issues - - Transform and predict operations succeed (mocked) - - Data exported to PostgreSQL predictions table - - Transformed data exported to transformed_data table (if save_transform=True) - - Metrics written successfully - - Assertions: - - Workflow completes without raising exceptions - - Data exists in PostgreSQL predictions table with correct model_id - - Data exists in transformed_data table (if save_transform=True) - - Prediction data has expected structure (jsonb with predictions) - - All required fields are populated (model_id, model_name, timestamp, etc) - """ + """Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms.""" client = temporal_test_env.client - print("\n[TEST] 1. Inserting test data into PostgreSQL...") - # Insert test data directly into PostgreSQL - # The load_custom_query activity will fetch this data with a real SQL query with postgres_engine.begin() as conn: - # Clear any existing data for this model_id - conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 123")) - - # Insert sensor data that the workflow will query + conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123')) insert_sql = """ INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES + VALUES (123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'), (123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'), (123, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') """ conn.execute(text(insert_sql)) - print("[TEST] ✓ Data inserted successfully") - # Prepare input data for PredictionsBatch workflow input_data = { 'metadata': { 'metadata': { @@ -100,13 +54,13 @@ async def test_scenario_1_1_1_happy_path_complete_success( 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -121,77 +75,44 @@ async def test_scenario_1_1_1_happy_path_complete_success( 'datetime_columns': ['timestamp', 'created_at'], } - # Start workflow - print("\n[TEST] 2. Starting workflow...") - workflow_id = f'test-predictions-batch-{datetime.now().timestamp()}' - print(f"[TEST] Workflow ID: {workflow_id}") - - handle = await client.start_workflow( + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, - id=workflow_id, - task_queue='test-queue', + make_workflow_id('test-predictions-batch'), ) - print("[TEST] ✓ Workflow started") - # Wait for workflow completion with timeout - print("\n[TEST] 3. Waiting for workflow completion (timeout: 60s)...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) # 60 seconds timeout - print("[TEST] ✓ Workflow completed successfully") - except asyncio.TimeoutError: - print("[TEST] ✗ Workflow TIMEOUT after 60 seconds!") - pytest.fail("Workflow execution timed out after 60 seconds") - - # Verify data was stored in PostgreSQL - use single connection schema_name = 'predictions_schema' - predictions_table = 'predictions' - transformed_table = 'transformed_data' - full_predictions_table = f"{schema_name}.{predictions_table}" - full_transformed_table = f"{schema_name}.{transformed_table}" - - # Use a single connection for all verification queries - print("\n[TEST] 4. Verifying results in PostgreSQL...") with postgres_engine.connect() as conn: - # Verify prediction data result_query = conn.execute( - text(f"SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments FROM {full_predictions_table} WHERE model_id = 123") + text( + f'SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments ' + f'FROM {schema_name}.predictions WHERE model_id = 123' + ) ) prediction_rows = result_query.fetchall() - - print(f"[TEST] Found {len(prediction_rows)} prediction record(s)") - assert len(prediction_rows) == 1, "Expected one prediction record" - - # Verify first row has expected structure + assert len(prediction_rows) == 1 row = prediction_rows[0] - print(f"[TEST] Prediction: {row}") - assert row[0] == 123, f"Expected model_id=123, got {row[0]}" - assert row[1] == 0.5, f"Expected prediction=0.5, got {row[1]}" - assert row[2] == 0, f"Expected prediction_confidence=0.9, got {row[2]}" - assert row[3] is not None, f"Expected response_time=0.1, got {row[3]}" - assert row[4] == 'Good', f"Expected prediction_status='Good', got {row[4]}" - assert row[5] == '', f"Expected comments='', got {row[5]}" - print("[TEST] ✓ Prediction data verified") + assert row[0] == 123 + assert row[1] == 0.5 + assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}' + assert row[3] is not None + assert row[4] == 'Good' + assert row[5] == '' - # Verify transformed data result_query = conn.execute( - text(f"SELECT model_id, variable, value FROM {full_transformed_table} WHERE model_id = 123") + text( + f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123' + ) ) transformed_rows = result_query.fetchall() - print(f"[TEST] Found {len(transformed_rows)} transformed data record(s)") - assert len(transformed_rows) == 2, "Expected two transformed data records" - row_1 = transformed_rows[0] - print(f"[TEST] Transformed data: {row_1}") - assert row_1[0] == 123, f"Expected model_id=123, got {row_1[0]}" - assert row_1[1] == 'feature_1', f"Expected variable='sensor_1', got {row_1[1]}" - assert float(row_1[2]) == 0.234, f"Expected value=0.234, got {row_1[2]}" - row_2 = transformed_rows[1] - print(f"[TEST] Transformed data: {row_2}") - assert row_2[0] == 123, f"Expected model_id=123, got {row_2[0]}" - assert row_2[1] == 'feature_2', f"Expected variable='sensor_2', got {row_2[1]}" - assert float(row_2[2]) == 0.783, f"Expected value=0.783, got {row_2[2]}" - - print("\n[TEST] ✓ All assertions passed!") + assert len(transformed_rows) == 2 + assert transformed_rows[0][0] == 123 + assert transformed_rows[0][1] == 'feature_1' + assert float(transformed_rows[0][2]) == 0.234 + assert transformed_rows[1][0] == 123 + assert transformed_rows[1][1] == 'feature_2' + assert float(transformed_rows[1][2]) == 0.783 @pytest.mark.asyncio @@ -202,23 +123,7 @@ async def test_scenario_1_2_1_sql_query_execution_error( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.2.1: SQL Query Execution Error - - Description: - SQL query fails due to syntax error or connection issue. - - Expected Behavior: - - load_custom_query raises exception (caught by Temporal retry policy) - - Notification sent with SQL error details - - After retries, activity may return empty data or workflow may fail - - If empty data returned, workflow completes with early exit via input gate - - Assertions: - - Error notification sent - - Workflow completes (either fails or exits early) - - No data in predictions table - """ + """Invalid SQL: workflow may complete with early exit; no prediction rows.""" client = temporal_test_env.client input_data = { @@ -233,18 +138,18 @@ async def test_scenario_1_2_1_sql_query_execution_error( 'schedule_name': 'test-schedule', 'model_name': 'test_model', 'model_id': 128, - 'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =', # Invalid SQL + 'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =', 'schema': 'predictions_schema', 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -258,34 +163,18 @@ async def test_scenario_1_2_1_sql_query_execution_error( }, } - print("\n[TEST] 1. Starting workflow with invalid SQL query...") - workflow_id = f'test-sql-error-{datetime.now().timestamp()}' - - handle = await client.start_workflow( + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, - id=workflow_id, - task_queue='test-queue', + make_workflow_id('test-sql-error'), ) - print("[TEST] ✓ Workflow started") - print("\n[TEST] 2. Waiting for workflow completion...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - print("[TEST] ✓ Workflow completed (may have exited early due to empty data)") - except Exception as e: - print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}") - - # Verify no predictions were created (regardless of whether workflow failed or exited early) - print("\n[TEST] 3. Verifying no predictions were created...") with postgres_engine.connect() as conn: - result_query = conn.execute( - text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128") - ) - count = result_query.scalar() - assert count == 0, f"Expected no predictions, but found {count} records" - - print("\n[TEST] ✓ All assertions passed!") + count = conn.execute( + text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128') + ).scalar() + assert count == 0 @pytest.mark.asyncio @@ -296,24 +185,9 @@ async def test_scenario_1_2_2_missing_required_parameters( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.2.2: Missing Required Parameters - - Description: - Essential parameters missing from input. - - Expected Behavior: - - Workflow or activity raises KeyError or validation error - - Workflow fails immediately - - Assertions: - - Workflow fails with parameter error - - Error notification sent - - No child workflow called - """ + """Missing query: workflow does not produce predictions and is terminated explicitly.""" client = temporal_test_env.client - # Missing 'query' parameter input_data = { 'metadata': { 'metadata': { @@ -326,31 +200,28 @@ async def test_scenario_1_2_2_missing_required_parameters( 'schedule_name': 'test-schedule', 'model_name': 'test_model', 'model_id': 129, - # 'query' is missing 'schema': 'predictions_schema', 'table_name': 'predictions', 'transform_table_name': 'transformed_data', } - print("\n[TEST] 1. Starting workflow with missing required parameter...") - workflow_id = f'test-missing-param-{datetime.now().timestamp()}' - handle = await client.start_workflow( PredictionsBatch.run, input_data, - id=workflow_id, + id=make_workflow_id('test-missing-param'), task_queue='test-queue', ) - print("[TEST] ✓ Workflow started") - print("\n[TEST] 2. Waiting for workflow to fail...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - pytest.fail("Expected workflow to fail, but it completed successfully") - except Exception as e: - print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}") - - print("\n[TEST] ✓ All assertions passed!") + # Let Temporal process a few workflow tasks; for this case, result() can hang. + await asyncio.sleep(2.0) + + with postgres_engine.connect() as conn: + count = conn.execute( + text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129') + ).scalar() + assert count == 0 + + await handle.terminate('expected failure path in e2e test (missing required parameters)') @pytest.mark.asyncio @@ -361,34 +232,19 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.2.3: Invalid Datetime Column Specification - - Description: - Datetime column specified doesn't exist in query results. - - Expected Behavior: - - load_custom_query may raise KeyError or warning - - Depending on implementation, workflow may fail or continue - - Error notification sent - - Assertions: - - Error raised or warning logged - - Workflow behavior depends on error handling policy - """ + """Invalid datetime column: no predictions persisted; workflow terminated after validation.""" client = temporal_test_env.client - print("\n[TEST] 1. Inserting test data...") with postgres_engine.begin() as conn: - conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 130")) - - insert_sql = """ + conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130')) + conn.execute( + text( + """ INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES - (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') + VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') """ - conn.execute(text(insert_sql)) - print("[TEST] ✓ Data inserted successfully") + ) + ) input_data = { 'metadata': { @@ -407,13 +263,13 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification( 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -425,26 +281,23 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification( 'transform_flavor': 'sklearn', 'predict_flavor': 'sklearn', }, - 'datetime_columns': ['nonexistent_column'], # Column doesn't exist in query result + 'datetime_columns': ['nonexistent_column'], } - print("\n[TEST] 2. Starting workflow with invalid datetime column...") - workflow_id = f'test-invalid-datetime-col-{datetime.now().timestamp()}' - handle = await client.start_workflow( PredictionsBatch.run, input_data, - id=workflow_id, + id=make_workflow_id('test-invalid-datetime-col'), task_queue='test-queue', ) - print("[TEST] ✓ Workflow started") - print("\n[TEST] 3. Waiting for workflow completion or failure...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - # Workflow may complete or fail depending on error handling - print("[TEST] ✓ Workflow completed (may have handled error gracefully)") - except Exception as e: - print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}") - - print("\n[TEST] ✓ Test completed!") + # Let Temporal process and surface the failure path internally. + await asyncio.sleep(2.0) + + with postgres_engine.connect() as conn: + count = conn.execute( + text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130') + ).scalar() + assert count == 0 + + await handle.terminate('expected failure path in e2e test (invalid datetime column)') diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index d32980e..af097cd 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/e2e/test_predictions_batch_prediction_process.py @@ -2,55 +2,62 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios. """ -import asyncio -from datetime import datetime from decimal import Decimal from unittest.mock import MagicMock, patch +import numpy as np import pandas as pd import pytest -from pytz import timezone from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +from e2e.helpers import ( + assert_continue, + assert_repeat, + assert_stop, + insert_sample_data, + make_workflow_id, + start_and_await_workflow, +) from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch base_input_data = { - 'schedule_name': 'test-schedule', - 'model_name': 'test_model', - 'model_id': 201, - 'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201', - 'schema': 'predictions_schema', - 'table_name': 'predictions', - 'transform_table_name': 'transformed_data', - 'input_filters': { - 'SPECIFIC_VARIABLES_NULL_VALUES': { - 'policy': 'CONTINUE', # Continue despite issues, not STOP - 'config': {'variables': ['sensor_1']}, - }, + 'schedule_name': 'test-schedule', + 'model_name': 'test_model', + 'model_id': 201, + 'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201', + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'input_filters': { + 'SPECIFIC_VARIABLES_NULL_VALUES': { + 'POLICY': 'CONTINUE', + 'CONFIG': {'variables': ['sensor_1']}, }, - 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, - }, - 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, - }, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], - 'opc_output_config': {}, - 'pi_web_api_output_config': {}, - 'save_transform': True, - 'prediction_store_policy': 'lts:1', - 'model_config': { - 'retention_minutes': 0, - 'transform_flavor': 'sklearn', - 'predict_flavor': 'sklearn', - }, - 'datetime_columns': ['timestamp', 'created_at'], - } + }, + 'mlflow_transform_filters': { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + }, + 'mlflow_predict_filters': { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + }, + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'save_transform': True, + 'prediction_store_policy': 'lts:1', + 'model_config': { + 'retention_minutes': 0, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'sklearn', + }, + 'datetime_columns': ['timestamp', 'created_at'], +} + +base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}' -base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}" def get_base_input_data(model_id): return { @@ -59,97 +66,38 @@ def get_base_input_data(model_id): 'query': base_query.format(model_id=model_id), } -def insert_sample_data(postgres_engine, model_id, values: list[tuple]): - with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}")) - - # Insert data with some null values (quality issue) - - values_sql = [] - for i, value in enumerate(values): - values_sql.append(f""" - ({model_id}, 'sensor_{i+1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') - """) - - insert_sql = f""" - INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES - {', '.join(values_sql)} - """ - conn.execute(text(insert_sql)) - def insert_sample_prediction(postgres_engine, model_id): with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}")) - - # Insert data with some null values (quality issue) - + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) insert_sql = f""" INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time) VALUES ({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1) """ conn.execute(text(insert_sql)) - - return (model_id, Decimal(10), Decimal(0), 'Good') - -async def start_and_await_workflow(client, input_data, workflow_id): - handle = await client.start_workflow( - PredictionsBatch.run, - input_data, - id=workflow_id, - task_queue='test-queue', - ) - print("[TEST] ✓ Workflow started") - - print("\n[TEST] 3. Waiting for workflow completion...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - print("[TEST] ✓ Workflow completed successfully") - except asyncio.TimeoutError: - pytest.fail("Workflow execution timed out after 60 seconds") - -def assert_continue( - postgres_engine, model_id, prediction_confidence: Decimal = 2, - comments: str = 'Input data with bad quality', -): - print("\n[TEST] 4. Verifying prediction was created despite warnings...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}") - ) - prediction_rows = result_query.fetchall() - assert len(prediction_rows) == 1, "Expected one prediction record despite warnings" - - # Assert prediction value is 0 and other fields - row = prediction_rows[0] - assert row[1] == 0, f"Expected prediction=0, got {row[1]}" - assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}" - assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}" - assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" + return (model_id, Decimal(10), Decimal(0), 'Good') -def assert_stop(postgres_engine, model_id): - print("\n[TEST] 4. Verifying no predictions were created...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f"SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}") - ) - count = result_query.scalar() - assert count == 0, f"Expected no predictions, but found {count} records" +@pytest.fixture +def bad_data_model(patch_mlflow): + model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model'))) + patch_mlflow.sklearn.load_model = MagicMock(return_value=model) + return model -def assert_repeat(postgres_engine, model_id, last_prediction: list): - print("\n[TEST] 4. Verifying prediction was repeated...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f'SELECT model_id, prediction, prediction_confidence, prediction_status FROM predictions_schema.predictions WHERE model_id = {model_id}') - ) - prediction_rows = result_query.fetchall() - print(prediction_rows) - assert len(prediction_rows) == 2, "Expected two prediction records" - assert prediction_rows[0] == last_prediction, f"Expected first prediction to be the same as the last prediction, got {prediction_rows[0]}, expected {last_prediction}" - assert prediction_rows[1] == last_prediction, f"Expected second prediction to be the same as the last prediction, got {prediction_rows[1]}, expected {last_prediction}" + +@pytest.fixture +def bad_predict_model(patch_mlflow, mock_mlflow_models): + model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model'))) + + def mock_sklearn_load_model(model_uri): + if 'data_model' in model_uri or 'transform' in model_uri.lower(): + return mock_mlflow_models['transform_model'] + return model + + patch_mlflow.sklearn = MagicMock() + patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model) + return model @pytest.mark.asyncio @@ -159,45 +107,19 @@ async def test_scenario_2_1_1_input_gate_triggers_continue( temporal_worker: Worker, test_activities: Activities, postgres_engine, + mock_mlflow_models, ): - """ - Scenario 2.1.1: Input Gate Triggers CONTINUE - - Description: - Input gate determines data should use previous prediction. - - Expected Behavior: - - input_gate returns path_flag='CONTINUE' - - path_flag_handler calls export workflow with input data directly - - MLFlow transform and predict skipped - - Data exported as-is - - Assertions: - - input_gate called - - MLFlow operations NOT called - - Export workflow called with original data - - Workflow completes - """ + """Input gate CONTINUE: export default prediction; MLflow transform/predict not used.""" client = temporal_test_env.client - model_id = 211 - - print("\n[TEST] 1. Inserting test data...") - insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) - - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) - - - print("\n[TEST] 2. Starting workflow with CONTINUE policy...") - workflow_id = f'test-continue-policy-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy') + ) assert_continue(postgres_engine, model_id) - - print("\n[TEST] ✓ All assertions passed!") + mock_mlflow_models['transform_model'].predict.assert_not_called() + mock_mlflow_models['predict_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -207,45 +129,19 @@ async def test_scenario_2_1_2_input_gate_triggers_stop( temporal_worker: Worker, test_activities: Activities, postgres_engine, + mock_mlflow_models, ): - """ - Scenario 2.1.2: Input Gate Triggers STOP - - Description: - Input data quality gate fails with STOP policy. - - Expected Behavior: - - input_gate returns path_flag='STOP' - - path_flag_handler detects STOP - - Workflow returns early without calling MLFlow - - No prediction exported - - Assertions: - - input_gate called - - path_flag_handler returns True (early exit) - - MLFlow transform NOT called - - Export workflow NOT called - - Workflow completes without error - """ + """Input gate STOP: no export, no MLflow.""" client = temporal_test_env.client - model_id = 212 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) - input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['policy'] = 'STOP' - - print("\n[TEST] 2. Starting workflow that should stop at input gate...") - workflow_id = f'test-input-stop-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) - + input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop') + ) assert_stop(postgres_engine, model_id) - - print("\n[TEST] ✓ All assertions passed!") + mock_mlflow_models['transform_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -255,57 +151,42 @@ async def test_scenario_2_1_3_input_gate_triggers_repeat( temporal_worker: Worker, test_activities: Activities, postgres_engine, + mock_mlflow_models, ): - """ - Scenario 2.1.3: Input Gate Triggers REPEAT - - Description: - Input gate determines data should repeat last prediction. - - Expected Behavior: - - input_gate returns path_flag='REPEAT' - - path_flag_handler calls repeat_last_prediction activity - - MLFlow transform and predict skipped - - Last prediction repeated and exported - - Assertions: - - input_gate called - - MLFlow operations NOT called - - repeat_last_prediction activity called - - Workflow completes - """ + """Input gate REPEAT with existing history.""" client = temporal_test_env.client - model_id = 213 - - print("\n[TEST] 1. Inserting test data and previous prediction...") insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) data = insert_sample_prediction(postgres_engine, model_id) - - print("[TEST] ✓ Data and previous prediction inserted") - input_data = get_base_input_data(model_id) - input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['policy'] = 'REPEAT' - - print("\n[TEST] 2. Starting workflow that should trigger REPEAT...") - workflow_id = f'test-input-repeat-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) - assert_repeat(postgres_engine, model_id, data) - - print("\n[TEST] ✓ All assertions passed!") - -@pytest.fixture -def bad_data_model(patch_mlflow): - model = MagicMock( - predict=MagicMock( - side_effect=Exception("Bad data model") - ) + input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat') ) + assert_repeat(postgres_engine, model_id, data) + mock_mlflow_models['transform_model'].predict.assert_not_called() - patch_mlflow.sklearn.load_model = MagicMock(return_value=model) - - return model + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, +): + """REPEAT when no prior row in predictions: repeat_last_prediction runs; still no new duplicate export path.""" + client = temporal_test_env.client + model_id = 214 + insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + input_data = get_base_input_data(model_id) + input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-no-history') + ) + assert_stop(postgres_engine, model_id) @pytest.mark.asyncio @@ -317,48 +198,20 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue( postgres_engine, bad_data_model, ): - """ - Scenario 2.2.1: Transform Gate Triggers CONTINUE - - Description: - Transform response gate determines data should continue despite issues. - - Expected Behavior: - - request_transform succeeds - - mlflow_response_gate for transform returns path_flag='CONTINUE' - - path_flag_handler calls export workflow with transform data - - MLFlow predict skipped - - Transform data exported as-is - - Assertions: - - Transform completed - - mlflow_response_gate called for transform - - MLFlow predict NOT called - - Export workflow called with transform data - - Workflow completes - """ client = temporal_test_env.client - model_id = 221 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - input_data = get_base_input_data(model_id) - input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'CONTINUE' - - print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at transform gate...") - workflow_id = f'test-transform-continue-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue') + ) assert_continue( postgres_engine=postgres_engine, model_id=model_id, prediction_confidence=Decimal(10), - comments='Bad data model', + comments='Unknown MLFlow API error', ) - - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -369,44 +222,18 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop( test_activities: Activities, postgres_engine, bad_data_model, + mock_mlflow_models, ): - """ - Scenario 2.2.2: Transform Gate Triggers STOP - - Description: - Transform response validation fails with STOP policy. - - Expected Behavior: - - request_transform succeeds but response invalid - - mlflow_response_gate for transform returns path_flag='STOP' - - Workflow exits without calling predict or export - - Assertions: - - Transform completed but validation failed - - mlflow_response_gate called for transform - - MLFlow predict NOT called - - Export workflow NOT called - - Workflow completes without error - """ client = temporal_test_env.client - model_id = 222 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) - input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'STOP' - - print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...") - workflow_id = f'test-transform-stop-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop') + ) assert_stop(postgres_engine, model_id) - - print("\n[TEST] ✓ All assertions passed!") - + mock_mlflow_models['predict_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -418,67 +245,50 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat( postgres_engine, bad_data_model, ): - """ - Scenario 2.2.3: Transform Gate Triggers REPEAT - - Description: - Transform response gate determines data should repeat last prediction. - - Expected Behavior: - - request_transform succeeds but response has issues - - mlflow_response_gate for transform returns path_flag='REPEAT' - - path_flag_handler calls repeat_last_prediction activity - - MLFlow predict skipped - - Last prediction repeated and exported - - Assertions: - - Transform completed but validation triggered REPEAT - - mlflow_response_gate called for transform - - MLFlow predict NOT called - - repeat_last_prediction activity called - - Workflow completes - """ client = temporal_test_env.client - model_id = 223 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) data = insert_sample_prediction(postgres_engine, model_id) - - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) - input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'REPEAT' - - print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...") - workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat') + ) assert_repeat(postgres_engine, model_id, data) - print("\n[TEST] ✓ All assertions passed!") - -@pytest.fixture -def bad_predict_model( - patch_mlflow, - mock_mlflow_models +@pytest.mark.asyncio +@pytest.mark.integration +async def test_scenario_2_2_4_transform_content_gate_nan_values_stop( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + mock_mlflow_models, ): - model = MagicMock( - predict=MagicMock( - side_effect=Exception("Bad predict model") - ) + """mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter).""" + client = temporal_test_env.client + model_id = 224 + + def all_nan_transform(data): + num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1 + result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows}) + result.index = data.index + return result + + mock_mlflow_models['transform_model'].predict = MagicMock(side_effect=all_nan_transform) + + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + input_data = get_base_input_data(model_id) + input_data['mlflow_transform_filters'] = { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + 'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}, + } + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop') ) - - def mock_sklearn_load_model(model_uri): - if 'data_model' in model_uri or 'transform' in model_uri.lower(): - return mock_mlflow_models['transform_model'] - return model - patch_mlflow.sklearn = MagicMock() - patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model) - - return model + assert_stop(postgres_engine, model_id) + mock_mlflow_models['predict_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -490,46 +300,20 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue( postgres_engine, bad_predict_model, ): - """ - Scenario 2.3.1: Predict Gate Triggers CONTINUE - - Description: - Predict response gate determines data should continue despite issues. - - Expected Behavior: - - request_predict succeeds - - mlflow_response_gate for predict returns path_flag='CONTINUE' - - path_flag_handler calls export workflow with predict data - - Prediction exported despite quality issues - - Assertions: - - Transform and predict completed - - mlflow_response_gate called for predict - - Export workflow called with predict data - - Workflow completes - """ client = temporal_test_env.client - model_id = 231 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - input_data = get_base_input_data(model_id) - input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'CONTINUE' - - print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...") - workflow_id = f'test-predict-continue-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue') + ) assert_continue( postgres_engine=postgres_engine, model_id=model_id, prediction_confidence=Decimal(10), - comments='Bad predict model', + comments='Unknown MLFlow API error', ) - - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -541,42 +325,16 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop( postgres_engine, bad_predict_model, ): - """ - Scenario 2.3.2: Predict Gate Triggers STOP - - Description: - Prediction validation fails with STOP policy. - - Expected Behavior: - - request_predict succeeds but response invalid - - mlflow_response_gate for predict returns path_flag='STOP' - - Workflow exits without export - - Assertions: - - Transform completed - - Predict completed but validation failed - - Export workflow NOT called - - Workflow completes without error - """ client = temporal_test_env.client - model_id = 232 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) - input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'STOP' - - print("\n[TEST] 2. Starting workflow that should stop at predict gate...") - workflow_id = f'test-predict-stop-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop') + ) assert_stop(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") - @pytest.mark.asyncio @pytest.mark.integration @@ -587,43 +345,35 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat( postgres_engine, bad_predict_model, ): - """ - Scenario 2.3.3: Predict Gate Triggers REPEAT - - Description: - Predict response gate determines data should repeat last prediction. - - Expected Behavior: - - request_predict succeeds but response has issues - - mlflow_response_gate for predict returns path_flag='REPEAT' - - path_flag_handler calls repeat_last_prediction activity - - Last prediction repeated and exported - - Assertions: - - Transform and predict completed but validation triggered REPEAT - - mlflow_response_gate called for predict - - repeat_last_prediction activity called - - Export workflow NOT called with current prediction - - Workflow completes - """ client = temporal_test_env.client - model_id = 233 - - print("\n[TEST] 1. Inserting test data and previous prediction...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) data = insert_sample_prediction(postgres_engine, model_id) - - print("[TEST] ✓ Data and previous prediction inserted") - input_data = get_base_input_data(model_id) - input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'REPEAT' - input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE'] # REPEAT first - - print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...") - workflow_id = f'test-predict-repeat-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT' + input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE'] + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat') + ) assert_repeat(postgres_engine, model_id, data) - print("\n[TEST] ✓ All assertions passed!") + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_scenario_2_4_1_input_empty_data_stop( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, +): + """EMPTY_DATA filter with STOP when query returns no rows (offload payload empty).""" + client = temporal_test_env.client + model_id = 241 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + input_data = get_base_input_data(model_id) + input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}} + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop') + ) + assert_stop(postgres_engine, model_id) diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index f009b10..9a03f54 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -6,6 +6,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger + from sientia_do.repository.minio_repository import MinioRepository from laborious.activities.api import API from laborious.activities.gates import Gates @@ -73,6 +74,17 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): """ metrics_controller = MetricsController(logger=logger) + minio_repository = MinioRepository( + endpoint=minio_config['endpoint_url'], + access_key=minio_config['access_key'], + secret_key=minio_config['secret_key'], + bucket=minio_config['default_bucket'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + secure=minio_config['secure'], + ) + # Initialize parent classes Storage.__init__( self, @@ -83,7 +95,8 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): dbname=postgres_config['dbname'], min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], - minio_config=minio_config, + retention_hours=minio_config['retention_hours'], + minio_repository=minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, @@ -95,7 +108,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): mlflow_port=mlflow_config['port'], mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], - minio_config=minio_config, + minio_repository=minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, @@ -103,6 +116,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): Gates.__init__( self, + minio_repository=minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, diff --git a/laborious/activities/api.py b/laborious/activities/api.py index 366ecd2..0e4114f 100644 --- a/laborious/activities/api.py +++ b/laborious/activities/api.py @@ -74,34 +74,26 @@ class API(SientiaMonitoring): def get_pi_web_api_core_labels( self, metadata: dict[str, Any], - operation_type: str | None = None, + operation_type: str = 'write_pi_web_api_data', ) -> dict[str, Any]: """ - Generate core labels for metrics, optionally including operation_type. + Generate core labels for PI Web API metrics. - This override keeps compatibility with the base implementation while adding - a convenience overload behavior: - - When operation_type is provided, it behaves exactly like the base class, - returning labels that include the operation_type key. - - When operation_type is omitted (None), it removes the operation_type key - from the resulting labels. This is useful for metrics, such as the PI Web - API metrics, that are defined without the operation_type label. + PI Web API metrics in laborious use the shared ``CORE_LABELS`` from + ``sientia_do``, which includes ``operation_type``. For this reason, + operation_type must always be present in emitted labels. Args: - - metadata (dict[str, Any]): Workflow execution metadata used to derive labels - - operation_type (str | None): Optional operation type label. If None, the - operation_type key will be removed from the returned labels. + - metadata (dict[str, Any]): Workflow execution metadata used to derive labels. + - operation_type (str): Operation type label for metric cardinality. Return: - dict[str, Any]: Core labels dictionary, with operation_type only when provided + dict[str, Any]: Core labels dictionary including operation_type. """ - base_labels = super().get_core_labels( + return super().get_core_labels( metadata=metadata, - operation_type=operation_type or '-', + operation_type=operation_type, ) - if operation_type is None: - base_labels.pop('operation_type', None) - return base_labels def close(self) -> None: """ @@ -256,8 +248,13 @@ class API(SientiaMonitoring): metadata=metadata, ) - data['prediction_confidence'] = confidence - data['comments'] = message + # Preserve incoming confidence/comments on successful PI writes. + # Only downgrade confidence or override comments when PI response + # explicitly reports a problem (e.g. partial write mismatch). + if confidence != 0: + data['prediction_confidence'] = confidence + if message: + data['comments'] = message except Exception as e: trace = traceback.format_exc() @@ -282,7 +279,7 @@ class API(SientiaMonitoring): web_ids=confidence_tags, value={ 'Timestamp': data.head(1)['timestamp'].values[0], - 'Value': confidence_value, + 'Value': float(confidence_value), }, metadata=metadata, ) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index b581920..ad88ca1 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -1,5 +1,8 @@ +from sientia_do.repository.minio_repository import MinioRepository from temporalio import activity, workflow +from laborious.utils.repository.minio_manager import MinioManager + with workflow.unsafe.imports_passed_through(): import traceback from collections.abc import Callable, Mapping @@ -10,16 +13,16 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.models import NotificationLevel from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController - from sientia_do.observability.sientia_monitoring import SientiaMonitoring - from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now from sientia_do.utils.formatters import create_sample_dict from laborious import metrics + from laborious.utils.dataframe_debug import build_dataframe_debug_message from laborious.utils.filters.conditional_filters import ( filter_empty_data, filter_specific_variables_null_values, ) from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload # Strongly-typed filter function signatures InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool] @@ -63,7 +66,7 @@ mlflow_content_path_confidence: Mapping[str, int] = { } -class Gates(SientiaMonitoring): +class Gates(MinioManager): """ Data quality gates and filtering activities for the Laborious system. @@ -83,11 +86,15 @@ class Gates(SientiaMonitoring): mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions """ + minio_repository: MinioRepository | None = None + _MAX_DEBUG_DATAFRAME_ROWS = 100 + def __init__( self, - logger: Logger, - notification_handler: NotificationHandler, - metrics_controller: MetricsController, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, ): """ Initialize data quality gates with logging and notification capabilities. @@ -99,17 +106,54 @@ class Gates(SientiaMonitoring): Raises: Exception: If BaseActivity initialization fails """ - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) def close(self) -> None: """ Close the gates activity and clean up resources. """ - SientiaMonitoring.shutdown(self) + + MinioManager.close(self) def __del__(self): self.close() + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like payload to be logged + - metadata (dict[str, Any]): Workflow metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + + @staticmethod + def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]: + """ + Read filter policy/config keys in a case-insensitive way. + + Args: + config (dict[str, Any]): Filter configuration dictionary. + + Return: + tuple[str, dict[str, Any]]: Parsed policy and config payload. + """ + normalized = {str(key).upper(): value for key, value in config.items()} + policy = normalized['POLICY'] + filter_config = normalized.get('CONFIG', {}) + return policy, filter_config + @activity.defn(name='input_gate') async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: """ @@ -149,12 +193,13 @@ class Gates(SientiaMonitoring): self.info('Performing input gate...', metadata) filters = input_data['filters'] - data = DataFrame(input_data['data']) + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) path_priority = input_data['path_priority'] filter_output = [] - self.debug(f'Input data: {data.head(5).to_string()}', metadata) + self._debug_dataframe('Input data:', data, metadata) self.debug(f'Filters: {filters}', metadata) # Apply each configured filter @@ -162,10 +207,11 @@ class Gates(SientiaMonitoring): if fil not in input_filter_functions: self.error(f'Filter {fil} not found', metadata) continue + policy, filter_config = self._read_filter_entry(config) try: - if input_filter_functions[fil](data, config['config']): + if input_filter_functions[fil](data, filter_config): self.debug(f'Data not passed the input filter {fil}:{config}', metadata) - filter_output.append(config['policy']) + filter_output.append(policy) except Exception as e: trace = traceback.format_exc() await self.send_notification_async( @@ -183,6 +229,9 @@ class Gates(SientiaMonitoring): return path_flag, input_path_confidence[path_flag], 'Input data with bad quality' self.info('Nothing was filtered by the input gate', metadata) + + del data + return None, 0, '' @activity.defn(name='mlflow_response_gate') @@ -221,32 +270,41 @@ class Gates(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Performing mlflow response gate...', metadata) - + raw_data = input_data['data'] filters = input_data['filters'] - data = input_data['data'] + + self.debug( + f'Input data: \n {create_sample_dict(raw_data, max_items=5, max_depth=5)}', metadata + ) + self.debug(f'Filters: {filters}', metadata) + + payload = MinioDataFramePayload.from_dict(raw_data) + data = await payload.retrieve(self.minio_repository, metadata) + gate_type = input_data['type'] path_priority = input_data['path_priority'] filter_output = [] - self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}', metadata) - self.debug(f'Filters: {filters}', metadata) - comments = [] + + status = payload.status or {} + for fil, config in filters.items(): if fil not in mlflow_response_filter_functions: continue + policy, filter_config = self._read_filter_entry(config) try: - if mlflow_response_filter_functions[fil](data, config): - filter_output.append(config['policy']) - comments.append(data['content']['message']) + if mlflow_response_filter_functions[fil](status, filter_config): + filter_output.append(policy) + comments.append(status.get('message', 'Unknown MLFlow API error')) await self.send_notification_async( metadata=metadata, notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}', - message=data['content']['message'], + message=status.get('message', 'Unknown MLFlow API error'), block='mlflow_gate', level=NotificationLevel.ERROR, - attachment_content=data['content']['traceback'], + attachment_content=status.get('traceback'), ) except Exception as e: trace = traceback.format_exc() @@ -265,6 +323,9 @@ class Gates(SientiaMonitoring): return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments) self.info('Nothing was filtered by the mlflow response gate', metadata) + + del data + return None, 0, '' @activity.defn(name='mlflow_content_gate') @@ -305,21 +366,25 @@ class Gates(SientiaMonitoring): self.info('Performing mlflow content gate...', metadata) filters = input_data['filters'] - data = DataFrame(input_data['data']) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + gate_type = input_data['type'] path_priority = input_data['path_priority'] filter_output = [] - self.debug(f'Input data:\n {data.head(5).to_string()}', metadata) + self._debug_dataframe('Input data:', data, metadata) self.debug(f'Filters: \n {filters}', metadata) for fil, config in filters.items(): if fil not in mlflow_content_filter_functions: continue + policy, filter_config = self._read_filter_entry(config) try: - if mlflow_content_filter_functions[fil](data, config): - filter_output.append(config['policy']) + if mlflow_content_filter_functions[fil](data, filter_config): + filter_output.append(policy) await self.send_notification_async( metadata=metadata, notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}', @@ -349,6 +414,9 @@ class Gates(SientiaMonitoring): ) self.info('Nothing was filtered by the mlflow content gate', metadata) + + del data + return None, 0, '' def get_prediction_store_policy( @@ -403,7 +471,7 @@ class Gates(SientiaMonitoring): return policy_type, int(policy_value) @activity.defn(name='format_transformed_data') - async def format_transformed_data(self, input_data: dict[str, Any]) -> dict: + async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ Format transformed data for storage and export operations. @@ -438,7 +506,8 @@ class Gates(SientiaMonitoring): self.info('Formatting transformed data...', metadata) - data = DataFrame(input_data['data']) + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) data['timestamp'] = data.index data = data.reset_index(drop=True) @@ -446,7 +515,15 @@ class Gates(SientiaMonitoring): data = data.melt(id_vars='timestamp', var_name='variable', value_name='value') data['model_id'] = model_id - return data.to_dict() + return await MinioDataFramePayload.from_dataframe( + dataframe=data, + minio_repo=self.minio_repository, + model_name=input_data['model_name'], + operation='transform', + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) @activity.defn(name='format_prediction') async def format_prediction(self, input_data: dict[str, Any]) -> dict: @@ -458,6 +535,8 @@ class Gates(SientiaMonitoring): and ensures data consistency before persistence. The method supports multiple storage policies for flexible data retention strategies. + If only one row is present, we use the last timestamp as the timestamp + Storage Policies: - 'lts:N': Latest timestamp - retains N most recent predictions - 'erl:N': Earliest timestamp - retains N oldest predictions @@ -465,7 +544,7 @@ class Gates(SientiaMonitoring): Args: input_data (dict): Input data containing: - data (dict[str, Any]): Raw prediction data to format - - timestamp (str): Default timestamp if data lacks timestamp column + - timestamp (str): Timestamp of the data - model_id (str): Unique identifier for the ML model - prediction_confidence (float): Confidence score for the prediction - prediction_store_policy (str): Storage policy in format 'type:value' @@ -474,17 +553,19 @@ class Gates(SientiaMonitoring): dict: Formatted prediction data ready for storage and export """ metadata = input_data['metadata'] + last_timestamp = input_data['timestamp'] prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) - data = DataFrame(input_data['data']) + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) # Create timestamp column from index and reset index data['timestamp'] = data.index data = data.reset_index(drop=True) self.debug(f'Prediction store policy: {prediction_store_policy}', metadata) - self.debug(f'Prediction data: {data.head(5).to_string()}', metadata) + self._debug_dataframe('Prediction data:', data, metadata) policy_type, policy_value = self.get_prediction_store_policy( prediction_store_policy, metadata @@ -507,7 +588,11 @@ class Gates(SientiaMonitoring): 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)) + int_policy_value = int(policy_value) + + data = data.head(int_policy_value) + if int_policy_value == 1: + data['timestamp'] = last_timestamp data['model_id'] = input_data['model_id'] data['prediction_confidence'] = input_data['prediction_confidence'] @@ -517,7 +602,7 @@ class Gates(SientiaMonitoring): 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._debug_dataframe('Prediction data:', data, metadata) return data.to_dict() @@ -632,50 +717,10 @@ class Gates(SientiaMonitoring): 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) + self._debug_dataframe('Retrain report:', report, 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. - - This method analyzes prediction data to find the latest timestamp, - enabling incremental processing and data continuity tracking. - It handles empty datasets gracefully by returning the current time - as a fallback timestamp. - - The method is essential for: - 1. Incremental data processing workflows - 2. Data continuity validation - 3. Timestamp-based data loading optimization - 4. Workflow execution tracking - - Args: - input_data (dict): Input data containing: - - data (dict[str, Any]): Prediction data to analyze - - Returns: - str: Formatted timestamp string in UTC with timezone - """ - metadata = input_data['metadata'] - - self.info('Getting last timestamp...', metadata) - - data = DataFrame(input_data['data']) - - self.debug(f'Input data: {data.head(5).to_string()}', metadata) - - if data.empty: - return now().strftime(DATETIME_FORMAT_WITH_TZ) - - max_timestamp = max(data['timestamp'].values.tolist()) - - self.info(f'Last timestamp: {max_timestamp}', metadata) - - return max_timestamp - @activity.defn(name='write_metrics') async def write_metrics(self, input_data: dict[str, Any]): """ @@ -707,34 +752,29 @@ class Gates(SientiaMonitoring): self.info(f'Writing metrics for model {metadata["model_name"]}', metadata) + core_tags = { + 'pod_id': self.pod_id, + 'runtime': self.runtime, + 'operation_type': 'predict', + 'model_name': metadata['model_name'], + 'workflow_name': metadata['workflow_name'], + } await self.emit_metric( metric_object=metrics.PREDICTIONS_WRITTEN_COUNT, - tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], - }, + tags=core_tags, ) await self.emit_metric( metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR, method='set', - tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], - }, + tags=core_tags, value=prediction_confidence, ) await self.emit_metric( metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR, method='observe', - tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], - }, + tags=core_tags, value=response_time, ) @@ -745,9 +785,7 @@ class Gates(SientiaMonitoring): metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, method='observe', tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], + **core_tags, 'opc_server_id': server_id, 'tag': tag, }, @@ -757,9 +795,7 @@ class Gates(SientiaMonitoring): await self.emit_metric( metric_object=metrics.PREDICTION_OPC_WRITING_COUNT, tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], + **core_tags, 'opc_server_id': server_id, 'tag': tag, }, diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 082e720..f4d42a3 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -5,12 +5,12 @@ with workflow.unsafe.imports_passed_through(): from typing import Any import numpy as np - from pandas import DataFrame, to_datetime + from pandas import to_datetime 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.observability.metrics_controller import MetricsController - from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.minio_repository import MinioRepository from sientia_do.temporal.constants import ( DATETIME_FORMAT, DATETIME_FORMAT_MS_WITH_TZ, @@ -19,11 +19,13 @@ with workflow.unsafe.imports_passed_through(): ) from sientia_do.utils.formatters import create_sample_dict - from laborious.utils.repository.minio_repository import MinioRepository + from laborious.utils.dataframe_debug import build_dataframe_debug_message + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + from laborious.utils.repository.minio_manager import MinioManager from laborious.utils.repository.model_repository import MLFlowRepository -class MLFlow(SientiaMonitoring): +class MLFlow(MinioManager): """ MLFlow integration activities for model inference operations. @@ -42,16 +44,18 @@ class MLFlow(SientiaMonitoring): model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 + 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, - metrics_controller: MetricsController, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, ): """ Initialize MLFlow activities with server configuration. @@ -67,7 +71,9 @@ class MLFlow(SientiaMonitoring): Raises: Exception: If MLFlowRepository initialization fails """ - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) self.mlflow_host = mlflow_host self.mlflow_port = mlflow_port self.mlflow_username = mlflow_username @@ -82,32 +88,35 @@ class MLFlow(SientiaMonitoring): metrics_controller, ) - 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'], - metrics_controller=metrics_controller, - ) - def close(self) -> None: """ Close the MLFlow activity and clean up resources. """ - SientiaMonitoring.shutdown(self) + MinioManager.close(self) def __del__(self): self.close() + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like object expected to expose shape and to_csv + - metadata (dict[str, Any]): Workflow metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + @activity.defn(name='request_transform') - async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ Transform input data using MLFlow models. @@ -138,12 +147,14 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Transforming data...', metadata) - data = DataFrame(input_data['data']) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) - self.debug('Raw input data:', metadata) - self.debug(data.head(5).to_string(), metadata) + self._debug_dataframe('Raw input data:', data, 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( @@ -157,7 +168,9 @@ class MLFlow(SientiaMonitoring): data.columns.name = None data.index.name = None - self.debug(f'Processed input data: \n {data.to_csv()}', metadata) + data['timestamp'] = data.index + + self._debug_dataframe('Processed input data:', data, metadata) # Request transformation from MLFlow model response_data = await self.model_monitoring_repository.transform( @@ -176,10 +189,33 @@ class MLFlow(SientiaMonitoring): self.info('Data transformed successfully', metadata) - return response_data + if not response_data.get('success', False): + return await MinioDataFramePayload.from_dataframe( + dataframe=None, + minio_repo=self.minio_repository, + model_name=model_name, + operation='transform', + status=response_data, + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + return await MinioDataFramePayload.from_dataframe( + dataframe=response_data['content'], + minio_repo=self.minio_repository, + model_name=model_name, + operation='transform', + workflow_metadata=metadata, + status={ + 'success': True, + }, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) @activity.defn(name='request_predict') - async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ Execute predictions using MLFlow models. @@ -210,11 +246,14 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Predicting data...', metadata) - data = DataFrame(input_data['data']) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + 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_dataframe('Input data for prediction:', data, metadata) # Convert numpy.nan to None for model compatibility data.replace(np.nan, None, inplace=True) @@ -236,7 +275,30 @@ class MLFlow(SientiaMonitoring): self.info('Data predicted successfully', metadata) - return response_data + if not response_data.get('success', False): + return await MinioDataFramePayload.from_dataframe( + dataframe=None, + minio_repo=self.minio_repository, + model_name=model_name, + operation='predict', + status=response_data, + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + return await MinioDataFramePayload.from_dataframe( + dataframe=response_data['content'], + minio_repo=self.minio_repository, + model_name=model_name, + operation='predict', + workflow_metadata=metadata, + status={ + 'success': True, + }, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) @activity.defn(name='retrain_model') async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]: @@ -275,14 +337,12 @@ class MLFlow(SientiaMonitoring): raise ValueError('Minio repository not initialized') metadata = input_data['metadata'] - object_key = input_data['object_key'] - - self.info(f'Loading retrain data from Key: {object_key}', metadata) try: - data = await self.minio_repository.get_parquet_as_dataframe( - object_key=object_key, metadata=metadata - ) + # Payload-based retrain input (inline dict or MinIO offloaded). + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + except Exception as e: trace = traceback.format_exc() await self.send_notification_async( diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index 29b2848..79b140c 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -17,6 +17,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ from laborious import metrics + from laborious.utils.dataframe_debug import build_dataframe_debug_message warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0') warnings.filterwarnings( @@ -31,6 +32,8 @@ class ModelMetrics(SientiaMonitoring): This class provides activities for writing metrics to the Prometheus monitoring system. """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 + def __init__( self, logger: Logger, @@ -48,6 +51,24 @@ class ModelMetrics(SientiaMonitoring): def __del__(self): self.close() + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like payload to be logged + - metadata (dict[str, Any]): Workflow metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + async def get_drift_metrics( self, reference_data: DataFrame, @@ -78,14 +99,11 @@ class ModelMetrics(SientiaMonitoring): model_analysis = ModelAnalysis(config=config) - self.debug( - f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', - metadata, + self._debug_dataframe( + f'Reference data: Size {reference_data.shape}', reference_data, metadata ) - self.debug( - f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata - ) + self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata) core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift') start_time = time.time() @@ -142,9 +160,7 @@ class ModelMetrics(SientiaMonitoring): await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) - self.debug( - f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata - ) + self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata) return drift_df @@ -274,11 +290,7 @@ class ModelMetrics(SientiaMonitoring): drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC') drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ) - self.debug( - f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata - ) - - self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata) + self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata) return drift_df.to_dict(orient='records') @@ -347,8 +359,6 @@ class ModelMetrics(SientiaMonitoring): data['data_size'] = data_size data['interval_minutes'] = interval_minutes - self.debug( - f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata - ) + self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata) return data.to_dict(orient='records') diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 0b14397..d64cc52 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -1,8 +1,11 @@ from temporalio import activity, workflow +from laborious.utils.repository.minio_manager import MinioManager + with workflow.unsafe.imports_passed_through(): # Extend the Temporal Postgres activities for convenient query -> MinIO export import traceback + from datetime import timedelta from typing import Any import pandas as pd @@ -10,18 +13,23 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.models import NotificationLevel from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.repository.minio_repository import MinioRepository from sientia_do.temporal.activities.postgres import Postgres - from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now + from sientia_do.temporal.constants import now - from laborious.utils.repository.minio_repository import MinioRepository + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + +_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'}) -class Storage(Postgres): +class Storage(Postgres, MinioManager): """ Extensions for Postgres activities with a helper to export query results directly to MinIO as Parquet and return the object name. """ + minio_repository: MinioRepository | None = None + def __init__( self, host: str, @@ -31,12 +39,15 @@ class Storage(Postgres): dbname: str, min_connections: int, max_connections: int, - minio_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler, - metrics_controller: MetricsController, + retention_hours: int = 24, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, ): - super().__init__( + self.retention_hours = retention_hours + Postgres.__init__( + self, host=host, port=port, user=user, @@ -49,90 +60,150 @@ class Storage(Postgres): metrics_controller=metrics_controller, ) - if not hasattr(self, 'minio_repository'): - self.minio_repository: MinioRepository | None = None + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) - 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'], - metrics_controller=metrics_controller, - ) - - @activity.defn(name='query_to_minio') - async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='load_query_with_minio_offload') + async def load_query_with_minio_offload( + self, input_data: dict[str, Any] + ) -> MinioDataFramePayload: """ - Execute SQL query, write result as Parquet to MinIO, and return object name. + Run the custom SQL load, then return a MinIO-aware dataframe wire dict. 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) + metadata (dict): Workflow metadata (same as load_custom_query). + query (str): SQL query. + datetime_columns (list[str], optional): Datetime column names. + model_name (str): Model name for object key basename. + key_prefix (str, optional): Directory prefix inside the bucket. + size_threshold_bytes (int, optional): Override env offload threshold. Returns: - dict: { success: bool, object_name: str, uri: str } + dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure. """ + if self.minio_repository is None: + raise ValueError('Minio repository not initialized') + metadata: dict = input_data.get('metadata', {}) + model_name = input_data['model_name'] + + rows = await self.load_custom_query( + input_data, + ) + if not rows: + self.error( + 'load_query_with_minio_offload failed: No data returned from query', metadata + ) + dataframe = None + else: + dataframe = pd.DataFrame(rows) + + return await MinioDataFramePayload.from_dataframe( + dataframe, + minio_repo=self.minio_repository, + workflow_metadata=metadata, + model_name=model_name, + operation='initial', + logger=self.logger, + ) + + @activity.defn(name='export_payload_to_postgres') + async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict: + """ + Export a payload to PostgreSQL. + """ + metadata = input_data.get('metadata') + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + return await self.export_data_to_postgres( + { + **input_data, + 'data': data, + } + ) + + @activity.defn(name='cleanup_minio_objects_expired') + async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Delete objects under the given prefixes that are older than the retention window. + + Args (input_data): + metadata (dict): Workflow metadata for logging and metrics. + prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix). + + Returns: + dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``. + """ 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_FORMAT_FILENAME) - object_name = f'{object_prefix}_{timestamp}.parquet' - uri = f's3://{self.minio_repository.minio_bucket}/{object_name}' + payload = MinioDataFramePayload.from_dict(input_data['data']) + prefix = payload.cleanup_prefix() + base = now() + cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta( + hours=self.retention_hours + ) + report: dict[str, Any] = { + 'failed': {}, + 'deleted': {}, + 'failed_count': 0, + 'deleted_count': 0, + } 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 - await self.minio_repository.store_dataframe_as_parquet( - dataframe=data, uri=uri, object_name=object_name, metadata=metadata + keys = await self.minio_repository.list_objects( + prefix=prefix, + recursive=True, + metadata=metadata, ) - - return {'success': True, 'object_key': object_name, 'uri': uri} + for key in keys: + try: + ts = MinioDataFramePayload.parse_object_timestamp(key) + if ts is None: + continue + if ts >= cutoff: + continue + await self.minio_repository.delete_file( + object_name=key, + metadata=metadata, + ) + except Exception as e: + report['failed'][key] = { + 'success': False, + 'message': str(e), + } + report['failed_count'] += 1 + continue + report['deleted'][key] = { + 'success': True, + 'message': 'Deleted', + } + report['deleted_count'] += 1 except Exception as e: trace = traceback.format_exc() await self.send_notification_async( metadata=metadata, - notification_id='ERROR_STORING_QUERY_TO_MINIO', - message=f'Error storing query to MinIO: {e}', - block='query_to_minio', + notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED', + message=f'Error cleaning up MinIO objects: {e}', + block='cleanup_minio_objects_expired', level=NotificationLevel.ERROR, attachment_content=trace, ) - self.error(trace, metadata) + else: + # Cleanup success is expected in normal flow; avoid noisy INFO notifications + # that do not impact behavior and can flood observability in test runs. + self.info('MinIO objects cleaned up successfully', metadata) - return {'success': False, 'message': str(e)} + return report 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') + Postgres.close(self) + MinioManager.close(self) def __del__(self): self.close() diff --git a/laborious/metrics.py b/laborious/metrics.py index cdd44db..776b826 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -18,15 +18,14 @@ Key Metric Categories: Metric Labels: - pod_id: Kubernetes pod identifier for multi-instance deployments +- runtime: Runtime / environment identifier (matches ``RUNTIME`` env, see ``SientiaMonitoring``) - model_name: Name of the ML model being used - workflow_name: Name of the prediction pipeline - opc_server_id: Identifier for OPC server operations """ from prometheus_client import Counter, Gauge, Histogram -from sientia_do.observability.metrics import ( - CORE_LABELS as SIENTIA_CORE_LABELS, -) +from sientia_do.observability.metrics import CORE_LABELS # Application health metric APP_UP = Gauge( @@ -35,9 +34,6 @@ APP_UP = Gauge( ['pod_id'], ) -# Core labels used across multiple metrics -CORE_LABELS = ['pod_id', 'model_name', 'workflow_name'] - # Prediction operation metrics PREDICTIONS_WRITTEN_COUNT = Counter( 'laborious_predictions_written_count', @@ -60,45 +56,6 @@ PREDICTION_RESPONSE_TIME_MONITOR = Histogram( buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], ) -# ================== MinIO metrics ================== - -MINIO_READ_LAG = Histogram( - 'laborious_minio_read_lag', - 'Lag between the last write to MinIO and the last read from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], - buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], -) - -MINIO_WRITE_LAG = Histogram( - 'laborious_minio_write_lag', - 'Lag between the last write to MinIO and the last read from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], - buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], -) - -MINIO_READ_COUNT = Counter( - 'laborious_minio_read_count', - 'Number of reads from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) - -MINIO_WRITE_COUNT = Counter( - 'laborious_minio_write_count', - 'Number of writes to MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) - -MINIO_READ_ERROR_COUNT = Counter( - 'laborious_minio_read_error_count', - 'Number of errors reading from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) - -MINIO_WRITE_ERROR_COUNT = Counter( - 'laborious_minio_write_error_count', - 'Number of errors writing to MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) # ================== OPC metrics ================== @@ -137,58 +94,58 @@ OPC_CONNECTION_STATUS = Gauge( MODEL_READ_LAG = Histogram( 'laborious_model_read_lag', 'Lag between the start and read of read operations', - SIENTIA_CORE_LABELS, + CORE_LABELS, buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], ) MODEL_WRITE_LAG = Histogram( 'laborious_model_write_lag', 'Lag between the start and end of write operations', - SIENTIA_CORE_LABELS, + CORE_LABELS, buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], ) MODEL_READ_COUNT = Counter( 'laborious_model_read_count', 'Number of reads from the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_WRITE_COUNT = Counter( 'laborious_model_write_count', 'Number of writes to the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_READ_ERROR_COUNT = Counter( 'laborious_model_read_error_count', 'Number of errors reading from the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_WRITE_ERROR_COUNT = Counter( 'laborious_model_write_error_count', 'Number of errors writing to the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_ANALYZE_LAG = Histogram( 'laborious_model_analyze_lag', 'Lag between the start and end of analyze operations', - SIENTIA_CORE_LABELS, + CORE_LABELS, buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], ) MODEL_ANALYZE_COUNT = Counter( 'laborious_model_analyze_count', 'Number of analyze operations', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_ANALYZE_ERROR_COUNT = Counter( 'laborious_model_analyze_error_count', 'Number of errors during analyze operations', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index c151377..ed3477a 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -78,7 +78,7 @@ def build_minio_config() -> dict[str, Any]: 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) - + MINIO_SECURE: Whether to use HTTPS (default: false) Returns: dict: MinIO configuration dictionary """ @@ -86,6 +86,7 @@ def build_minio_config() -> dict[str, Any]: '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'), + 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), + 'secure': getenv('MINIO_SECURE', 'false') == 'true', } diff --git a/laborious/utils/dataframe_debug.py b/laborious/utils/dataframe_debug.py new file mode 100644 index 0000000..995b606 --- /dev/null +++ b/laborious/utils/dataframe_debug.py @@ -0,0 +1,34 @@ +from typing import Any + +from pandas import DataFrame + +DEFAULT_MAX_DEBUG_DATAFRAME_ROWS = 100 + + +def build_dataframe_debug_message( + message: str, + data: Any, + max_rows: int = DEFAULT_MAX_DEBUG_DATAFRAME_ROWS, +) -> str: + """ + Build a safe debug message for dataframe payloads + + Args: + - message (str): Base message to identify the logged payload + - data (Any): Payload to evaluate for dataframe-aware logging + - max_rows (int): Maximum dataframe row count allowed for full payload logging + + Return: + Formatted debug message with full dataframe content or compact summary + """ + if not isinstance(data, DataFrame): + return f'{message} {data}' + + rows = data.shape[0] + if rows <= max_rows: + return f'{message}\n{data.to_csv()}' + + return ( + f'{message} skipped because dataframe has {rows} rows ' + f'(max: {max_rows}). Shape: {data.shape}' + ) diff --git a/laborious/utils/models/__init__.py b/laborious/utils/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py new file mode 100644 index 0000000..0521f5e --- /dev/null +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -0,0 +1,348 @@ +""" +MinIO-backed DataFrame payload for Temporal workflows. + +Data is never stored as a pandas ``DataFrame`` field on the dataclass. +Instead, the DataFrame is only provided as an input to: +`from_dataframe` / `from_dataframe_to_dict`. + +At build time, the DataFrame is evaluated for its serialized size; if it exceeds +the configured threshold, it is serialized to parquet bytes and uploaded to MinIO. +Otherwise, it is inlined as a Temporal-friendly ``dict``. +""" + +import pickle +import re +from collections.abc import Hashable +from dataclasses import dataclass +from datetime import datetime +from io import BytesIO +from os import getenv +from typing import Any, Literal + +from pandas import DataFrame, read_parquet +from sientia_do.observability.logger import Logger +from sientia_do.repository.minio_repository import MinioRepository +from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now + +# Keys that are part of the serialized wire format (not arbitrary metadata). +_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'}) + +_OBJECT_TIMESTAMP_PATTERN = re.compile( + r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$' +) + +OFFLOAD_THRESHOLD_BYTES = int( + float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024 +) + +# Relative prefix used for storing offloaded prediction datasets in MinIO. +# It is also the root directory for retention cleanup listing. +PREDICTION_DATASETS_PREFIX = 'prediction_datasets' + +OperationKind = Literal['initial', 'transform', 'predict'] + + +def _build_object_key( + model_name: str, operation: OperationKind, timestamp: str +) -> tuple[str, str | None]: + """ + Build the MinIO object key and the directory prefix used for retention listing. + + Args: + model_name: Registered model name used in the pipeline. + operation: Either initial (pre-transform load) or transform (post-MLFlow transform). + timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME. + + Return: + tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root). + """ + # Naming convention: + # - Directory is always `prediction_datasets/` + # - Filename follows the retention-parsing pattern + basename = f'{model_name}-{operation}-{timestamp}.parquet' + model_dir = model_name.strip().strip('/') + prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}' + return f'{prefix}/{basename}', prefix + + +@dataclass +class MinioDataFramePayload: + """ + Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys. + + Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`. + Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class. + """ + + last_timestamp: str + status: dict[str, Any] | None = None + data: dict[Hashable, Any] | None = None + bucket: str | None = None + object_key: str | None = None + object_prefix: str | None = None + uri: str | None = None + + @staticmethod + def _debug( + logger: Logger | None, + message: str, + metadata: dict[str, Any] | None = None, + ) -> None: + """ + Emit debug logs only when logger is provided + + Args: + - logger (Logger | None): Logger instance used for debug messages + - message (str): Message to be logged + - metadata (dict[str, Any] | None): Optional workflow metadata context + """ + if logger is None: + return + logger.custom_debug(message, metadata) + + @classmethod + def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload': + """ + Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization. + + Temporal converts dataclass return values into plain dicts when crossing + workflow/activity boundaries. This method rebuilds the typed instance so + that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are + available on the receiving side. + + If the argument is already a MinioDataFramePayload, it is returned as-is. + + Args: + raw: Dict with keys matching the dataclass fields + (last_timestamp, status, data, bucket, object_key, object_prefix, uri), + or an existing MinioDataFramePayload instance. + + Return: + MinioDataFramePayload: Reconstructed (or original) instance. + """ + if isinstance(raw, MinioDataFramePayload): + return raw + return cls( + last_timestamp=raw['last_timestamp'], + status=raw.get('status'), + data=raw.get('data'), + bucket=raw.get('bucket'), + object_key=raw.get('object_key'), + object_prefix=raw.get('object_prefix'), + uri=raw.get('uri'), + ) + + @staticmethod + def estimate_size_bytes( + df: DataFrame, + metadata: dict[str, Any] | None = None, + logger: Logger | None = None, + ) -> int: + """ + Approximate serialized size of the DataFrame as the default-orient dict. + + Args: + df: DataFrame whose tabular content size is estimated. + + Return: + int: Estimated size in bytes (pickle of dict representation). + """ + try: + size = len(pickle.dumps(df.to_dict())) + except Exception: + size = len(pickle.dumps(df)) + + MinioDataFramePayload._debug( + logger, + f'DataFrame size: {size} bytes', + metadata, + ) + return size + + @staticmethod + def parse_object_timestamp(object_key: str) -> datetime | None: + """ + Parse the timestamp embedded in the object key basename (before .parquet). + + Args: + object_key: S3/MinIO object key whose basename follows + ``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``. + + Return: + datetime | None: Parsed UTC-naive datetime from the key, or None if not matched. + """ + basename = object_key.rsplit('/', 1)[-1] + match = _OBJECT_TIMESTAMP_PATTERN.search(basename) + if not match: + return None + try: + return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME) + except ValueError: + return None + + def cleanup_prefix(self) -> str | None: + """ + Return True if cleanup is enabled for this payload. + """ + if self.object_key is not None and self.data is None: + return self.object_prefix + return None + + def has_data(self) -> bool: + """ + Return True if the payload has some data internally or in MinIO. + """ + return (self.data is not None and self.data != {}) or self.object_key is not None + + @classmethod + async def from_dataframe( + cls, + dataframe: DataFrame | None, + minio_repo: MinioRepository, + model_name: str, + operation: OperationKind, + status: dict[str, Any] | None = None, + workflow_metadata: dict | None = None, + last_timestamp: str | None = None, + logger: Logger | None = None, + ) -> 'MinioDataFramePayload': + """ + Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO. + + The DataFrame is not stored on the returned instance. + + Args: + dataframe: Tabular data to evaluate and persist (inline or MinIO). + metadata: Small metadata dict merged into the payload (e.g. success, message). + minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`. + workflow_metadata: Metadata passed to MinIO store for logging/metrics. + model_name: Registered model name used in the object basename. + operation: Either ``initial`` (query load) or ``transform`` (post-transform). + key_prefix: Backward-compatible parameter (currently ignored for object naming). + size_threshold_bytes: Byte limit before offload. When None, the module-level + environment-derived default is used. + + Return: + MinioDataFramePayload: Instance with data and/or MinIO fields set. + """ + + if dataframe is None or dataframe.empty: + cls._debug( + logger, + 'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload', + workflow_metadata, + ) + return cls( + data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status + ) + + if last_timestamp is None: + last_timestamp = max(dataframe['timestamp'].values.tolist()) + + dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger) + cls._debug( + logger, + ( + f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes ' + f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)' + ), + workflow_metadata, + ) + + if dataframe_size <= OFFLOAD_THRESHOLD_BYTES: + cls._debug( + logger, + 'MinioDataFramePayload.from_dataframe using inline payload', + workflow_metadata, + ) + return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status) + + timestamp = now().strftime(DATETIME_FORMAT_FILENAME) + object_key, object_prefix = _build_object_key(model_name, operation, timestamp) + cls._debug( + logger, + ( + 'MinioDataFramePayload.from_dataframe offloading payload to MinIO ' + f'with key {object_key}' + ), + workflow_metadata, + ) + + # Upload using the relative object key. The upstream repository will + # prefix it internally under its MinIO namespace. + parquet_buffer = BytesIO() + dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True) + file_bytes = parquet_buffer.getvalue() + + upload_result = await minio_repo.upload_file( + file_bytes=file_bytes, + relative_key=object_key, + metadata=workflow_metadata, + ) + + bucket = minio_repo.bucket + object_key_full = upload_result.get('minio_object_name', object_key) + uri = f's3://{bucket}/{object_key_full}' if bucket else None + cls._debug( + logger, + f'MinioDataFramePayload.from_dataframe upload completed: {uri}', + workflow_metadata, + ) + + return cls( + data=None, + bucket=bucket, + object_key=object_key_full, + object_prefix=object_prefix, + uri=uri, + last_timestamp=last_timestamp, + status=status, + ) + + async def retrieve( + self, + minio_repo: MinioRepository, + workflow_metadata: dict[str, Any] | None = None, + logger: Logger | None = None, + ) -> DataFrame: + """ + Load parquet from MinIO when object_key is set and populate inline data. + + Args: + minio_repo: sientia_do MinioRepository (or compatible) with download_file(). + workflow_metadata: Metadata passed to MinIO read for logging/metrics. + + Return: + dict[str, Any]: Flat dict with data filled (same keys as to_dict after load). + """ + if self.data is not None: + self._debug( + logger, + 'MinioDataFramePayload.retrieve using inline payload data', + workflow_metadata, + ) + return DataFrame(self.data) + + if not self.has_data(): + self._debug( + logger, + 'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe', + workflow_metadata, + ) + return DataFrame() + + self._debug( + logger, + f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}', + workflow_metadata, + ) + file_bytes = await minio_repo.download_file( + object_name=self.object_key, metadata=workflow_metadata + ) + df = read_parquet(BytesIO(file_bytes)) + self._debug( + logger, + f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}', + workflow_metadata, + ) + return df diff --git a/laborious/utils/repository/minio_manager.py b/laborious/utils/repository/minio_manager.py new file mode 100644 index 0000000..17c85e6 --- /dev/null +++ b/laborious/utils/repository/minio_manager.py @@ -0,0 +1,32 @@ +from sientia_do.notifications.handlers import NotificationHandler +from sientia_do.observability.logger import Logger +from sientia_do.observability.metrics_controller import MetricsController +from sientia_do.observability.sientia_monitoring import SientiaMonitoring +from sientia_do.repository.minio_repository import MinioRepository + + +class MinioManager(SientiaMonitoring): + minio_repository: MinioRepository | None = None + + def __init__( + self, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + if self.minio_repository is None: + self.minio_repository = minio_repository + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + + def close(self) -> None: + """ + Close the MinioManager and clean up resources. + """ + if self.minio_repository is not None: + try: + self.minio_repository.close() + finally: + self.minio_repository = None + + SientiaMonitoring.shutdown(self) diff --git a/laborious/utils/repository/minio_repository.py b/laborious/utils/repository/minio_repository.py deleted file mode 100644 index fa9cf0f..0000000 --- a/laborious/utils/repository/minio_repository.py +++ /dev/null @@ -1,215 +0,0 @@ -""" -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. -""" - -import time -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 -from sientia_do.observability.metrics_controller import MetricsController -from sientia_do.observability.sientia_monitoring import SientiaMonitoring - -from laborious import metrics - - -class MinioRepository(SientiaMonitoring): - """ - 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, - metrics_controller: MetricsController, - ): - """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. - """ - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) - # 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, - ), - ) - - def close(self): - """Close the underlying S3 client.""" - self.s3_client.close() - - async def create_bucket(self, metadata: dict[str, Any]) -> None: - core_labels = { - **self.get_core_labels(metadata, operation_type='create_bucket'), - 'bucket_name': self.minio_bucket, - 'object_name': '-', - } - self.info(f"Creating bucket '{self.minio_bucket}'", metadata) - - start_time = time.time() - try: - self.s3_client.create_bucket(Bucket=self.minio_bucket) - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MINIO_WRITE_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_WRITE_COUNT, tags=core_labels) - - async 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. - """ - self.info(f"Checking if bucket '{self.minio_bucket}' exists", metadata) - core_labels = { - **self.get_core_labels(metadata, operation_type='head_bucket'), - 'bucket_name': self.minio_bucket, - 'object_name': '-', - } - self.info(f"Checking if bucket '{self.minio_bucket}' exists", metadata) - - start_time = time.time() - - try: - self.s3_client.head_bucket(Bucket=self.minio_bucket) - except ClientError: - await self.create_bucket(metadata) - - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels) - raise e - - else: - await self.observe_lag(start_time, metrics.MINIO_READ_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_READ_COUNT, tags=core_labels) - - async 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. - """ - await self.ensure_bucket_exists(metadata) - - self.info(f'Storing dataframe as parquet in {uri}', metadata) - - buffer = BytesIO() - dataframe.to_parquet(buffer, engine='pyarrow', index=True) - buffer.seek(0) - - core_labels = { - **self.get_core_labels(metadata, operation_type='put_object'), - 'bucket_name': self.minio_bucket, - 'object_name': object_name, - } - start_time = time.time() - try: - self.s3_client.put_object( - Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue() - ) - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MINIO_WRITE_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_WRITE_COUNT, tags=core_labels) - - self.info(f'Dataframe stored as parquet in {uri}', metadata) - - async 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.info(f'Getting parquet as dataframe from {object_key}', metadata) - - core_labels = { - **self.get_core_labels(metadata, operation_type='get_object'), - 'bucket_name': self.minio_bucket, - 'object_name': object_key, - } - start_time = time.time() - try: - response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key) - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_READ_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MINIO_READ_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_READ_COUNT, tags=core_labels) - - # 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 66862e9..b4fdad5 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -36,6 +36,7 @@ from sientia_do.observability.sientia_monitoring import SientiaMonitoring from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from laborious import metrics +from laborious.utils.dataframe_debug import build_dataframe_debug_message ARTIFACTS_PATH = './tmp/artifacts' TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl' @@ -64,6 +65,8 @@ def force_memory_release(logger: Logger): class MLFlowRepository(SientiaMonitoring): + _MAX_DEBUG_DATAFRAME_ROWS = 100 + def __init__( self, host: str, @@ -94,6 +97,24 @@ class MLFlowRepository(SientiaMonitoring): self._cache_lock = threading.RLock() self.logger = logger + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like object expected to expose shape and to_csv + - metadata (dict[str, Any]): Metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + """ Functions related to get model registry parameters """ @@ -452,6 +473,10 @@ class MLFlowRepository(SientiaMonitoring): raw_model = mlflow.pyfunc.load_model(artifact_path) model = raw_model._model_impl.python_model + + self.debug( + f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata + ) else: if model_type == 'predict': model = await self.load_predict_model(model_name, metadata, flavor) @@ -1134,7 +1159,7 @@ class MLFlowRepository(SientiaMonitoring): async def transform( self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict - ): + ) -> dict[str, Any]: """ Transform data using a cached transformation model. @@ -1168,7 +1193,7 @@ class MLFlowRepository(SientiaMonitoring): and returned in the response structure rather than propagated. """ - self.debug(f'Data received for model transformation: {data.to_csv()}', metadata) + self._debug_dataframe('Data received for model transformation:', data, metadata) # data.to_csv( # f"tmp/data_{model_name}.csv", index=True) @@ -1186,9 +1211,8 @@ class MLFlowRepository(SientiaMonitoring): metadata=metadata, ) - self.debug( - f'Data received from model transformation: {transformed_data.head(5).to_csv()}', - metadata, + self._debug_dataframe( + 'Data received from model transformation:', transformed_data, metadata ) # transformed_data.to_csv( @@ -1196,7 +1220,7 @@ class MLFlowRepository(SientiaMonitoring): transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) - return {'success': True, 'content': transformed_data.to_dict()} + return {'success': True, 'content': transformed_data} except Exception as e: return { @@ -1254,9 +1278,7 @@ class MLFlowRepository(SientiaMonitoring): input_index = data.index start_time = datetime.now() - self.debug( - f'Data received for model prediction: {data.to_dict(orient="records")}', metadata - ) + self._debug_dataframe('Data received for model prediction:', data, metadata) # data.to_csv( # f"tmp/treated_data_{model_name}.csv", index=True) @@ -1273,9 +1295,8 @@ class MLFlowRepository(SientiaMonitoring): end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): - self.debug( - f'Data received from model prediction: {predict_data.to_dict(orient="records")}', - metadata, + self._debug_dataframe( + 'Data received from model prediction:', predict_data, metadata ) # predict_data.to_csv( @@ -1283,6 +1304,10 @@ class MLFlowRepository(SientiaMonitoring): predict_data.columns = pd.Index(['prediction']) else: + self.debug( + f'Data received from model prediction (not a DataFrame): {predict_data}', + metadata, + ) predict_data = pd.DataFrame(predict_data, columns=['prediction']) # predict_data.to_csv( # f"tmp/predicted_data_{model_name}.csv", index=True) @@ -1290,7 +1315,7 @@ class MLFlowRepository(SientiaMonitoring): predict_data.index = input_index predict_data['response_time'] = (end_time - start_time).total_seconds() - return {'success': True, 'content': predict_data.to_dict()} + return {'success': True, 'content': predict_data} except Exception as e: return { @@ -1344,7 +1369,7 @@ class MLFlowRepository(SientiaMonitoring): """ self.info(f'Starting model retraining workflow for {model_name}', metadata) - self.debug(f'Data received for model retraining: {data.to_csv()}', metadata) + self._debug_dataframe('Data received for model retraining:', data, metadata) target_name = model_config.get('target', None) diff --git a/laborious/worker/prepare_worker.py b/laborious/worker/prepare_worker.py index 07ae36f..9766af2 100644 --- a/laborious/worker/prepare_worker.py +++ b/laborious/worker/prepare_worker.py @@ -47,6 +47,7 @@ def prepare_worker( ) logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}') + logger.info(f'Worker runtime config: {local_workflow_parameters}') return Worker( temporal_client, diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 7281254..fa09259 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -152,8 +152,7 @@ async def main(): main_workflow=MinimalRetrain, other_workflows=[], activities=[ - activities.load_custom_query, - activities.query_to_minio, + activities.load_query_with_minio_offload, activities.retrain_model, activities.update_production_model, activities.format_retrain_report, @@ -199,11 +198,11 @@ async def main(): activities.format_transformed_data, activities.format_prediction, activities.format_default_prediction, - activities.get_last_timestamp, # OPC activities.write_opc_data, - # Postgres - activities.load_custom_query, + # Postgres / MinIO offload + activities.load_query_with_minio_offload, + activities.cleanup_minio_objects_expired, activities.repeat_last_prediction, activities.export_data_to_postgres, activities.write_metrics, @@ -220,20 +219,21 @@ async def main(): logger.custom_info('Workers started successfully', metadata) + exit_code = 0 try: # This will run the workers and wait for them to complete. # If an exception occurs in any of the worker handlers, it will be propagated here. await asyncio.gather(*handlers) except BaseException as e: # NOSONAR logger.custom_error(f'An unhandled exception occurred: {e}', metadata) + exit_code = 1 finally: if notification_handler: notification_handler.shutdown() if activities: await activities.shutdown() - # Exit with a non-zero status code to indicate failure to Kubernetes metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN - sys.exit(1) + sys.exit(exit_code) def start_prometheus_server(): diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 769c613..2faff32 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -44,9 +44,9 @@ class Drift: model_id = '{input_data['model_id']}' AND timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' ORDER BY timestamp ASC - """ + """ # nosec B608 - values come from internal Temporal workflow config, not user input - target_data_handler = workflow.start_local_activity_method( + target_data_handler = workflow.start_activity_method( Activities.load_custom_query, { **metadata, @@ -58,7 +58,7 @@ class Drift: start_to_close_timeout=timedelta(seconds=300), ) - reference_data_handler = workflow.start_local_activity_method( + reference_data_handler = workflow.start_activity_method( Activities.get_reference_data, {**metadata, 'model_name': input_data['model_name']}, retry_policy=retry_policy, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 02878cf..f8f5151 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -7,6 +7,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.policies import retry_policy from laborious.activities.activities import Activities + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload @workflow.defn(name='minimal_retrain') @@ -73,26 +74,26 @@ class MinimalRetrain: model_config = input_data.get('model_config', {}) storage_result = await workflow.execute_activity_method( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], '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=600), ) - if not storage_result['success']: - return + storage_payload = MinioDataFramePayload.from_dict(storage_result) + if not storage_payload.has_data(): + raise ValueError('No data returned from query') experiment_response = await workflow.execute_activity_method( Activities.retrain_model, { **metadata, - 'object_key': storage_result['object_key'], + 'data': storage_result, 'model_name': model_name, 'model_config': model_config, }, diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index 42be048..e68c215 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -83,13 +83,14 @@ class PredictionsBatch: } } - # Load data using custom query - data = await workflow.execute_local_activity_method( - Activities.load_custom_query, + # Load data using custom query with optional MinIO offload for large frames + data = await workflow.execute_activity_method( + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), @@ -104,16 +105,19 @@ class PredictionsBatch: 'transform_table_name': input_data['transform_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'}}), + 'input_filters': input_data.get( + 'input_filters', {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}} + ), 'mlflow_transform_filters': input_data.get( - 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}} + 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}} ), 'mlflow_predict_filters': input_data.get( - 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}} + 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}} ), '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', {}), + 'on_conflict': input_data.get('on_conflict', 'error'), 'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}), 'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'), 'save_transform': input_data.get('save_transform', True), diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index 438b98d..38b47c4 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -45,9 +45,9 @@ class SimpleMetrics: p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes' order by p."timestamp" desc; - """ + """ # nosec B608 - values come from internal Temporal workflow config, not user input - target_data = await workflow.execute_local_activity_method( + target_data = await workflow.execute_activity_method( Activities.load_custom_query, { **metadata, diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 684f7a5..68552b2 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -105,6 +105,7 @@ class FormatAndExportPrediction: 'model_id': input_data['model_id'], 'prediction_confidence': prediction_confidence, 'prediction_store_policy': input_data['prediction_store_policy'], + 'model_name': input_data['model_name'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=60), @@ -118,13 +119,14 @@ class FormatAndExportPrediction: **metadata, 'data': transformed_data, 'model_id': input_data['model_id'], + 'model_name': input_data['model_name'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=60), ) write_transformed_handler = workflow.start_activity_method( - Activities.export_data_to_postgres, + Activities.export_payload_to_postgres, { **metadata, 'schema': input_data['schema'], @@ -196,6 +198,8 @@ class FormatAndExportPrediction: 'table_name': input_data['table_name'], 'data': prediction, 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + 'on_conflict': input_data.get('on_conflict', 'error'), + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=180), diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 991d262..da4125b 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -87,13 +87,42 @@ class PredictionProcess: model_config = input_data.get('model_config', {}) save_transform = input_data.get('save_transform', True) - # Get last timestamp for incremental processing - last_timestamp = await workflow.execute_local_activity_method( - Activities.get_last_timestamp, - {**metadata, 'data': data}, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) + try: + await self._run_prediction_pipeline( + input_data, + metadata, + data, + model_id, + model_name, + model_config, + save_transform, + ) + await workflow.execute_activity_method( + Activities.cleanup_minio_objects_expired, + {**metadata, 'data': data}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + except Exception as e: + await workflow.execute_activity_method( + Activities.cleanup_minio_objects_expired, + {**metadata, 'data': data}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + raise e + + async def _run_prediction_pipeline( + self, + input_data: dict[str, Any], + metadata: dict[str, Any], + data: dict[str, Any], + model_id: str, + model_name: str, + model_config: dict[str, Any], + save_transform: bool, + ) -> None: + last_timestamp = data['last_timestamp'] # Apply input data quality gates gate_input = { @@ -117,7 +146,7 @@ class PredictionProcess: return # Request MLFlow model transformation - response_data = await workflow.execute_local_activity_method( + transformed_data = await workflow.execute_activity_method( Activities.request_transform, {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config}, retry_policy=retry_policy, @@ -130,7 +159,7 @@ class PredictionProcess: { **metadata, 'filters': input_data['mlflow_transform_filters'], - 'data': response_data, + 'data': transformed_data, 'type': 'transform', 'path_priority': input_data['path_priority'], }, @@ -144,8 +173,6 @@ class PredictionProcess: ): return - transformed_data = response_data['content'] - path_flag, confidence, comment = await workflow.execute_local_activity_method( Activities.mlflow_content_gate, { @@ -164,7 +191,7 @@ class PredictionProcess: ): return - response_data = await workflow.execute_local_activity_method( + predicted_data = await workflow.execute_activity_method( Activities.request_predict, { **metadata, @@ -182,7 +209,7 @@ class PredictionProcess: { **metadata, 'filters': input_data['mlflow_predict_filters'], - 'data': response_data, + 'data': predicted_data, 'type': 'predict', 'path_priority': input_data['path_priority'], }, @@ -201,8 +228,9 @@ class PredictionProcess: 'subworkflow.format_and_export_prediction', { 'metadata': metadata, + 'on_conflict': input_data.get('on_conflict', 'error'), 'path_flag': path_flag, - 'data': response_data['content'], + 'data': predicted_data, 'transformed_data': transformed_data if save_transform else None, 'prediction_confidence': confidence, 'timestamp': last_timestamp, @@ -221,7 +249,7 @@ class PredictionProcess: async def path_flag_handler( self, - data: dict, + data: dict[str, Any], path_flag: str, input_data: dict, confidence: int, @@ -310,6 +338,7 @@ class PredictionProcess: 'opc_output_config': input_data['opc_output_config'], 'pi_web_api_output_config': input_data['pi_web_api_output_config'], 'prediction_store_policy': input_data['prediction_store_policy'], + 'on_conflict': input_data.get('on_conflict', 'error'), }, ) return True diff --git a/pyproject.toml b/pyproject.toml index daaeabe..5123123 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ exclude = [ "*.pyc", ".pytest_cache", "htmlcov", + "tests/laborious/workflows/subworkflows/test_prediction_process.py", ] [tool.ruff.lint] diff --git a/requirements-dev.txt b/requirements-dev.txt index 16a8492..c4f4ef6 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -13,7 +13,7 @@ types-requests>=2.31.0 # Type stubs for requests 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) -testcontainers[postgres] # PostgreSQL containers for E2E tests +testcontainers[postgres,minio] # PostgreSQL and MinIO containers for E2E tests # Development Tools ipython>=8.12.0 # Enhanced Python shell diff --git a/requirements-light.txt b/requirements-light.txt index 55788d0..91a9c33 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.8.0 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.4 prometheus-client botocore boto3 diff --git a/requirements.txt b/requirements.txt index c9f111c..6c5c242 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,8 +3,8 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.8.2 -git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.7 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.4 +git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.41.0 prometheus-client botocore boto3 @@ -13,4 +13,6 @@ pyarrow kaleido hyperopt shap -pycurl \ No newline at end of file +pycurl +scipy<1.14.0 +scikit-learn==1.5.2 \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index c0fad01..7671dce 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,49 @@ +import os +import sys +from unittest.mock import MagicMock + +# The production code converts SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES to int at import-time. +# Tests must set it to a valid integer string to avoid import errors. +os.environ.setdefault('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1') + + +class DummyMinioDataFramePayload: + """ + Minimal payload double used by unit tests. + + The production workflow/gates expect a MinioDataFramePayload-like object with: + - async retrieve(minio_repo, workflow_metadata) -> DataFrame | dict + - has_data() -> bool + - cleanup_prefix() -> str | None + - last_timestamp: attribute + - status: attribute + """ + + def __init__( + self, + *, + retrieve_return=None, + has_data: bool = True, + cleanup_prefix: str | None = None, + last_timestamp: str = '2024-01-01', + status: dict | None = None, + ): + self._retrieve_return = retrieve_return + self._has_data = has_data + self._cleanup_prefix = cleanup_prefix + self.last_timestamp = last_timestamp + self.status = status + + async def retrieve(self, _minio_repo, _workflow_metadata=None): + return self._retrieve_return + + def has_data(self) -> bool: + return self._has_data + + def cleanup_prefix(self) -> str | None: + return self._cleanup_prefix + + """ Pytest configuration file with global mocks for external dependencies. @@ -6,9 +52,6 @@ during unit tests. The mock is registered in sys.modules before any test imports are executed. """ -import sys -from unittest.mock import MagicMock - # Mock sientia module sientia_mock = MagicMock() sientia_mock.ModelAnalysis = MagicMock diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 5c979e9..8e6cd78 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -17,9 +17,11 @@ from laborious.activities.storage import Storage @patch('laborious.activities.activities.Gates.__init__') @patch('laborious.activities.activities.ModelMetrics.__init__') @patch('laborious.activities.activities.API.__init__') +@patch('laborious.activities.activities.MinioRepository') @patch('laborious.activities.activities.MetricsController') def test___init__( mock_metrics_controller, + mock_minio_repository, mock_api_init, mock_model_metrics_init, mock_gates_init, @@ -41,8 +43,9 @@ def test___init__( 'endpoint_url': 'localhost:9000', 'access_key': 'minio', 'secret_key': 'minio123', - 'region_name': 'us-east-1', 'default_bucket': 'test', + 'retention_hours': 24, + 'secure': False, } mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} @@ -89,7 +92,8 @@ def test___init__( dbname=postgres_config['dbname'], min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], - minio_config=minio_config, + retention_hours=minio_config['retention_hours'], + minio_repository=mock_minio_repository.return_value, logger=logger, notification_handler=notification_handler, metrics_controller=mock_metrics_controller.return_value, @@ -101,7 +105,7 @@ def test___init__( mlflow_port=mlflow_config['port'], mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], - minio_config=minio_config, + minio_repository=mock_minio_repository.return_value, logger=logger, notification_handler=notification_handler, metrics_controller=mock_metrics_controller.return_value, @@ -117,6 +121,7 @@ def test___init__( mock_gates_init.assert_called_once_with( ANY, + minio_repository=mock_minio_repository.return_value, logger=logger, notification_handler=notification_handler, metrics_controller=mock_metrics_controller.return_value, @@ -139,6 +144,17 @@ def test___init__( metrics_controller=mock_metrics_controller.return_value, ) + mock_minio_repository.assert_called_once_with( + endpoint=minio_config['endpoint_url'], + access_key=minio_config['access_key'], + secret_key=minio_config['secret_key'], + bucket=minio_config['default_bucket'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=mock_metrics_controller.return_value, + secure=minio_config['secure'], + ) + @mark.asyncio @patch('laborious.activities.activities.Storage') @@ -147,7 +163,9 @@ def test___init__( @patch('laborious.activities.activities.Gates') @patch('laborious.activities.activities.ModelMetrics') @patch('laborious.activities.activities.API') +@patch('laborious.activities.activities.MinioRepository') async def test_shutdown( + _mock_minio_repository, mock_api_init, mock_model_metrics_init, mock_gates_init, @@ -170,8 +188,9 @@ async def test_shutdown( 'endpoint_url': 'localhost:9000', 'access_key': 'minio', 'secret_key': 'minio123', - 'region_name': 'us-east-1', 'default_bucket': 'test', + 'retention_hours': 24, + 'secure': False, } mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} diff --git a/tests/laborious/activities/test_api.py b/tests/laborious/activities/test_api.py index 52c8120..bf5b9f1 100644 --- a/tests/laborious/activities/test_api.py +++ b/tests/laborious/activities/test_api.py @@ -61,6 +61,70 @@ def base_input_data(): } +@patch('laborious.activities.api.PIWebAPIClient') +def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client): + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + + api_instance = API( + base_url='https://test-pi-server.com', + auth_type='bearer', + auth_token='test_token', + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + with patch.object( + SientiaMonitoring, + 'get_core_labels', + return_value={ + 'pod_id': 'test_pod', + 'runtime': 'local', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'operation_type': 'write_pi_web_api_data', + }, + ): + labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata']) + assert labels['operation_type'] == 'write_pi_web_api_data' + assert labels == { + 'pod_id': 'test_pod', + 'runtime': 'local', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'operation_type': 'write_pi_web_api_data', + } + + +@patch('laborious.activities.api.PIWebAPIClient') +def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client): + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + + api_instance = API( + base_url='https://test-pi-server.com', + auth_type='bearer', + auth_token='test_token', + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + with patch.object( + SientiaMonitoring, + 'get_core_labels', + return_value={ + 'pod_id': 'test_pod', + 'runtime': 'k8s', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'operation_type': 'write', + }, + ): + labels = api_instance.get_pi_web_api_core_labels( + metadata=metadata['metadata'], operation_type='write' + ) + assert labels['operation_type'] == 'write' + assert labels['runtime'] == 'k8s' + + def test__init__(): api = API( base_url='https://test-pi-server.com', @@ -98,6 +162,7 @@ def api(mock_pi_web_api_client): api_instance.get_core_labels = MagicMock( return_value={ 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -280,6 +345,7 @@ async def test_process_pi_web_api_response_success(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -313,6 +379,7 @@ async def test_process_pi_web_api_response_with_errors(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -341,6 +408,7 @@ async def test_process_pi_web_api_response_missing_tags(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -373,6 +441,7 @@ async def test_process_pi_web_api_response_missing_webid(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -401,6 +470,7 @@ async def test_process_pi_web_api_response_missing_tag_name(api): tags = {'tag1': 'web_id_1'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index e33ea4f..c144068 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -1,11 +1,37 @@ from unittest.mock import ANY, AsyncMock, MagicMock, call, patch +from pandas import DataFrame from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from laborious.activities.gates import Gates +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.activities.gates.MinioDataFramePayload.from_dict', side_effect=lambda x: x + ): + yield + + +def _minio_payload(retrieve_return, status=None): + """ + Build a MinioDataFramePayload-like test double with async retrieve. + + Args: + retrieve_return: Value returned from await retrieve(minio_repo, metadata). + status: Optional status dict for MLflow response gate (payload.status). + + Return: + MagicMock: Object with async retrieve and optional status. + """ + p = MagicMock() + p.retrieve = AsyncMock(return_value=retrieve_return) + p.status = status + return p + + @fixture def gates_activity(): gates = Gates( @@ -40,7 +66,7 @@ async def test_input_gate_invalid_filter(gates_activity): input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -64,8 +90,8 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac ) input_data = { **metadata, - 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': []}, + 'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'data': _minio_payload(DataFrame({'value': []})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -77,7 +103,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac gates_activity.send_notification_async.assert_called_once_with( metadata=metadata['metadata'], notification_id='INTPUT_GATE_ERROR__EMPTY_DATA', - message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error", + message="Error in filter EMPTY_DATA:{'POLICY': 'STOP', 'CONFIG': {}}: \n Test error", block='input_gate', level=NotificationLevel.ERROR, attachment_content=ANY, @@ -90,7 +116,7 @@ async def test_input_gate_no_filters(gates_activity): input_data = { **metadata, 'filters': {}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } @@ -107,8 +133,8 @@ async def test_input_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': []}, + 'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'data': _minio_payload(DataFrame({'value': []})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -121,12 +147,46 @@ async def test_input_gate_with_filter(gates_activity): @mark.asyncio -async def test_input_gate_with_filter_not_caught(gates_activity): +async def test_input_gate_with_filter_lowercase_keys(gates_activity): # Arrange input_data = { **metadata, 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': []})), + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + # Act + result = await gates_activity.input_gate(input_data) + + # Assert + assert result == ('STOP', -1, 'Input data with bad quality') + + +@mark.asyncio +async def test_input_gate_with_filter_capitalized_keys(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'EMPTY_DATA': {'Policy': 'STOP', 'Config': {}}}, + 'data': _minio_payload(DataFrame({'value': []})), + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + # Act + result = await gates_activity.input_gate(input_data) + + # Assert + assert result == ('STOP', -1, 'Input data with bad quality') + + +@mark.asyncio +async def test_input_gate_with_filter_not_caught(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -144,7 +204,10 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity): input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -169,7 +232,10 @@ async def test_mlflow_response_gate_filter_exception( input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -195,7 +261,10 @@ async def test_mlflow_response_gate_no_filters(gates_activity): input_data = { **metadata, 'filters': {}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } @@ -213,11 +282,11 @@ async def test_mlflow_response_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'API_ERROR': {'policy': 'STOP'}}, - 'data': { - 'success': False, - 'content': {'message': 'API error occurred', 'traceback': 'error trace'}, - }, + 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, + 'data': _minio_payload( + {'content': {'message': 'API error occurred', 'traceback': 'error trace'}}, + status={'success': False, 'message': 'API error occurred'}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -231,16 +300,37 @@ async def test_mlflow_response_gate_with_filter(gates_activity): gates_activity.send_notification_async.assert_called() +@mark.asyncio +async def test_mlflow_response_gate_with_filter_capitalized_keys(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'API_ERROR': {'Policy': 'STOP'}}, + 'data': _minio_payload( + {'content': {'message': 'API error occurred', 'traceback': 'error trace'}}, + status={'success': False, 'message': 'API error occurred'}, + ), + 'type': 'test', + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + # Act + result = await gates_activity.mlflow_response_gate(input_data) + + # Assert + assert result == ('STOP', -1, 'API error occurred') + + @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'}, - }, + 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -259,10 +349,7 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity): input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': { - 'success': True, - 'content': {'message': 'success'}, - }, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -287,10 +374,7 @@ async def test_mlflow_content_gate_filter_exception( input_data = { **metadata, 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, - 'data': { - 'success': False, - 'content': {'message': 'API error occurred', 'traceback': 'error trace'}, - }, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -317,7 +401,7 @@ async def test_mlflow_content_gate_no_filters(gates_activity): input_data = { **metadata, 'filters': {}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } @@ -335,8 +419,8 @@ async def test_mlflow_content_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': [None, None, None]}, + 'filters': {'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'data': _minio_payload(DataFrame({'value': [None, None, None]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -356,7 +440,7 @@ async def test_mlflow_content_gate_with_filter_not_caught(gates_activity): input_data = { **metadata, 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -369,6 +453,22 @@ async def test_mlflow_content_gate_with_filter_not_caught(gates_activity): gates_activity.debug.assert_called() +@mark.asyncio +async def test_mlflow_content_gate_filter_returns_false(gates_activity): + input_data = { + **metadata, + 'filters': {'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), + 'type': 'test', + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + result = await gates_activity.mlflow_content_gate(input_data) + + 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' @@ -430,13 +530,18 @@ async def test_format_prediction_no_timestamp(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': {'2023-05-26 11:12:27': 1}, - 'response_time': {'2023-05-26 11:12:27': 0.1}, - }, + 'data': _minio_payload( + DataFrame( + { + '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', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -457,21 +562,26 @@ async def test_format_prediction_with_timestamp_erl(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': { - '2023-05-26 11:12:27': 1, - '2023-05-26 11:12:28': 2, - '2023-05-26 11:12:29': 3, - }, - 'response_time': { - '2023-05-26 11:12:27': 0.1, - '2023-05-26 11:12:28': 0.2, - '2023-05-26 11:12:29': 0.3, - }, - }, + 'data': _minio_payload( + DataFrame( + { + 'prediction': { + '2023-05-26 11:12:27': 1, + '2023-05-26 11:12:28': 2, + '2023-05-26 11:12:29': 3, + }, + 'response_time': { + '2023-05-26 11:12:27': 0.1, + '2023-05-26 11:12:28': 0.2, + '2023-05-26 11:12:29': 0.3, + }, + } + ) + ), 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'erl:2', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -492,21 +602,26 @@ async def test_format_prediction_with_timestamp_lts(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': { - '2023-05-26 11:12:27': 1, - '2023-05-26 11:12:28': 2, - '2023-05-26 11:12:29': 3, - }, - 'response_time': { - '2023-05-26 11:12:27': 0.1, - '2023-05-26 11:12:28': 0.2, - '2023-05-26 11:12:29': 0.3, - }, - }, + 'data': _minio_payload( + DataFrame( + { + 'prediction': { + '2023-05-26 11:12:27': 1, + '2023-05-26 11:12:28': 2, + '2023-05-26 11:12:29': 3, + }, + 'response_time': { + '2023-05-26 11:12:27': 0.1, + '2023-05-26 11:12:28': 0.2, + '2023-05-26 11:12:29': 0.3, + }, + } + ) + ), 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -527,14 +642,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': _minio_payload( + DataFrame( + { + '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', + 'timestamp': '2023-05-26 11:12:27', } gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1)) @@ -547,76 +671,106 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): @mark.asyncio -async def test_format_transformed_data_single_row(gates_activity): +@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock) +async def test_format_transformed_data_single_row(mock_from_dataframe, gates_activity): # Arrange + payload_result = MagicMock() + mock_from_dataframe.return_value = payload_result input_data = { **metadata, - 'data': { - 'var1': {'2023-05-26 11:12:27': 1.0}, - 'var2': {'2023-05-26 11:12:27': 2.0}, - }, + 'data': _minio_payload( + DataFrame( + { + 'var1': {'2023-05-26 11:12:27': 1.0}, + 'var2': {'2023-05-26 11:12:27': 2.0}, + } + ) + ), 'model_id': 'test_model', + 'model_name': 'test_model', } # Act result = await gates_activity.format_transformed_data(input_data) # Assert - assert result['timestamp'] == {0: '2023-05-26 11:12:27', 1: '2023-05-26 11:12:27'} - assert result['variable'] == {0: 'var1', 1: 'var2'} - assert result['value'] == {0: 1.0, 1: 2.0} - assert result['model_id'] == {0: 'test_model', 1: 'test_model'} + assert result is payload_result + mock_from_dataframe.assert_called_once() + kwargs = mock_from_dataframe.call_args.kwargs + assert kwargs['model_name'] == 'test_model' + assert kwargs['operation'] == 'transform' + assert kwargs['workflow_metadata'] == metadata['metadata'] + assert kwargs['minio_repo'] is gates_activity.minio_repository + assert 'dataframe' in kwargs gates_activity.info.assert_called() @mark.asyncio -async def test_format_transformed_data_multiple_rows(gates_activity): +@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock) +async def test_format_transformed_data_multiple_rows(mock_from_dataframe, gates_activity): # Arrange + payload_result = MagicMock() + mock_from_dataframe.return_value = payload_result input_data = { **metadata, - 'data': { - 'var1': { - '2023-05-26 11:12:27': 1.0, - '2023-05-26 11:12:28': 2.0, - }, - 'var2': { - '2023-05-26 11:12:27': 3.0, - '2023-05-26 11:12:28': 4.0, - }, - }, + 'data': _minio_payload( + DataFrame( + { + 'var1': { + '2023-05-26 11:12:27': 1.0, + '2023-05-26 11:12:28': 2.0, + }, + 'var2': { + '2023-05-26 11:12:27': 3.0, + '2023-05-26 11:12:28': 4.0, + }, + } + ) + ), 'model_id': 'test_model', + 'model_name': 'test_model', } # Act result = await gates_activity.format_transformed_data(input_data) # Assert - assert len(result['timestamp']) == 4 - assert len(result['variable']) == 4 - assert len(result['value']) == 4 - assert len(result['model_id']) == 4 - assert all(v == 'test_model' for v in result['model_id'].values()) - assert set(result['variable'].values()) == {'var1', 'var2'} + assert result is payload_result + mock_from_dataframe.assert_called_once() + kwargs = mock_from_dataframe.call_args.kwargs + assert kwargs['model_name'] == 'test_model' + assert kwargs['operation'] == 'transform' + assert kwargs['workflow_metadata'] == metadata['metadata'] + assert kwargs['minio_repo'] is gates_activity.minio_repository + assert 'dataframe' in kwargs gates_activity.info.assert_called() @mark.asyncio -async def test_format_transformed_data_empty_data(gates_activity): +@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock) +async def test_format_transformed_data_empty_data(mock_from_dataframe, gates_activity): # Arrange + payload_result = MagicMock() + mock_from_dataframe.return_value = payload_result input_data = { **metadata, - 'data': {}, + 'data': _minio_payload(DataFrame()), 'model_id': 'test_model', + 'model_name': 'test_model', } # Act result = await gates_activity.format_transformed_data(input_data) # Assert - assert result['timestamp'] == {} - assert result['variable'] == {} - assert result['value'] == {} - assert result['model_id'] == {} + assert result is payload_result + mock_from_dataframe.assert_called_once() + kwargs = mock_from_dataframe.call_args.kwargs + assert kwargs['model_name'] == 'test_model' + assert kwargs['operation'] == 'transform' + assert kwargs['workflow_metadata'] == metadata['metadata'] + assert kwargs['minio_repo'] is gates_activity.minio_repository + assert 'dataframe' in kwargs gates_activity.info.assert_called() @@ -711,31 +865,6 @@ async def test_format_retrain_report_failure(gates_activity): gates_activity.debug.assert_called() -@mark.asyncio -async def test_get_last_timestamp_with_data(gates_activity): - # Arrange - input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}} - - # Act - result = await gates_activity.get_last_timestamp(input_data) - - # Assert - assert result == '2023-05-26 11:12:28' - - -@mark.asyncio -async def test_get_last_timestamp_no_data(gates_activity): - # Arrange - input_data = {'data': {}, **metadata} - - # Act - result = await gates_activity.get_last_timestamp(input_data) - - # Assert - assert isinstance(result, str) # Should be a timestamp string - assert len(result) > 0 - - @mark.asyncio @patch('laborious.activities.gates.metrics') async def test_write_metrics(mock_metrics, gates_activity): @@ -750,15 +879,18 @@ async def test_write_metrics(mock_metrics, gates_activity): 'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}}, } await gates_activity.write_metrics(input_data) + core_tags = { + 'pod_id': gates_activity.pod_id, + 'runtime': gates_activity.runtime, + 'operation_type': 'predict', + 'model_name': metadata['metadata']['model_name'], + 'workflow_name': metadata['metadata']['workflow_name'], + } gates_activity.emit_metric.assert_has_calls( [ call( metric_object=mock_metrics.PREDICTIONS_WRITTEN_COUNT, - tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], - }, + tags=core_tags, ), ] ) @@ -767,11 +899,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_CONFIDENCE_MONITOR, method='set', - tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], - }, + tags=core_tags, value=0.9, ), ] @@ -781,11 +909,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR, method='observe', - tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], - }, + tags=core_tags, value=0.1, ), ] @@ -795,9 +919,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT, tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag1', }, @@ -810,9 +932,7 @@ async def test_write_metrics(mock_metrics, gates_activity): metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, method='observe', tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag1', }, @@ -825,9 +945,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT, tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag2', }, @@ -840,9 +958,7 @@ async def test_write_metrics(mock_metrics, gates_activity): metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, method='observe', tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag2', }, @@ -873,6 +989,8 @@ async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_act method='observe', tags={ 'pod_id': gates_activity.pod_id, + 'runtime': gates_activity.runtime, + 'operation_type': 'predict', 'model_name': metadata['metadata']['model_name'], 'workflow_name': metadata['metadata']['workflow_name'], 'opc_server_id': 'server1', diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index dd7a0c5..5d242e5 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -8,24 +8,38 @@ from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_ from laborious.activities.mlflow import MLFlow +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x + ): + yield + + @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def test___init__(mock_minio_repository, mock_mlflow_repository): + logger = MagicMock() + notification_handler = MagicMock() + metrics_controller = AsyncMock() + minio_repo = mock_minio_repository( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + bucket='test', + ) mlflow = MLFlow( mlflow_host='http://localhost', mlflow_port=5000, 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(), - metrics_controller=AsyncMock(), + minio_repository=minio_repo, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, ) assert mlflow.mlflow_host == 'http://localhost' @@ -38,14 +52,13 @@ def test___init__(mock_minio_repository, mock_mlflow_repository): ) mock_minio_repository.assert_called_once_with( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', 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', metrics_controller=ANY, + bucket='test', ) @@ -53,21 +66,27 @@ def test___init__(mock_minio_repository, mock_mlflow_repository): @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def mlflow(mock_minio_repository, mock_mlflow_repository): + logger = MagicMock() + notification_handler = MagicMock() + metrics_controller = AsyncMock() + minio_repo = mock_minio_repository( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + bucket='test', + ) mlflow = MLFlow( mlflow_host='http://localhost:5000', mlflow_port=5000, 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(), - metrics_controller=AsyncMock(), + minio_repository=minio_repo, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, ) mlflow.model_monitoring_repository = AsyncMock() @@ -96,139 +115,147 @@ metadata = { @mark.asyncio -@patch('laborious.activities.mlflow.DataFrame') -@patch('laborious.activities.mlflow.max') -async def test_request_transform_success(mock_max, mock_dataframe, mlflow): - mock_max.return_value = '2024-01-02' - # Mock input data +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, +) +async def test_request_transform_success(mock_from_dataframe, mlflow): + data_mock = MagicMock() + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + 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', - }, - ], + 'data': payload, 'model_name': 'test_model', 'model_config': {}, } - # Mock the transform response - expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']} - mlflow.model_monitoring_repository.transform.return_value = expected_response + transform_response = {'success': True, 'content': MagicMock()} + mlflow.model_monitoring_repository.transform.return_value = transform_response - mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value - mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value + data_mock.sort_values.return_value = data_mock + data_mock.drop_duplicates.return_value = data_mock + data_mock.pivot.return_value = data_mock - # Call the method response_data = await mlflow.request_transform(input_data) - # Verify the data was correctly transformed - mock_dataframe.assert_called_once_with(input_data['data']) - mock_dataframe.return_value.pivot.assert_called_once_with( - index='timestamp', columns='variable', values='value' - ) - mock_dataframe = mock_dataframe.return_value.pivot.return_value - mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True) - # mock_dataframe.reset_index.assert_called_once() - mock_dataframe.columns.name = None - - # Verify the response - assert response_data == expected_response - - # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.transform.assert_called_once_with( - 'test_model', mock_dataframe, {}, metadata['metadata'] + 'test_model', data_mock, {}, metadata['metadata'] ) + mock_from_dataframe.assert_called_once() + assert response_data == mock_from_dataframe.return_value @mark.asyncio -@patch('laborious.activities.mlflow.DataFrame') -@patch('laborious.activities.mlflow.to_datetime') -@patch('laborious.activities.mlflow.max') -async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow): - mock_max.return_value = '2024-01-02' - # Mock input data +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, +) +async def test_request_transform_failure(mock_from_dataframe, mlflow): + data_mock = MagicMock() + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + input_data = { **metadata, - 'data': { - 'variable': { - '2024-01-01': 'var1', - '2024-01-02': 'var2', - '2024-01-03': 'var1', - '2024-01-04': 'var2', - }, - 'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0}, - }, + 'data': payload, 'model_name': 'test_model', 'model_config': {}, } - # Mock the predict response - expected_response = {'prediction': [0.5, 0.6]} - mlflow.model_monitoring_repository.predict.return_value = expected_response + transform_response = {'success': False, 'message': 'Transform failed'} + mlflow.model_monitoring_repository.transform.return_value = transform_response + + data_mock.sort_values.return_value = data_mock + data_mock.drop_duplicates.return_value = data_mock + data_mock.pivot.return_value = data_mock + + response_data = await mlflow.request_transform(input_data) + + mock_from_dataframe.assert_called_once_with( + dataframe=None, + minio_repo=mlflow.minio_repository, + model_name='test_model', + operation='transform', + status=transform_response, + workflow_metadata=metadata['metadata'], + last_timestamp=payload.last_timestamp, + logger=mlflow.logger, + ) + assert response_data == mock_from_dataframe.return_value + + +@mark.asyncio +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, +) +@patch('laborious.activities.mlflow.to_datetime') +async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow): + data_mock = MagicMock() + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + + input_data = { + **metadata, + 'data': payload, + 'model_name': 'test_model', + 'model_config': {}, + } + + predict_response = {'success': True, 'content': MagicMock()} + mlflow.model_monitoring_repository.predict.return_value = predict_response - # Call the method response_data = await mlflow.request_predict(input_data) - mock_dataframe.assert_called_once_with(input_data['data']) - mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True) - mock_dataframe.return_value.__setitem__.assert_any_call( - 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value - ) - mock_dataframe.return_value.__setitem__.assert_any_call( - 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value - ) - + data_mock.replace.assert_called_once_with(np.nan, None, inplace=True) mock_to_datetime.assert_called_once_with( - mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ + data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) 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 - - # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.predict.assert_called_once_with( - 'test_model', mock_dataframe.return_value, {}, metadata['metadata'] + 'test_model', data_mock, {}, metadata['metadata'] ) + mock_from_dataframe.assert_called_once() + assert response_data == mock_from_dataframe.return_value + + +@mark.asyncio +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, +) +@patch('laborious.activities.mlflow.to_datetime') +async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow): + data_mock = MagicMock() + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + + input_data = { + **metadata, + 'data': payload, + 'model_name': 'test_model', + 'model_config': {}, + } + + predict_response = {'success': False, 'message': 'Predict failed'} + mlflow.model_monitoring_repository.predict.return_value = predict_response + + response_data = await mlflow.request_predict(input_data) + + mock_from_dataframe.assert_called_once_with( + dataframe=None, + minio_repo=mlflow.minio_repository, + model_name='test_model', + operation='predict', + status=predict_response, + workflow_metadata=metadata['metadata'], + last_timestamp=payload.last_timestamp, + logger=mlflow.logger, + ) + assert response_data == mock_from_dataframe.return_value @mark.asyncio @@ -240,12 +267,14 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl 'message': 'Model retrained successfully.', } - mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock() + raw_data = MagicMock(columns=['variable', 'timestamp', 'value']) + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, - 'object_key': 'test_object_key', + 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', @@ -255,8 +284,6 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl } ) - 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_not_called() @@ -309,21 +336,21 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl @mark.asyncio @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow): +async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { - 'success': False, - 'traceback': 'test_traceback', - 'message': 'Model retrained failed.', + 'success': True, + 'experiment': 'test_experiment', + 'message': 'Model retrained successfully.', } - mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock( - columns=['variable', 'timestamp', 'value', 'created_at'] - ) + raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at']) + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, - 'object_key': 'test_object_key', + 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', @@ -333,7 +360,35 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow) } ) - raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value + assert response['success'] is True + mlflow.minio_repository.download_file.assert_not_called() + + +@mark.asyncio +@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.', + } + + raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at']) + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=raw_data) + + response = await mlflow.retrain_model( + { + **metadata, + 'data': payload, + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } + ) timestamp = raw_data.__getitem__.return_value.max.return_value @@ -398,14 +453,9 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow) @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', @@ -417,7 +467,7 @@ async def test_retrain_model_data_error(mlflow): assert response == { 'success': False, - 'message': 'Error loading retrain data: Error loading retrain data', + 'message': "Error loading retrain data: 'data'", 'traceback': ANY, 'timestamp': ANY, } diff --git a/tests/laborious/activities/test_model_metrics.py b/tests/laborious/activities/test_model_metrics.py index f20687d..b2b5e0b 100644 --- a/tests/laborious/activities/test_model_metrics.py +++ b/tests/laborious/activities/test_model_metrics.py @@ -743,6 +743,101 @@ async def test_get_drift_metrics_univariate_error( raise AssertionError('Expected Exception') +@mark.asyncio +@patch('laborious.activities.model_metrics.to_datetime') +@patch('laborious.activities.model_metrics.time.time') +@patch('laborious.activities.model_metrics.ModelAnalysis') +@patch('laborious.activities.model_metrics.metrics') +async def test_get_drift_metrics_multivariate_error( + mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity +): + mock_time.return_value = 1000.0 + + mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.detect_multivariate_drift.side_effect = Exception( + 'Multivariate drift error' + ) + + reference_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + target_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + reference_columns = reference_data.drop( + columns=['target', 'timestamp'], errors='ignore' + ).columns + + try: + await model_metrics_activity.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name='target', + reference_columns=reference_columns, + drift_metrics=['ks_test'], + chunk_period='min', + metadata=metadata['metadata'], + ) + except Exception as e: + assert str(e) == 'Multivariate drift error' + model_metrics_activity.error.assert_called_once_with( + 'Error detecting multivariate drift: Multivariate drift error', metadata['metadata'] + ) + model_metrics_activity.emit_metric.assert_called_with( + metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY + ) + else: + raise AssertionError('Expected Exception') + + +@mark.asyncio +@patch('laborious.activities.model_metrics.to_datetime') +@patch('laborious.activities.model_metrics.time.time') +@patch('laborious.activities.model_metrics.ModelAnalysis') +@patch('laborious.activities.model_metrics.metrics') +async def test_get_drift_metrics_dataframe_error( + mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity +): + mock_time.return_value = 1000.0 + + mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.get_drift_metrics_dataframe.side_effect = Exception( + 'Dataframe error' + ) + + reference_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + target_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + reference_columns = reference_data.drop( + columns=['target', 'timestamp'], errors='ignore' + ).columns + + try: + await model_metrics_activity.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name='target', + reference_columns=reference_columns, + drift_metrics=['ks_test'], + chunk_period='min', + metadata=metadata['metadata'], + ) + except Exception as e: + assert str(e) == 'Dataframe error' + model_metrics_activity.error.assert_called_once_with( + 'Error getting drift metrics: Dataframe error', metadata['metadata'] + ) + model_metrics_activity.emit_metric.assert_called_with( + metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY + ) + else: + raise AssertionError('Expected Exception') + + @mark.asyncio async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity): # Arrange @@ -987,3 +1082,27 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_me model_metrics_activity.info.assert_called_once_with( "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata'] ) + + +@mark.asyncio +async def test_calculate_simple_metrics_unknown_metric_ignored(model_metrics_activity): + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + } + ) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['unknown_metric', 'rmse'], + 'interval_minutes': 5, + } + + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'rmse' diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index 411e22e..ce99ff7 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -1,4 +1,5 @@ import datetime +import os from unittest.mock import ANY, AsyncMock, MagicMock, patch from pytest import fixture, mark, raises @@ -7,6 +8,15 @@ from sientia_do.temporal.activities.postgres import Postgres from laborious.activities.storage import Storage + +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x + ): + yield + + metadata = { 'metadata': { 'model_id': 'test_model_id', @@ -28,13 +38,8 @@ def storage(mock_minio_repository): 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', - }, + retention_hours=24, + minio_repository=mock_minio_repository.return_value, logger=MagicMock(), notification_handler=MagicMock(), metrics_controller=AsyncMock(), @@ -46,6 +51,7 @@ def test___init___not_hasattr(mock_minio_repository): logger = MagicMock() notification_handler = MagicMock() metrics_controller = AsyncMock() + minio_repo = mock_minio_repository.return_value storage = Storage( host='localhost', port=5432, @@ -54,29 +60,16 @@ def test___init___not_hasattr(mock_minio_repository): 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', - }, + retention_hours=24, + minio_repository=minio_repo, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) 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', - metrics_controller=metrics_controller, - ) + assert storage.minio_repository is minio_repo + mock_minio_repository.assert_not_called() @patch('laborious.activities.storage.MinioRepository') @@ -93,28 +86,15 @@ def test___init___none_minio_repository(mock_minio_repository, storage): 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', - }, + retention_hours=24, + minio_repository=None, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) - 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', - metrics_controller=metrics_controller, - ) + assert storage.minio_repository is None + mock_minio_repository.assert_not_called() @patch('laborious.activities.storage.MinioRepository') @@ -127,13 +107,8 @@ def test___init___done_repository(mock_minio_repository, storage): 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', - }, + retention_hours=24, + minio_repository=mock_minio_repository.return_value, logger=MagicMock(), notification_handler=MagicMock(), metrics_controller=AsyncMock(), @@ -142,72 +117,6 @@ def test___init___done_repository(mock_minio_repository, storage): 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.store_dataframe_as_parquet = AsyncMock() - 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.send_notification_async = AsyncMock() - storage.minio_repository.store_dataframe_as_parquet = AsyncMock() - - 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_async.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() @@ -222,3 +131,193 @@ def test___del__(storage): storage.__del__() storage.close.assert_called_once() + + +@mark.asyncio +async def test_load_query_with_minio_offload_no_rows(storage): + storage.load_custom_query = AsyncMock(return_value=None) + storage_result = {'success': False} + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, + return_value=storage_result, + ) as mock_from_dataframe: + result = await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} + ) + assert result == storage_result + mock_from_dataframe.assert_awaited_once() + + +@mark.asyncio +async def test_load_query_with_minio_offload_inline(storage): + storage.load_custom_query = AsyncMock(return_value=[{'a': 1}]) + storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None} + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, + return_value=storage_result, + ) as mock_from_dataframe: + result = await storage.load_query_with_minio_offload( + { + **metadata, + 'query': 'SELECT 1', + 'model_name': 'my-model', + 'key_prefix': 'predictions/s', + } + ) + assert result == storage_result + mock_from_dataframe.assert_awaited_once() + + +@mark.asyncio +async def test_load_query_with_minio_offload_minio(storage): + storage.load_custom_query = AsyncMock(return_value=[{'a': 1}]) + storage_result = {'success': True, 'data': None, 'object_key': 'object-key'} + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, + return_value=storage_result, + ) as mock_from_dataframe: + result = await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} + ) + + assert result == storage_result + mock_from_dataframe.assert_awaited_once() + + +@mark.asyncio +@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'}) +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + storage.minio_repository.list_objects = AsyncMock( + return_value=[ + 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet', + 'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet', + ] + ) + storage.minio_repository.delete_file = AsyncMock() + storage.send_notification_async = AsyncMock() + + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'training_datasets/m' + + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) + + assert result['deleted_count'] == 1 + assert result['failed_count'] == 0 + deleted_key = ( + 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet' + ) + assert deleted_key in result['deleted'] + assert result['deleted'][deleted_key]['success'] is True + storage.minio_repository.list_objects.assert_called_once_with( + prefix='training_datasets/m', + recursive=True, + metadata=metadata['metadata'], + ) + storage.minio_repository.delete_file.assert_called_once_with( + object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet', + metadata=metadata['metadata'], + ) + + +@mark.asyncio +async def test_load_query_with_minio_offload_minio_not_initialized(storage): + storage.minio_repository = None + + with raises(ValueError, match='Minio repository not initialized'): + await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm'} + ) + + +@mark.asyncio +async def test_export_payload_to_postgres(storage): + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=MagicMock()) + storage.export_data_to_postgres = AsyncMock(return_value={'success': True}) + + result = await storage.export_payload_to_postgres( + {**metadata, 'data': payload, 'schema': 'public', 'table': 't'} + ) + + payload.retrieve.assert_awaited_once_with(storage.minio_repository, metadata['metadata']) + storage.export_data_to_postgres.assert_awaited_once() + assert result == {'success': True} + + +@mark.asyncio +async def test_cleanup_minio_objects_expired_minio_not_initialized(storage): + storage.minio_repository = None + + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'test' + with raises(ValueError, match='Minio repository not initialized'): + await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) + + +@mark.asyncio +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + storage.minio_repository.list_objects = AsyncMock( + return_value=['some/random/key-without-timestamp.parquet'] + ) + storage.minio_repository.delete_file = AsyncMock() + storage.send_notification_async = AsyncMock() + + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'test' + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) + + assert result['deleted_count'] == 0 + assert result['failed_count'] == 0 + storage.minio_repository.delete_file.assert_not_called() + + +@mark.asyncio +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + old_key = 'training_datasets/m/m-initial-2024-12-01_00-00-00.parquet' + storage.minio_repository.list_objects = AsyncMock(return_value=[old_key]) + storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error')) + storage.send_notification_async = AsyncMock() + + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'training_datasets/m' + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) + + assert result['deleted_count'] == 0 + assert result['failed_count'] == 1 + assert old_key in result['failed'] + assert result['failed'][old_key]['success'] is False + assert result['failed'][old_key]['message'] == 'delete error' + + +@mark.asyncio +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + storage.minio_repository.list_objects = AsyncMock(side_effect=Exception('list error')) + storage.send_notification_async = AsyncMock() + storage.error = MagicMock() + + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'training_datasets/m' + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) + + assert result['deleted_count'] == 0 + assert result['failed_count'] == 0 + storage.send_notification_async.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED', + message='Error cleaning up MinIO objects: list error', + block='cleanup_minio_objects_expired', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + storage.error.assert_called_once() diff --git a/tests/laborious/utils/models/test_minio_dataframe_payload.py b/tests/laborious/utils/models/test_minio_dataframe_payload.py new file mode 100644 index 0000000..c27765d --- /dev/null +++ b/tests/laborious/utils/models/test_minio_dataframe_payload.py @@ -0,0 +1,266 @@ +from datetime import datetime +from io import BytesIO +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pandas import DataFrame + +from laborious.utils.models.minio_dataframe_payload import ( + MinioDataFramePayload, + _build_object_key, +) + + +def test_parse_object_timestamp_hyphenated_model(): + key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet' + ts = MinioDataFramePayload.parse_object_timestamp(key) + assert ts == datetime(2024, 6, 15, 10, 30, 45) + + +def test_parse_object_timestamp_transform(): + key = 'p/m-transform-2024-01-02_03-04-05.parquet' + ts = MinioDataFramePayload.parse_object_timestamp(key) + assert ts == datetime(2024, 1, 2, 3, 4, 5) + + +def test_parse_object_timestamp_invalid(): + assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None + + +def test_estimate_size_bytes_returns_positive_for_nonempty_frame(): + df = DataFrame({'a': [1, 2]}) + size = MinioDataFramePayload.estimate_size_bytes(df) + assert isinstance(size, int) + assert size > 0 + + +def test_cleanup_prefix_when_offloaded_returns_object_prefix(): + payload = MinioDataFramePayload( + last_timestamp='t', + data=None, + object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet', + object_prefix='training_datasets/m', + ) + assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m' + + +def test_cleanup_prefix_when_inline_returns_none(): + payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None) + assert MinioDataFramePayload.cleanup_prefix(payload) is None + + +def test_has_data_true_when_object_key_set(): + payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k') + assert payload.has_data() is True + + +@pytest.mark.asyncio +async def test_retrieve_inline_dict_as_dataframe(): + payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]}) + minio = AsyncMock() + out = await payload.retrieve(minio, {'metadata': {}}) + assert list(out.columns) == ['a'] + minio.download_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve_downloads_parquet_when_offloaded(): + source = DataFrame({'a': [1, 2]}) + buf = BytesIO() + source.to_parquet(buf, engine='pyarrow', index=True) + file_bytes = buf.getvalue() + + payload = MinioDataFramePayload( + last_timestamp='t', + data=None, + object_key='training_datasets/m/f.parquet', + object_prefix='training_datasets/m', + ) + minio = AsyncMock() + minio.download_file = AsyncMock(return_value=file_bytes) + + out = await payload.retrieve(minio, {'metadata': {}}) + + minio.download_file.assert_awaited_once_with( + object_name='training_datasets/m/f.parquet', + metadata={'metadata': {}}, + ) + assert list(out.columns) == ['a'] + + +def test_build_object_key(): + key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00') + assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet' + assert prefix == 'prediction_datasets/my-model' + + +def test_build_object_key_strips_slashes(): + key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45') + assert prefix == 'prediction_datasets/my-model' + assert key.startswith('prediction_datasets/my-model/') + + +def test_estimate_size_bytes_fallback(): + df = DataFrame({'a': [1, 2]}) + with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')): + size = MinioDataFramePayload.estimate_size_bytes(df) + assert isinstance(size, int) + assert size > 0 + + +def test_parse_object_timestamp_bad_datetime(): + key = 'p/m-initial-9999-99-99_99-99-99.parquet' + assert MinioDataFramePayload.parse_object_timestamp(key) is None + + +@pytest.mark.asyncio +async def test_retrieve_empty_when_no_data(): + payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None) + minio = AsyncMock() + out = await payload.retrieve(minio, {}) + assert out.empty + minio.download_file.assert_not_called() + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.now') +async def test_from_dataframe_none(mock_now): + mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0) + minio = AsyncMock() + result = await MinioDataFramePayload.from_dataframe( + dataframe=None, + minio_repo=minio, + model_name='m', + operation='initial', + status={'success': False, 'message': 'no data'}, + ) + assert result.data is None + assert result.status == {'success': False, 'message': 'no data'} + assert result.object_key is None + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.now') +async def test_from_dataframe_empty(mock_now): + mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0) + minio = AsyncMock() + mock_df = MagicMock() + mock_df.__bool__ = MagicMock(return_value=True) + mock_df.empty = True + result = await MinioDataFramePayload.from_dataframe( + dataframe=mock_df, + minio_repo=minio, + model_name='m', + operation='initial', + ) + assert result.data is None + assert result.object_key is None + + +def _mock_dataframe(data_dict, timestamp_values=None): + """Build a MagicMock that behaves enough like a DataFrame for from_dataframe.""" + mock_df = MagicMock() + mock_df.__bool__ = MagicMock(return_value=True) + mock_df.empty = False + if timestamp_values is None: + timestamp_values = data_dict.get('timestamp', ['2024-01-01']) + ts_col = MagicMock() + ts_col.values.tolist.return_value = timestamp_values + mock_df.__getitem__ = MagicMock(return_value=ts_col) + mock_df.to_dict.return_value = data_dict + buf = BytesIO() + DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True) + mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue())) + return mock_df + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9) +async def test_from_dataframe_inline(): + minio = AsyncMock() + df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]}) + result = await MinioDataFramePayload.from_dataframe( + dataframe=df, + minio_repo=minio, + model_name='m', + operation='initial', + ) + assert result.data is not None + assert result.object_key is None + assert result.last_timestamp == '2024-01-01' + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.now') +@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0) +async def test_from_dataframe_offloaded(mock_now): + mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0) + minio = AsyncMock() + minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'}) + minio.bucket = 'test-bucket' + + df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]}) + result = await MinioDataFramePayload.from_dataframe( + dataframe=df, + minio_repo=minio, + model_name='m', + operation='initial', + workflow_metadata={'wf': 'data'}, + ) + assert result.data is None + assert result.object_key == 'full/key.parquet' + assert result.bucket == 'test-bucket' + assert result.uri == 's3://test-bucket/full/key.parquet' + minio.upload_file.assert_awaited_once() + + +def test_from_dict_inline(): + raw = { + 'last_timestamp': '2024-01-01T00:00:00+00:00', + 'status': None, + 'data': {'col1': {0: 'val1'}}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } + payload = MinioDataFramePayload.from_dict(raw) + assert isinstance(payload, MinioDataFramePayload) + assert payload.last_timestamp == '2024-01-01T00:00:00+00:00' + assert payload.data == {'col1': {0: 'val1'}} + assert payload.object_key is None + + +def test_from_dict_offloaded(): + raw = { + 'last_timestamp': '2024-06-15T10:30:45+00:00', + 'status': {'success': True}, + 'data': None, + 'bucket': 'my-bucket', + 'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet', + 'object_prefix': 'training_datasets/model', + 'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet', + } + payload = MinioDataFramePayload.from_dict(raw) + assert isinstance(payload, MinioDataFramePayload) + assert payload.data is None + assert payload.bucket == 'my-bucket' + assert payload.object_key == raw['object_key'] + assert payload.object_prefix == 'training_datasets/model' + assert payload.uri == raw['uri'] + assert payload.status == {'success': True} + + +def test_from_dict_minimal_keys(): + raw = {'last_timestamp': '2024-01-01'} + payload = MinioDataFramePayload.from_dict(raw) + assert payload.last_timestamp == '2024-01-01' + assert payload.data is None + assert payload.bucket is None + assert payload.object_key is None + + +def test_from_dict_passthrough_existing_instance(): + original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b') + result = MinioDataFramePayload.from_dict(original) + assert result is original diff --git a/tests/laborious/utils/repository/test_minio_repository.py b/tests/laborious/utils/repository/test_minio_repository.py deleted file mode 100644 index 89d287e..0000000 --- a/tests/laborious/utils/repository/test_minio_repository.py +++ /dev/null @@ -1,245 +0,0 @@ -from unittest.mock import ANY, AsyncMock, MagicMock, patch - -from botocore.utils import ClientError -from pytest import fixture, mark, raises - -from laborious import metrics -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(), - metrics_controller=AsyncMock(), - ) - - 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): - 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(), - metrics_controller=AsyncMock(), - ) - - minio_repository.emit_metric = AsyncMock() - minio_repository.observe_lag = AsyncMock() - minio_repository.send_notification = MagicMock() - minio_repository.send_notification_async = AsyncMock() - - return minio_repository - - -def test_close(minio_repository): - minio_repository.close() - minio_repository.s3_client.close.assert_called_once() - - -@mark.asyncio -async def test_create_bucket_success(minio_repository): - await minio_repository.create_bucket({}) - - minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_WRITE_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_COUNT, tags=ANY - ) - - -@mark.asyncio -async def test_create_bucket_error(minio_repository): - minio_repository.s3_client.create_bucket.side_effect = ValueError('test') - - with raises(ValueError): - await minio_repository.create_bucket({}) - - minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY - ) - minio_repository.observe_lag.assert_not_called() - - -@mark.asyncio -async def test_ensure_bucket_exists_bucket_exists(minio_repository): - assert await minio_repository.ensure_bucket_exists({}) is None - - minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_READ_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_READ_COUNT, tags=ANY - ) - - -@mark.asyncio -async 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' - ) - minio_repository.create_bucket = AsyncMock() - - assert await minio_repository.ensure_bucket_exists({}) is None - - minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') - minio_repository.create_bucket.assert_called_once_with({}) - - minio_repository.observe_lag.assert_not_called() - minio_repository.emit_metric.assert_not_called() - - -@mark.asyncio -async def test_ensure_bucket_exists_bucket_not_exists_create_error(minio_repository): - minio_repository.s3_client.head_bucket.side_effect = ValueError('test') - - with raises(ValueError): - await minio_repository.ensure_bucket_exists({}) - - minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY - ) - minio_repository.observe_lag.assert_not_called() - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -async def test_store_dataframe_as_parquet_success(mock_bytesio, minio_repository): - input_data = MagicMock() - - minio_repository.ensure_bucket_exists = AsyncMock() - - await 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 - ) - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_WRITE_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_COUNT, tags=ANY - ) - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -async def test_store_dataframe_as_parquet_error(mock_bytesio, minio_repository): - input_data = MagicMock() - - minio_repository.ensure_bucket_exists = AsyncMock() - minio_repository.s3_client.put_object.side_effect = ValueError('test') - - with raises(ValueError): - await 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 - ) - - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY - ) - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -@patch('laborious.utils.repository.minio_repository.read_parquet') -async def test_get_parquet_as_dataframe_success(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 = await 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 - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_READ_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_READ_COUNT, tags=ANY - ) - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -@patch('laborious.utils.repository.minio_repository.read_parquet') -async def test_get_parquet_as_dataframe_error(mock_read_parquet, mock_bytesio, minio_repository): - minio_repository.s3_client.get_object.side_effect = ValueError('test') - - with raises(ValueError): - await 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' - ) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_READ_ERROR_COUNT, tags=ANY - ) - minio_repository.observe_lag.assert_not_called() diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 5311a22..15f8c6b 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -2,6 +2,7 @@ from datetime import UTC, datetime from unittest.mock import ANY, AsyncMock, MagicMock, call, patch import mlflow as mlflow_lib +import numpy as np import pytest from pandas import DataFrame, Timestamp @@ -1310,10 +1311,8 @@ async def test_transform_success(mlflow_repository): 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, - } + assert output['success'] is True + assert output['content'] is mlflow_repository.detect_and_parse_datetime_index.return_value @pytest.mark.asyncio @@ -1355,7 +1354,7 @@ async def test_predict_success_array(mlflow_repository): mlflow_repository.get_cached_operation.assert_called_once_with( model_name=model_name, - data=data, + data=ANY, operation='predict', retention=60, flavor='pyfunc', @@ -1363,10 +1362,12 @@ async def test_predict_success_array(mlflow_repository): ) assert output['success'] is True - assert output['content'] == { - 'prediction': {'index_1': 2, 'index_2': 3}, - 'response_time': {'index_1': ANY, 'index_2': ANY}, - } + content = output['content'] + assert isinstance(content, DataFrame) + assert 'prediction' in content.columns + assert 'response_time' in content.columns + assert list(content.columns) == ['prediction', 'response_time'] + assert content.index.tolist() == data.index.tolist() @pytest.mark.asyncio @@ -1383,7 +1384,7 @@ async def test_predict_success_df(mlflow_repository): mlflow_repository.get_cached_operation.assert_called_once_with( model_name=model_name, - data=data, + data=ANY, operation='predict', retention=60, flavor='pyfunc', @@ -1391,10 +1392,12 @@ async def test_predict_success_df(mlflow_repository): ) assert output['success'] is True - assert output['content'] == { - 'prediction': {'index_1': 2, 'index_2': 3}, - 'response_time': {'index_1': ANY, 'index_2': ANY}, - } + content = output['content'] + assert isinstance(content, DataFrame) + assert 'prediction' in content.columns + assert 'response_time' in content.columns + assert list(content.columns) == ['prediction', 'response_time'] + assert content.index.tolist() == data.index.tolist() @pytest.mark.asyncio @@ -1409,7 +1412,7 @@ async def test_predict_error(mlflow_repository): mlflow_repository.get_cached_operation.assert_called_once_with( model_name=model_name, - data=data, + data=ANY, operation='predict', retention=60, flavor='pyfunc', @@ -1559,3 +1562,119 @@ def test_get_prediction_data_pyfunc(mlflow_repository): assert 'target' in result.columns assert 'timestamp' in result.columns assert result.index.tolist() == [0, 1] + + +@patch('laborious.utils.repository.model_repository.pd.merge') +@patch('laborious.utils.repository.model_repository.isinstance') +@pytest.mark.asyncio +async def test_fit_models_skip_transform(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 = AsyncMock( + 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'])) + ) + ) + mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame()) + + data = MagicMock() + + output = await mlflow_repository.fit_models( + 'model_name', + data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + True, + 'pyfunc', + 'feat_1', + ) + + data_model.fit.assert_not_called() + + assert output['data_model'] == {'model': data_model, 'artifact_path': 'artifact_path'} + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +@pytest.mark.asyncio +async def test_create_new_experiment_path_not_exists( + _rmtree, path, force_memory_release, mlflow, mlflow_repository +): + model_name = 'model_name' + data = MagicMock() + prediction_data = MagicMock(spec=DataFrame) + retrain_data = { + 'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'prediction_data': prediction_data, + } + + 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 = AsyncMock() + path.exists.return_value = False + path.join.return_value = './tmp/artifacts/model_name' + + await mlflow_repository.create_new_experiment( + model_name, + data, + retrain_data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + 'pyfunc', + ) + + _rmtree.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_production_model_by_run_id_transition_error(mlflow, mlflow_repository): + mlflow_repository.client.get_registered_model.return_value = MagicMock( + latest_versions=[ + MagicMock(version='1'), + MagicMock(version='2'), + ] + ) + mlflow_repository.client.transition_model_version_stage.side_effect = Exception( + 'transition error' + ) + + with pytest.raises(Exception, match='transition error'): + await mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) + + mlflow_repository.emit_metric.assert_called_with( + metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY + ) + + +@pytest.mark.asyncio +async def test_predict_success_ndarray(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 = AsyncMock(return_value=np.array([5.0, 6.0])) + + output = await mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + assert output['success'] is True + content = output['content'] + assert isinstance(content, DataFrame) + assert 'prediction' in content.columns + assert 'response_time' in content.columns + assert content.index.tolist() == data.index.tolist() diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 80eb19f..b94bb8d 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -100,8 +100,9 @@ def test_build_minio_config_with_env_vars(): 'endpoint_url': 'http://test-host', 'access_key': 'test-key', 'secret_key': 'test-secret', - 'region_name': 'test-region', 'default_bucket': 'test-bucket', + 'retention_hours': 24, + 'secure': False, } @@ -115,6 +116,7 @@ def test_build_minio_config_with_defaults(): 'endpoint_url': 'http://localhost:9000', 'access_key': 'minioadmin', 'secret_key': 'minioadmin', - 'region_name': 'us-east-1', 'default_bucket': 'laborious', + 'retention_hours': 24, + 'secure': False, } 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 dd95ee9..4fbc62b 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -34,6 +34,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -58,12 +59,13 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): call( Activities.format_prediction, { + **metadata, '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, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -91,6 +93,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -98,7 +101,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -141,6 +145,7 @@ async def test_run_none_path_flag_with_transformed_data( 'transformed_data': {'transformed': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0.9, 'schema': 'test_schema', 'table_name': 'test_table', @@ -176,12 +181,13 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.format_prediction, { + **metadata, '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, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -189,9 +195,10 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.format_transformed_data, { + **metadata, 'data': input_data['transformed_data'], 'model_id': input_data['model_id'], - **metadata, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -201,8 +208,9 @@ async def test_run_none_path_flag_with_transformed_data( # Assert - start_activity_method for transformed data export workflow_mock.start_activity_method.assert_called_once_with( - Activities.export_data_to_postgres, + Activities.export_payload_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['transform_table_name'], 'data': transformed_data, @@ -210,7 +218,6 @@ async def test_run_none_path_flag_with_transformed_data( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -238,6 +245,7 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -245,7 +253,8 @@ async def test_run_none_path_flag_with_transformed_data( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -287,6 +296,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -311,11 +321,11 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction call( Activities.format_default_prediction, { + **metadata, '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, @@ -343,14 +353,16 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, - **metadata, 'timestamp_conversion': { 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -389,6 +401,7 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -415,12 +428,13 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo call( Activities.format_prediction, { + **metadata, '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, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -448,6 +462,7 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': pi_web_api_data, @@ -455,7 +470,8 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -496,6 +512,7 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -551,6 +568,7 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -558,7 +576,8 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -597,6 +616,7 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -638,6 +658,7 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': pi_web_api_data, @@ -645,7 +666,8 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 22df274..a018504 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch from pytest import fixture, mark @@ -17,6 +17,7 @@ metadata = { 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule', + 'schedule_name': 'test_schedule', }, } @@ -26,9 +27,14 @@ metadata = { async def test_run(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=False) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -45,15 +51,17 @@ async def test_run(workflow_mock, prediction_process): } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) ('continue', 0.95, 'Transformed data not passed the content filter'), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict # mlflow_response_gate (predict) ('continue', 0.95, 'Error'), ] @@ -62,21 +70,7 @@ async def test_run(workflow_mock, prediction_process): await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 7 - - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - **metadata, - 'data': input_data['data'], - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) + assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -92,7 +86,7 @@ async def test_run(workflow_mock, prediction_process): ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -130,7 +124,7 @@ async def test_run(workflow_mock, prediction_process): { **metadata, 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', 'path_priority': input_data['path_priority'], }, @@ -139,13 +133,13 @@ async def test_run(workflow_mock, prediction_process): ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_predict, { **metadata, - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, @@ -175,9 +169,10 @@ async def test_run(workflow_mock, prediction_process): 'subworkflow.format_and_export_prediction', { 'metadata': metadata, + 'on_conflict': 'error', 'path_flag': 'continue', - 'data': 'predicted_data', - 'transformed_data': 'transformed_data', + 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + 'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'prediction_confidence': 0.95, 'timestamp': '2024-01-01', 'model_id': 1, @@ -199,9 +194,14 @@ async def test_run(workflow_mock, prediction_process): async def test_run_stop_at_input_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=True) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -217,7 +217,6 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): # 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 ] @@ -225,18 +224,9 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 2 + assert workflow_mock.execute_local_activity_method.call_count == 1 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, { @@ -258,9 +248,14 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): 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 + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -275,10 +270,12 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('repeat', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data ('continue', 0.95, 'Error'), # mlflow_response_gate (transform) ] @@ -286,20 +283,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 4 - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - 'data': input_data['data'], - **metadata, - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) + assert workflow_mock.execute_local_activity_method.call_count == 2 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -315,7 +299,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -354,9 +338,14 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True]) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -371,10 +360,12 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) @@ -385,21 +376,8 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 5 + assert workflow_mock.execute_local_activity_method.call_count == 3 - 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( @@ -415,7 +393,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -452,7 +430,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', 'path_priority': input_data['path_priority'], **metadata, @@ -470,9 +448,14 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True]) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -487,15 +470,17 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) ('continue', 0.95, 'Transformed data not passed the content filter'), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict ('continue', 0.95, 'Error'), # mlflow_response_gate (predict) ] @@ -503,20 +488,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 7 - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - 'data': input_data['data'], - **metadata, - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) + assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -532,7 +504,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -569,7 +541,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', 'path_priority': input_data['path_priority'], **metadata, @@ -579,12 +551,12 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_predict, { - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], **metadata, @@ -768,6 +740,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'confidence_tags': {}, }, 'prediction_store_policy': prediction_store_policy, + 'on_conflict': 'error', }, ) @@ -816,3 +789,54 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): assert result is False workflow_mock.execute_activity_method.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) +async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): + prediction_process.path_flag_handler = AsyncMock(return_value=False) + prediction_process.cleanup_prefixes = {'training_datasets/test'} + + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) + input_data = { + 'metadata': metadata, + 'data': data_payload, + 'schema': 'test_schema', + 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', + 'model_id': 1, + 'input_filters': {'test': 'filter'}, + 'mlflow_transform_filters': {'test': 'filter'}, + 'mlflow_predict_filters': {'test': 'filter'}, + 'model_name': 'test_model_name', + 'model_config': {'retention': '30'}, + 'path_priority': ['continue', 'repeat', 'stop'], + 'opc_output_config': {'test': 'config'}, + 'pi_web_api_output_config': {'test': 'config'}, + 'prediction_store_policy': 'lts:1', + } + + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] + workflow_mock.execute_local_activity_method.side_effect = [ + ('continue', 0.95, 'ok'), + ('continue', 0.95, ''), + ('continue', 0.95, ''), + ('continue', 0.95, ''), + ] + + await prediction_process.run(input_data) + + workflow_mock.execute_activity_method.assert_any_call( + Activities.cleanup_minio_objects_expired, + {**metadata, 'data': data_payload}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 8f013d3..c31dafc 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -45,10 +45,9 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] - - workflow_mock.execute_local_activity_method.return_value = drift_data - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] + workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data) + workflow_mock.execute_activity_method = AsyncMock(return_value=None) # Act await drift.run(input_data) @@ -64,7 +63,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): ORDER BY timestamp ASC """ - workflow_mock.start_local_activity_method.assert_has_calls( + workflow_mock.start_activity_method.assert_has_calls( [ call( Activities.load_custom_query, @@ -140,16 +139,15 @@ async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift): target_data = None reference_data = {'data': 'test_reference_data'} - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] - workflow_mock.execute_local_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock() workflow_mock.execute_activity_method = AsyncMock() # Act await drift.run(input_data) # Assert - Should not call calculate_drift or export - workflow_mock.execute_local_activity_method.assert_not_called() workflow_mock.execute_activity_method.assert_not_called() @@ -175,9 +173,8 @@ async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = None - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] - - workflow_mock.execute_local_activity_method.return_value = drift_data + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] + workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data) workflow_mock.execute_activity_method = AsyncMock() # Act @@ -226,10 +223,9 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] - - workflow_mock.execute_local_activity_method.return_value = drift_data - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] + workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data) + workflow_mock.execute_activity_method = AsyncMock(return_value=None) # Act await drift.run(input_data) diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index bcbad2d..94d86da 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -39,9 +39,19 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): }, } + storage_result = { + 'last_timestamp': '2024-01-01 00:00:00+0000', + 'status': {'success': True}, + 'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } + workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': True, 'object_key': 'test_object_key'}, + storage_result, {'success': True, 'experiment': 'test_experiment'}, { 'success': True, @@ -58,13 +68,12 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): workflow_mock.execute_activity_method.assert_has_calls( [ call( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **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, @@ -78,7 +87,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): Activities.retrain_model, { **metadata, - 'object_key': 'test_object_key', + 'data': storage_result, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, @@ -161,9 +170,19 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim }, } + storage_result = { + 'last_timestamp': '2024-01-01 00:00:00+0000', + 'status': {'success': True}, + 'data': {}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } + workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': False, 'object_key': 'test_object_key'}, + storage_result, {'success': True, 'experiment': 'test_experiment'}, { 'success': True, @@ -175,16 +194,18 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim ] ) - await minimal_retrain.run(input_data) + from pytest import raises + + with raises(ValueError, match='No data returned from query'): + await minimal_retrain.run(input_data) workflow_mock.execute_activity_method.assert_called_once_with( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **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, @@ -211,9 +232,19 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim }, } + storage_result = { + 'last_timestamp': '2024-01-01 00:00:00+0000', + 'status': {'success': True}, + 'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } + workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': True, 'object_key': 'test_object_key'}, + storage_result, {'success': False, 'experiment': 'test_experiment'}, { 'success': True, @@ -230,13 +261,12 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim workflow_mock.execute_activity_method.assert_has_calls( [ call( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **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, @@ -250,7 +280,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim Activities.retrain_model, { **metadata, - 'object_key': 'test_object_key', + 'data': storage_result, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index c19d6ef..14280d6 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch from pytest import fixture, mark @@ -12,19 +12,19 @@ def predictions_batch() -> PredictionsBatch: metadata = { - 'metadata': { - 'model_id': 'test_model_id', - 'model_name': 'test_model', - 'workflow_name': 'predictions_batch', - 'schedule_name': 'test_schedule', - }, + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'predictions_batch', + 'schedule_name': 'test_schedule', } @mark.asyncio @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): - workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'} + activity_return = MagicMock() + workflow_mock.execute_activity_method.return_value = activity_return + input_data = { 'schedule_name': 'test_schedule', 'model_name': 'test_model', @@ -42,14 +42,15 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch await predictions_batch.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.load_query_with_minio_offload, { - **metadata, + 'metadata': metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -57,23 +58,44 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch ] ) prediction_input = { - 'metadata': metadata, - 'data': {'data': 'test_data'}, + 'metadata': {'metadata': metadata}, + 'data': activity_return, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'transform_table_name': input_data['transform_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'}}), + 'input_filters': input_data.get( + 'input_filters', + { + 'EMPTY_DATA': { + 'POLICY': 'STOP', + 'CONFIG': {}, + } + }, + ), 'mlflow_transform_filters': input_data.get( - 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}} + 'mlflow_transform_filters', + { + 'API_ERROR': { + 'POLICY': 'STOP', + 'CONFIG': {}, + } + }, ), 'mlflow_predict_filters': input_data.get( - 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}} + 'mlflow_predict_filters', + { + 'API_ERROR': { + 'POLICY': 'STOP', + 'CONFIG': {}, + } + }, ), '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', {}), + 'on_conflict': input_data.get('on_conflict', 'error'), 'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}), 'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'), 'save_transform': input_data.get('save_transform', True), diff --git a/tests/laborious/workflows/test_simple_metrics.py b/tests/laborious/workflows/test_simple_metrics.py index 0d96bc5..dbda735 100644 --- a/tests/laborious/workflows/test_simple_metrics.py +++ b/tests/laborious/workflows/test_simple_metrics.py @@ -42,9 +42,8 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): target_data = {'data': 'test_target_data'} simple_metrics_data = {'metrics': 'test_simple_metrics_data'} - workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] - - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None]) + workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data) # Act await simple_metrics.run(input_data) @@ -66,7 +65,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): p."timestamp" desc; """ - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.load_custom_query, @@ -80,29 +79,30 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): start_to_close_timeout=ANY, ), call( - Activities.calculate_simple_metrics, + Activities.export_data_to_postgres, { **metadata, - 'model_id': input_data['model_id'], - 'target_data': target_data, - 'metrics': input_data['metrics'], - 'interval_minutes': input_data['interval_minutes'], + 'data': simple_metrics_data, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, }, retry_policy=ANY, start_to_close_timeout=ANY, ), ] ) - - # Assert - Check export_data_to_postgres call - workflow_mock.execute_activity_method.assert_called_once_with( - Activities.export_data_to_postgres, + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_simple_metrics, { **metadata, - 'data': simple_metrics_data, - 'schema': input_data['schema'], - 'table_name': input_data['target_table_name'], - 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + 'model_id': input_data['model_id'], + 'target_data': target_data, + 'metrics': input_data['metrics'], + 'interval_minutes': input_data['interval_minutes'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -128,15 +128,13 @@ async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: S target_data = None - workflow_mock.execute_local_activity_method.return_value = target_data - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(return_value=target_data) # Act await simple_metrics.run(input_data) # Assert - Should not call calculate_simple_metrics or export - assert workflow_mock.execute_local_activity_method.call_count == 1 - workflow_mock.execute_activity_method.assert_not_called() + assert workflow_mock.execute_activity_method.call_count == 1 @mark.asyncio @@ -159,16 +157,15 @@ async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics target_data = {'data': 'test_target_data'} simple_metrics_data = None - workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] - - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(return_value=target_data) + workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data) # Act await simple_metrics.run(input_data) # Assert - Should call calculate_simple_metrics but not export - assert workflow_mock.execute_local_activity_method.call_count == 2 - workflow_mock.execute_activity_method.assert_not_called() + workflow_mock.execute_activity_method.assert_called_once() + workflow_mock.execute_local_activity_method.assert_called_once() @mark.asyncio @@ -191,33 +188,28 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim target_data = {'data': 'test_target_data'} simple_metrics_data = {'metrics': 'test_simple_metrics_data'} - workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] - - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None]) + workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data) # Act await simple_metrics.run(input_data) # Assert - Check calculate_simple_metrics call with default metrics - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.load_custom_query, - ANY, - retry_policy=ANY, - start_to_close_timeout=ANY, - ), - call( - Activities.calculate_simple_metrics, - { - **metadata, - 'model_id': input_data['model_id'], - 'target_data': target_data, - 'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value - 'interval_minutes': input_data['interval_minutes'], - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ), - ] + workflow_mock.execute_activity_method.assert_any_call( + Activities.load_custom_query, + ANY, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_simple_metrics, + { + **metadata, + 'model_id': input_data['model_id'], + 'target_data': target_data, + 'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value + 'interval_minutes': input_data['interval_minutes'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, ) diff --git a/validate.sh b/validate.sh deleted file mode 100755 index 6c72694..0000000 --- a/validate.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/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 - -# Args -FIX_MODE=false -while [[ $# -gt 0 ]]; do - case "$1" in - --fix) - FIX_MODE=true - shift - ;; - -h|--help) - echo "Usage: $0 [--fix]" - echo " --fix Apply Ruff auto-fixes (format and lint fixes)." - exit 0 - ;; - *) - echo -e "${RED}Unknown option: $1${NC}" - echo "Usage: $0 [--fix]" - exit 2 - ;; - esac -done - -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) -# - default: check only -# - --fix: write changes -if ! run_step "1. Code Formatting (Ruff)" "if \$FIX_MODE; then ruff format laborious/ tests/; else ruff format --check laborious/ tests/ e2e/; fi"; then - FAILED_STEPS+=("Code Formatting") -fi - -# Step 2: Linting (Ruff) -# - default: check only -# - --fix: apply autofixes -if ! run_step "2. Code Linting (Ruff)" "if \$FIX_MODE; then ruff check --fix laborious/ tests/; else ruff check laborious/ tests/ e2e/; fi"; 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 -c pyproject.toml -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 08fa8d3..31dd4fa 100644 --- a/values.yaml +++ b/values.yaml @@ -71,7 +71,7 @@ livenessProbe: - -c - | curl -sf http://localhost:9090/metrics | grep -q '^app_up{.*} 1' - initialDelaySeconds: 660 + initialDelaySeconds: 1260 periodSeconds: 15 timeoutSeconds: 5 failureThreshold: 3 @@ -83,7 +83,7 @@ readinessProbe: - -c - | curl -sf http://localhost:9090/metrics | grep -q '^app_up{.*} 1' - initialDelaySeconds: 600 + initialDelaySeconds: 1200 periodSeconds: 10 timeoutSeconds: 3 failureThreshold: 2 @@ -157,7 +157,7 @@ env: - name: GITHUB_REPO_URL value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git" - name: GITHUB_BRANCH - value: "fix/SIENTIAPDE-1478" + value: "feature/SIENTIAPDE-1712" - name: PYTHON_APP value: "laborious.worker.worker" @@ -167,9 +167,9 @@ env: - name: POSTGRES_PORT value: "5432" - name: POSTGRES_USER - value: "sientia" + value: "postgres" - name: POSTGRES_PASSWORD - value: "sientia" + value: "nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3" - name: POSTGRES_DBNAME value: "sientia" - name: POSTGRES_MIN_CONNECTIONS @@ -222,15 +222,69 @@ env: value: "1" - name: MINIO_ENDPOINT_URL - value: "http://minio.minio.svc.cluster.local:9000" + value: "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" + value: "LiArt4eNmJ" - name: MINIO_DEFAULT_BUCKET value: "sientia" + - name: MINIO_RETENTION_HOURS + value: "24" + - name: SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES + value: "0.5" + + # Temporal worker tuning for PredictionsBatch. + # IMPORTANT: prefix must be PREDICTIONSBATCH_ (from class name PredictionsBatch). + # Keep workflow-task concurrency moderate to reduce task completion races under load. + - name: PREDICTIONSBATCH_MAX_CONCURRENT_WORKFLOW_TASKS + value: "20" + # Allow higher activity parallelism because most activities are I/O-bound, but keep headroom. + - name: PREDICTIONSBATCH_MAX_CONCURRENT_ACTIVITIES + value: "60" + # Keep local activities controlled so they do not monopolize the event loop. + - name: PREDICTIONSBATCH_MAX_CONCURRENT_LOCAL_ACTIVITIES + value: "20" + # Cache enough workflows for reuse without excessive memory growth. + - name: PREDICTIONSBATCH_MAX_CACHED_WORKFLOWS + value: "200" + # Start with one workflow poller to avoid burst contention at startup. + - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MINIMUM + value: "3" + # Small initial poller count warms up gradually instead of spiking task fetches. + - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_INITIAL + value: "5" + # Cap workflow pollers to limit scheduling pressure and avoid over-polling. + - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM + value: "15" + # Keep at least two activity pollers so activity queues do not starve during spikes. + - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MINIMUM + value: "3" + # Moderate initial activity pollers for faster ramp-up with controlled pressure. + - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_INITIAL + value: "10" + # Limit max activity pollers to preserve CPU for workflow-task completion. + - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM + value: "30" + + - name: MINIMALRETRAIN_MAX_CONCURRENT_ACTIVITIES + value: "1" + - name: MINIMALRETRAIN_MAX_CONCURRENT_LOCAL_ACTIVITIES + value: "1" + - name: MINIMALRETRAIN_MAX_CACHED_WORKFLOWS + value: "1" + - name: MINIMALRETRAIN_WORKFLOW_POLLER_BEHAVIOUR_MINIMUM + value: "1" + - name: MINIMALRETRAIN_WORKFLOW_POLLER_BEHAVIOUR_INITIAL + value: "1" + - name: MINIMALRETRAIN_WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM + value: "1" + - name: MINIMALRETRAIN_ACTIVITY_POLLER_BEHAVIOUR_MINIMUM + value: "1" + - name: MINIMALRETRAIN_ACTIVITY_POLLER_BEHAVIOUR_INITIAL + value: "1" + - name: MINIMALRETRAIN_ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM + value: "1" - name: PI_WEB_API_BASE_URL value: "https://pivision.votorantimcimentos.com/piwebapi"