gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge pull request #38 from Aignosi/feature/SIENTIAPDE-1712

Browse Source 
Feat/Refactor: MinIO Offload, Data Tracking, and Drift Detection Workflows
This commit is contained in:
vitor-aignosi
2026-04-02 16:20:48 -03:00
committed by GitHub
parent 9dc3cb3ba0 25430b5b26
commit 473bd0b03f
54 changed files with 3854 additions and 2580 deletions
Download Patch File Download Diff File Expand all files Collapse all files

281
README.md
View File

@@ -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://<bucket>/<...>` 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.
**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments.

195
e2e/conftest.py
View File

@@ -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

174
e2e/helpers.py Normal file
View File

@@ -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()}'

39
e2e/scenarios.md
View File

@@ -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

74
e2e/test_child_workflows_e2e.py Normal file
View File

@@ -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'

124
e2e/test_minio_offload.py Normal file
View File

@@ -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

362
e2e/test_predictions_batch_format_export.py
View File

@@ -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!")

307
e2e/test_predictions_batch_main_workflow.py
View File

@@ -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)')

614
e2e/test_predictions_batch_prediction_process.py
View File

@@ -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)

18
laborious/activities/activities.py
View File

@@ -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,

39
laborious/activities/api.py
View File

@@ -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,
)

228
laborious/activities/gates.py
View File

@@ -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,
},

142
laborious/activities/mlflow.py
View File

@@ -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(

44
laborious/activities/model_metrics.py
View File

@@ -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')

201
laborious/activities/storage.py
View File

@@ -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()

65
laborious/metrics.py
View File

@@ -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,
)

5
laborious/utils/connectors_config.py
View File

@@ -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',
}

34
laborious/utils/dataframe_debug.py Normal file
View File

@@ -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}'
)

0
laborious/utils/models/__init__.py Normal file
View File

348
laborious/utils/models/minio_dataframe_payload.py Normal file
View File

@@ -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/<model_name>`
# - 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

32
laborious/utils/repository/minio_manager.py Normal file
View File

@@ -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)

215
laborious/utils/repository/minio_repository.py
View File

@@ -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)

53
laborious/utils/repository/model_repository.py
View File

@@ -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)

1
laborious/worker/prepare_worker.py
View File

@@ -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,

14
laborious/worker/worker.py
View File

@@ -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():

6
laborious/workflows/drift.py
View File

@@ -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,

11
laborious/workflows/minimal_retrain.py
View File

@@ -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,
},

16
laborious/workflows/predictions_batch.py
View File

@@ -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),

4
laborious/workflows/simple_metrics.py
View File

@@ -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,

6
laborious/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -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),

59
laborious/workflows/sub_workflows/prediction_process.py
View File

@@ -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

1
pyproject.toml
View File

@@ -23,6 +23,7 @@ exclude = [
"*.pyc",
".pytest_cache",
"htmlcov",
"tests/laborious/workflows/subworkflows/test_prediction_process.py",
]
[tool.ruff.lint]

2
requirements-dev.txt
View File

@@ -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

2
requirements-light.txt
View File

@@ -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

8
requirements.txt
View File

@@ -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
pycurl
scipy<1.14.0
scikit-learn==1.5.2

49
tests/conftest.py
View File

@@ -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

27
tests/laborious/activities/test_activities.py
View File

@@ -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'}

70
tests/laborious/activities/test_api.py
View File

@@ -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',
}

420
tests/laborious/activities/test_gates.py
View File

@@ -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',

344
tests/laborious/activities/test_mlflow.py
View File

@@ -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,
}

119
tests/laborious/activities/test_model_metrics.py
View File

@@ -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'

327
tests/laborious/activities/test_storage.py
View File

@@ -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()

266
tests/laborious/utils/models/test_minio_dataframe_payload.py Normal file
View File

@@ -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

245
tests/laborious/utils/repository/test_minio_repository.py
View File

@@ -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()

149
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -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()

6
tests/laborious/utils/test_connectors_config.py
View File

@@ -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,
}

48
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -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,

218
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -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,
)

26
tests/laborious/workflows/test_drift.py
View File

@@ -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)

54
tests/laborious/workflows/test_minimal_retrain.py
View File

@@ -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'],
},

54
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -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),

92
tests/laborious/workflows/test_simple_metrics.py
View File

@@ -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,
)

124
validate.sh
View File

@@ -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

72
values.yaml
View File

@@ -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"