Snapshot of fix/QTZPOC-13 source tree
Code-only import without upstream history. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
35
.env.example
Normal file
35
.env.example
Normal file
@@ -0,0 +1,35 @@
|
|||||||
|
POSTGRES_HOST="paradedb-rw.paradedb.svc.cluster.local"
|
||||||
|
POSTGRES_PORT="5432"
|
||||||
|
POSTGRES_USER="sientia"
|
||||||
|
POSTGRES_PASSWORD="password"
|
||||||
|
POSTGRES_DBNAME="sientia"
|
||||||
|
POSTGRES_MIN_CONNECTIONS="10"
|
||||||
|
POSTGRES_MAX_CONNECTIONS="30"
|
||||||
|
|
||||||
|
MLFLOW_HOST="http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
|
||||||
|
MLFLOW_PORT="80"
|
||||||
|
MLFLOW_USERNAME="aignosi"
|
||||||
|
MLFLOW_PASSWORD="mlflow_password"
|
||||||
|
|
||||||
|
OPC_ID="1"
|
||||||
|
OPC_URL="opc.tcp://sientia-opc-simulator-opc.sientia.svc.cluster.local:4840"
|
||||||
|
|
||||||
|
LOG_LEVEL="DEBUG"
|
||||||
|
HTTP_METRICS_PORT="9090"
|
||||||
|
HTTP_SDK_METRICS_PORT="9091"
|
||||||
|
PROJECT_NAME="sientia-laborious"
|
||||||
|
|
||||||
|
TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233"
|
||||||
|
TEMPORAL_NAMESPACE="laborious"
|
||||||
|
|
||||||
|
MONGODB_USERNAME="mongo_user"
|
||||||
|
MONGODB_PASSWORD="mongo_db_password"
|
||||||
|
MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"
|
||||||
|
MONGODB_DATABASE="sientia"
|
||||||
|
MONGODB_TTL_INDEX_HOURS="1"
|
||||||
|
|
||||||
|
MINIO_ENDPOINT_URL="http://localhost:9000"
|
||||||
|
MINIO_ACCESS_KEY="sientia"
|
||||||
|
MINIO_SECRET_KEY="sientia"
|
||||||
|
MINIO_REGION_NAME="sa-east-1"
|
||||||
|
MINIO_DEFAULT_BUCKET="sientia"
|
||||||
31
.github/workflows/quality-gate.yml
vendored
Normal file
31
.github/workflows/quality-gate.yml
vendored
Normal file
@@ -0,0 +1,31 @@
|
|||||||
|
name: Quality gate
|
||||||
|
|
||||||
|
on:
|
||||||
|
push:
|
||||||
|
branches:
|
||||||
|
- main
|
||||||
|
- 'release/**'
|
||||||
|
- 'feature/**'
|
||||||
|
pull_request:
|
||||||
|
branches:
|
||||||
|
- main
|
||||||
|
- 'release/**'
|
||||||
|
- 'feature/**'
|
||||||
|
types: [ opened, synchronize, reopened ]
|
||||||
|
|
||||||
|
concurrency:
|
||||||
|
group: ${{ github.workflow }}-${{ github.ref }}
|
||||||
|
cancel-in-progress: true
|
||||||
|
|
||||||
|
jobs:
|
||||||
|
quality-gate:
|
||||||
|
uses: Aignosi/github_workflow_templates/.github/workflows/python-quality-gate.yml@main
|
||||||
|
permissions:
|
||||||
|
contents: read
|
||||||
|
pull-requests: write
|
||||||
|
issues: write
|
||||||
|
with:
|
||||||
|
project_name: 'laborious'
|
||||||
|
repositories: 'sientia-dataops-library, sientia-mlops-library'
|
||||||
|
requirements_file: 'requirements-light.txt'
|
||||||
|
secrets: inherit
|
||||||
25
.github/workflows/release.yml
vendored
Normal file
25
.github/workflows/release.yml
vendored
Normal file
@@ -0,0 +1,25 @@
|
|||||||
|
name: Create Release on Merge to Main
|
||||||
|
|
||||||
|
on:
|
||||||
|
pull_request:
|
||||||
|
types: [closed]
|
||||||
|
branches:
|
||||||
|
- main
|
||||||
|
workflow_dispatch:
|
||||||
|
inputs:
|
||||||
|
version:
|
||||||
|
description: 'Version to release'
|
||||||
|
required: false
|
||||||
|
type: string
|
||||||
|
|
||||||
|
jobs:
|
||||||
|
release:
|
||||||
|
if: |
|
||||||
|
(github.event_name == 'pull_request' && github.event.pull_request.merged == true) ||
|
||||||
|
github.event_name == 'workflow_dispatch'
|
||||||
|
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-release.yml@main
|
||||||
|
permissions: write-all
|
||||||
|
with:
|
||||||
|
project_name: 'laborious'
|
||||||
|
release_version: ${{ github.event.inputs.version || '' }}
|
||||||
|
secrets: inherit
|
||||||
59
.gitignore
vendored
Normal file
59
.gitignore
vendored
Normal file
@@ -0,0 +1,59 @@
|
|||||||
|
# Ignorar volumes do Docker
|
||||||
|
docker-compose.override.yml
|
||||||
|
**/db_data/
|
||||||
|
**/kafka-volume/
|
||||||
|
**/zookeeper-volume/
|
||||||
|
**/mage_data/
|
||||||
|
**/minio_data/
|
||||||
|
**/venv/
|
||||||
|
**/certs/*.pem
|
||||||
|
**/certs/*.der
|
||||||
|
**/certs/*.csr
|
||||||
|
**/deploy/*.yaml
|
||||||
|
scouter/.file_versions/
|
||||||
|
scouter/pipelines/**/triggers.yaml
|
||||||
|
**/postgres_data/**
|
||||||
|
# Ignorar arquivos e diretórios de cache do Python
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
*.pyd
|
||||||
|
|
||||||
|
# Ignorar logs
|
||||||
|
*.log
|
||||||
|
|
||||||
|
# Ignorar arquivos de configuração locais
|
||||||
|
.vscode/
|
||||||
|
.pytest_cache/
|
||||||
|
.idea/
|
||||||
|
*.swp
|
||||||
|
|
||||||
|
# Ignorar arquivos temporários
|
||||||
|
*.tmp
|
||||||
|
*.bak
|
||||||
|
*.old
|
||||||
|
.secret
|
||||||
|
|
||||||
|
# Ignorar coverage
|
||||||
|
htmlcov/
|
||||||
|
.coverage
|
||||||
|
coverage.xml
|
||||||
|
|
||||||
|
# git keys
|
||||||
|
git_key*
|
||||||
|
|
||||||
|
git_log
|
||||||
|
|
||||||
|
.env
|
||||||
|
|
||||||
|
tmp/
|
||||||
|
catboost_info/
|
||||||
|
|
||||||
|
.ruff_cache/
|
||||||
|
.mypy_cache/
|
||||||
|
mlruns/
|
||||||
|
|
||||||
|
relatorio*
|
||||||
|
openspec/
|
||||||
|
.cursor/
|
||||||
|
.claude/
|
||||||
7
Makefile
Normal file
7
Makefile
Normal file
@@ -0,0 +1,7 @@
|
|||||||
|
VERSION = 1.0.8
|
||||||
|
name = sientia-laborious
|
||||||
|
# ENVIRONMENT = production
|
||||||
|
|
||||||
|
docker-hub:
|
||||||
|
@docker build --no-cache -t aignosi.azurecr.io/$(name):$(VERSION) .
|
||||||
|
@docker push aignosi.azurecr.io/$(name):$(VERSION)
|
||||||
170
docs/opc-communication.md
Normal file
170
docs/opc-communication.md
Normal file
@@ -0,0 +1,170 @@
|
|||||||
|
# OPC UA communication (Laborious)
|
||||||
|
|
||||||
|
Laborious exports predictions to OPC UA servers through `OpcRepository` ([`laborious/utils/repository/opc_repository.py`](../laborious/utils/repository/opc_repository.py)) and the Temporal activity layer in [`laborious/activities/opc.py`](../laborious/activities/opc.py).
|
||||||
|
|
||||||
|
Implementation plan for session/channel recovery on Tier-1 `Bad*` errors: [`.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md`](../.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md).
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
```text
|
||||||
|
Worker (long-lived)
|
||||||
|
└── OpcRepository per OPC server id (from OPC_CONFIG / env)
|
||||||
|
├── connect / disconnect / validate_connection (read-only)
|
||||||
|
├── _connect_locked / _reconnect_locked (under _connection_lock)
|
||||||
|
├── write_data (single attempt per call)
|
||||||
|
└── background reconnect on Tier-1 Bad*, protocol closed, or stale session
|
||||||
|
|
||||||
|
Temporal activity write_opc_data
|
||||||
|
└── OPC.manage_output_tags → write_data per tag (sequential per activity)
|
||||||
|
```
|
||||||
|
|
||||||
|
One worker process holds one `OpcRepository` instance per configured server. Multiple Temporal activities can call `write_data` concurrently on the same repository.
|
||||||
|
|
||||||
|
## Connection lifecycle
|
||||||
|
|
||||||
|
| Phase | Behavior |
|
||||||
|
|-------|----------|
|
||||||
|
| Startup | `init_opc()` creates repositories and calls `connect()` → `_connect_locked()` |
|
||||||
|
| Steady state | `validate_connection()` is read-only (`protocol.state` only); `_session_ready` is checked in `write_data` |
|
||||||
|
| Tier-1 Bad* / protocol closed | `_start_reconnect` → `_run_reconnect` → `_reconnect_locked()` (respects `reconnection_interval`) |
|
||||||
|
| Write | `write_data()` checks reconnect task, `_session_ready`, validates protocol, then one `get_node` + `write_value` |
|
||||||
|
| Shutdown | `close()` disconnects all repositories |
|
||||||
|
|
||||||
|
### Session and channel timeouts
|
||||||
|
|
||||||
|
Requested session and secure-channel lifetime: **10 minutes** (`OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS` in `opc_repository.py`). The server may revise these values; negotiated values are logged after connect and exposed as `opc_session_revised_timeout_milliseconds`.
|
||||||
|
|
||||||
|
### Reconnection interval
|
||||||
|
|
||||||
|
`OPC_RECONNECTION_INTERVAL` is in **seconds** (default `120`). It gates **background** reconnect after Tier-1 `Bad*`, closed protocol, or stale session (`last_reconnection_time` is updated only in `_reconnect_locked()`). It limits load on the OPC server when many workflows fail at once.
|
||||||
|
|
||||||
|
## Concurrency: connection lock and session readiness
|
||||||
|
|
||||||
|
To allow **multiple concurrent writes** when the session is healthy, but **block all writes** while the connection is being torn down or re-established:
|
||||||
|
|
||||||
|
| Primitive | Role |
|
||||||
|
|-----------|------|
|
||||||
|
| `_connection_lock` (`asyncio.Lock`) | Held for the entire `disconnect` → `connect` path. Only one connection-maintenance task at a time. |
|
||||||
|
| `_session_ready` (`asyncio.Event`) | Set when a session is ready for writes; cleared before reconnect starts and set again after a successful connect. |
|
||||||
|
|
||||||
|
**Connection methods (caller holds `_connection_lock` for `_*_locked` helpers):**
|
||||||
|
|
||||||
|
| Method | Role |
|
||||||
|
|--------|------|
|
||||||
|
| `_create_client()` | Create asyncua `Client` + optional `set_security`; raises if `client` already exists |
|
||||||
|
| `_open_session()` | `client.connect()` + metrics; raises if session already open or client missing |
|
||||||
|
| `_connect_locked()` | `_create_client()` (when needed) + `_open_session()`; raises if already connected |
|
||||||
|
| `_disconnect_locked()` | Teardown session and clear `client` |
|
||||||
|
| `_reconnect_locked()` | `_disconnect_locked()` + `_connect_locked()`; sets `last_reconnection_time` |
|
||||||
|
|
||||||
|
Public `connect()` / `disconnect()` acquire the lock and call `_connect_locked()` / `_disconnect_locked()`.
|
||||||
|
|
||||||
|
**Write path (`write_data`):**
|
||||||
|
|
||||||
|
1. If a reconnect task is **in flight** → **fail immediately** (`opc_error_kind=reconnect_in_progress`).
|
||||||
|
2. If `_session_ready` is cleared and no task is running → schedule reconnect (`SessionNotReady`); fail with `connection_lost` or `reconnect_in_progress` if a task started.
|
||||||
|
3. `validate_connection()` checks `protocol.state` only (read-only). If closed → schedule reconnect (`ProtocolClosed`) and fail with `opc_error_kind=connection_lost`.
|
||||||
|
4. Single `get_node` + `write_value` (no retry). Tier-1 `Bad*` on write also schedules reconnect.
|
||||||
|
|
||||||
|
**Reconnect path (`_run_reconnect`):**
|
||||||
|
|
||||||
|
1. `_start_reconnect` clears `_session_ready` and schedules the task when the interval allows and `_allow_reconnect` is true.
|
||||||
|
2. `async with _connection_lock:` → `_reconnect_locked()`.
|
||||||
|
3. `_session_ready` is set on successful `_open_session()`.
|
||||||
|
4. `disconnect()` sets `_allow_reconnect=False` so shutdown does not respawn sessions.
|
||||||
|
|
||||||
|
A second `_connect_locked()` while a session is already open raises `OpcSessionAlreadyConnectedError` (disconnect first).
|
||||||
|
|
||||||
|
**asyncua note:** Concurrent `write_value` on the same session is only safe if the stack tolerates it. If production shows issues, serialize writes with an optional `asyncio.Semaphore(1)` while keeping the connection lock semantics above.
|
||||||
|
|
||||||
|
**Future threads:** replace `asyncio.Lock` / `Event` with `threading` primitives or route all OPC I/O through one dedicated loop.
|
||||||
|
|
||||||
|
## Reconnect triggers
|
||||||
|
|
||||||
|
Background reconnect is scheduled when:
|
||||||
|
|
||||||
|
- `validate_connection()` sees a closed or missing protocol (`ProtocolClosed`).
|
||||||
|
- `_session_ready` is clear after a failed reconnect (`SessionNotReady`).
|
||||||
|
- A write raises a Tier-1 `UaStatusCodeError` in `RECONNECTABLE_OPC_BAD_NAMES`.
|
||||||
|
|
||||||
|
For Tier-1 `Bad*` when the server invalidates the session (e.g. `BadSessionIdInvalid`) but the client still sees transport as open, `write_data` fails once, records the OPC status in metrics, and **schedules** reconnect if:
|
||||||
|
|
||||||
|
- The exception is a `UaStatusCodeError` whose name is in `RECONNECTABLE_OPC_BAD_NAMES` (see plan), and
|
||||||
|
- `reconnection_interval` has elapsed since `last_reconnection_time`, and
|
||||||
|
- No reconnect task is already running.
|
||||||
|
|
||||||
|
There is **no write retry**: the failed export is not sent again in the same activity.
|
||||||
|
|
||||||
|
## Prediction confidence and PostgreSQL comments
|
||||||
|
|
||||||
|
| `prediction_confidence` | Meaning |
|
||||||
|
|-------------------------|---------|
|
||||||
|
| (unchanged) | Successful OPC export |
|
||||||
|
| **12** | Generic OPC write failure (`OPC_WRITTING_ERROR_CONFIDENCE`) |
|
||||||
|
| **14** | Tier-1 session/channel `Bad*` on export (`OPC_SESSION_BAD_CONFIDENCE`) |
|
||||||
|
| **14** | Write while reconnect in progress (`OPC_SESSION_BAD_CONFIDENCE`, comment `OPC UA reconnect in progress`) |
|
||||||
|
| **13** | PI Web API write failure (separate path) |
|
||||||
|
|
||||||
|
Session/channel errors use a stable comment for counting:
|
||||||
|
|
||||||
|
```text
|
||||||
|
OPC UA session/channel error: BadSessionIdInvalid
|
||||||
|
```
|
||||||
|
|
||||||
|
Reconnect-in-progress exports use:
|
||||||
|
|
||||||
|
```text
|
||||||
|
OPC UA reconnect in progress
|
||||||
|
```
|
||||||
|
|
||||||
|
Example SQL:
|
||||||
|
|
||||||
|
```sql
|
||||||
|
SELECT count(*) FROM predictions WHERE prediction_confidence = 14;
|
||||||
|
SELECT count(*) FROM predictions WHERE comments LIKE 'OPC UA session/channel error:%';
|
||||||
|
```
|
||||||
|
|
||||||
|
## Prometheus metrics (`opc_*`)
|
||||||
|
|
||||||
|
Defined in [`laborious/metrics.py`](../laborious/metrics.py). Do not rename in production without a dashboard migration.
|
||||||
|
|
||||||
|
| Metric | Purpose |
|
||||||
|
|--------|---------|
|
||||||
|
| `opc_connections_initiated_total` | Connection attempts |
|
||||||
|
| `opc_connections_failed_total` | Failed connects |
|
||||||
|
| `opc_connection_status` | Gauge 1=connected, 0=disconnected |
|
||||||
|
| `opc_session_created_total` | Session established after connect |
|
||||||
|
| `opc_session_closed_total` | Disconnect initiated |
|
||||||
|
| `opc_session_revised_timeout_milliseconds` | Negotiated session timeout (ms) |
|
||||||
|
| `opc_write_attempts_total` | Per write; label `result` = `OK` or exception name |
|
||||||
|
| `opc_write_inter_arrival_over_session_timeout_total` | Successful writes spaced longer than revised session timeout |
|
||||||
|
|
||||||
|
Legacy activity metrics: `laborious_prediction_opc_writing_count`, `laborious_prediction_opc_writing_response_time_monitor`.
|
||||||
|
|
||||||
|
## Environment variables
|
||||||
|
|
||||||
|
| Variable | Default | Description |
|
||||||
|
|----------|---------|-------------|
|
||||||
|
| `OPC_CONFIG` | — | JSON map of server configs (overrides single-server env) |
|
||||||
|
| `OPC_ID` | `1` | Server id |
|
||||||
|
| `OPC_URL` | `opc.tcp://localhost:4840` | Endpoint |
|
||||||
|
| `OPC_SERVER_NAME` | `default_server` | Label for metrics/logs |
|
||||||
|
| `OPC_SERVER_URI` | same as URL | Application URI / cert SAN |
|
||||||
|
| `OPC_CERT_PATH` | — | Client certificate (secure mode) |
|
||||||
|
| `OPC_PRIVATE_KEY_PATH` | — | Client private key |
|
||||||
|
| `OPC_SERVER_CERT_PATH` | — | Server certificate |
|
||||||
|
| `OPC_RECONNECTION_INTERVAL` | `120` | Minimum seconds between reconnects |
|
||||||
|
|
||||||
|
## Operations checklist
|
||||||
|
|
||||||
|
- Correlate `BadSessionIdInvalid` in `opc_write_attempts_total` with `opc_session_closed_total` / `opc_session_created_total` (reconnect may finish after the row is stored with confidence 14).
|
||||||
|
- Use confidence **14** and comment prefix for session invalidation rates; use **12** for other OPC failures.
|
||||||
|
- Respect `OPC_RECONNECTION_INTERVAL` under parallel load; bursts of confidence 14 are expected until the next successful cycle.
|
||||||
|
|
||||||
|
## Related tests
|
||||||
|
|
||||||
|
- Unit: [`tests/laborious/utils/repository/test_opc_repository.py`](../tests/laborious/utils/repository/test_opc_repository.py)
|
||||||
|
- Unit: [`tests/laborious/activities/test_opc.py`](../tests/laborious/activities/test_opc.py)
|
||||||
|
- E2E (mock OPC): [`e2e/test_predictions_batch_format_export.py`](../e2e/test_predictions_batch_format_export.py)
|
||||||
|
- E2E (in-process asyncua server + real `OpcRepository`): [`e2e/test_opc_real_server.py`](../e2e/test_opc_real_server.py) — scenarios 3.1.2, 3.2.2, 3.2.4, 3.2.5
|
||||||
|
- Scenarios: [`e2e/scenarios.md`](../e2e/scenarios.md)
|
||||||
3
e2e/__init__.py
Normal file
3
e2e/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
"""
|
||||||
|
End-to-end tests for laborious temporal workflows.
|
||||||
|
"""
|
||||||
767
e2e/conftest.py
Normal file
767
e2e/conftest.py
Normal file
@@ -0,0 +1,767 @@
|
|||||||
|
"""
|
||||||
|
Pytest configuration and fixtures for E2E tests.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
|
# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
|
||||||
|
_model_analysis_module = MagicMock()
|
||||||
|
_model_analysis_module.ModelAnalysis = MagicMock
|
||||||
|
sys.modules.setdefault('sientia', MagicMock())
|
||||||
|
sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
|
||||||
|
from io import BytesIO
|
||||||
|
|
||||||
|
import joblib
|
||||||
|
import pandas as pd
|
||||||
|
import pytest
|
||||||
|
import pytest_asyncio
|
||||||
|
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||||||
|
from sientia_do.observability.metrics_controller import MetricsController
|
||||||
|
from sklearn.compose import ColumnTransformer
|
||||||
|
from sklearn.linear_model import LinearRegression
|
||||||
|
from sklearn.pipeline import Pipeline
|
||||||
|
from sqlalchemy import create_engine, text
|
||||||
|
from temporalio.testing import WorkflowEnvironment
|
||||||
|
from temporalio.worker import Worker
|
||||||
|
from testcontainers.minio import MinioContainer
|
||||||
|
from testcontainers.postgres import PostgresContainer
|
||||||
|
|
||||||
|
from e2e.opc_test_server import OpcE2ETestServer
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||||
|
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
||||||
|
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||||
|
|
||||||
|
# Test constants
|
||||||
|
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
|
||||||
|
TEST_DATABASE_NAME = 'test_db'
|
||||||
|
|
||||||
|
# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised.
|
||||||
|
JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model'
|
||||||
|
JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact'
|
||||||
|
|
||||||
|
|
||||||
|
@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():
|
||||||
|
"""
|
||||||
|
Create a PostgreSQL container using testcontainers.
|
||||||
|
|
||||||
|
This fixture creates a real PostgreSQL database in a Docker container
|
||||||
|
that will be used for all tests in the session.
|
||||||
|
"""
|
||||||
|
postgres = PostgresContainer('postgres:15')
|
||||||
|
postgres.start()
|
||||||
|
yield postgres
|
||||||
|
postgres.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def postgres_engine(postgres_container):
|
||||||
|
"""
|
||||||
|
Create SQLAlchemy engine for PostgreSQL test database.
|
||||||
|
|
||||||
|
This fixture creates a connection to the PostgreSQL container
|
||||||
|
created by the postgres_container fixture.
|
||||||
|
"""
|
||||||
|
engine = create_engine(postgres_container.get_connection_url())
|
||||||
|
|
||||||
|
yield engine
|
||||||
|
|
||||||
|
engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
def _create_schema_and_tables(engine):
|
||||||
|
"""
|
||||||
|
Helper function to create schema and tables in the given engine.
|
||||||
|
|
||||||
|
Creates predictions_schema with:
|
||||||
|
- laborious_data: Input data table for queries
|
||||||
|
- predictions: Output predictions table
|
||||||
|
- transformed_data: Output transformed data table
|
||||||
|
"""
|
||||||
|
# Use begin() to ensure transaction is properly committed
|
||||||
|
with engine.begin() as conn:
|
||||||
|
# Create predictions_schema
|
||||||
|
conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema'))
|
||||||
|
|
||||||
|
# Create laborious_data table (input data from sensors)
|
||||||
|
create_laborious_data_sql = """
|
||||||
|
CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
|
||||||
|
id SERIAL NOT NULL,
|
||||||
|
model_id int4 NOT NULL,
|
||||||
|
variable text NOT NULL,
|
||||||
|
value numeric NULL,
|
||||||
|
"timestamp" timestamptz NOT NULL,
|
||||||
|
created_at timestamptz NOT NULL,
|
||||||
|
PRIMARY KEY (id)
|
||||||
|
);
|
||||||
|
"""
|
||||||
|
conn.execute(text(create_laborious_data_sql))
|
||||||
|
|
||||||
|
# Create predictions table
|
||||||
|
create_predictions_sql = """
|
||||||
|
CREATE TABLE if not exists predictions_schema.predictions (
|
||||||
|
id SERIAL NOT NULL ,
|
||||||
|
model_id int4 NOT NULL,
|
||||||
|
prediction numeric NULL,
|
||||||
|
prediction_confidence numeric NOT NULL,
|
||||||
|
response_time numeric NOT NULL,
|
||||||
|
prediction_status text NOT NULL,
|
||||||
|
"timestamp" timestamptz NOT NULL,
|
||||||
|
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||||
|
"comments" text NULL,
|
||||||
|
PRIMARY KEY (id, created_at)
|
||||||
|
);
|
||||||
|
"""
|
||||||
|
conn.execute(text(create_predictions_sql))
|
||||||
|
|
||||||
|
# Create transformed_data table
|
||||||
|
create_transformed_sql = """
|
||||||
|
CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
|
||||||
|
id SERIAL NOT NULL,
|
||||||
|
model_id int4 NOT NULL,
|
||||||
|
variable text NOT NULL,
|
||||||
|
value numeric NULL,
|
||||||
|
"timestamp" timestamptz NOT NULL,
|
||||||
|
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||||
|
PRIMARY KEY (id)
|
||||||
|
);
|
||||||
|
"""
|
||||||
|
conn.execute(text(create_transformed_sql))
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(autouse=True)
|
||||||
|
def setup_postgres_schema_and_tables(postgres_engine):
|
||||||
|
"""
|
||||||
|
Automatically create necessary schema and tables before each test.
|
||||||
|
|
||||||
|
This fixture runs automatically (autouse=True) and ensures
|
||||||
|
that the predictions_schema and tables exist with the correct structure.
|
||||||
|
"""
|
||||||
|
_create_schema_and_tables(postgres_engine)
|
||||||
|
yield
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def mock_logger():
|
||||||
|
"""Mock logger for testing."""
|
||||||
|
|
||||||
|
def message(message):
|
||||||
|
print(f'[LOG] {message}')
|
||||||
|
|
||||||
|
def custom_message(message, _metadata=None):
|
||||||
|
print(f'[LOG] {message}')
|
||||||
|
|
||||||
|
logger = MagicMock()
|
||||||
|
logger.info = MagicMock(side_effect=message)
|
||||||
|
logger.debug = MagicMock(side_effect=message)
|
||||||
|
logger.error = MagicMock(side_effect=message)
|
||||||
|
logger.warning = MagicMock(side_effect=message)
|
||||||
|
logger.custom_info = MagicMock(side_effect=custom_message)
|
||||||
|
logger.custom_debug = MagicMock(side_effect=custom_message)
|
||||||
|
logger.custom_error = MagicMock(side_effect=custom_message)
|
||||||
|
logger.custom_warning = MagicMock(side_effect=custom_message)
|
||||||
|
return logger
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def mock_mongo_client():
|
||||||
|
"""
|
||||||
|
Mock MongoDB client to avoid real connections.
|
||||||
|
|
||||||
|
This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
|
||||||
|
allowing us to use a real NotificationHandler instance without connecting to MongoDB.
|
||||||
|
"""
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_db = MagicMock()
|
||||||
|
mock_collection = MagicMock()
|
||||||
|
|
||||||
|
# Configure the mock chain: client[database] -> db[collection] -> collection
|
||||||
|
mock_client.__getitem__.return_value = mock_db
|
||||||
|
mock_db.__getitem__.return_value = mock_collection
|
||||||
|
|
||||||
|
# Mock server_info() to avoid connection attempts
|
||||||
|
mock_client.server_info = MagicMock()
|
||||||
|
|
||||||
|
# Mock insert_one for notifications
|
||||||
|
mock_collection.insert_one = MagicMock()
|
||||||
|
|
||||||
|
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):
|
||||||
|
"""
|
||||||
|
Create a real NotificationHandler instance with mocked MongoDB client.
|
||||||
|
|
||||||
|
This fixture creates a real CoreNotificationHandler instance but mocks
|
||||||
|
the underlying MongoDB connection to avoid real database connections.
|
||||||
|
"""
|
||||||
|
# Patch MongoClient where it's imported in the handlers module
|
||||||
|
with patch('sientia_do.notifications.handlers.MongoClient', return_value=mock_mongo_client):
|
||||||
|
handler = CoreNotificationHandler(
|
||||||
|
connection_string=TEST_MONGODB_CONNECTION_STRING,
|
||||||
|
database=TEST_DATABASE_NAME,
|
||||||
|
logger=mock_logger,
|
||||||
|
project_name='laborious',
|
||||||
|
)
|
||||||
|
yield handler
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def metrics_controller(mock_logger):
|
||||||
|
"""Create a real MetricsController instance."""
|
||||||
|
return MetricsController(logger=mock_logger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
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()
|
||||||
|
|
||||||
|
# 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
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def mock_pi_web_api_repository():
|
||||||
|
"""Mock PI Web API repository for PI Web API operations."""
|
||||||
|
mock_repo = MagicMock()
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def mock_opc_repository():
|
||||||
|
"""Mock OPC repository for OPC operations."""
|
||||||
|
mock_repo = MagicMock()
|
||||||
|
mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1}))
|
||||||
|
mock_repo.disconnect = AsyncMock()
|
||||||
|
return mock_repo
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def opc_e2e_server():
|
||||||
|
"""
|
||||||
|
In-process asyncua OPC UA server for E2E tests against OpcRepository.
|
||||||
|
"""
|
||||||
|
server = OpcE2ETestServer()
|
||||||
|
await server.start()
|
||||||
|
try:
|
||||||
|
yield server
|
||||||
|
finally:
|
||||||
|
await server.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def patch_create_engine(postgres_engine):
|
||||||
|
"""Patch create_engine to return test postgres_engine."""
|
||||||
|
with patch(
|
||||||
|
'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine
|
||||||
|
):
|
||||||
|
yield
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def patch_minio_repository(mock_minio_repository):
|
||||||
|
"""Patch MinioRepository to return mock."""
|
||||||
|
# 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
|
||||||
|
def patch_pi_web_api_repository(mock_pi_web_api_repository):
|
||||||
|
"""Patch MLflowRepository to return mock."""
|
||||||
|
with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
|
||||||
|
yield
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def mock_mlflow_models():
|
||||||
|
"""Create mock models for MLflow load_model methods."""
|
||||||
|
# Mock transform model - returns DataFrame with same index as input
|
||||||
|
mock_transform_model = MagicMock()
|
||||||
|
|
||||||
|
def mock_transform_predict(data):
|
||||||
|
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
||||||
|
print(data.to_csv())
|
||||||
|
print(data.index)
|
||||||
|
result = pd.DataFrame(
|
||||||
|
{
|
||||||
|
'feature_1': [0.234] * num_rows,
|
||||||
|
'feature_2': [0.783] * num_rows,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
result.index = data.index
|
||||||
|
return result
|
||||||
|
|
||||||
|
mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
|
||||||
|
|
||||||
|
# Mock predict model - returns array/list of predictions
|
||||||
|
mock_predict_model = MagicMock()
|
||||||
|
|
||||||
|
def mock_predict_predict(data):
|
||||||
|
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
||||||
|
return [0.5] * num_rows
|
||||||
|
|
||||||
|
mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
|
||||||
|
|
||||||
|
# Mock PyFuncModel for compressed models
|
||||||
|
mock_pyfunc_model = MagicMock()
|
||||||
|
mock_pyfunc_model._model_impl = MagicMock()
|
||||||
|
mock_pyfunc_model._model_impl.python_model = mock_transform_model
|
||||||
|
|
||||||
|
return {
|
||||||
|
'transform_model': mock_transform_model,
|
||||||
|
'predict_model': mock_predict_model,
|
||||||
|
'pyfunc_model': mock_pyfunc_model,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _build_joblib_fallback_estimator():
|
||||||
|
"""
|
||||||
|
Build a real sklearn estimator for the joblib fallback artifact.
|
||||||
|
|
||||||
|
`request_predict` appends a string `timestamp` column to the transformed
|
||||||
|
data, so the estimator selects the transform output columns explicitly and
|
||||||
|
drops the rest.
|
||||||
|
"""
|
||||||
|
features = ['feature_1', 'feature_2']
|
||||||
|
estimator = Pipeline(
|
||||||
|
steps=[
|
||||||
|
(
|
||||||
|
'select',
|
||||||
|
ColumnTransformer([('features', 'passthrough', features)], remainder='drop'),
|
||||||
|
),
|
||||||
|
('regressor', LinearRegression()),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
train = pd.DataFrame(
|
||||||
|
{
|
||||||
|
'feature_1': [0.0, 1.0, 2.0, 3.0],
|
||||||
|
'feature_2': [0.0, 1.0, 0.0, 1.0],
|
||||||
|
'timestamp': ['2024-01-01 12:00:00'] * 4,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3)
|
||||||
|
return estimator
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def joblib_fallback_artifact(tmp_path):
|
||||||
|
"""
|
||||||
|
Real artifact directory holding a joblib-compressed sklearn estimator.
|
||||||
|
|
||||||
|
Mirrors what the joblib fallback finds after downloading `prediction_model`:
|
||||||
|
a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an
|
||||||
|
`MLmodel` placeholder.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict with 'path' (the artifact directory) and 'estimator' (the object
|
||||||
|
that was serialized, so tests can compute the expected prediction).
|
||||||
|
"""
|
||||||
|
artifact_dir = tmp_path / 'prediction_model'
|
||||||
|
artifact_dir.mkdir()
|
||||||
|
|
||||||
|
estimator = _build_joblib_fallback_estimator()
|
||||||
|
joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3)
|
||||||
|
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
|
||||||
|
|
||||||
|
return {'path': str(artifact_dir), 'estimator': estimator}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def joblib_fallback_artifact_without_pickle(tmp_path):
|
||||||
|
"""Artifact directory with no top-level `.pkl`, so the joblib fallback also fails."""
|
||||||
|
artifact_dir = tmp_path / 'prediction_model_without_pickle'
|
||||||
|
artifact_dir.mkdir()
|
||||||
|
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
|
||||||
|
return str(artifact_dir)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
|
||||||
|
"""Patch mlflow module in repository with load_model mocks."""
|
||||||
|
mock_mlflow = MagicMock()
|
||||||
|
|
||||||
|
# Mock sklearn.load_model
|
||||||
|
def mock_sklearn_load_model(model_uri):
|
||||||
|
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
|
||||||
|
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
|
||||||
|
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
||||||
|
return mock_mlflow_models['transform_model']
|
||||||
|
return mock_mlflow_models['predict_model']
|
||||||
|
|
||||||
|
mock_mlflow.sklearn = MagicMock()
|
||||||
|
mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
|
||||||
|
|
||||||
|
# Mock pyfunc.load_model
|
||||||
|
def mock_pyfunc_load_model(model_uri):
|
||||||
|
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
|
||||||
|
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
|
||||||
|
if 'artifacts' in model_uri or 'tmp' in model_uri:
|
||||||
|
return mock_mlflow_models['pyfunc_model']
|
||||||
|
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
||||||
|
return mock_mlflow_models['transform_model']
|
||||||
|
return mock_mlflow_models['predict_model']
|
||||||
|
|
||||||
|
mock_mlflow.pyfunc = MagicMock()
|
||||||
|
mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
|
||||||
|
|
||||||
|
# Mock pytorch.load_model
|
||||||
|
mock_mlflow.pytorch = MagicMock()
|
||||||
|
mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
|
||||||
|
|
||||||
|
# Mock other mlflow methods that might be called
|
||||||
|
mock_mlflow.set_tracking_uri = MagicMock()
|
||||||
|
mock_mlflow.get_run = MagicMock(
|
||||||
|
return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))
|
||||||
|
)
|
||||||
|
mock_mlflow.tracking = MagicMock()
|
||||||
|
mock_mlflow_client = MagicMock(
|
||||||
|
search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
|
||||||
|
search_model_versions=MagicMock(
|
||||||
|
return_value=[
|
||||||
|
MagicMock(
|
||||||
|
current_stage='Production', version='1', source='runs:/artifacts/test_run_id'
|
||||||
|
)
|
||||||
|
]
|
||||||
|
),
|
||||||
|
)
|
||||||
|
# Only the joblib fallback calls download_artifacts directly on the client.
|
||||||
|
mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path'])
|
||||||
|
mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)
|
||||||
|
|
||||||
|
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
|
||||||
|
yield mock_mlflow
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mlflow_client(patch_mlflow):
|
||||||
|
"""The MlflowClient mock the repository holds as `self.client`."""
|
||||||
|
return patch_mlflow.tracking.MlflowClient.return_value
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(scope='function')
|
||||||
|
async def test_activities(
|
||||||
|
postgres_engine,
|
||||||
|
postgres_container,
|
||||||
|
mock_logger,
|
||||||
|
notification_handler,
|
||||||
|
metrics_controller,
|
||||||
|
mock_minio_repository,
|
||||||
|
patch_create_engine,
|
||||||
|
patch_minio_repository,
|
||||||
|
patch_mlflow,
|
||||||
|
patch_pi_web_api_repository,
|
||||||
|
mock_opc_repository,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Create Activities instance with test dependencies.
|
||||||
|
|
||||||
|
This fixture creates a real Activities instance with:
|
||||||
|
- PostgreSQL database (via testcontainers)
|
||||||
|
- Mocked MinIO client
|
||||||
|
- Real NotificationHandler and MetricsController (with mocked underlying services)
|
||||||
|
"""
|
||||||
|
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={
|
||||||
|
# Host:port only; Minio() prepends http(s):// from the secure flag.
|
||||||
|
'endpoint_url': 'localhost:9000',
|
||||||
|
'access_key': 'test',
|
||||||
|
'secret_key': 'test',
|
||||||
|
'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:
|
||||||
|
# Cleanup - ALWAYS runs, even if test fails
|
||||||
|
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."""
|
||||||
|
env = await WorkflowEnvironment.start_time_skipping()
|
||||||
|
async with env:
|
||||||
|
yield env
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(scope='function')
|
||||||
|
async def temporal_worker(temporal_test_env, test_activities):
|
||||||
|
"""Create Temporal worker with test activities."""
|
||||||
|
async with Worker(
|
||||||
|
temporal_test_env.client,
|
||||||
|
task_queue='test-queue',
|
||||||
|
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(scope='function')
|
||||||
|
async def test_activities_real_opc(
|
||||||
|
postgres_engine,
|
||||||
|
postgres_container,
|
||||||
|
opc_e2e_server: OpcE2ETestServer,
|
||||||
|
mock_logger,
|
||||||
|
notification_handler,
|
||||||
|
metrics_controller,
|
||||||
|
patch_create_engine,
|
||||||
|
patch_minio_repository,
|
||||||
|
patch_mlflow,
|
||||||
|
patch_pi_web_api_repository,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Activities with a real OpcRepository connected to the in-process OPC UA server.
|
||||||
|
"""
|
||||||
|
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': 'localhost:9000',
|
||||||
|
'access_key': 'test',
|
||||||
|
'secret_key': 'test',
|
||||||
|
'default_bucket': 'test-bucket',
|
||||||
|
'retention_hours': 24,
|
||||||
|
'secure': False,
|
||||||
|
},
|
||||||
|
opc_config={
|
||||||
|
'1': {
|
||||||
|
'id': '1',
|
||||||
|
'server_name': 'e2e-opc',
|
||||||
|
'url': opc_e2e_server.url,
|
||||||
|
'server_uri': opc_e2e_server.url,
|
||||||
|
'cert_path': None,
|
||||||
|
'private_key_path': None,
|
||||||
|
'server_cert_path': None,
|
||||||
|
'reconnection_interval': 0,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
pi_web_api_config={
|
||||||
|
'base_url': 'http://localhost:8080',
|
||||||
|
'auth_type': 'bearer',
|
||||||
|
'auth_token': 'test_token',
|
||||||
|
},
|
||||||
|
logger=mock_logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
)
|
||||||
|
await activities.init_opc()
|
||||||
|
repo = activities.opc_repository['1']
|
||||||
|
assert repo._session_ready.is_set(), 'OPC E2E server connection failed during init_opc'
|
||||||
|
try:
|
||||||
|
yield activities
|
||||||
|
finally:
|
||||||
|
await activities.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(scope='function')
|
||||||
|
async def temporal_worker_real_opc(temporal_test_env, test_activities_real_opc):
|
||||||
|
"""Temporal worker backed by Activities using the in-process OPC UA server."""
|
||||||
|
async with Worker(
|
||||||
|
temporal_test_env.client,
|
||||||
|
task_queue='test-queue',
|
||||||
|
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
|
||||||
|
activities=_worker_activity_list(test_activities_real_opc),
|
||||||
|
) as worker:
|
||||||
|
yield worker
|
||||||
186
e2e/helpers.py
Normal file
186
e2e/helpers.py
Normal file
@@ -0,0 +1,186 @@
|
|||||||
|
"""
|
||||||
|
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 = None,
|
||||||
|
comments_contains: str | None = None,
|
||||||
|
) -> 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 exact comments string (optional).
|
||||||
|
comments_contains: Substring expected in comments when queued (optional).
|
||||||
|
"""
|
||||||
|
|
||||||
|
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]}"
|
||||||
|
)
|
||||||
|
if comments is not None:
|
||||||
|
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
||||||
|
if comments_contains is not None:
|
||||||
|
assert comments_contains in row[4], (
|
||||||
|
f"Expected comments to contain '{comments_contains}', 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."""
|
||||||
|
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
|
||||||
|
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()}'
|
||||||
192
e2e/opc_test_server.py
Normal file
192
e2e/opc_test_server.py
Normal file
@@ -0,0 +1,192 @@
|
|||||||
|
"""
|
||||||
|
In-process OPC UA server for E2E tests (asyncua).
|
||||||
|
|
||||||
|
Provides writable prediction/confidence nodes and optional write faults
|
||||||
|
(Tier-1 BadSessionIdInvalid via PreWrite callback).
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import socket
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from asyncua import Server, ua
|
||||||
|
from asyncua.common.callback import CallbackType
|
||||||
|
from asyncua.common.utils import ServiceError
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from asyncua.common.node import Node
|
||||||
|
|
||||||
|
|
||||||
|
UNKNOWN_NODE_ID = 'ns=99;i=9999'
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class OpcE2ENodeIds:
|
||||||
|
"""NodeId strings used in opc_output_config for E2E workflows."""
|
||||||
|
|
||||||
|
prediction: str
|
||||||
|
confidence: str
|
||||||
|
unknown: str = UNKNOWN_NODE_ID
|
||||||
|
|
||||||
|
|
||||||
|
class OpcE2ETestServer:
|
||||||
|
"""
|
||||||
|
Ephemeral asyncua server with Laborious E2E variables and controllable faults.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
host: Bind address (default 127.0.0.1).
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, host: str = '127.0.0.1') -> None:
|
||||||
|
self._host = host
|
||||||
|
self._server: Server | None = None
|
||||||
|
self._prediction_node: Node | None = None
|
||||||
|
self._confidence_node: Node | None = None
|
||||||
|
self._session_bad_on_write = False
|
||||||
|
self._url: str | None = None
|
||||||
|
self._node_ids: OpcE2ENodeIds | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def url(self) -> str:
|
||||||
|
if self._url is None:
|
||||||
|
raise RuntimeError('OPC E2E server is not started')
|
||||||
|
return self._url
|
||||||
|
|
||||||
|
@property
|
||||||
|
def node_ids(self) -> OpcE2ENodeIds:
|
||||||
|
if self._node_ids is None:
|
||||||
|
raise RuntimeError('OPC E2E server is not started')
|
||||||
|
return self._node_ids
|
||||||
|
|
||||||
|
def set_session_bad_on_write(self, enabled: bool) -> None:
|
||||||
|
"""
|
||||||
|
When enabled, every client Write is rejected with BadSessionIdInvalid.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
enabled (bool): Turn Tier-1 session fault injection on or off.
|
||||||
|
"""
|
||||||
|
self._session_bad_on_write = enabled
|
||||||
|
|
||||||
|
async def start(self) -> OpcE2ENodeIds:
|
||||||
|
"""
|
||||||
|
Start the OPC UA server on a free TCP port.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
OpcE2ENodeIds: NodeId strings for prediction and confidence tags.
|
||||||
|
"""
|
||||||
|
port = _free_port(self._host)
|
||||||
|
self._url = f'opc.tcp://{self._host}:{port}/freeopcua/server/'
|
||||||
|
|
||||||
|
server = Server()
|
||||||
|
server.set_endpoint(self._url)
|
||||||
|
await server.init()
|
||||||
|
server.iserver.callback_service.addListener(
|
||||||
|
CallbackType.PreWrite,
|
||||||
|
self._pre_write_callback,
|
||||||
|
)
|
||||||
|
|
||||||
|
idx = await server.register_namespace('http://sientia.test/laborious-e2e')
|
||||||
|
e2e_object = await server.nodes.objects.add_object(idx, 'LaboriousE2E')
|
||||||
|
prediction = await e2e_object.add_variable(
|
||||||
|
idx,
|
||||||
|
'Prediction',
|
||||||
|
ua.Variant(0.0, ua.VariantType.Float),
|
||||||
|
)
|
||||||
|
confidence = await e2e_object.add_variable(
|
||||||
|
idx,
|
||||||
|
'Confidence',
|
||||||
|
ua.Variant(0.0, ua.VariantType.Float),
|
||||||
|
)
|
||||||
|
await prediction.set_writable()
|
||||||
|
await confidence.set_writable()
|
||||||
|
|
||||||
|
await server.start()
|
||||||
|
self._server = server
|
||||||
|
self._prediction_node = prediction
|
||||||
|
self._confidence_node = confidence
|
||||||
|
self._node_ids = OpcE2ENodeIds(
|
||||||
|
prediction=prediction.nodeid.to_string(),
|
||||||
|
confidence=confidence.nodeid.to_string(),
|
||||||
|
)
|
||||||
|
return self._node_ids
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""Stop the OPC UA server and release the listening port."""
|
||||||
|
if self._server is not None:
|
||||||
|
await self._server.stop()
|
||||||
|
self._server = None
|
||||||
|
self._prediction_node = None
|
||||||
|
self._confidence_node = None
|
||||||
|
self._url = None
|
||||||
|
self._node_ids = None
|
||||||
|
self._session_bad_on_write = False
|
||||||
|
|
||||||
|
async def read_prediction(self) -> float:
|
||||||
|
"""
|
||||||
|
Read the current prediction variable value from the address space.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
float: Stored prediction value.
|
||||||
|
"""
|
||||||
|
if self._prediction_node is None:
|
||||||
|
raise RuntimeError('OPC E2E server is not started')
|
||||||
|
value = await self._prediction_node.read_value()
|
||||||
|
return float(value)
|
||||||
|
|
||||||
|
async def read_confidence(self) -> float:
|
||||||
|
"""
|
||||||
|
Read the current confidence variable value from the address space.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
float: Stored confidence value.
|
||||||
|
"""
|
||||||
|
if self._confidence_node is None:
|
||||||
|
raise RuntimeError('OPC E2E server is not started')
|
||||||
|
value = await self._confidence_node.read_value()
|
||||||
|
return float(value)
|
||||||
|
|
||||||
|
async def _pre_write_callback(self, event, _service) -> None:
|
||||||
|
# Only client writes are faulted. asyncua's own clock task writes CurrentTime
|
||||||
|
# through this same PreWrite dispatch, and raising there kills its time_task;
|
||||||
|
# the stored exception then resurfaces as a teardown failure in server.stop().
|
||||||
|
if self._session_bad_on_write and getattr(event, 'is_external', False):
|
||||||
|
raise ServiceError(ua.StatusCodes.BadSessionIdInvalid)
|
||||||
|
|
||||||
|
|
||||||
|
def _free_port(host: str) -> int:
|
||||||
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||||
|
sock.bind((host, 0))
|
||||||
|
return int(sock.getsockname()[1])
|
||||||
|
|
||||||
|
|
||||||
|
def build_opc_output_config(
|
||||||
|
node_ids: OpcE2ENodeIds,
|
||||||
|
*,
|
||||||
|
prediction_tag: str | None = None,
|
||||||
|
confidence_tag: str | None = None,
|
||||||
|
prediction_only: bool = False,
|
||||||
|
server_key: str = '1',
|
||||||
|
) -> dict[str, dict]:
|
||||||
|
"""
|
||||||
|
Build opc_output_config for PredictionsBatch using real server NodeIds.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
node_ids (OpcE2ENodeIds): Node ids from OpcE2ETestServer.
|
||||||
|
prediction_tag (str | None): Override prediction NodeId (default: node_ids.prediction).
|
||||||
|
confidence_tag (str | None): Override confidence NodeId (default: node_ids.confidence).
|
||||||
|
prediction_only (bool): When True, omit confidence_tags (single write per activity).
|
||||||
|
server_key (str): OPC server id key in opc_output_config.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict: opc_output_config payload for workflow input.
|
||||||
|
"""
|
||||||
|
pred = prediction_tag if prediction_tag is not None else node_ids.prediction
|
||||||
|
conf = confidence_tag if confidence_tag is not None else node_ids.confidence
|
||||||
|
server_config: dict = {
|
||||||
|
'prediction_tags': {pred: {'data_type': 'float'}},
|
||||||
|
}
|
||||||
|
if not prediction_only:
|
||||||
|
server_config['confidence_tags'] = {conf: {'data_type': 'float'}}
|
||||||
|
return {server_key: server_config}
|
||||||
564
e2e/scenarios.md
Normal file
564
e2e/scenarios.md
Normal file
@@ -0,0 +1,564 @@
|
|||||||
|
# Test Scenarios for Predictions Batch Workflow
|
||||||
|
|
||||||
|
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.
|
||||||
|
- **OPC tests (real server)**: `e2e/test_opc_real_server.py` uses an in-process **asyncua** server and real `OpcRepository` (`test_activities_real_opc`). Scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 are covered there. Other E2E modules keep the OPC mock.
|
||||||
|
- Run only OPC real-server tests: `pytest e2e/test_opc_real_server.py -m "integration and opc"`.
|
||||||
|
|
||||||
|
## Workflow Overview
|
||||||
|
|
||||||
|
The `predictions_batch` workflow:
|
||||||
|
1. Loads data using a custom SQL query
|
||||||
|
2. Prepares prediction configuration
|
||||||
|
3. Delegates to `prediction_process` child workflow which:
|
||||||
|
- Retrieves last timestamp for incremental processing
|
||||||
|
- Applies input data quality gates
|
||||||
|
- Executes MLFlow transform operation
|
||||||
|
- Validates transform response
|
||||||
|
- Executes MLFlow predict operation
|
||||||
|
- Validates predict response
|
||||||
|
- Delegates to `format_and_export_prediction` child workflow
|
||||||
|
4. The `format_and_export_prediction` workflow:
|
||||||
|
- Formats prediction data (normal or default)
|
||||||
|
- Exports to PI Web API (optional)
|
||||||
|
- Exports to OPC server (optional)
|
||||||
|
- Exports to PostgreSQL
|
||||||
|
- Writes metrics
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 1. Predictions Batch - Main Workflow Scenarios
|
||||||
|
|
||||||
|
### 1.1 Success Scenarios
|
||||||
|
|
||||||
|
#### Scenario 1.1.1: Happy Path - Complete Success
|
||||||
|
**Description**: Workflow completes successfully with valid SQL query and all activities succeed
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid `schedule_name`, `model_name`, `model_id`
|
||||||
|
- Valid `query` returning non-empty DataFrame
|
||||||
|
- Valid `schema`, `table_name`, `transform_table_name`
|
||||||
|
- Optional `datetime_columns` for timestamp parsing
|
||||||
|
- Optional `input_filters`, `mlflow_transform_filters`, `mlflow_predict_filters`
|
||||||
|
- Optional `path_priority`, `opc_output_config`, `pi_web_api_output_config`
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- `load_custom_query` returns DataFrame with data
|
||||||
|
- Workflow prepares prediction input with all configurations
|
||||||
|
- `prediction_process` child workflow executes successfully
|
||||||
|
- All gates pass with no issues
|
||||||
|
- Transform and predict operations succeed
|
||||||
|
- Data exported to PostgreSQL
|
||||||
|
- Metrics written
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- SQL query executed once
|
||||||
|
- `prediction_process` workflow called with correct parameters
|
||||||
|
- Data exists in PostgreSQL (predictions table)
|
||||||
|
- Metrics recorded
|
||||||
|
- No errors raised
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 1.2 Error Scenarios
|
||||||
|
|
||||||
|
#### Scenario 1.2.1: SQL Query Execution Error
|
||||||
|
**Description**: SQL query fails due to syntax error or connection issue
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Invalid SQL query (syntax error)
|
||||||
|
- Or database connection unavailable
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 1.2.2: Missing Required Parameters
|
||||||
|
**Description**: Essential parameters missing from input
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Missing `query` or `model_id` or `schema` or `table_name`
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 1.2.3: Invalid Datetime Column Specification
|
||||||
|
**Description**: Datetime column specified doesn't exist in query results
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- `datetime_columns: ['nonexistent_column']`
|
||||||
|
- Query results don't have this column
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 2. Prediction Process - Child Workflow Scenarios
|
||||||
|
|
||||||
|
### 2.1 Input gate Early Exit Scenarios
|
||||||
|
|
||||||
|
#### Scenario 2.1.1: Input Gate Triggers CONTINUE
|
||||||
|
**Description**: Input gate determines data should use previous prediction
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Data that should continue with input data as prediction
|
||||||
|
- `input_filters` configured with `POLICY: 'CONTINUE'`
|
||||||
|
- `path_priority` includes CONTINUE
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
|
||||||
|
#### Scenario 2.1.2: Input Gate Triggers STOP
|
||||||
|
**Description**: Input data quality gate fails with STOP policy
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Data with EMPTY_DATA or other critical issues
|
||||||
|
- `input_filters` configured with `POLICY: 'STOP'`
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
|
||||||
|
#### Scenario 2.1.3: Input Gate Triggers REPEAT
|
||||||
|
**Description**: Input gate determines data should repeat last prediction
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Data with quality issues that require using previous prediction
|
||||||
|
- `input_filters` configured with `POLICY: 'REPEAT'`
|
||||||
|
- `path_priority` includes REPEAT
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 2.2 Transform gate Early Exit Scenarios
|
||||||
|
|
||||||
|
#### Scenario 2.2.1: Transform Gate Triggers CONTINUE
|
||||||
|
**Description**: Transform response gate determines data should continue despite issues
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid input data
|
||||||
|
- Transform response has quality issues but policy is CONTINUE
|
||||||
|
- `mlflow_transform_filters` configured with `POLICY: 'CONTINUE'`
|
||||||
|
- `path_priority` includes CONTINUE
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 2.2.2: Transform Gate Triggers STOP
|
||||||
|
**Description**: Transform response validation fails with STOP policy
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid input data
|
||||||
|
- Transform response has critical errors
|
||||||
|
- `mlflow_transform_filters` configured with `POLICY: 'STOP'`
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 2.2.3: Transform Gate Triggers REPEAT
|
||||||
|
**Description**: Transform response gate determines data should repeat last prediction
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid input data
|
||||||
|
- Transform response has quality issues that require using previous prediction
|
||||||
|
- `mlflow_transform_filters` configured with `POLICY: 'REPEAT'`
|
||||||
|
- `path_priority` includes REPEAT
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 2.3 Predict gate Early Exit Scenarios
|
||||||
|
|
||||||
|
#### Scenario 2.3.1: Predict Gate Triggers CONTINUE
|
||||||
|
**Description**: Predict response gate determines data should continue despite issues
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid input and transform data
|
||||||
|
- Predict response has quality issues but policy is CONTINUE
|
||||||
|
- `mlflow_predict_filters` configured with `POLICY: 'CONTINUE'`
|
||||||
|
- `path_priority` includes CONTINUE
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 2.3.2: Predict Gate Triggers STOP
|
||||||
|
**Description**: Prediction validation fails with STOP policy
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid input and transform
|
||||||
|
- Predict response has critical errors
|
||||||
|
- `mlflow_predict_filters` configured with `POLICY: 'STOP'`
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 2.3.3: Predict Gate Triggers REPEAT
|
||||||
|
**Description**: Predict response gate determines data should repeat last prediction
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid input and transform data
|
||||||
|
- Predict response has quality issues that require using previous prediction
|
||||||
|
- `mlflow_predict_filters` configured with `POLICY: 'REPEAT'`
|
||||||
|
- `path_priority` includes REPEAT
|
||||||
|
|
||||||
|
**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
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 3. Format and Export Prediction - Child Workflow Scenarios
|
||||||
|
|
||||||
|
### 3.1 Success Scenarios
|
||||||
|
|
||||||
|
#### Scenario 3.1.1: Default Prediction Export
|
||||||
|
**Description**: Error prediction path creates default prediction
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- `path_flag: 'ERROR'` or other non-None value (not STOP/CONTINUE/REPEAT)
|
||||||
|
- `comment` provided with error details
|
||||||
|
|
||||||
|
**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.2: Export with OPC only
|
||||||
|
**Description**: Export to PostgreSQL and OPC server only (no PI Web API)
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- `path_flag: None`
|
||||||
|
- `opc_output_config` configured with valid OPC settings
|
||||||
|
- `pi_web_api_output_config: None` or `{}`
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- Normal formatting
|
||||||
|
- PostgreSQL export executed
|
||||||
|
- OPC export executed
|
||||||
|
- PI Web API activity skipped
|
||||||
|
- Metrics written with OPC metrics
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- PI Web API activity NOT called
|
||||||
|
- OPC activity called
|
||||||
|
- PostgreSQL export called
|
||||||
|
- Metrics written with `opc_metrics` populated
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 3.1.3: Export with PI Web API only
|
||||||
|
**Description**: Export to PostgreSQL and PI Web API only (no OPC)
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- `path_flag: None`
|
||||||
|
- `pi_web_api_output_config` configured with valid PI Web API settings
|
||||||
|
- `opc_output_config: None` or `{}`
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- Normal formatting
|
||||||
|
- PostgreSQL export executed
|
||||||
|
- PI Web API export executed
|
||||||
|
- OPC activity skipped
|
||||||
|
- Metrics written without OPC metrics
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- OPC activity NOT called
|
||||||
|
- PI Web API activity called
|
||||||
|
- PostgreSQL export called
|
||||||
|
- Metrics written with empty `opc_metrics`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 3.1.4: Export Without Optional Outputs
|
||||||
|
**Description**: Export only to PostgreSQL (no OPC or PI Web API)
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- `path_flag: None`
|
||||||
|
- `opc_output_config: None` or `{}`
|
||||||
|
- `pi_web_api_output_config: None` or `{}`
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- Normal formatting
|
||||||
|
- Only PostgreSQL export executed
|
||||||
|
- OPC and PI Web API activities skipped
|
||||||
|
- Metrics written without OPC metrics
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- PI Web API activity NOT called
|
||||||
|
- OPC activity NOT called
|
||||||
|
- PostgreSQL export called
|
||||||
|
- Metrics written with empty `opc_metrics`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 3.1.5: Export Without Transformed Data
|
||||||
|
**Description**: Only prediction exported, no transform table
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- `path_flag: None`
|
||||||
|
- `transformed_data: None` or `save_transform: False`
|
||||||
|
- `opc_output_config: None` or `{}`
|
||||||
|
- `pi_web_api_output_config: None` or `{}`
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- Only prediction formatted and exported
|
||||||
|
- Transform export skipped
|
||||||
|
- Single PostgreSQL write
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- `format_transformed_data` NOT called
|
||||||
|
- One PostgreSQL export
|
||||||
|
- Transform table remains empty
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### 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
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid prediction
|
||||||
|
- PI Web API service unavailable or invalid config
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- `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 (when applicable)
|
||||||
|
- Workflow completes
|
||||||
|
- PostgreSQL contains the prediction with `prediction_confidence` 13 and expected `comments`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 3.2.2: OPC Write Error
|
||||||
|
**Description**: OPC server write fails
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid prediction
|
||||||
|
- OPC server unavailable or invalid configuration
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- `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 (when applicable)
|
||||||
|
- Workflow completes
|
||||||
|
- PostgreSQL contains the prediction with `prediction_confidence` 12 and expected `comments`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 3.2.4: OPC Session / Channel Bad* (Tier-1)
|
||||||
|
**Description**: OPC write fails with a Tier-1 session or channel status (e.g. `BadSessionIdInvalid`) while transport may still appear open on the client
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid prediction and OPC output config
|
||||||
|
- Mock or server returning Tier-1 `UaStatusCodeError` on write (no write retry in the same activity)
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- `write_opc_data` fails forward for affected tags; background reconnect may be scheduled if `OPC_RECONNECTION_INTERVAL` allows
|
||||||
|
- Workflow **completes**
|
||||||
|
- PostgreSQL row uses **`prediction_confidence` 14** and comment prefix `OPC UA session/channel error:` (including OPC status name)
|
||||||
|
- `opc_write_attempts_total` records `result=BadSessionIdInvalid` (or matching status); no second write attempt in the same activity
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- Workflow completes
|
||||||
|
- `prediction_confidence = 14`
|
||||||
|
- `comments` matches `OPC UA session/channel error:%`
|
||||||
|
- Generic OPC error confidence **12** is not used for this case
|
||||||
|
|
||||||
|
**Reference**: [docs/opc-communication.md](../docs/opc-communication.md), plan `.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### Scenario 3.2.5: OPC Write Blocked During Reconnect
|
||||||
|
**Description**: A write is attempted while the repository is reconnecting (session not ready)
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid prediction
|
||||||
|
- Simulated slow reconnect (e.g. delayed `connect`) or concurrent writes where the first triggers reconnect
|
||||||
|
|
||||||
|
**Expected Behavior**:
|
||||||
|
- Second write (or parallel write) is rejected **immediately** when reconnect is in progress or `_session_ready` is cleared — **without** calling `write_value`
|
||||||
|
- No wait/sleep on the write path; no duplicate `connect` from parallel writers (connection lock)
|
||||||
|
- `prediction_confidence = 14`, `comments = OPC UA reconnect in progress` (distinguish from Tier-1 `Bad*` via comment prefix in SQL)
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- At most one reconnect sequence (`disconnect` + `connect`) for the overlapping window
|
||||||
|
- No write retry after failure
|
||||||
|
- Tests in `test_opc_repository` (unit) and optional e2e in `test_predictions_batch_format_export.py`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
#### 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
|
||||||
|
|
||||||
|
**Input**:
|
||||||
|
- Valid prediction
|
||||||
|
- Two prediction tags configured
|
||||||
|
- PI Web API returns partial success (one tag succeeds, one fails)
|
||||||
|
|
||||||
|
**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 (single activity attempt; no retry loop)
|
||||||
|
|
||||||
|
**Assertions**:
|
||||||
|
- One tag written successfully
|
||||||
|
- One tag failed
|
||||||
|
- Error confidence set in prediction
|
||||||
|
- Error notification sent
|
||||||
|
- Workflow completes
|
||||||
|
|
||||||
|
---
|
||||||
76
e2e/test_child_workflows_e2e.py
Normal file
76
e2e/test_child_workflows_e2e.py
Normal file
@@ -0,0 +1,76 @@
|
|||||||
|
"""
|
||||||
|
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'
|
||||||
136
e2e/test_minio_offload.py
Normal file
136
e2e/test_minio_offload.py
Normal file
@@ -0,0 +1,136 @@
|
|||||||
|
"""
|
||||||
|
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
|
||||||
211
e2e/test_opc_real_server.py
Normal file
211
e2e/test_opc_real_server.py
Normal file
@@ -0,0 +1,211 @@
|
|||||||
|
"""
|
||||||
|
E2E tests for OPC export using an in-process asyncua server and real OpcRepository.
|
||||||
|
|
||||||
|
Covers scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 from e2e/scenarios.md.
|
||||||
|
Mock-based OPC tests remain in test_predictions_batch_format_export.py.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
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 e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config
|
||||||
|
from e2e.test_predictions_batch_format_export import get_base_input_data
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.activities.opc import OPC_RECONNECT_IN_PROGRESS_COMMENT
|
||||||
|
from laborious.utils.repository.opc_repository import OpcRepository
|
||||||
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||||
|
|
||||||
|
|
||||||
|
async def _reconnect_under_held_lock(
|
||||||
|
repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Hold the connection lock until released, then reconnect.
|
||||||
|
|
||||||
|
The lock must stay held for the whole workflow, and how long the workflow takes to
|
||||||
|
reach the OPC write is not bounded — a fixed sleep here lets `_reconnect_locked`
|
||||||
|
finish first, which sets `_session_ready` again and the write then succeeds.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
repo (OpcRepository): Connected repository.
|
||||||
|
acquired (asyncio.Event): Set once the lock is held, before any write can race.
|
||||||
|
release (asyncio.Event): Awaited while holding the lock; set by the caller.
|
||||||
|
"""
|
||||||
|
async with repo._connection_lock:
|
||||||
|
acquired.set()
|
||||||
|
await release.wait()
|
||||||
|
await repo._reconnect_locked()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
@pytest.mark.opc
|
||||||
|
async def test_scenario_3_1_2_export_with_opc_only_real_server(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker_real_opc: Worker,
|
||||||
|
test_activities_real_opc: Activities,
|
||||||
|
opc_e2e_server: OpcE2ETestServer,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.1.2 (real OPC): connect, write prediction and confidence, verify server values.
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 412
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
|
||||||
|
input_data['pi_web_api_output_config'] = None
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client,
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
make_workflow_id('test-opc-real-happy'),
|
||||||
|
)
|
||||||
|
|
||||||
|
test_activities_real_opc.pi_web_api_client.write_value.assert_not_called()
|
||||||
|
assert await opc_e2e_server.read_prediction() == pytest.approx(0.5)
|
||||||
|
assert await opc_e2e_server.read_confidence() == pytest.approx(0.0)
|
||||||
|
assert_prediction(postgres_engine, model_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
@pytest.mark.opc
|
||||||
|
async def test_scenario_3_2_2_opc_write_error_real_server(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker_real_opc: Worker,
|
||||||
|
test_activities_real_opc: Activities,
|
||||||
|
opc_e2e_server: OpcE2ETestServer,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.2 (real OPC): unknown NodeId yields generic write failure (confidence 12).
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 422
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
node_ids = opc_e2e_server.node_ids
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = build_opc_output_config(
|
||||||
|
node_ids,
|
||||||
|
prediction_tag=UNKNOWN_NODE_ID,
|
||||||
|
confidence_tag=UNKNOWN_NODE_ID,
|
||||||
|
)
|
||||||
|
input_data['pi_web_api_output_config'] = None
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client,
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
make_workflow_id('test-opc-real-bad-node'),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_prediction(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
prediction_confidence=12,
|
||||||
|
comments='Some data could not be written to OPC servers',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
@pytest.mark.opc
|
||||||
|
async def test_scenario_3_2_4_opc_session_bad_real_server(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker_real_opc: Worker,
|
||||||
|
test_activities_real_opc: Activities,
|
||||||
|
opc_e2e_server: OpcE2ETestServer,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.4 (real OPC): server PreWrite fault injects BadSessionIdInvalid (confidence 14).
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 424
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
opc_e2e_server.set_session_bad_on_write(True)
|
||||||
|
try:
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = build_opc_output_config(
|
||||||
|
opc_e2e_server.node_ids,
|
||||||
|
prediction_only=True,
|
||||||
|
)
|
||||||
|
input_data['pi_web_api_output_config'] = None
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client,
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
make_workflow_id('test-opc-real-session-bad'),
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
opc_e2e_server.set_session_bad_on_write(False)
|
||||||
|
|
||||||
|
assert_prediction(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
prediction_confidence=14,
|
||||||
|
comments_contains='OPC UA session/channel error: BadSessionIdInvalid',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
@pytest.mark.opc
|
||||||
|
async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker_real_opc: Worker,
|
||||||
|
test_activities_real_opc: Activities,
|
||||||
|
opc_e2e_server: OpcE2ETestServer,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.5 (real OPC): writes rejected while reconnect holds the connection lock.
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 425
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
repo = test_activities_real_opc.opc_repository['1']
|
||||||
|
repo._session_ready.clear()
|
||||||
|
acquired = asyncio.Event()
|
||||||
|
release = asyncio.Event()
|
||||||
|
reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release))
|
||||||
|
await acquired.wait()
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
|
||||||
|
input_data['pi_web_api_output_config'] = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client,
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
make_workflow_id('test-opc-real-reconnect-block'),
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
release.set()
|
||||||
|
await reconnect_task
|
||||||
|
|
||||||
|
assert_prediction(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
prediction_confidence=14,
|
||||||
|
comments_contains=OPC_RECONNECT_IN_PROGRESS_COMMENT,
|
||||||
|
)
|
||||||
800
e2e/test_predictions_batch_format_export.py
Normal file
800
e2e/test_predictions_batch_format_export.py
Normal file
@@ -0,0 +1,800 @@
|
|||||||
|
"""
|
||||||
|
End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from decimal import Decimal
|
||||||
|
from typing import Any, cast
|
||||||
|
from unittest.mock import AsyncMock, call
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
base_input_data = {
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 301,
|
||||||
|
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 301',
|
||||||
|
'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'],
|
||||||
|
}
|
||||||
|
|
||||||
|
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 {
|
||||||
|
**base_input_data,
|
||||||
|
'model_id': model_id,
|
||||||
|
'query': base_query.format(model_id=model_id),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_1_1_default_prediction_export(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
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
|
||||||
|
|
||||||
|
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'},
|
||||||
|
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||||
|
}
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_1'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01 12:00:00+0000',
|
||||||
|
'Value': 0,
|
||||||
|
},
|
||||||
|
metadata={
|
||||||
|
'model_id': 311,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_2'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01 12:00:00+0000',
|
||||||
|
'Value': 2,
|
||||||
|
},
|
||||||
|
metadata={
|
||||||
|
'model_id': 311,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
],
|
||||||
|
any_order=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||||
|
opc_write_data.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
'addr_1',
|
||||||
|
0,
|
||||||
|
'float',
|
||||||
|
{
|
||||||
|
'model_id': 311,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
'addr_2',
|
||||||
|
2,
|
||||||
|
'float',
|
||||||
|
{
|
||||||
|
'model_id': 311,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
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',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_1_2_export_with_opc_only(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.1.2: Export with OPC only
|
||||||
|
|
||||||
|
Description:
|
||||||
|
Export to PostgreSQL and OPC server only (no PI Web API).
|
||||||
|
|
||||||
|
Expected Behavior:
|
||||||
|
- Normal formatting
|
||||||
|
- PostgreSQL export executed
|
||||||
|
- OPC export executed
|
||||||
|
- PI Web API activity skipped
|
||||||
|
- Metrics written with OPC metrics
|
||||||
|
|
||||||
|
Assertions:
|
||||||
|
- PI Web API activity NOT called
|
||||||
|
- OPC activity called
|
||||||
|
- PostgreSQL export called
|
||||||
|
- Metrics written with opc_metrics populated
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
model_id = 312
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
input_data['pi_web_api_output_config'] = None # No PI Web API config
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-only')
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||||
|
opc_write_data.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
'addr_1',
|
||||||
|
0.5,
|
||||||
|
'float',
|
||||||
|
{
|
||||||
|
'model_id': 312,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
'addr_2',
|
||||||
|
0,
|
||||||
|
'float',
|
||||||
|
{
|
||||||
|
'model_id': 312,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
test_activities.pi_web_api_client.write_value.assert_not_called()
|
||||||
|
|
||||||
|
assert_prediction(postgres_engine, model_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.1.3: Export with PI Web API only
|
||||||
|
|
||||||
|
Description:
|
||||||
|
Export to PostgreSQL and PI Web API only (no OPC).
|
||||||
|
|
||||||
|
Expected Behavior:
|
||||||
|
- Normal formatting
|
||||||
|
- PostgreSQL export executed
|
||||||
|
- PI Web API export executed
|
||||||
|
- OPC activity skipped
|
||||||
|
- Metrics written without OPC metrics
|
||||||
|
|
||||||
|
Assertions:
|
||||||
|
- OPC activity NOT called
|
||||||
|
- PI Web API activity called
|
||||||
|
- PostgreSQL export called
|
||||||
|
- Metrics written with empty opc_metrics
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
model_id = 313
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['pi_web_api_output_config'] = {
|
||||||
|
'endpoint': 'test_endpoint',
|
||||||
|
'prediction_tags': {'tag_1': 'web_id_1'},
|
||||||
|
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||||
|
}
|
||||||
|
input_data['opc_output_config'] = None # No OPC config
|
||||||
|
|
||||||
|
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(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_1'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01 12:00:00+0000',
|
||||||
|
'Value': 0.5,
|
||||||
|
},
|
||||||
|
metadata={
|
||||||
|
'model_id': 313,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_2'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01 12:00:00+0000',
|
||||||
|
'Value': 0,
|
||||||
|
},
|
||||||
|
metadata={
|
||||||
|
'model_id': 313,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
],
|
||||||
|
any_order=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||||
|
opc_write_data.assert_not_called()
|
||||||
|
|
||||||
|
assert_prediction(postgres_engine, model_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.1.4: Export Without Optional Outputs
|
||||||
|
|
||||||
|
Description:
|
||||||
|
Export only to PostgreSQL (no OPC or PI Web API).
|
||||||
|
|
||||||
|
Expected Behavior:
|
||||||
|
- Normal formatting
|
||||||
|
- Only PostgreSQL export executed
|
||||||
|
- OPC and PI Web API activities skipped
|
||||||
|
- Metrics written without OPC metrics
|
||||||
|
|
||||||
|
Assertions:
|
||||||
|
- PI Web API activity NOT called
|
||||||
|
- OPC activity NOT called
|
||||||
|
- PostgreSQL export called
|
||||||
|
- Metrics written with empty opc_metrics
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
model_id = 314
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
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()
|
||||||
|
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||||
|
opc_write_data.assert_not_called()
|
||||||
|
|
||||||
|
assert_prediction(postgres_engine, model_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_1_5_export_without_transformed_data(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.1.5: Export Without Transformed Data
|
||||||
|
|
||||||
|
Description:
|
||||||
|
Only prediction exported, no transform table.
|
||||||
|
|
||||||
|
Expected Behavior:
|
||||||
|
- Only prediction formatted and exported
|
||||||
|
- Transform export skipped
|
||||||
|
- Single PostgreSQL write
|
||||||
|
|
||||||
|
Assertions:
|
||||||
|
- format_transformed_data NOT called
|
||||||
|
- One PostgreSQL export
|
||||||
|
- Transform table remains empty
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
model_id = 315
|
||||||
|
|
||||||
|
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}')
|
||||||
|
)
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['save_transform'] = False # Don't save transformed data
|
||||||
|
input_data['pi_web_api_output_config'] = {
|
||||||
|
'endpoint': 'test_endpoint',
|
||||||
|
'prediction_tags': {'tag_1': 'web_id_1'},
|
||||||
|
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||||
|
}
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_1'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01 12:00:00+0000',
|
||||||
|
'Value': 0.5,
|
||||||
|
},
|
||||||
|
metadata={
|
||||||
|
'model_id': 315,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_2'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01 12:00:00+0000',
|
||||||
|
'Value': 0,
|
||||||
|
},
|
||||||
|
metadata={
|
||||||
|
'model_id': 315,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
],
|
||||||
|
any_order=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||||
|
opc_write_data.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
'addr_1',
|
||||||
|
0.5,
|
||||||
|
'float',
|
||||||
|
{
|
||||||
|
'model_id': 315,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
'addr_2',
|
||||||
|
0,
|
||||||
|
'float',
|
||||||
|
{
|
||||||
|
'model_id': 315,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
with postgres_engine.connect() as conn:
|
||||||
|
result_query = conn.execute(
|
||||||
|
text(
|
||||||
|
f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
|
||||||
|
)
|
||||||
|
)
|
||||||
|
count = result_query.scalar()
|
||||||
|
assert count == 0, f'Expected transform table to be empty, but found {count} records'
|
||||||
|
|
||||||
|
assert_prediction(postgres_engine, model_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
notification_inserts,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.1: PI Web API Write Error
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
test_activities.pi_web_api_client.write_value.side_effect = Exception(
|
||||||
|
'PI Web API service unavailable'
|
||||||
|
)
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['pi_web_api_output_config'] = {
|
||||||
|
'endpoint': 'test_endpoint',
|
||||||
|
'prediction_tags': {'tag_1': 'web_id_1'},
|
||||||
|
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||||
|
}
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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',
|
||||||
|
)
|
||||||
|
assert notification_inserts.call_count >= 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_2_2_opc_write_error(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.2: OPC Write Error
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
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',
|
||||||
|
'level': NotificationLevel.ERROR,
|
||||||
|
'attachment_content': 'OPC server unavailable',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
input_data['pi_web_api_output_config'] = {
|
||||||
|
'endpoint': 'test_endpoint',
|
||||||
|
'prediction_tags': {'tag_1': 'web_id_1'},
|
||||||
|
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||||
|
}
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error')
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_prediction(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
prediction_confidence=12,
|
||||||
|
comments='Some data could not be written to OPC servers',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_2_4_opc_session_bad_error(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.4: OPC session/channel Tier-1 Bad* (e.g. BadSessionIdInvalid).
|
||||||
|
|
||||||
|
PostgreSQL stores prediction_confidence 14 and a stable session error comment.
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
model_id = 324
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
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': 'BadSessionIdInvalid',
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': NotificationLevel.ERROR,
|
||||||
|
'attachment_content': 'BadSessionIdInvalid',
|
||||||
|
'opc_error_kind': 'session_bad',
|
||||||
|
'opc_status': 'BadSessionIdInvalid',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
input_data['pi_web_api_output_config'] = None
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-session-bad')
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_prediction(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
prediction_confidence=14,
|
||||||
|
comments='OPC UA session/channel error: BadSessionIdInvalid',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_3_2_3_pi_web_api_partial_write_error(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Scenario 3.2.3: PI Web API Partial Write Error
|
||||||
|
|
||||||
|
Partial PI write: confidence 13, descriptive comments, workflow completes
|
||||||
|
without an activity retry loop.
|
||||||
|
"""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
model_id = 323
|
||||||
|
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
|
||||||
|
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'] = {
|
||||||
|
'endpoint': 'test_endpoint',
|
||||||
|
'prediction_tags': {'tag_1': 'web_id_1', 'tag_3': 'web_id_3'},
|
||||||
|
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||||
|
}
|
||||||
|
input_data['opc_output_config'] = {
|
||||||
|
'1': {
|
||||||
|
'prediction_tags': {
|
||||||
|
'addr_1': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'confidence_tags': {
|
||||||
|
'addr_2': {
|
||||||
|
'data_type': 'float',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error')
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_prediction(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
prediction_confidence=13,
|
||||||
|
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.",
|
||||||
|
)
|
||||||
186
e2e/test_predictions_batch_joblib_fallback.py
Normal file
186
e2e/test_predictions_batch_joblib_fallback.py
Normal file
@@ -0,0 +1,186 @@
|
|||||||
|
"""
|
||||||
|
End-to-end tests for the joblib fallback on the MLflow model load paths.
|
||||||
|
|
||||||
|
`joblib.load` is NOT mocked here: the artifact fixture holds a real sklearn
|
||||||
|
estimator serialized with `joblib.dump(..., compress=3)`, so a successful run
|
||||||
|
proves the fallback really deserializes a compressed pickle and serves it.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pandas as pd
|
||||||
|
import pytest
|
||||||
|
from sqlalchemy import text
|
||||||
|
from temporalio.testing import WorkflowEnvironment
|
||||||
|
from temporalio.worker import Worker
|
||||||
|
|
||||||
|
from e2e.conftest import JOBLIB_FALLBACK_FLAVOR_ERROR, JOBLIB_FALLBACK_MODEL_NAME
|
||||||
|
from e2e.helpers import (
|
||||||
|
assert_prediction,
|
||||||
|
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_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||||
|
|
||||||
|
|
||||||
|
def build_input_data(model_id: int, model_name: str) -> dict:
|
||||||
|
"""Workflow input for a plain transform + predict run of a given model."""
|
||||||
|
return {
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_id': model_id,
|
||||||
|
'query': base_query.format(model_id=model_id),
|
||||||
|
'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'],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def expected_fallback_prediction(estimator) -> float:
|
||||||
|
"""Prediction the fixture estimator produces for the mocked transform output."""
|
||||||
|
features = pd.DataFrame(
|
||||||
|
{
|
||||||
|
'feature_1': [0.234],
|
||||||
|
'feature_2': [0.783],
|
||||||
|
'timestamp': ['2024-01-01 12:00:00'],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return float(estimator.predict(features)[0])
|
||||||
|
|
||||||
|
|
||||||
|
def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
|
||||||
|
"""Assert exactly one Good prediction row whose value matches `expected`."""
|
||||||
|
with postgres_engine.connect() as conn:
|
||||||
|
rows = conn.execute(
|
||||||
|
text(
|
||||||
|
'SELECT prediction, prediction_status '
|
||||||
|
'FROM predictions_schema.predictions '
|
||||||
|
f'WHERE model_id = {model_id}'
|
||||||
|
)
|
||||||
|
).fetchall()
|
||||||
|
|
||||||
|
assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
|
||||||
|
assert abs(float(rows[0][0]) - expected) < tolerance, (
|
||||||
|
f'Expected prediction close to {expected}, got {rows[0][0]}'
|
||||||
|
)
|
||||||
|
assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"
|
||||||
|
|
||||||
|
|
||||||
|
def logged_messages(mock_logger, level: str) -> str:
|
||||||
|
"""Concatenate every message logged at a given custom level."""
|
||||||
|
calls = getattr(mock_logger, f'custom_{level}').call_args_list
|
||||||
|
return ' '.join(str(logged.args[0]) for logged in calls)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_prediction_workflow_recovered_by_joblib_fallback(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
joblib_fallback_artifact,
|
||||||
|
mlflow_client,
|
||||||
|
mock_logger,
|
||||||
|
):
|
||||||
|
"""Flavor loader fails, joblib fallback loads the real compressed pickle, prediction lands."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 271
|
||||||
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||||
|
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback')
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_single_prediction_close(
|
||||||
|
postgres_engine,
|
||||||
|
model_id,
|
||||||
|
expected_fallback_prediction(joblib_fallback_artifact['estimator']),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert mlflow_client.download_artifacts.call_count == 1
|
||||||
|
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
|
||||||
|
|
||||||
|
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged_messages(mock_logger, 'warning')
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_prediction_workflow_without_fallback_is_unchanged(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
mlflow_client,
|
||||||
|
mock_mlflow_models,
|
||||||
|
):
|
||||||
|
"""A model whose flavor loader works keeps the current result and downloads nothing."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 272
|
||||||
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||||
|
input_data = build_input_data(model_id, 'test_model')
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-no-fallback')
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||||
|
mock_mlflow_models['predict_model'].predict.assert_called()
|
||||||
|
mlflow_client.download_artifacts.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_prediction_workflow_reports_both_causes_when_fallback_fails(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
joblib_fallback_artifact_without_pickle,
|
||||||
|
mlflow_client,
|
||||||
|
mock_logger,
|
||||||
|
):
|
||||||
|
"""Both attempts fail: workflow stops and the logged traceback carries both causes."""
|
||||||
|
mlflow_client.download_artifacts.return_value = joblib_fallback_artifact_without_pickle
|
||||||
|
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 273
|
||||||
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||||
|
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback-failed')
|
||||||
|
)
|
||||||
|
|
||||||
|
assert_stop(postgres_engine, model_id)
|
||||||
|
|
||||||
|
logged = logged_messages(mock_logger, 'debug')
|
||||||
|
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged
|
||||||
|
assert 'No .pkl file found in prediction_model artifact' in logged
|
||||||
302
e2e/test_predictions_batch_main_workflow.py
Normal file
302
e2e/test_predictions_batch_main_workflow.py
Normal file
@@ -0,0 +1,302 @@
|
|||||||
|
"""
|
||||||
|
End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
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.activities.activities import Activities
|
||||||
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_1_1_1_happy_path_complete_success(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
with postgres_engine.begin() as conn:
|
||||||
|
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
|
||||||
|
(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))
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'metadata': {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 123,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 123,
|
||||||
|
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 123',
|
||||||
|
'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'],
|
||||||
|
}
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client,
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
make_workflow_id('test-predictions-batch'),
|
||||||
|
)
|
||||||
|
|
||||||
|
schema_name = 'predictions_schema'
|
||||||
|
with postgres_engine.connect() as conn:
|
||||||
|
result_query = conn.execute(
|
||||||
|
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()
|
||||||
|
assert len(prediction_rows) == 1
|
||||||
|
row = prediction_rows[0]
|
||||||
|
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] == ''
|
||||||
|
|
||||||
|
result_query = conn.execute(
|
||||||
|
text(
|
||||||
|
f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123'
|
||||||
|
)
|
||||||
|
)
|
||||||
|
transformed_rows = result_query.fetchall()
|
||||||
|
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
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_1_2_1_sql_query_execution_error(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""Invalid SQL: workflow may complete with early exit; no prediction rows."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'metadata': {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 128,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 128,
|
||||||
|
'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': {}},
|
||||||
|
},
|
||||||
|
'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',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client,
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
make_workflow_id('test-sql-error'),
|
||||||
|
)
|
||||||
|
|
||||||
|
with postgres_engine.connect() as conn:
|
||||||
|
count = conn.execute(
|
||||||
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
|
||||||
|
).scalar()
|
||||||
|
assert count == 0
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_1_2_2_missing_required_parameters(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""Missing query: workflow does not produce predictions and is terminated explicitly."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'metadata': {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 129,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 129,
|
||||||
|
'schema': 'predictions_schema',
|
||||||
|
'table_name': 'predictions',
|
||||||
|
'transform_table_name': 'transformed_data',
|
||||||
|
}
|
||||||
|
|
||||||
|
handle = await client.start_workflow(
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
id=make_workflow_id('test-missing-param'),
|
||||||
|
task_queue='test-queue',
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_1_2_3_invalid_datetime_column_specification(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
):
|
||||||
|
"""Invalid datetime column: no predictions persisted; workflow terminated after validation."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
|
||||||
|
with postgres_engine.begin() as conn:
|
||||||
|
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')
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'metadata': {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 130,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'schedule_name': 'test-schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 130,
|
||||||
|
'query': 'SELECT timestamp, variable, value FROM predictions_schema.laborious_data WHERE model_id = 130',
|
||||||
|
'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': ['nonexistent_column'],
|
||||||
|
}
|
||||||
|
|
||||||
|
handle = await client.start_workflow(
|
||||||
|
PredictionsBatch.run,
|
||||||
|
input_data,
|
||||||
|
id=make_workflow_id('test-invalid-datetime-col'),
|
||||||
|
task_queue='test-queue',
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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)')
|
||||||
385
e2e/test_predictions_batch_prediction_process.py
Normal file
385
e2e/test_predictions_batch_prediction_process.py
Normal file
@@ -0,0 +1,385 @@
|
|||||||
|
"""
|
||||||
|
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from decimal import Decimal
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pandas as pd
|
||||||
|
import pytest
|
||||||
|
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',
|
||||||
|
'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'],
|
||||||
|
}
|
||||||
|
|
||||||
|
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 {
|
||||||
|
**base_input_data,
|
||||||
|
'model_id': model_id,
|
||||||
|
'query': base_query.format(model_id=model_id),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
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_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')
|
||||||
|
|
||||||
|
|
||||||
|
@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
|
||||||
|
|
||||||
|
|
||||||
|
@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
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_1_1_input_gate_triggers_continue(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
mock_mlflow_models,
|
||||||
|
):
|
||||||
|
"""Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 211
|
||||||
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||||
|
input_data = get_base_input_data(model_id)
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
|
||||||
|
)
|
||||||
|
assert_continue(postgres_engine, model_id)
|
||||||
|
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
||||||
|
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_1_2_input_gate_triggers_stop(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
mock_mlflow_models,
|
||||||
|
):
|
||||||
|
"""Input gate STOP: no export, no MLflow."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 212
|
||||||
|
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'] = 'STOP'
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
|
||||||
|
)
|
||||||
|
assert_stop(postgres_engine, model_id)
|
||||||
|
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_1_3_input_gate_triggers_repeat(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
mock_mlflow_models,
|
||||||
|
):
|
||||||
|
"""Input gate REPEAT with existing history."""
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 213
|
||||||
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||||
|
data = insert_sample_prediction(postgres_engine, 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')
|
||||||
|
)
|
||||||
|
assert_repeat(postgres_engine, model_id, data)
|
||||||
|
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@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
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_2_1_transform_gate_triggers_continue(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
bad_data_model,
|
||||||
|
):
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 221
|
||||||
|
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'
|
||||||
|
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='Unknown MLFlow API error',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_2_2_transform_gate_triggers_stop(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
bad_data_model,
|
||||||
|
mock_mlflow_models,
|
||||||
|
):
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 222
|
||||||
|
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'
|
||||||
|
await start_and_await_workflow(
|
||||||
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
|
||||||
|
)
|
||||||
|
assert_stop(postgres_engine, model_id)
|
||||||
|
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_2_3_transform_gate_triggers_repeat(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
bad_data_model,
|
||||||
|
):
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 223
|
||||||
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||||
|
data = insert_sample_prediction(postgres_engine, model_id)
|
||||||
|
input_data = get_base_input_data(model_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)
|
||||||
|
|
||||||
|
|
||||||
|
@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,
|
||||||
|
):
|
||||||
|
"""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')
|
||||||
|
)
|
||||||
|
assert_stop(postgres_engine, model_id)
|
||||||
|
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_3_1_predict_gate_triggers_continue(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
bad_predict_model,
|
||||||
|
):
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 231
|
||||||
|
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'
|
||||||
|
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='Unknown MLFlow API error',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_3_2_predict_gate_triggers_stop(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
bad_predict_model,
|
||||||
|
):
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 232
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
input_data = get_base_input_data(model_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)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_scenario_2_3_3_predict_gate_triggers_repeat(
|
||||||
|
temporal_test_env: WorkflowEnvironment,
|
||||||
|
temporal_worker: Worker,
|
||||||
|
test_activities: Activities,
|
||||||
|
postgres_engine,
|
||||||
|
bad_predict_model,
|
||||||
|
):
|
||||||
|
client = temporal_test_env.client
|
||||||
|
model_id = 233
|
||||||
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||||
|
data = insert_sample_prediction(postgres_engine, model_id)
|
||||||
|
input_data = get_base_input_data(model_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)
|
||||||
|
|
||||||
|
|
||||||
|
@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)
|
||||||
13
encode.sh
Executable file
13
encode.sh
Executable file
@@ -0,0 +1,13 @@
|
|||||||
|
source ./venv/bin/activate
|
||||||
|
|
||||||
|
pip install pathspec
|
||||||
|
pip install pyyaml
|
||||||
|
|
||||||
|
echo "
|
||||||
|
.git" >> .gitignore
|
||||||
|
|
||||||
|
python encrypt.py ./ code --ignore .gitignore --chunk-size 100000
|
||||||
|
|
||||||
|
sed -i '/.git/d' .gitignore
|
||||||
|
|
||||||
|
xdg-open .
|
||||||
113
encrypt.py
Normal file
113
encrypt.py
Normal file
@@ -0,0 +1,113 @@
|
|||||||
|
import os
|
||||||
|
import argparse
|
||||||
|
from pathspec import PathSpec
|
||||||
|
import yaml # type: ignore
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
'''
|
||||||
|
Usage:
|
||||||
|
python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000
|
||||||
|
'''
|
||||||
|
|
||||||
|
|
||||||
|
def load_ignore_patterns(ignore_file, include_library):
|
||||||
|
# Ensure the .gitignore file exists
|
||||||
|
if not os.path.exists(ignore_file):
|
||||||
|
raise FileNotFoundError(f"Ignore file not found at {ignore_file}")
|
||||||
|
|
||||||
|
# Load and parse the .gitignore patterns
|
||||||
|
with open(ignore_file, 'r') as file:
|
||||||
|
patterns = file.readlines()
|
||||||
|
if not include_library:
|
||||||
|
patterns.append('**/deploy/library/')
|
||||||
|
|
||||||
|
spec = PathSpec.from_lines('gitwildmatch', patterns)
|
||||||
|
return spec
|
||||||
|
|
||||||
|
|
||||||
|
def is_ignored(file_path, spec):
|
||||||
|
"""Check if a file should be ignored based on the ignore patterns."""
|
||||||
|
return spec.match_file(file_path) if spec else False
|
||||||
|
|
||||||
|
|
||||||
|
def encode_file_tree_to_yaml(directory, ignore_file, include_library):
|
||||||
|
"""Encode the file tree into a single YAML file."""
|
||||||
|
ignore_patterns = load_ignore_patterns(
|
||||||
|
ignore_file, include_library) if ignore_file else None
|
||||||
|
file_tree: dict[str, Any] = {}
|
||||||
|
|
||||||
|
for root, dirs, files in os.walk(directory):
|
||||||
|
# Skip ignored directories
|
||||||
|
dirs[:] = [d for d in dirs if not is_ignored(
|
||||||
|
os.path.join(root, d), ignore_patterns)]
|
||||||
|
|
||||||
|
for file in files:
|
||||||
|
file_path = os.path.join(root, file)
|
||||||
|
|
||||||
|
# Skip ignored files
|
||||||
|
if is_ignored(file_path, ignore_patterns):
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Read file content
|
||||||
|
try:
|
||||||
|
with open(file_path, 'r', encoding='utf-8') as f:
|
||||||
|
content = f.read()
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error reading file {file_path}: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
# Create nested dictionary structure
|
||||||
|
path_parts = os.path.relpath(file_path, directory).split(os.sep)
|
||||||
|
current_level = file_tree
|
||||||
|
|
||||||
|
# all except the last part (the file name)
|
||||||
|
for part in path_parts[:-1]:
|
||||||
|
current_level = current_level.setdefault(part, {})
|
||||||
|
|
||||||
|
# Add the file and its content
|
||||||
|
current_level[path_parts[-1]] = content
|
||||||
|
return yaml.dump(file_tree, default_flow_style=False)
|
||||||
|
|
||||||
|
|
||||||
|
def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None):
|
||||||
|
"""Chunk the YAML content and write it to the output file."""
|
||||||
|
|
||||||
|
chunks = [yaml_content] if chunk_size is None else [
|
||||||
|
yaml_content[i:i + chunk_size] for i in range(0, len(yaml_content), chunk_size)]
|
||||||
|
|
||||||
|
for i, chunk in enumerate(chunks):
|
||||||
|
chunk_file = f"{output_file}_{i}.yaml"
|
||||||
|
# Write the file tree to the output YAML file
|
||||||
|
with open(chunk_file, 'w', encoding='utf-8') as yaml_file:
|
||||||
|
yaml_file.write(chunk)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Encrypts file tree to yaml file")
|
||||||
|
parser.add_argument("input_directory", help="Directory to encode")
|
||||||
|
parser.add_argument("output_yaml_file", help="Output YAML file")
|
||||||
|
parser.add_argument("--ignore", default=None,
|
||||||
|
help="Path to the ignore file")
|
||||||
|
parser.add_argument("--chunk-size", type=int, default=None,
|
||||||
|
help="Chunk size for the output YAML file")
|
||||||
|
parser.add_argument("--library", type=bool, default=False,
|
||||||
|
help="Incude the library in the output YAML file")
|
||||||
|
|
||||||
|
# Parse arguments
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# Example usage
|
||||||
|
directory_to_encode = args.input_directory
|
||||||
|
ignore_file_path = args.ignore
|
||||||
|
output_yaml_file = args.output_yaml_file
|
||||||
|
include_library = args.library
|
||||||
|
|
||||||
|
content = encode_file_tree_to_yaml(
|
||||||
|
directory_to_encode, ignore_file_path, include_library)
|
||||||
|
chunk_and_write_file_tree_to_yaml(
|
||||||
|
content, output_yaml_file, args.chunk_size)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
2
git-requirements-mapping.txt
Normal file
2
git-requirements-mapping.txt
Normal file
@@ -0,0 +1,2 @@
|
|||||||
|
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git:sientia-do
|
||||||
|
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git:sientia
|
||||||
25
inter_arrival.py
Normal file
25
inter_arrival.py
Normal file
@@ -0,0 +1,25 @@
|
|||||||
|
# %%
|
||||||
|
|
||||||
|
# Load logs.txt
|
||||||
|
with open('logs.txt', 'r') as file:
|
||||||
|
lines = file.readlines()
|
||||||
|
|
||||||
|
# %%
|
||||||
|
import re
|
||||||
|
# Grep "inter-arrival_s=number" with regex
|
||||||
|
intervals = []
|
||||||
|
for line in lines:
|
||||||
|
match = re.search(r'inter-arrival_s=([0-9.]+)', line)
|
||||||
|
if match:
|
||||||
|
intervals.append(float(match.group(1)))
|
||||||
|
# %%
|
||||||
|
|
||||||
|
print(intervals)
|
||||||
|
# %%
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
plt.plot(intervals)
|
||||||
|
plt.ylabel('Inter-arrival time (s)')
|
||||||
|
plt.xlabel('Sample')
|
||||||
|
plt.title('Inter-arrival time distribution')
|
||||||
|
plt.show()
|
||||||
|
# %%
|
||||||
0
laborious/__init__.py
Normal file
0
laborious/__init__.py
Normal file
0
laborious/activities/__init__.py
Normal file
0
laborious/activities/__init__.py
Normal file
169
laborious/activities/activities.py
Normal file
169
laborious/activities/activities.py
Normal file
@@ -0,0 +1,169 @@
|
|||||||
|
from sientia_do.observability.metrics_controller import MetricsController
|
||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
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
|
||||||
|
from laborious.activities.mlflow import MLFlow
|
||||||
|
from laborious.activities.model_metrics import ModelMetrics
|
||||||
|
from laborious.activities.opc import OPC
|
||||||
|
from laborious.activities.storage import Storage
|
||||||
|
|
||||||
|
|
||||||
|
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||||
|
"""
|
||||||
|
Main activities orchestrator for the Laborious system.
|
||||||
|
|
||||||
|
This class combines functionality from multiple activity classes to provide
|
||||||
|
a unified interface for all workflow operations. It manages database connections,
|
||||||
|
MLFlow model interactions, data quality validation, and OPC server communications.
|
||||||
|
|
||||||
|
The class implements multiple inheritance to combine specialized functionality:
|
||||||
|
- Storage: Database operations and data persistence
|
||||||
|
- MLFlow: Model inference and transformation operations
|
||||||
|
- Gates: Data quality validation and filtering mechanisms
|
||||||
|
- OPC: Real-time data export to OPC servers
|
||||||
|
- ModelMetrics: Model performance metrics and drift detection
|
||||||
|
- API: PI Web API export operations for industrial systems
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
postgres_config (dict): PostgreSQL connection configuration
|
||||||
|
mlflow_config (dict): MLFlow server configuration
|
||||||
|
opc_config (dict): OPC server configuration
|
||||||
|
pi_web_api_config (dict): PI Web API server configuration
|
||||||
|
logger (Logger): Logging and observability instance
|
||||||
|
notification_handler (NotificationHandler): Notification management instance
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
postgres_config: dict[str, Any],
|
||||||
|
mlflow_config: dict[str, Any],
|
||||||
|
minio_config: dict[str, Any],
|
||||||
|
opc_config: dict[str, Any],
|
||||||
|
pi_web_api_config: dict[str, Any],
|
||||||
|
logger: Logger,
|
||||||
|
notification_handler: NotificationHandler,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize the Activities orchestrator with all required configurations.
|
||||||
|
|
||||||
|
This constructor initializes all parent classes with their respective
|
||||||
|
configurations and sets up the foundation for all activity operations.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
postgres_config: PostgreSQL connection configuration dictionary
|
||||||
|
Required keys: host, port, user, password, dbname, min_connections, max_connections
|
||||||
|
mlflow_config: MLFlow server configuration dictionary
|
||||||
|
Required keys: host, port, username, password
|
||||||
|
opc_config: OPC server configuration dictionary
|
||||||
|
Can contain multiple server configurations
|
||||||
|
pi_web_api_config: PI Web API server configuration dictionary
|
||||||
|
Required keys: base_url, auth_type, auth_token
|
||||||
|
logger: Logger instance for observability and debugging
|
||||||
|
notification_handler: Notification handler for alerts and monitoring
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If any parent class initialization fails
|
||||||
|
"""
|
||||||
|
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,
|
||||||
|
host=postgres_config['host'],
|
||||||
|
port=postgres_config['port'],
|
||||||
|
user=postgres_config['user'],
|
||||||
|
password=postgres_config['password'],
|
||||||
|
dbname=postgres_config['dbname'],
|
||||||
|
min_connections=postgres_config['min_connections'],
|
||||||
|
max_connections=postgres_config['max_connections'],
|
||||||
|
retention_hours=minio_config['retention_hours'],
|
||||||
|
minio_repository=minio_repository,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
MLFlow.__init__(
|
||||||
|
self,
|
||||||
|
mlflow_host=mlflow_config['host'],
|
||||||
|
mlflow_port=mlflow_config['port'],
|
||||||
|
mlflow_username=mlflow_config['username'],
|
||||||
|
mlflow_password=mlflow_config['password'],
|
||||||
|
minio_repository=minio_repository,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
Gates.__init__(
|
||||||
|
self,
|
||||||
|
minio_repository=minio_repository,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
OPC.__init__(
|
||||||
|
self,
|
||||||
|
opc_servers=opc_config,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
ModelMetrics.__init__(
|
||||||
|
self,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
API.__init__(
|
||||||
|
self,
|
||||||
|
base_url=pi_web_api_config['base_url'],
|
||||||
|
auth_type=pi_web_api_config['auth_type'],
|
||||||
|
auth_token=pi_web_api_config['auth_token'],
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def shutdown(self):
|
||||||
|
"""
|
||||||
|
Gracefully shutdown all activities and clean up resources.
|
||||||
|
|
||||||
|
This method ensures proper cleanup of all resources including:
|
||||||
|
- PostgreSQL connection pools
|
||||||
|
- OPC server connections
|
||||||
|
- PI Web API client connections
|
||||||
|
- MLFlow model repositories
|
||||||
|
- Any other resources that need explicit cleanup
|
||||||
|
|
||||||
|
The method should be called before the application terminates to ensure
|
||||||
|
proper resource cleanup and prevent resource leaks.
|
||||||
|
"""
|
||||||
|
Storage.close(self)
|
||||||
|
MLFlow.close(self)
|
||||||
|
Gates.close(self)
|
||||||
|
await OPC.aclose(self)
|
||||||
|
ModelMetrics.close(self)
|
||||||
|
API.close(self)
|
||||||
305
laborious/activities/api.py
Normal file
305
laborious/activities/api.py
Normal file
@@ -0,0 +1,305 @@
|
|||||||
|
from temporalio import activity, workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
import json
|
||||||
|
import traceback
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from pandas import DataFrame
|
||||||
|
from sientia_do.notifications.handlers import 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.pi_web_api_client import PIWebAPIClient
|
||||||
|
|
||||||
|
from laborious import metrics
|
||||||
|
|
||||||
|
|
||||||
|
PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13
|
||||||
|
|
||||||
|
|
||||||
|
class API(SientiaMonitoring):
|
||||||
|
"""
|
||||||
|
PI Web API operations for writing prediction data to PI Web API.
|
||||||
|
|
||||||
|
This class provides Temporal activities for interacting with the PI Web API
|
||||||
|
to write prediction and confidence values to industrial systems. It handles
|
||||||
|
error scenarios gracefully by setting error confidence values and sending
|
||||||
|
notifications when write operations fail.
|
||||||
|
|
||||||
|
The class implements comprehensive error handling for both prediction and
|
||||||
|
confidence value writes, ensuring that partial failures are properly
|
||||||
|
reported and handled.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
base_url: str,
|
||||||
|
auth_type: str,
|
||||||
|
auth_token: str,
|
||||||
|
logger: Logger,
|
||||||
|
notification_handler: NotificationHandler,
|
||||||
|
metrics_controller: MetricsController,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Initialize API activity with PI Web API client.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
base_url (str): Base URL of the PI Web API server
|
||||||
|
auth_type (str): Authentication type ('basic' or 'bearer')
|
||||||
|
auth_token (str): Authentication token
|
||||||
|
logger (Logger): Logger instance for operation logging
|
||||||
|
notification_handler (NotificationHandler): Handler for system notifications
|
||||||
|
metrics_controller (MetricsController): Controller for metrics collection
|
||||||
|
"""
|
||||||
|
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||||
|
self.pi_web_api_client = PIWebAPIClient(
|
||||||
|
base_url=base_url,
|
||||||
|
auth_config={
|
||||||
|
'type': auth_type,
|
||||||
|
'token': auth_token,
|
||||||
|
},
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
headers_config={
|
||||||
|
'Content-Type': 'application/json',
|
||||||
|
'Accept': 'application/json',
|
||||||
|
'x-requested-with': 'piwebapistreams',
|
||||||
|
'User-Agent': 'Aig-Laborious-Agent/1.0',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_pi_web_api_core_labels(
|
||||||
|
self,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
operation_type: str = 'write_pi_web_api_data',
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Generate core labels for PI Web API metrics.
|
||||||
|
|
||||||
|
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): Operation type label for metric cardinality.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict[str, Any]: Core labels dictionary including operation_type.
|
||||||
|
"""
|
||||||
|
return super().get_core_labels(
|
||||||
|
metadata=metadata,
|
||||||
|
operation_type=operation_type,
|
||||||
|
)
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""
|
||||||
|
Close the PI Web API client and shutdown monitoring services.
|
||||||
|
|
||||||
|
This method properly closes all connections and resources associated
|
||||||
|
with the PI Web API client and monitoring services.
|
||||||
|
"""
|
||||||
|
self.pi_web_api_client.close()
|
||||||
|
SientiaMonitoring.shutdown(self)
|
||||||
|
|
||||||
|
async def process_pi_web_api_response(
|
||||||
|
self,
|
||||||
|
response_data: list[dict[str, Any]],
|
||||||
|
tags: dict[str, str],
|
||||||
|
core_labels: dict[str, str],
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
) -> tuple[int, str]:
|
||||||
|
"""
|
||||||
|
Process the response data from PI Web API write operation.
|
||||||
|
|
||||||
|
Validates that all tags were successfully written, emits metrics for each tag
|
||||||
|
(success or error), and returns the appropriate prediction confidence value.
|
||||||
|
Sets error confidence if any tag write fails or if the number of written tags
|
||||||
|
doesn't match the expected count.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- response_data (dict[str, Any]): The response data from the PI Web API write operation.
|
||||||
|
- tags (dict[str, str]): The tags that were written to the PI Web API.
|
||||||
|
- core_labels (dict[str, str]): The core labels of the workflow execution.
|
||||||
|
- metadata (dict[str, Any]): The metadata of the workflow execution.
|
||||||
|
Returns:
|
||||||
|
int: Prediction confidence value (0 for success, 13 for errors)
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Convert tags from name:webid to webid:name
|
||||||
|
tags = {w: t for t, w in tags.items()}
|
||||||
|
|
||||||
|
tag_names = list[str](tags.values())
|
||||||
|
|
||||||
|
confidence = 0
|
||||||
|
|
||||||
|
message = ''
|
||||||
|
|
||||||
|
# Evaluate response for each tag
|
||||||
|
written_tags = []
|
||||||
|
for item in response_data:
|
||||||
|
web_id = item.get('WebId')
|
||||||
|
if not web_id:
|
||||||
|
self.error('The response did not contain some WebIds', metadata)
|
||||||
|
continue
|
||||||
|
errors = item.get('Errors', [])
|
||||||
|
tag_name = tags.get(web_id)
|
||||||
|
if not tag_name:
|
||||||
|
self.error(
|
||||||
|
f'The response did not contain the tag name for WebId {web_id}', metadata
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
if errors:
|
||||||
|
self.error(
|
||||||
|
f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata
|
||||||
|
)
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT,
|
||||||
|
tags={
|
||||||
|
**core_labels,
|
||||||
|
'tag_name': tag_name,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
else:
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT,
|
||||||
|
tags={
|
||||||
|
**core_labels,
|
||||||
|
'tag_name': tag_name,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
written_tags.append(tag_name)
|
||||||
|
|
||||||
|
if len(written_tags) != len(tag_names):
|
||||||
|
message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.'
|
||||||
|
|
||||||
|
self.error(
|
||||||
|
f'{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||||
|
message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
|
||||||
|
block='write_pi_web_api_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
)
|
||||||
|
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
|
||||||
|
return confidence, message
|
||||||
|
|
||||||
|
@activity.defn(name='write_pi_web_api_data')
|
||||||
|
async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||||
|
"""
|
||||||
|
Write prediction and confidence data to PI Web API.
|
||||||
|
|
||||||
|
Writes prediction values and confidence scores to PI Web API using configured
|
||||||
|
web IDs. Processes responses to validate writes and emit metrics. Handles errors
|
||||||
|
gracefully by setting error confidence values when writes fail and sending
|
||||||
|
notifications for both prediction and confidence write errors.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict[str, Any]): The input data containing:
|
||||||
|
- metadata (dict[str, Any]): Workflow execution metadata
|
||||||
|
- pi_web_api_output_config (dict[str, Any]): PI Web API configuration with:
|
||||||
|
- endpoint (str): PI Web API endpoint URL
|
||||||
|
- prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions
|
||||||
|
- confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence
|
||||||
|
- data (dict[str, Any]): Prediction data, its a dataframe converted to dict.
|
||||||
|
Returns:
|
||||||
|
dict[Any, Any]: Data dictionary with potentially modified confidence values
|
||||||
|
If prediction write fails, prediction_confidence is set to error value (13)
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
data = DataFrame(input_data['data'])
|
||||||
|
pi_web_api_output_config = input_data['pi_web_api_output_config']
|
||||||
|
|
||||||
|
self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata)
|
||||||
|
|
||||||
|
raw_prediction_tags = pi_web_api_output_config['prediction_tags']
|
||||||
|
raw_confidence_tags = pi_web_api_output_config['confidence_tags']
|
||||||
|
prediction_tags = list[str](raw_prediction_tags.values())
|
||||||
|
confidence_tags = list(raw_confidence_tags.values())
|
||||||
|
|
||||||
|
core_labels = self.get_pi_web_api_core_labels(metadata)
|
||||||
|
|
||||||
|
prediction_value = data.head(1)['prediction'].values[0]
|
||||||
|
confidence_value = data.head(1)['prediction_confidence'].values[0]
|
||||||
|
|
||||||
|
try:
|
||||||
|
prediction_response = await self.pi_web_api_client.write_value(
|
||||||
|
web_ids=prediction_tags,
|
||||||
|
value={
|
||||||
|
'Timestamp': data.head(1)['timestamp'].values[0],
|
||||||
|
'Value': prediction_value,
|
||||||
|
},
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
confidence, message = await self.process_pi_web_api_response(
|
||||||
|
response_data=prediction_response,
|
||||||
|
tags=raw_prediction_tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||||
|
message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
|
||||||
|
block='write_pi_web_api_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.error(trace, metadata)
|
||||||
|
|
||||||
|
data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
data['comments'] = str(e)
|
||||||
|
|
||||||
|
return data.to_dict()
|
||||||
|
|
||||||
|
try:
|
||||||
|
confidence_response = await self.pi_web_api_client.write_value(
|
||||||
|
web_ids=confidence_tags,
|
||||||
|
value={
|
||||||
|
'Timestamp': data.head(1)['timestamp'].values[0],
|
||||||
|
'Value': float(confidence_value),
|
||||||
|
},
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.process_pi_web_api_response(
|
||||||
|
response_data=confidence_response,
|
||||||
|
tags=raw_confidence_tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
|
||||||
|
message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
|
||||||
|
block='write_pi_web_api_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
return data.to_dict()
|
||||||
804
laborious/activities/gates.py
Normal file
804
laborious/activities/gates.py
Normal file
@@ -0,0 +1,804 @@
|
|||||||
|
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
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from pandas import DataFrame
|
||||||
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
from sientia_do.observability.logger import Logger
|
||||||
|
from sientia_do.observability.metrics_controller import MetricsController
|
||||||
|
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]
|
||||||
|
ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
|
||||||
|
ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||||
|
|
||||||
|
# Input filter function mappings
|
||||||
|
input_filter_functions: dict[str, InputFilterFunc] = {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
|
||||||
|
'EMPTY_DATA': filter_empty_data,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Confidence mappings kept separate from function maps to avoid Union types
|
||||||
|
input_path_confidence: Mapping[str, int] = {
|
||||||
|
'STOP': -1,
|
||||||
|
'CONTINUE': 2,
|
||||||
|
'REPEAT': -1,
|
||||||
|
}
|
||||||
|
|
||||||
|
# MLFlow response filter function mappings
|
||||||
|
mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
|
||||||
|
'API_ERROR': api_error_filter,
|
||||||
|
}
|
||||||
|
|
||||||
|
mlflow_response_path_confidence: Mapping[str, int] = {
|
||||||
|
'STOP': -1,
|
||||||
|
'CONTINUE': 10,
|
||||||
|
'REPEAT': -1,
|
||||||
|
}
|
||||||
|
|
||||||
|
# MLFlow content filter function mappings
|
||||||
|
mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
|
||||||
|
'NAN_VALUES': nan_values_filter,
|
||||||
|
'EMPTY_DATA': filter_empty_data,
|
||||||
|
}
|
||||||
|
|
||||||
|
mlflow_content_path_confidence: Mapping[str, int] = {
|
||||||
|
'STOP': -1,
|
||||||
|
'CONTINUE': 18,
|
||||||
|
'REPEAT': -1,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class Gates(MinioManager):
|
||||||
|
"""
|
||||||
|
Data quality gates and filtering activities for the Laborious system.
|
||||||
|
|
||||||
|
This class implements comprehensive data quality validation and filtering
|
||||||
|
mechanisms that can be applied at different stages of the prediction pipeline.
|
||||||
|
It provides configurable filters with policy-based decision making to ensure
|
||||||
|
data integrity and quality throughout the ML workflow.
|
||||||
|
|
||||||
|
The class supports multiple filter types and implements a flexible policy
|
||||||
|
system that can be configured for different validation requirements. Each
|
||||||
|
filter returns a path decision (STOP, CONTINUE, REPEAT) along with confidence
|
||||||
|
scores and detailed comments for monitoring and debugging.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
input_filter_functions (dict): Mapping of input filter names to functions
|
||||||
|
mlflow_response_filter_functions (dict): Mapping of MLFlow response filter names to functions
|
||||||
|
mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
|
||||||
|
"""
|
||||||
|
|
||||||
|
minio_repository: MinioRepository | None = None
|
||||||
|
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
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.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
logger: Logger instance for observability and debugging
|
||||||
|
notification_handler: Notification handler for alerts and monitoring
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If BaseActivity initialization fails
|
||||||
|
"""
|
||||||
|
MinioManager.__init__(
|
||||||
|
self, minio_repository, logger, notification_handler, metrics_controller
|
||||||
|
)
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""
|
||||||
|
Close the gates activity and clean up resources.
|
||||||
|
"""
|
||||||
|
|
||||||
|
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]:
|
||||||
|
"""
|
||||||
|
Apply input data quality filters and validation.
|
||||||
|
|
||||||
|
This activity validates input data quality using configurable filters
|
||||||
|
before proceeding with ML operations. It applies multiple filter types
|
||||||
|
and returns a path decision based on the filter results and configured
|
||||||
|
policies.
|
||||||
|
|
||||||
|
The method implements a comprehensive filtering system that:
|
||||||
|
1. Applies configured filters to input data
|
||||||
|
2. Evaluates filter results against policy configurations
|
||||||
|
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
|
||||||
|
4. Provides confidence scores and detailed comments
|
||||||
|
5. Handles errors gracefully with notification integration
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Configuration and data for input validation
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- filters (dict): Filter configuration and policies
|
||||||
|
- data (dict): Input data to validate
|
||||||
|
- path_priority (list[str]): Priority order for path decisions
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple: (path_flag, confidence, comment)
|
||||||
|
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||||
|
- confidence (int): Confidence score for the decision
|
||||||
|
- comment (str): Detailed explanation of the decision
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If filter execution fails or configuration is invalid
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
|
||||||
|
self.info('Performing input gate...', metadata)
|
||||||
|
|
||||||
|
filters = input_data['filters']
|
||||||
|
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_dataframe('Input data:', data, metadata)
|
||||||
|
self.debug(f'Filters: {filters}', metadata)
|
||||||
|
|
||||||
|
# Apply each configured filter
|
||||||
|
for fil, config in filters.items():
|
||||||
|
if fil not in input_filter_functions:
|
||||||
|
self.error(f'Filter {fil} not found', metadata)
|
||||||
|
continue
|
||||||
|
policy, filter_config = self._read_filter_entry(config)
|
||||||
|
try:
|
||||||
|
if input_filter_functions[fil](data, filter_config):
|
||||||
|
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
|
||||||
|
filter_output.append(policy)
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id=f'INTPUT_GATE_ERROR__{fil}',
|
||||||
|
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||||
|
block='input_gate',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
for path_flag in path_priority:
|
||||||
|
if path_flag in filter_output:
|
||||||
|
self.info(f'Input gate result: {path_flag}', metadata)
|
||||||
|
return path_flag, input_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')
|
||||||
|
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||||
|
"""
|
||||||
|
Validate MLFlow API response quality and integrity.
|
||||||
|
|
||||||
|
This activity validates MLFlow API responses to ensure they meet quality
|
||||||
|
standards before proceeding with further processing. It applies response-specific
|
||||||
|
filters and determines appropriate path decisions based on response quality.
|
||||||
|
|
||||||
|
The method implements response validation that:
|
||||||
|
1. Applies MLFlow response-specific filters
|
||||||
|
2. Evaluates API response quality and integrity
|
||||||
|
3. Determines path decisions based on response validation results
|
||||||
|
4. Provides confidence scores and detailed validation comments
|
||||||
|
5. Handles API errors and response validation failures
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Configuration and data for response validation
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- filters (dict): Response filter configuration and policies
|
||||||
|
- data (dict): MLFlow API response data to validate
|
||||||
|
- type (str): Type of MLFlow operation (transform, predict)
|
||||||
|
- path_priority (list[str]): Priority order for path decisions
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple: (path_flag, confidence, comment)
|
||||||
|
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||||
|
- confidence (int): Confidence score for the decision
|
||||||
|
- comment (str): Detailed explanation of the decision
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If response validation fails or configuration is invalid
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.info('Performing mlflow response gate...', metadata)
|
||||||
|
raw_data = input_data['data']
|
||||||
|
filters = input_data['filters']
|
||||||
|
|
||||||
|
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 = []
|
||||||
|
|
||||||
|
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](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=status.get('message', 'Unknown MLFlow API error'),
|
||||||
|
block='mlflow_gate',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=status.get('traceback'),
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
|
||||||
|
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||||
|
block='mlflow_gate',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
for path_flag in path_priority:
|
||||||
|
if path_flag in filter_output:
|
||||||
|
self.info(f'Mlflow response gate result: {path_flag}', metadata)
|
||||||
|
return path_flag, mlflow_response_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')
|
||||||
|
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||||
|
"""
|
||||||
|
Validate MLFlow prediction content quality and integrity.
|
||||||
|
|
||||||
|
This activity validates the content of MLFlow predictions to ensure they
|
||||||
|
meet quality standards before export and persistence. It applies content-specific
|
||||||
|
filters and determines appropriate path decisions based on content quality.
|
||||||
|
|
||||||
|
The method implements content validation that:
|
||||||
|
1. Applies MLFlow content-specific filters
|
||||||
|
2. Evaluates prediction content quality and integrity
|
||||||
|
3. Determines path decisions based on content validation results
|
||||||
|
4. Provides confidence scores and detailed validation comments
|
||||||
|
5. Handles content validation failures and quality issues
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Configuration and data for content validation
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- filters (dict): Content filter configuration and policies
|
||||||
|
- data (dict): MLFlow prediction content to validate
|
||||||
|
- type (str): Type of MLFlow operation (transform, predict)
|
||||||
|
- path_priority (list[str]): Priority order for path decisions
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple: (path_flag, confidence, comment)
|
||||||
|
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||||
|
- confidence (int): Confidence score for the decision
|
||||||
|
- comment (str): Detailed explanation of the decision
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If content validation fails or configuration is invalid
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.info('Performing mlflow content gate...', metadata)
|
||||||
|
|
||||||
|
filters = input_data['filters']
|
||||||
|
|
||||||
|
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_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, filter_config):
|
||||||
|
filter_output.append(policy)
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
|
||||||
|
message=f'Data not passed the content filter {fil}:{config}',
|
||||||
|
block='mlflow_gate',
|
||||||
|
level=NotificationLevel.WARNING,
|
||||||
|
attachment_content=data.to_string(),
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
|
||||||
|
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||||
|
block='mlflow_gate',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
for path_flag in path_priority:
|
||||||
|
if path_flag in filter_output:
|
||||||
|
self.info(f'Mlflow content gate result: {path_flag}', metadata)
|
||||||
|
return (
|
||||||
|
path_flag,
|
||||||
|
mlflow_content_path_confidence[path_flag],
|
||||||
|
'Transformed data not passed the content filter',
|
||||||
|
)
|
||||||
|
|
||||||
|
self.info('Nothing was filtered by the mlflow content gate', metadata)
|
||||||
|
|
||||||
|
del data
|
||||||
|
|
||||||
|
return None, 0, ''
|
||||||
|
|
||||||
|
def get_prediction_store_policy(
|
||||||
|
self, prediction_store_policy: str, metadata: dict[str, Any]
|
||||||
|
) -> tuple[str, int]:
|
||||||
|
"""
|
||||||
|
Parse and validate prediction store policy configuration.
|
||||||
|
|
||||||
|
This method parses prediction store policy strings in the format 'type:value'
|
||||||
|
and validates them against allowed policy types and values. It provides
|
||||||
|
sensible defaults for invalid configurations and logs policy validation
|
||||||
|
failures for operational monitoring.
|
||||||
|
|
||||||
|
Supported Policy Types:
|
||||||
|
- 'lts': Latest timestamp - sorts data by timestamp descending
|
||||||
|
- 'erl': Earliest timestamp - sorts data by timestamp ascending
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prediction_store_policy (str): Policy string in format 'type:value'
|
||||||
|
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple[str, int]: (policy_type, policy_value)
|
||||||
|
- policy_type (str): Validated policy type ('lts' or 'erl')
|
||||||
|
- policy_value (int): Number of rows to retain
|
||||||
|
"""
|
||||||
|
policy_elements = prediction_store_policy.split(':')
|
||||||
|
|
||||||
|
if len(policy_elements) < 2:
|
||||||
|
self.error(
|
||||||
|
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
return 'lts', 1
|
||||||
|
|
||||||
|
policy_type = policy_elements[0]
|
||||||
|
policy_value = policy_elements[1]
|
||||||
|
|
||||||
|
# If the policy_type is not lts or erl, we use the default policy
|
||||||
|
# If the policty_value is not a number or 0, we use the default policy
|
||||||
|
if (
|
||||||
|
policy_type not in ['lts', 'erl']
|
||||||
|
or not policy_value.isdigit()
|
||||||
|
or int(policy_value) == 0
|
||||||
|
):
|
||||||
|
self.error(
|
||||||
|
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
return 'lts', 1
|
||||||
|
|
||||||
|
return policy_type, int(policy_value)
|
||||||
|
|
||||||
|
@activity.defn(name='format_transformed_data')
|
||||||
|
async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||||
|
"""
|
||||||
|
Format transformed data for storage and export operations.
|
||||||
|
|
||||||
|
This method formats transformed data from MLFlow model transformations
|
||||||
|
into a standardized format suitable for database storage. It converts
|
||||||
|
wide-format data (columns as variables) into long-format (melted)
|
||||||
|
with proper timestamp handling and model identification.
|
||||||
|
|
||||||
|
The formatting process includes:
|
||||||
|
1. Converting input data dictionary to DataFrame
|
||||||
|
2. Extracting timestamps from DataFrame index
|
||||||
|
3. Resetting index to create sequential row numbers
|
||||||
|
4. Melting data from wide format to long format (variable-value pairs)
|
||||||
|
5. Adding model_id for data lineage tracking
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict)
|
||||||
|
- model_id (str): Unique identifier for the ML model
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Formatted data dictionary with keys:
|
||||||
|
- timestamp (dict): Timestamp values indexed by row number
|
||||||
|
- variable (dict): Variable names indexed by row number
|
||||||
|
- value (dict): Variable values indexed by row number
|
||||||
|
- model_id (dict): Model identifiers indexed by row number
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
|
||||||
|
model_id = input_data['model_id']
|
||||||
|
|
||||||
|
self.info('Formatting transformed data...', metadata)
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
|
||||||
|
data['model_id'] = model_id
|
||||||
|
|
||||||
|
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:
|
||||||
|
"""
|
||||||
|
Format prediction data according to configured storage policies.
|
||||||
|
|
||||||
|
This method formats prediction data for storage and export operations.
|
||||||
|
It applies timestamp-based sorting policies, adds metadata fields,
|
||||||
|
and ensures data consistency before persistence. The method supports
|
||||||
|
multiple storage policies for flexible data retention strategies.
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- data (dict[str, Any]): Raw prediction data to format
|
||||||
|
- 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'
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
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)
|
||||||
|
|
||||||
|
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_dataframe('Prediction data:', data, metadata)
|
||||||
|
|
||||||
|
policy_type, policy_value = self.get_prediction_store_policy(
|
||||||
|
prediction_store_policy, metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
# If data has no timestamp, we use the default timestamp and not sort the data
|
||||||
|
self.info(
|
||||||
|
f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
|
||||||
|
if policy_type == 'lts':
|
||||||
|
self.debug('Sorting data by timestamp descending', metadata)
|
||||||
|
data = data.sort_values(by='timestamp', ascending=False)
|
||||||
|
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
|
||||||
|
elif policy_type == 'erl':
|
||||||
|
self.debug('Sorting data by timestamp ascending', metadata)
|
||||||
|
data = data.sort_values(by='timestamp', ascending=True)
|
||||||
|
else:
|
||||||
|
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
|
||||||
|
raise ValueError(f'Invalid policy type: {policy_type}')
|
||||||
|
|
||||||
|
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']
|
||||||
|
data['prediction_status'] = 'Good'
|
||||||
|
data['comments'] = ''
|
||||||
|
data = data.sort_values(by='timestamp', ascending=False)
|
||||||
|
data = data.reset_index(drop=True)
|
||||||
|
|
||||||
|
self.info(f'Prediction formatted: {len(data)} rows', metadata)
|
||||||
|
self._debug_dataframe('Prediction data:', data, metadata)
|
||||||
|
|
||||||
|
return data.to_dict()
|
||||||
|
|
||||||
|
@activity.defn(name='format_default_prediction')
|
||||||
|
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||||
|
"""
|
||||||
|
Create and format default prediction data for error conditions.
|
||||||
|
|
||||||
|
This method generates default prediction data when the main prediction
|
||||||
|
pipeline encounters errors or quality issues. It creates a standardized
|
||||||
|
data structure with zero values for predictions and useful metadata
|
||||||
|
for operational monitoring and debugging.
|
||||||
|
|
||||||
|
The default prediction serves as a fallback mechanism to:
|
||||||
|
1. Maintain data pipeline continuity during failures
|
||||||
|
2. Provide operational visibility into prediction quality issues
|
||||||
|
3. Enable downstream systems to handle error conditions gracefully
|
||||||
|
4. Support debugging and troubleshooting efforts
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- timestamp (str): Timestamp for the default prediction
|
||||||
|
- model_id (str): Unique identifier for the ML model
|
||||||
|
- prediction_confidence (float): Confidence score (typically low for errors)
|
||||||
|
- comment (str): Error description or operational comment
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Formatted default prediction data with error indicators
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.debug('Formatting default prediction...', metadata)
|
||||||
|
|
||||||
|
data = DataFrame(
|
||||||
|
{
|
||||||
|
'prediction': [0],
|
||||||
|
'response_time': [0],
|
||||||
|
'timestamp': [input_data['timestamp']],
|
||||||
|
'model_id': [input_data['model_id']],
|
||||||
|
'prediction_confidence': [input_data['prediction_confidence']],
|
||||||
|
'prediction_status': ['Bad'],
|
||||||
|
'comments': [input_data['comment']],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
self.info(f'Default prediction formatted: {data.size} rows', metadata)
|
||||||
|
return data.to_dict()
|
||||||
|
|
||||||
|
@activity.defn(name='format_retrain_report')
|
||||||
|
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
|
||||||
|
"""
|
||||||
|
Format retrain report data for storage and audit trail maintenance.
|
||||||
|
|
||||||
|
This method formats model retraining operation results into a standardized
|
||||||
|
report format suitable for database storage and operational monitoring.
|
||||||
|
It captures retraining status, timestamps, and model version information
|
||||||
|
for comprehensive audit trails and operational visibility.
|
||||||
|
|
||||||
|
The formatting process includes:
|
||||||
|
1. Extracting retraining experiment response data
|
||||||
|
2. Capturing model update report information (version, MLflow IDs)
|
||||||
|
3. Formatting timestamps and status information
|
||||||
|
4. Conditionally including version information for successful retrains
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- experiment_response (dict): Retraining experiment response containing:
|
||||||
|
- success (bool): Retraining operation success status
|
||||||
|
- timestamp (str): Timestamp of the retraining operation
|
||||||
|
- message (str): Status message or error description
|
||||||
|
- update_report (dict): Model update report containing:
|
||||||
|
- version (str): New model version identifier
|
||||||
|
- mlflow_run_id (str): MLflow run identifier
|
||||||
|
- mlflow_experiment_id (str): MLflow experiment identifier
|
||||||
|
- model_id (str): Unique identifier for the ML model
|
||||||
|
- model_name (str): Name of the ML model
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Formatted retrain report dictionary with keys:
|
||||||
|
- model_id (dict): Model identifiers indexed by row number
|
||||||
|
- model_name (dict): Model names indexed by row number
|
||||||
|
- timestamp (dict): Retraining timestamps indexed by row number
|
||||||
|
- status (dict): Retraining status messages indexed by row number
|
||||||
|
- version (dict, optional): Model versions indexed by row number
|
||||||
|
Only included if experiment_response['success'] is True
|
||||||
|
- mlflow_run_id (dict, optional): MLflow run IDs indexed by row number
|
||||||
|
Only included if experiment_response['success'] is True
|
||||||
|
- mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number
|
||||||
|
Only included if experiment_response['success'] is True
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.info('Formatting retrain report...', metadata)
|
||||||
|
|
||||||
|
experiment_response = input_data['experiment_response']
|
||||||
|
update_report = input_data['update_report']
|
||||||
|
model_id = input_data['model_id']
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
|
||||||
|
report = DataFrame(
|
||||||
|
{
|
||||||
|
'model_id': [model_id],
|
||||||
|
'model_name': [model_name],
|
||||||
|
'timestamp': [experiment_response['timestamp']],
|
||||||
|
'status': [experiment_response['message']],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
if experiment_response['success']:
|
||||||
|
# Retrain was successfull
|
||||||
|
report['version'] = update_report['version']
|
||||||
|
report['mlflow_run_id'] = update_report['mlflow_run_id']
|
||||||
|
report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
|
||||||
|
|
||||||
|
self._debug_dataframe('Retrain report:', report, metadata)
|
||||||
|
|
||||||
|
return report.to_dict()
|
||||||
|
|
||||||
|
@activity.defn(name='write_metrics')
|
||||||
|
async def write_metrics(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Write prediction performance metrics to Prometheus monitoring system.
|
||||||
|
|
||||||
|
This method records comprehensive metrics for prediction operations,
|
||||||
|
enabling operational monitoring, performance analysis, and alerting.
|
||||||
|
It tracks prediction counts, confidence levels, and response times
|
||||||
|
for each model and pipeline combination.
|
||||||
|
|
||||||
|
Metrics Recorded:
|
||||||
|
1. Prediction Count: Incremental counter for successful predictions
|
||||||
|
2. Confidence Monitor: Current confidence level for predictions
|
||||||
|
3. Response Time Monitor: Histogram of prediction response times
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- metadata (dict[str, Any]): Workflow execution metadata
|
||||||
|
- prediction (dict[str, Any]): Prediction data with metrics
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If metrics writing fails or configuration is invalid
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
prediction = DataFrame(input_data['prediction'])
|
||||||
|
prediction_confidence = prediction['prediction_confidence'].values[0]
|
||||||
|
response_time = prediction['response_time'].values[0]
|
||||||
|
opc_metrics = input_data['opc_metrics']
|
||||||
|
|
||||||
|
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=core_tags,
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
|
||||||
|
method='set',
|
||||||
|
tags=core_tags,
|
||||||
|
value=prediction_confidence,
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
|
||||||
|
method='observe',
|
||||||
|
tags=core_tags,
|
||||||
|
value=response_time,
|
||||||
|
)
|
||||||
|
|
||||||
|
for server_id, tags in opc_metrics.items():
|
||||||
|
for tag, response_time in tags.items():
|
||||||
|
if response_time is not None:
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
|
||||||
|
method='observe',
|
||||||
|
tags={
|
||||||
|
**core_tags,
|
||||||
|
'opc_server_id': server_id,
|
||||||
|
'tag': tag,
|
||||||
|
},
|
||||||
|
value=response_time,
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
|
||||||
|
tags={
|
||||||
|
**core_tags,
|
||||||
|
'opc_server_id': server_id,
|
||||||
|
'tag': tag,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)
|
||||||
528
laborious/activities/mlflow.py
Normal file
528
laborious/activities/mlflow.py
Normal file
@@ -0,0 +1,528 @@
|
|||||||
|
from temporalio import activity, workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
import traceback
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
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.repository.minio_repository import MinioRepository
|
||||||
|
from sientia_do.temporal.constants import (
|
||||||
|
DATETIME_FORMAT,
|
||||||
|
DATETIME_FORMAT_MS_WITH_TZ,
|
||||||
|
DATETIME_FORMAT_WITH_TZ,
|
||||||
|
now,
|
||||||
|
)
|
||||||
|
from sientia_do.utils.formatters import create_sample_dict
|
||||||
|
|
||||||
|
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(MinioManager):
|
||||||
|
"""
|
||||||
|
MLFlow integration activities for model inference operations.
|
||||||
|
|
||||||
|
This class provides activities for interacting with MLFlow models, including
|
||||||
|
data transformation and prediction operations. It handles authentication,
|
||||||
|
data preprocessing, and model management with configurable retention policies.
|
||||||
|
|
||||||
|
The class implements comprehensive error handling and logging for all
|
||||||
|
MLFlow operations, ensuring reliable model inference in production environments.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
mlflow_host (str): MLFlow server hostname
|
||||||
|
mlflow_port (int): MLFlow server port
|
||||||
|
mlflow_username (str): MLFlow authentication username
|
||||||
|
mlflow_password (str): MLFlow authentication password
|
||||||
|
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
|
||||||
|
"""
|
||||||
|
|
||||||
|
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
mlflow_host: str,
|
||||||
|
mlflow_port: int,
|
||||||
|
mlflow_username: str,
|
||||||
|
mlflow_password: str,
|
||||||
|
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.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
mlflow_host: MLFlow server hostname or IP address
|
||||||
|
mlflow_port: MLFlow server port number
|
||||||
|
mlflow_username: Username for MLFlow authentication
|
||||||
|
mlflow_password: Password for MLFlow authentication
|
||||||
|
logger: Logger instance for observability and debugging
|
||||||
|
notification_handler: Notification handler for alerts and monitoring
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If MLFlowRepository initialization fails
|
||||||
|
"""
|
||||||
|
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
|
||||||
|
self.mlflow_password = mlflow_password
|
||||||
|
|
||||||
|
self.model_monitoring_repository = MLFlowRepository(
|
||||||
|
f'{mlflow_host}:{mlflow_port}',
|
||||||
|
mlflow_username,
|
||||||
|
mlflow_password,
|
||||||
|
logger,
|
||||||
|
notification_handler,
|
||||||
|
metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""
|
||||||
|
Close the MLFlow activity and clean up resources.
|
||||||
|
"""
|
||||||
|
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]) -> MinioDataFramePayload:
|
||||||
|
"""
|
||||||
|
Transform input data using MLFlow models.
|
||||||
|
|
||||||
|
This activity processes input data through MLFlow model transformation,
|
||||||
|
including data preprocessing, format conversion, and validation. It handles
|
||||||
|
data deduplication, pivoting, and cleanup to ensure optimal model performance.
|
||||||
|
|
||||||
|
The transformation process includes:
|
||||||
|
1. Data deduplication based on variable and timestamp
|
||||||
|
2. Data pivoting for model input format
|
||||||
|
3. Null value handling and cleanup
|
||||||
|
4. MLFlow model transformation request
|
||||||
|
5. Response validation and logging
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Configuration and data for transformation
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- data (dict): Input data for transformation
|
||||||
|
- model_name (str): Name of the MLFlow model to use
|
||||||
|
- model_retention (int): Model retention period in minutes
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Transformed data from MLFlow model
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If transformation fails or MLFlow model is unavailable
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.info('Transforming data...', metadata)
|
||||||
|
|
||||||
|
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_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(
|
||||||
|
subset=['variable', 'timestamp'], keep='first'
|
||||||
|
)
|
||||||
|
|
||||||
|
# Pivot data for model input format
|
||||||
|
data = data.pivot(index='timestamp', columns='variable', values='value')
|
||||||
|
data.fillna(np.nan, inplace=True)
|
||||||
|
|
||||||
|
data.columns.name = None
|
||||||
|
data.index.name = None
|
||||||
|
|
||||||
|
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(
|
||||||
|
model_name, data, model_config, metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
self.debug(
|
||||||
|
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.debug(
|
||||||
|
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.info('Data transformed successfully', metadata)
|
||||||
|
|
||||||
|
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]) -> MinioDataFramePayload:
|
||||||
|
"""
|
||||||
|
Execute predictions using MLFlow models.
|
||||||
|
|
||||||
|
This activity performs ML model inference using MLFlow models with the
|
||||||
|
transformed data. It handles data format conversion, null value processing,
|
||||||
|
and model prediction requests with comprehensive error handling.
|
||||||
|
|
||||||
|
The prediction process includes:
|
||||||
|
1. Data format validation and cleanup
|
||||||
|
2. Null value handling for model compatibility
|
||||||
|
3. MLFlow model prediction request
|
||||||
|
4. Response validation and logging
|
||||||
|
5. Performance monitoring and metrics
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Configuration and data for prediction
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- data (dict): Transformed data for prediction
|
||||||
|
- model_name (str): Name of the MLFlow model to use
|
||||||
|
- model_retention (int): Model retention period in minutes
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Prediction results from MLFlow model
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If prediction fails or MLFlow model is unavailable
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.info('Predicting data...', metadata)
|
||||||
|
|
||||||
|
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_dataframe('Input data for prediction:', data, metadata)
|
||||||
|
|
||||||
|
# Convert numpy.nan to None for model compatibility
|
||||||
|
data.replace(np.nan, None, inplace=True)
|
||||||
|
|
||||||
|
data['timestamp'] = data.index
|
||||||
|
data['timestamp'] = to_datetime(
|
||||||
|
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||||
|
).dt.strftime(DATETIME_FORMAT)
|
||||||
|
|
||||||
|
# Request prediction from MLFlow model
|
||||||
|
response_data = await self.model_monitoring_repository.predict(
|
||||||
|
model_name, data, model_config, metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
self.debug(
|
||||||
|
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.info('Data predicted successfully', metadata)
|
||||||
|
|
||||||
|
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]:
|
||||||
|
"""
|
||||||
|
Retrain MLFlow models with updated training data.
|
||||||
|
|
||||||
|
This activity orchestrates the complete model retraining process,
|
||||||
|
including data preparation, model retraining execution, and result
|
||||||
|
validation. It handles data preprocessing, column cleanup, and
|
||||||
|
comprehensive error handling for production model management.
|
||||||
|
|
||||||
|
The retraining process includes:
|
||||||
|
1. Data timestamp extraction and validation
|
||||||
|
2. Column cleanup and data preparation
|
||||||
|
3. Data pivoting for model input format
|
||||||
|
4. MLFlow model retraining execution
|
||||||
|
5. Result validation and error handling
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- data (dict[str, Any]): Training data for model retraining
|
||||||
|
- model_name (str): Name of the MLFlow model to retrain
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Retraining results containing:
|
||||||
|
- status (str): Retraining operation status
|
||||||
|
- timestamp (str): Timestamp of the retraining operation
|
||||||
|
- experiment (str): MLFlow experiment identifier
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If retraining fails or encounters critical errors
|
||||||
|
"""
|
||||||
|
|
||||||
|
if self.minio_repository is None:
|
||||||
|
raise ValueError('Minio repository not initialized')
|
||||||
|
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
|
||||||
|
try:
|
||||||
|
# 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(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='ERROR_LOADING_RETRAIN_DATA',
|
||||||
|
message=f'Error loading retrain data: {e}',
|
||||||
|
block='retrain_model',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
self.error(trace, metadata)
|
||||||
|
return {
|
||||||
|
'success': False,
|
||||||
|
'message': f'Error loading retrain data: {e}',
|
||||||
|
'traceback': trace,
|
||||||
|
'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
||||||
|
}
|
||||||
|
|
||||||
|
self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
|
||||||
|
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
model_config = input_data.get('model_config', {})
|
||||||
|
|
||||||
|
self.info(f'Retraining model {model_name}...', metadata)
|
||||||
|
|
||||||
|
timestamp = data['timestamp'].max()
|
||||||
|
self.debug(f'Timestamp: {timestamp}', metadata)
|
||||||
|
|
||||||
|
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
||||||
|
if 'created_at' in data.columns:
|
||||||
|
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
||||||
|
subset=['variable', 'timestamp'], keep='first'
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')
|
||||||
|
|
||||||
|
data.drop(columns=['model_id'], inplace=True, errors='ignore')
|
||||||
|
data.drop(columns=['created_at'], inplace=True, errors='ignore')
|
||||||
|
|
||||||
|
# Pivot data for model input format
|
||||||
|
data = data.pivot(index='timestamp', columns='variable', values='value')
|
||||||
|
data.fillna(np.nan, inplace=True)
|
||||||
|
# data.reset_index(inplace=True)
|
||||||
|
data.columns.name = None
|
||||||
|
|
||||||
|
data['timestamp'] = data.index
|
||||||
|
data['timestamp'] = to_datetime(
|
||||||
|
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||||
|
).dt.strftime(DATETIME_FORMAT)
|
||||||
|
data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
|
||||||
|
|
||||||
|
data.columns.name = None
|
||||||
|
|
||||||
|
retrain_output = await self.model_monitoring_repository.retrain_model(
|
||||||
|
data=data, model_name=model_name, model_config=model_config, metadata=metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
if not retrain_output['success']:
|
||||||
|
trace = retrain_output['traceback']
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='RETRAIN_MODEL_ERROR',
|
||||||
|
message=f'Error retraining model {model_name}: {retrain_output["message"]}',
|
||||||
|
block='retrain_model',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
self.error(trace, metadata=metadata)
|
||||||
|
|
||||||
|
return {**retrain_output, 'timestamp': timestamp}
|
||||||
|
|
||||||
|
@activity.defn(name='update_production_model')
|
||||||
|
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||||
|
"""
|
||||||
|
Update production model with newly trained model version.
|
||||||
|
|
||||||
|
This activity manages the critical process of updating production
|
||||||
|
models with newly trained versions. It handles model deployment,
|
||||||
|
status tracking, and comprehensive reporting for operational
|
||||||
|
visibility and audit trails.
|
||||||
|
|
||||||
|
The update process includes:
|
||||||
|
1. Production model update execution
|
||||||
|
2. Status and metadata tracking
|
||||||
|
3. Comprehensive reporting and logging
|
||||||
|
4. Error handling and notification
|
||||||
|
5. Audit trail maintenance
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- model_name (str): Name of the MLFlow model to update
|
||||||
|
- experiment (str): MLFlow experiment identifier
|
||||||
|
- model_id (str): Unique identifier for the model version
|
||||||
|
- timestamp (str): Timestamp of the update operation
|
||||||
|
- status (str): Current status of the model update
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[Any, Any]: Comprehensive update report containing:
|
||||||
|
- model_id (str): Model version identifier
|
||||||
|
- model_name (str): Name of the updated model
|
||||||
|
- timestamp (str): Update operation timestamp
|
||||||
|
- status (str): Update operation status
|
||||||
|
- Additional MLFlow response metadata
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If production model update fails
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
experiment = input_data['experiment']
|
||||||
|
self.info(
|
||||||
|
f'Updating production model {model_name} from experiment {experiment}...', metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
response = await self.model_monitoring_repository.update_production_model(
|
||||||
|
experiment=experiment, model_name=model_name, metadata=metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
self.info(f'Production model {model_name} updated successfully', metadata)
|
||||||
|
return response
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
||||||
|
message=f'Error updating production model {model_name}: {e}',
|
||||||
|
block='update_production_model',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
self.error(trace, metadata=metadata)
|
||||||
|
raise e
|
||||||
|
|
||||||
|
@activity.defn(name='get_reference_data')
|
||||||
|
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
|
||||||
|
"""
|
||||||
|
Get reference data from the MLflow Model Registry.
|
||||||
|
|
||||||
|
This method retrieves evaluation reference data stored as artifacts in the
|
||||||
|
MLflow Model Registry. The reference data is typically used for model
|
||||||
|
drift detection, performance comparison, and quality validation. The method
|
||||||
|
loads the data from a CSV artifact file and formats timestamps for
|
||||||
|
consistent processing.
|
||||||
|
|
||||||
|
The method handles:
|
||||||
|
1. Loading evaluation data artifact from MLflow Model Registry
|
||||||
|
2. Timestamp parsing and formatting for consistency
|
||||||
|
3. Data conversion to dictionary format for workflow consumption
|
||||||
|
4. Graceful handling of missing reference data
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- model_name (str): Name of the MLFlow model to get reference data from
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
|
||||||
|
as a list of dictionaries. Returns None if reference data is not found
|
||||||
|
or if the artifact does not exist.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If artifact loading fails or encounters errors during processing
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
artifact = 'evaluation_data.csv'
|
||||||
|
|
||||||
|
reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
|
||||||
|
model_name=model_name, artifact_path=artifact, metadata=metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
if reference_data is None:
|
||||||
|
self.warning(f'Reference data not found for model {model_name}', metadata)
|
||||||
|
return None
|
||||||
|
|
||||||
|
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
||||||
|
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||||
|
|
||||||
|
return reference_data.to_dict(orient='records')
|
||||||
364
laborious/activities/model_metrics.py
Normal file
364
laborious/activities/model_metrics.py
Normal file
@@ -0,0 +1,364 @@
|
|||||||
|
from temporalio import activity, workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
import time
|
||||||
|
import traceback
|
||||||
|
import warnings
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from pandas import DataFrame, Index, to_datetime
|
||||||
|
from sientia.ModelAnalysis import ModelAnalysis
|
||||||
|
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.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(
|
||||||
|
'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class ModelMetrics(SientiaMonitoring):
|
||||||
|
"""
|
||||||
|
Metrics activities for the Laborious system.
|
||||||
|
|
||||||
|
This class provides activities for writing metrics to the Prometheus monitoring system.
|
||||||
|
"""
|
||||||
|
|
||||||
|
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
logger: Logger,
|
||||||
|
notification_handler: NotificationHandler,
|
||||||
|
metrics_controller: MetricsController,
|
||||||
|
):
|
||||||
|
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""
|
||||||
|
Close the model metrics activity and clean up resources.
|
||||||
|
"""
|
||||||
|
SientiaMonitoring.shutdown(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,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_drift_metrics(
|
||||||
|
self,
|
||||||
|
reference_data: DataFrame,
|
||||||
|
target_data: DataFrame,
|
||||||
|
target_name: str,
|
||||||
|
reference_columns: Index,
|
||||||
|
drift_metrics: list[str],
|
||||||
|
chunk_period: str,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
) -> DataFrame:
|
||||||
|
"""
|
||||||
|
Calculate univariate drift metrics for a model.
|
||||||
|
Args:
|
||||||
|
model_analysis (ModelAnalysis): Model analysis object
|
||||||
|
reference_data (DataFrame): Reference data
|
||||||
|
target_data (DataFrame): Target data
|
||||||
|
reference_columns (list[str]): Reference columns
|
||||||
|
drift_metrics (list[str]): Drift metrics
|
||||||
|
metadata (dict[str, Any]): Workflow execution metadata
|
||||||
|
"""
|
||||||
|
|
||||||
|
config = {
|
||||||
|
'target': target_name,
|
||||||
|
'prediction': 'prediction',
|
||||||
|
'timestamp': 'timestamp',
|
||||||
|
'features': reference_columns,
|
||||||
|
}
|
||||||
|
|
||||||
|
model_analysis = ModelAnalysis(config=config)
|
||||||
|
|
||||||
|
self._debug_dataframe(
|
||||||
|
f'Reference data: Size {reference_data.shape}', reference_data, 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()
|
||||||
|
try:
|
||||||
|
univariate_drift = model_analysis.detect_univariate_drift(
|
||||||
|
reference_df=reference_data,
|
||||||
|
analysis_df=target_data,
|
||||||
|
features=reference_columns,
|
||||||
|
timestamp_col=config['timestamp'],
|
||||||
|
methods=drift_metrics,
|
||||||
|
chunk_period=chunk_period,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self.error(f'Error detecting univariate drift: {e}', metadata)
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||||
|
)
|
||||||
|
raise e
|
||||||
|
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)
|
||||||
|
|
||||||
|
core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
|
||||||
|
start_time = time.time()
|
||||||
|
try:
|
||||||
|
multivariate_drift = model_analysis.detect_multivariate_drift(
|
||||||
|
reference_df=reference_data,
|
||||||
|
analysis_df=target_data,
|
||||||
|
features=reference_columns,
|
||||||
|
timestamp_col=config['timestamp'],
|
||||||
|
chunk_period=chunk_period,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self.error(f'Error detecting multivariate drift: {e}', metadata)
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||||
|
)
|
||||||
|
raise e
|
||||||
|
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)
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe')
|
||||||
|
try:
|
||||||
|
drift_df = model_analysis.get_drift_metrics_dataframe(
|
||||||
|
univariate_drift=univariate_drift,
|
||||||
|
multivariate_drift=multivariate_drift,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||||
|
)
|
||||||
|
raise e
|
||||||
|
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_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
|
||||||
|
|
||||||
|
return drift_df
|
||||||
|
|
||||||
|
@activity.defn(name='calculate_drift')
|
||||||
|
async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
|
||||||
|
"""
|
||||||
|
Calculate drift metrics for a model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict[str, Any]): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- model_name (str): Name of the MLFlow model to calculate drift for
|
||||||
|
- reference_data (pd.DataFrame): Reference data for the model
|
||||||
|
- target_data (pd.DataFrame): Target data for calculating drift
|
||||||
|
- target_name (str): Name of the target column
|
||||||
|
- drift_metrics (list[str]): List of drift metrics to calculate
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
model_id = input_data['model_id']
|
||||||
|
reference_raw_data = input_data['reference_data']
|
||||||
|
target_data = DataFrame(input_data['target_data'])
|
||||||
|
target_name = input_data['target_name']
|
||||||
|
drift_metrics = input_data['drift_metrics']
|
||||||
|
chunk_period = input_data['chunk_period']
|
||||||
|
|
||||||
|
if chunk_period not in ['min', 's']:
|
||||||
|
self.error(f'Invalid chunk period: {chunk_period}', metadata)
|
||||||
|
raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"')
|
||||||
|
|
||||||
|
self.info(f'Calculating drift for model {model_name}', metadata)
|
||||||
|
|
||||||
|
target_data = target_data.pivot(index='timestamp', columns='variable', values='value')
|
||||||
|
target_data['timestamp'] = target_data.index
|
||||||
|
target_data['timestamp'] = to_datetime(target_data['timestamp'])
|
||||||
|
target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||||
|
target_data = target_data.reset_index(drop=True)
|
||||||
|
target_data.dropna(inplace=True)
|
||||||
|
|
||||||
|
if reference_raw_data is not None:
|
||||||
|
self.info('Using reference data', metadata)
|
||||||
|
reference_data = DataFrame(reference_raw_data)
|
||||||
|
accurate = True
|
||||||
|
else:
|
||||||
|
# Get 30% first rows of target_data
|
||||||
|
self.warning('Using 30% first rows of target data as reference data', metadata)
|
||||||
|
target_data.sort_values(by='timestamp', ascending=True, inplace=True)
|
||||||
|
reference_data = target_data.head(int(len(target_data) * 0.3))
|
||||||
|
accurate = False
|
||||||
|
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
|
||||||
|
message='Using 30% first rows of target data as reference data',
|
||||||
|
block='model_metrics',
|
||||||
|
level=NotificationLevel.WARNING,
|
||||||
|
attachment_content=reference_data.to_csv(),
|
||||||
|
)
|
||||||
|
|
||||||
|
reference_columns = reference_data.drop(
|
||||||
|
columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
|
||||||
|
).columns
|
||||||
|
|
||||||
|
try:
|
||||||
|
drift_df = await self.get_drift_metrics(
|
||||||
|
reference_data=reference_data,
|
||||||
|
target_data=target_data,
|
||||||
|
target_name=target_name,
|
||||||
|
reference_columns=reference_columns,
|
||||||
|
drift_metrics=drift_metrics,
|
||||||
|
chunk_period=chunk_period,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
|
||||||
|
message=f'Error getting drift metrics: {e}',
|
||||||
|
block='model_metrics',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=traceback.format_exc(),
|
||||||
|
)
|
||||||
|
return []
|
||||||
|
|
||||||
|
if drift_df.empty:
|
||||||
|
self.warning('No drift metrics found', metadata)
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Drop unnecessary columns
|
||||||
|
drift_df.drop(columns=['p_value'], inplace=True)
|
||||||
|
|
||||||
|
# Extract timestamps only until minutes
|
||||||
|
if chunk_period == 'min':
|
||||||
|
target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
|
||||||
|
else:
|
||||||
|
target_timestamps = target_data['timestamp']
|
||||||
|
|
||||||
|
# Drop rows where timestamp is not in target data, to avoid save drift from reference
|
||||||
|
drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
|
||||||
|
|
||||||
|
if drift_df.empty:
|
||||||
|
self.warning(
|
||||||
|
'No drift metrics found after dropping rows where timestamp is not in target data',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Rename columns to match database columns
|
||||||
|
drift_df.rename(
|
||||||
|
columns={
|
||||||
|
'metric': 'method',
|
||||||
|
'statistic': 'value',
|
||||||
|
},
|
||||||
|
inplace=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Drop duplicates
|
||||||
|
drift_df.drop_duplicates(
|
||||||
|
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
|
||||||
|
)
|
||||||
|
|
||||||
|
drift_df['model_id'] = model_id
|
||||||
|
drift_df['accurate'] = accurate
|
||||||
|
|
||||||
|
drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
|
||||||
|
drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
|
||||||
|
drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
|
||||||
|
|
||||||
|
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
|
||||||
|
|
||||||
|
return drift_df.to_dict(orient='records')
|
||||||
|
|
||||||
|
@activity.defn(name='calculate_simple_metrics')
|
||||||
|
async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
|
||||||
|
"""
|
||||||
|
Calculate simple metrics for a model. Metrics available are:
|
||||||
|
- rmse
|
||||||
|
- mse
|
||||||
|
- mae
|
||||||
|
- r2
|
||||||
|
- accuracy
|
||||||
|
- precision
|
||||||
|
- recall
|
||||||
|
- f1
|
||||||
|
Args:
|
||||||
|
input_data (dict[str, Any]): Input data containing:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- model_id (str): ID of the MLFlow model
|
||||||
|
- target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns
|
||||||
|
- metrics (list[str]): List of metrics to calculate
|
||||||
|
Returns:
|
||||||
|
dict[Hashable, Any]: Dictionary containing the calculated metrics
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
model_id = input_data['model_id']
|
||||||
|
target_data = DataFrame(input_data['target_data'])
|
||||||
|
metrics = input_data['metrics']
|
||||||
|
interval_minutes = input_data['interval_minutes']
|
||||||
|
|
||||||
|
data_size = target_data.shape[0]
|
||||||
|
|
||||||
|
output_data = []
|
||||||
|
|
||||||
|
diff = target_data['target'] - target_data['prediction']
|
||||||
|
diff_squared = diff**2
|
||||||
|
|
||||||
|
self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
|
||||||
|
|
||||||
|
for metric in metrics:
|
||||||
|
if metric == 'rmse':
|
||||||
|
output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
|
||||||
|
elif metric == 'mse':
|
||||||
|
output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
|
||||||
|
elif metric == 'mae':
|
||||||
|
output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
|
||||||
|
elif metric == 'r2':
|
||||||
|
y_true = target_data['target']
|
||||||
|
y_mean = np.mean(y_true)
|
||||||
|
|
||||||
|
ss_res = np.sum(diff_squared)
|
||||||
|
ss_tot = np.sum((y_true - y_mean) ** 2)
|
||||||
|
|
||||||
|
# Evita divisão por zero
|
||||||
|
if ss_tot == 0:
|
||||||
|
r2_score = 0.0
|
||||||
|
else:
|
||||||
|
r2_score = 1 - (ss_res / ss_tot)
|
||||||
|
|
||||||
|
output_data.append({'metric': 'r2', 'value': r2_score})
|
||||||
|
|
||||||
|
data = DataFrame(output_data)
|
||||||
|
data['model_id'] = model_id
|
||||||
|
data['timestamp'] = target_data['timestamp'].max()
|
||||||
|
data['data_size'] = data_size
|
||||||
|
data['interval_minutes'] = interval_minutes
|
||||||
|
|
||||||
|
self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata)
|
||||||
|
|
||||||
|
return data.to_dict(orient='records')
|
||||||
515
laborious/activities/opc.py
Normal file
515
laborious/activities/opc.py
Normal file
@@ -0,0 +1,515 @@
|
|||||||
|
from temporalio import activity, workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
import traceback
|
||||||
|
from collections.abc import Hashable
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from pandas import DataFrame
|
||||||
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
from sientia_do.observability.logger import Logger
|
||||||
|
from sientia_do.observability.metrics_controller import MetricsController
|
||||||
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||||
|
|
||||||
|
from laborious.utils.repository.opc_repository import OpcRepository
|
||||||
|
|
||||||
|
OPC_WRITTING_ERROR_CONFIDENCE = 12
|
||||||
|
OPC_SESSION_BAD_CONFIDENCE = 14
|
||||||
|
OPC_SESSION_BAD_COMMENT_PREFIX = 'OPC UA session/channel error:'
|
||||||
|
OPC_WRITTING_ERROR_MESSAGE = 'Some data could not be written to OPC servers'
|
||||||
|
OPC_RECONNECT_IN_PROGRESS_COMMENT = 'OPC UA reconnect in progress'
|
||||||
|
OPC_COMMENT_SEPARATOR = ' | '
|
||||||
|
|
||||||
|
|
||||||
|
def _opc_session_bad_comment(opc_status: str | None) -> str:
|
||||||
|
status = opc_status or 'Unknown'
|
||||||
|
return f'{OPC_SESSION_BAD_COMMENT_PREFIX} {status}'
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_opc_write_error(
|
||||||
|
error_info: dict[str, Any] | None,
|
||||||
|
session_bad_seen: bool,
|
||||||
|
session_bad_status: str | None,
|
||||||
|
reconnect_in_progress_seen: bool,
|
||||||
|
) -> tuple[bool, str | None, bool]:
|
||||||
|
"""
|
||||||
|
Update session/reconnect flags from an OPC write error payload.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
error_info: Repository error details, or None when the write succeeded.
|
||||||
|
session_bad_seen: Whether a session_bad error was seen so far.
|
||||||
|
session_bad_status: Last known OPC status for session errors.
|
||||||
|
reconnect_in_progress_seen: Whether reconnect_in_progress was seen so far.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
Updated (session_bad_seen, session_bad_status, reconnect_in_progress_seen).
|
||||||
|
"""
|
||||||
|
if not error_info:
|
||||||
|
return session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||||
|
|
||||||
|
kind = error_info.get('opc_error_kind')
|
||||||
|
if kind == 'session_bad':
|
||||||
|
return True, error_info.get('opc_status', session_bad_status), reconnect_in_progress_seen
|
||||||
|
if kind == 'reconnect_in_progress':
|
||||||
|
return session_bad_seen, session_bad_status, True
|
||||||
|
return session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||||
|
|
||||||
|
|
||||||
|
class OPC(SientiaMonitoring):
|
||||||
|
"""
|
||||||
|
OPC server integration activities for real-time data export.
|
||||||
|
|
||||||
|
This class provides comprehensive OPC UA client functionality for connecting
|
||||||
|
to multiple OPC servers and writing prediction data in real-time. It implements
|
||||||
|
secure communication with certificate-based authentication and automatic
|
||||||
|
reconnection capabilities.
|
||||||
|
|
||||||
|
The class supports multiple OPC servers with individual configurations and
|
||||||
|
provides robust error handling and monitoring for production environments.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
opc_servers (dict): Configuration for multiple OPC servers
|
||||||
|
opc_repository (dict): Active OPC repository connections
|
||||||
|
logger (Logger): Logging and observability instance
|
||||||
|
notification_handler (NotificationHandler): Notification management instance
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
opc_servers: dict[str, dict[str, Any]],
|
||||||
|
logger: Logger,
|
||||||
|
notification_handler: NotificationHandler,
|
||||||
|
metrics_controller: MetricsController,
|
||||||
|
):
|
||||||
|
self.logger = logger
|
||||||
|
self.notification_handler = notification_handler
|
||||||
|
self.opc_servers = opc_servers
|
||||||
|
|
||||||
|
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||||
|
|
||||||
|
self.opc_repository: dict[str, OpcRepository] = {}
|
||||||
|
|
||||||
|
async def init_opc(self):
|
||||||
|
"""
|
||||||
|
Initialize OPC server connections and establish communication channels.
|
||||||
|
|
||||||
|
This method iterates through all configured OPC servers and attempts to
|
||||||
|
establish secure connections using certificate-based authentication.
|
||||||
|
Each server connection is managed independently, and connection failures
|
||||||
|
are reported through the notification system.
|
||||||
|
|
||||||
|
The method performs the following operations:
|
||||||
|
1. Creates OpcRepository instances for each configured server
|
||||||
|
2. Establishes secure connections with certificate validation
|
||||||
|
3. Reports connection success/failure through notifications
|
||||||
|
4. Logs connection status for operational visibility
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If OPC repository initialization fails or connection
|
||||||
|
establishment encounters critical errors
|
||||||
|
|
||||||
|
Note:
|
||||||
|
Connection failures are logged and reported but do not prevent
|
||||||
|
the initialization of other OPC servers. Each server is handled
|
||||||
|
independently to ensure maximum availability.
|
||||||
|
"""
|
||||||
|
self.logger.info('Initializing OPC servers...')
|
||||||
|
for opc_id, server in self.opc_servers.items():
|
||||||
|
self.opc_repository[opc_id] = OpcRepository(
|
||||||
|
opc_id=server['id'],
|
||||||
|
server_name=server['server_name'],
|
||||||
|
url=server['url'],
|
||||||
|
logger=self.logger,
|
||||||
|
server_uri=server['server_uri'],
|
||||||
|
cert_path=server['cert_path'],
|
||||||
|
private_key_path=server['private_key_path'],
|
||||||
|
server_cert_path=server['server_cert_path'],
|
||||||
|
notification_handler=self.notification_handler,
|
||||||
|
reconnection_interval=server['reconnection_interval'],
|
||||||
|
metrics_controller=self.metrics_controller,
|
||||||
|
)
|
||||||
|
is_connected, error_data = await self.opc_repository[opc_id].connect()
|
||||||
|
if not is_connected:
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata={
|
||||||
|
'model_id': '-',
|
||||||
|
'model_name': '-',
|
||||||
|
'workflow_name': '-',
|
||||||
|
'schedule_name': 'INITIALIZATION',
|
||||||
|
},
|
||||||
|
notification_id=error_data['notification_id'],
|
||||||
|
message=error_data['message'],
|
||||||
|
block=error_data['block'],
|
||||||
|
level=error_data.get('level', NotificationLevel.ERROR),
|
||||||
|
attachment_content=error_data.get('attachment_content', None),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.logger.info(
|
||||||
|
f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
|
||||||
|
)
|
||||||
|
|
||||||
|
async def write_data(
|
||||||
|
self,
|
||||||
|
server_id: str,
|
||||||
|
tag: str,
|
||||||
|
data: Any,
|
||||||
|
data_type: str,
|
||||||
|
tag_type: str,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
) -> tuple[float | None, dict[str, Any] | None]:
|
||||||
|
"""
|
||||||
|
Write data to a specific OPC server tag with comprehensive error handling.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[float | None, dict[str, Any] | None]: Response time on success, or
|
||||||
|
(None, error info_data) on repository failure.
|
||||||
|
"""
|
||||||
|
|
||||||
|
try:
|
||||||
|
is_success, info_data = await self.opc_repository[server_id].write_data(
|
||||||
|
tag, data, data_type, metadata
|
||||||
|
)
|
||||||
|
if not is_success:
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id=info_data['notification_id'],
|
||||||
|
message=info_data['message'],
|
||||||
|
block=info_data['block'],
|
||||||
|
level=info_data.get('level', NotificationLevel.ERROR),
|
||||||
|
attachment_content=info_data.get('attachment_content', None),
|
||||||
|
)
|
||||||
|
return None, info_data
|
||||||
|
return info_data['response_time'], None
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
|
||||||
|
message=f'Error writing data to OPC server: {e}',
|
||||||
|
block='write_opc_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
raise e
|
||||||
|
|
||||||
|
async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
|
||||||
|
"""
|
||||||
|
Validate that an OPC server is available and configured for write operations.
|
||||||
|
|
||||||
|
This method checks if the specified OPC server exists in the active
|
||||||
|
repository and is available for data writing operations. It provides
|
||||||
|
immediate feedback for server availability and logs validation failures
|
||||||
|
for operational monitoring.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
server_id (str): Unique identifier for the OPC server to validate
|
||||||
|
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if server is available, False otherwise
|
||||||
|
|
||||||
|
Note:
|
||||||
|
Server validation failures are automatically reported through the
|
||||||
|
notification system with detailed information about available servers.
|
||||||
|
This helps operators quickly identify configuration issues.
|
||||||
|
"""
|
||||||
|
if self.opc_repository.get(server_id) is None:
|
||||||
|
message = f'OPC server {server_id} not found to perform write operation.'
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='OPC_SERVER_NOT_FOUND',
|
||||||
|
message=message,
|
||||||
|
block='write_opc_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def _write_tags_from_config(
|
||||||
|
self,
|
||||||
|
server_id: str,
|
||||||
|
tags_config: dict[str, dict[str, Any]],
|
||||||
|
data: DataFrame,
|
||||||
|
data_column: str,
|
||||||
|
tag_type: str,
|
||||||
|
log_label: str,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
) -> tuple[dict[str, float | None], bool, str | None, bool]:
|
||||||
|
"""
|
||||||
|
Write a group of OPC tags and collect response times and error flags.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
server_id: Target OPC server identifier.
|
||||||
|
tags_config: Tag name to configuration mapping.
|
||||||
|
data: DataFrame with prediction/confidence columns.
|
||||||
|
data_column: Column name whose first row value is written.
|
||||||
|
tag_type: Tag category passed to write_data ('prediction' or 'confidence').
|
||||||
|
log_label: Human-readable label for success logs.
|
||||||
|
metadata: Context metadata for logging and notifications.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
(response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen)
|
||||||
|
"""
|
||||||
|
response_times: dict[str, float | None] = {}
|
||||||
|
session_bad_seen = False
|
||||||
|
session_bad_status: str | None = None
|
||||||
|
reconnect_in_progress_seen = False
|
||||||
|
|
||||||
|
for tag, tag_config in tags_config.items():
|
||||||
|
response_time, error_info = await self.write_data(
|
||||||
|
server_id=server_id,
|
||||||
|
tag=tag,
|
||||||
|
data=data.head(1)[data_column].values[0],
|
||||||
|
data_type=tag_config['data_type'],
|
||||||
|
tag_type=tag_type,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
session_bad_seen, session_bad_status, reconnect_in_progress_seen = (
|
||||||
|
_apply_opc_write_error(
|
||||||
|
error_info,
|
||||||
|
session_bad_seen,
|
||||||
|
session_bad_status,
|
||||||
|
reconnect_in_progress_seen,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if response_time is not None:
|
||||||
|
self.info(
|
||||||
|
f'{log_label} written to OPC server {server_id} for tag {tag}.',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
response_times[tag] = response_time
|
||||||
|
|
||||||
|
return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||||
|
|
||||||
|
async def manage_output_tags(
|
||||||
|
self,
|
||||||
|
server_id: str,
|
||||||
|
config: dict[str, Any],
|
||||||
|
data: DataFrame,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
) -> tuple[bool, dict[str, float | None], bool, str | None, bool]:
|
||||||
|
"""
|
||||||
|
Manage the writing of prediction and confidence data to OPC server tags.
|
||||||
|
|
||||||
|
This method orchestrates the writing of multiple data types to OPC servers
|
||||||
|
based on configuration. It handles both prediction data and confidence
|
||||||
|
values independently, allowing for flexible tag configuration and
|
||||||
|
comprehensive error handling.
|
||||||
|
|
||||||
|
The method supports two main tag types:
|
||||||
|
1. Prediction tags: Write actual prediction values to configured OPC tags
|
||||||
|
2. Confidence tags: Write confidence scores to separate OPC tags
|
||||||
|
|
||||||
|
Args:
|
||||||
|
server_id (str): Unique identifier for the target OPC server
|
||||||
|
config (dict[str, Any]): OPC tag configuration containing:
|
||||||
|
- prediction_tags (dict, optional): Prediction tag configurations
|
||||||
|
- confidence_tags (dict, optional): Confidence tag configurations
|
||||||
|
data (DataFrame): DataFrame containing prediction and confidence data
|
||||||
|
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||||
|
success (bool): Current success status to maintain across operations
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple[bool, int]: (overall_success, total_tags_written)
|
||||||
|
- overall_success: True if all configured tags were written successfully
|
||||||
|
- total_tags_written: Count of successfully written tags
|
||||||
|
"""
|
||||||
|
response_times: dict[str, float | None] = {}
|
||||||
|
session_bad_seen = False
|
||||||
|
session_bad_status: str | None = None
|
||||||
|
reconnect_in_progress_seen = False
|
||||||
|
|
||||||
|
tag_groups = (
|
||||||
|
('prediction_tags', 'prediction', 'prediction', 'Prediction data'),
|
||||||
|
('confidence_tags', 'prediction_confidence', 'confidence', 'Confidence data'),
|
||||||
|
)
|
||||||
|
for config_key, data_column, tag_type, log_label in tag_groups:
|
||||||
|
if config_key not in config:
|
||||||
|
continue
|
||||||
|
(
|
||||||
|
group_times,
|
||||||
|
group_session_bad,
|
||||||
|
group_status,
|
||||||
|
group_reconnect,
|
||||||
|
) = await self._write_tags_from_config(
|
||||||
|
server_id=server_id,
|
||||||
|
tags_config=config[config_key],
|
||||||
|
data=data,
|
||||||
|
data_column=data_column,
|
||||||
|
tag_type=tag_type,
|
||||||
|
log_label=log_label,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
response_times.update(group_times)
|
||||||
|
if group_session_bad:
|
||||||
|
session_bad_seen = True
|
||||||
|
session_bad_status = group_status or session_bad_status
|
||||||
|
if group_reconnect:
|
||||||
|
reconnect_in_progress_seen = True
|
||||||
|
|
||||||
|
success = None not in response_times.values()
|
||||||
|
return (
|
||||||
|
success,
|
||||||
|
response_times,
|
||||||
|
session_bad_seen,
|
||||||
|
session_bad_status,
|
||||||
|
reconnect_in_progress_seen,
|
||||||
|
)
|
||||||
|
|
||||||
|
@activity.defn(name='write_opc_data')
|
||||||
|
async def write_opc_data(
|
||||||
|
self, input_data: dict[str, Any]
|
||||||
|
) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
|
||||||
|
"""
|
||||||
|
Write prediction and confidence data to OPC servers. The two writing
|
||||||
|
operations are optional and independent of each other.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- input_data(dict[str, Any]): The input data. Contains the following keys:
|
||||||
|
- data(dict[str, Any]): The dataframe that contains the data to write
|
||||||
|
to the OPC servers.
|
||||||
|
- opc_output_config(dict[str, Any]): The OPC writing configuration.
|
||||||
|
The keys are the OPC server names and the values contain:
|
||||||
|
- prediction_tags(dict[str, Any]): The tags to write to the OPC servers.
|
||||||
|
- confidence_tags(dict[str, Any]): The tags to write to the OPC servers.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- dict[Any, Any]: The data that was written to the OPC servers.
|
||||||
|
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
self.info('Writing data to OPC servers...', metadata)
|
||||||
|
data = DataFrame(input_data['data'])
|
||||||
|
opc_output_config = input_data['opc_output_config']
|
||||||
|
self.info(f'Data to write: {data.size} rows', metadata)
|
||||||
|
|
||||||
|
success = True
|
||||||
|
session_bad_seen = False
|
||||||
|
session_bad_status: str | None = None
|
||||||
|
reconnect_in_progress_seen = False
|
||||||
|
|
||||||
|
metrics: dict[str, dict[str, float | None]] = {}
|
||||||
|
|
||||||
|
for server_id, config in opc_output_config.items():
|
||||||
|
if not await self.validate_server(server_id, metadata):
|
||||||
|
success = False
|
||||||
|
continue
|
||||||
|
|
||||||
|
(
|
||||||
|
local_success,
|
||||||
|
local_response_times,
|
||||||
|
local_session_bad,
|
||||||
|
local_status,
|
||||||
|
local_reconnect_in_progress,
|
||||||
|
) = await self.manage_output_tags(server_id, config, data, metadata)
|
||||||
|
metrics[server_id] = local_response_times
|
||||||
|
local_count = len(local_response_times)
|
||||||
|
success = success and local_success
|
||||||
|
if local_session_bad:
|
||||||
|
session_bad_seen = True
|
||||||
|
session_bad_status = local_status or session_bad_status
|
||||||
|
if local_reconnect_in_progress:
|
||||||
|
reconnect_in_progress_seen = True
|
||||||
|
|
||||||
|
self.info(
|
||||||
|
f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
return (
|
||||||
|
self.process_confidence(
|
||||||
|
data,
|
||||||
|
success,
|
||||||
|
metadata,
|
||||||
|
session_bad=session_bad_seen,
|
||||||
|
opc_status=session_bad_status,
|
||||||
|
reconnect_in_progress=reconnect_in_progress_seen,
|
||||||
|
),
|
||||||
|
metrics,
|
||||||
|
)
|
||||||
|
|
||||||
|
def process_confidence(
|
||||||
|
self,
|
||||||
|
data: DataFrame,
|
||||||
|
success: bool,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
*,
|
||||||
|
session_bad: bool = False,
|
||||||
|
opc_status: str | None = None,
|
||||||
|
reconnect_in_progress: bool = False,
|
||||||
|
) -> dict[Hashable, Any]:
|
||||||
|
"""
|
||||||
|
Process prediction confidence based on OPC write operation success.
|
||||||
|
|
||||||
|
This method updates the prediction confidence values in the DataFrame
|
||||||
|
based on the success status of OPC server write operations. If any
|
||||||
|
write operations failed, it sets the confidence to a predefined error
|
||||||
|
value to indicate data quality issues.
|
||||||
|
|
||||||
|
The method implements a confidence degradation strategy:
|
||||||
|
- Success: Maintains original confidence values
|
||||||
|
- Failure: Sets confidence to error value for operational awareness
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (DataFrame): DataFrame containing prediction and confidence data
|
||||||
|
success (bool): Overall success status of OPC write operations
|
||||||
|
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[Any, Any]: Processed data as a dictionary with updated confidence values
|
||||||
|
|
||||||
|
Note:
|
||||||
|
The error confidence value (OPC_WRITTING_ERROR_CONFIDENCE = 12) is
|
||||||
|
used to indicate that data was not successfully exported to OPC servers.
|
||||||
|
This allows downstream systems to handle data quality appropriately.
|
||||||
|
"""
|
||||||
|
|
||||||
|
if not success:
|
||||||
|
comment_parts: list[str] = []
|
||||||
|
confidence = OPC_WRITTING_ERROR_CONFIDENCE
|
||||||
|
|
||||||
|
if session_bad:
|
||||||
|
comment_parts.append(_opc_session_bad_comment(opc_status))
|
||||||
|
confidence = OPC_SESSION_BAD_CONFIDENCE
|
||||||
|
if reconnect_in_progress:
|
||||||
|
comment_parts.append(OPC_RECONNECT_IN_PROGRESS_COMMENT)
|
||||||
|
confidence = OPC_SESSION_BAD_CONFIDENCE
|
||||||
|
if not comment_parts:
|
||||||
|
comment_parts.append(OPC_WRITTING_ERROR_MESSAGE)
|
||||||
|
|
||||||
|
comments = OPC_COMMENT_SEPARATOR.join(comment_parts)
|
||||||
|
data['prediction_confidence'] = confidence
|
||||||
|
data['comments'] = comments
|
||||||
|
self.debug(
|
||||||
|
f'OPC write issues, confidence={confidence}, comments={comments}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.debug('Data written to OPC servers successfully.', metadata)
|
||||||
|
|
||||||
|
return data.to_dict()
|
||||||
|
|
||||||
|
async def aclose(self):
|
||||||
|
"""
|
||||||
|
Gracefully shutdown all OPC server connections and cleanup resources.
|
||||||
|
|
||||||
|
This method ensures proper cleanup of all active OPC server connections
|
||||||
|
by calling the disconnect method on each repository instance. It's
|
||||||
|
designed to be called during application shutdown to prevent resource
|
||||||
|
leaks and ensure clean termination.
|
||||||
|
|
||||||
|
The method performs the following cleanup operations:
|
||||||
|
1. Iterates through all active OPC repository connections
|
||||||
|
2. Calls disconnect() on each repository instance
|
||||||
|
3. Allows for graceful connection termination
|
||||||
|
4. Prevents resource leaks and connection hanging
|
||||||
|
|
||||||
|
Note:
|
||||||
|
This method should be called during application shutdown to ensure
|
||||||
|
proper cleanup. It handles all active connections regardless of
|
||||||
|
their current state and provides a clean shutdown experience.
|
||||||
|
"""
|
||||||
|
for opc in self.opc_repository.values():
|
||||||
|
await opc.disconnect()
|
||||||
210
laborious/activities/storage.py
Normal file
210
laborious/activities/storage.py
Normal file
@@ -0,0 +1,210 @@
|
|||||||
|
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
|
||||||
|
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.repository.minio_repository import MinioRepository
|
||||||
|
from sientia_do.temporal.activities.postgres import Postgres
|
||||||
|
from sientia_do.temporal.constants import now
|
||||||
|
|
||||||
|
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, 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,
|
||||||
|
port: int,
|
||||||
|
user: str,
|
||||||
|
password: str,
|
||||||
|
dbname: str,
|
||||||
|
min_connections: int,
|
||||||
|
max_connections: int,
|
||||||
|
retention_hours: int = 24,
|
||||||
|
minio_repository: MinioRepository | None = None,
|
||||||
|
logger: Logger | None = None,
|
||||||
|
notification_handler: NotificationHandler | None = None,
|
||||||
|
metrics_controller: MetricsController | None = None,
|
||||||
|
):
|
||||||
|
self.retention_hours = retention_hours
|
||||||
|
Postgres.__init__(
|
||||||
|
self,
|
||||||
|
host=host,
|
||||||
|
port=port,
|
||||||
|
user=user,
|
||||||
|
password=password,
|
||||||
|
dbname=dbname,
|
||||||
|
min_connections=min_connections,
|
||||||
|
max_connections=max_connections,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
MinioManager.__init__(
|
||||||
|
self, minio_repository, logger, notification_handler, metrics_controller
|
||||||
|
)
|
||||||
|
|
||||||
|
@activity.defn(name='load_query_with_minio_offload')
|
||||||
|
async def load_query_with_minio_offload(
|
||||||
|
self, input_data: dict[str, Any]
|
||||||
|
) -> MinioDataFramePayload:
|
||||||
|
"""
|
||||||
|
Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
|
||||||
|
|
||||||
|
Args (input_data):
|
||||||
|
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[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', {})
|
||||||
|
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:
|
||||||
|
keys = await self.minio_repository.list_objects(
|
||||||
|
prefix=prefix,
|
||||||
|
recursive=True,
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
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_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 report
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
"""Close Storage resources (MinIO client and Postgres engine)."""
|
||||||
|
if hasattr(self, 'engine'):
|
||||||
|
Postgres.close(self)
|
||||||
|
MinioManager.close(self)
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
self.close()
|
||||||
200
laborious/metrics.py
Normal file
200
laborious/metrics.py
Normal file
@@ -0,0 +1,200 @@
|
|||||||
|
"""
|
||||||
|
Laborious Metrics Module
|
||||||
|
|
||||||
|
This module defines all Prometheus metrics used by the Sientia DataOps Laborious system
|
||||||
|
for monitoring and observability. The metrics provide insights into system performance,
|
||||||
|
prediction quality, and operational health.
|
||||||
|
|
||||||
|
The metrics are designed to be scraped by Prometheus and can be visualized in
|
||||||
|
Grafana or other monitoring dashboards to provide real-time visibility into
|
||||||
|
the system's operation.
|
||||||
|
|
||||||
|
Key Metric Categories:
|
||||||
|
- Application Health: Overall system status and availability
|
||||||
|
- Prediction Operations: Count and performance of prediction operations
|
||||||
|
- Data Quality: Confidence levels and validation results
|
||||||
|
- Export Operations: Database and OPC export performance
|
||||||
|
- Response Times: Performance monitoring for various operations
|
||||||
|
|
||||||
|
Metric Labels:
|
||||||
|
- pod_id: Kubernetes pod identifier for multi-instance deployments
|
||||||
|
- 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
|
||||||
|
|
||||||
|
# Application health metric
|
||||||
|
APP_UP = Gauge(
|
||||||
|
'app_up',
|
||||||
|
'Indicates if the application is running (1) or shutting down (0)',
|
||||||
|
['pod_id'],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Prediction operation metrics
|
||||||
|
PREDICTIONS_WRITTEN_COUNT = Counter(
|
||||||
|
'laborious_predictions_written_count',
|
||||||
|
'Number of predictions written to the database table predictions',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Prediction quality metrics
|
||||||
|
PREDICTION_CONFIDENCE_MONITOR = Gauge(
|
||||||
|
'laborious_prediction_confidence_monitor',
|
||||||
|
'Current confidence of each prediction',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Prediction total response time
|
||||||
|
PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
|
||||||
|
'laborious_prediction_response_time_monitor',
|
||||||
|
'Current response time of each prediction',
|
||||||
|
CORE_LABELS,
|
||||||
|
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ================== OPC metrics ==================
|
||||||
|
|
||||||
|
|
||||||
|
PREDICTION_OPC_WRITING_COUNT = Counter(
|
||||||
|
'laborious_prediction_opc_writing_count',
|
||||||
|
'Number of predictions written to the OPC server',
|
||||||
|
[*CORE_LABELS, 'opc_server_id', 'tag'],
|
||||||
|
)
|
||||||
|
|
||||||
|
PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
|
||||||
|
'laborious_prediction_opc_writing_response_time_monitor',
|
||||||
|
'Current response time of each prediction written to the OPC server',
|
||||||
|
[*CORE_LABELS, 'opc_server_id', 'tag'],
|
||||||
|
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
|
||||||
|
)
|
||||||
|
|
||||||
|
OPC_CONNECTIONS_TOTAL = Counter(
|
||||||
|
'opc_connections_initiated_total',
|
||||||
|
'Total connection attempts to OPC servers',
|
||||||
|
['pod_id', 'server_name'],
|
||||||
|
)
|
||||||
|
OPC_CONNECTIONS_FAILED = Counter(
|
||||||
|
'opc_connections_failed_total',
|
||||||
|
'Total failed connection attempts to OPC servers',
|
||||||
|
['pod_id', 'server_name'],
|
||||||
|
)
|
||||||
|
OPC_CONNECTION_STATUS = Gauge(
|
||||||
|
'opc_connection_status',
|
||||||
|
'Connection status with the OPC server (1=connected, 0=disconnected)',
|
||||||
|
['pod_id', 'server_name', 'server_url'],
|
||||||
|
)
|
||||||
|
|
||||||
|
_OPC_SESSION_DEBUG_LABELS = ['pod_id', 'server_name', 'runtime', 'opc_server_id', 'session_id']
|
||||||
|
|
||||||
|
OPC_SESSION_CREATED_TOTAL = Counter(
|
||||||
|
'opc_session_created_total',
|
||||||
|
'OPC UA sessions established (after successful connect)',
|
||||||
|
_OPC_SESSION_DEBUG_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
OPC_SESSION_CLOSED_TOTAL = Counter(
|
||||||
|
'opc_session_closed_total',
|
||||||
|
'OPC UA client disconnects completed (session tear-down initiated)',
|
||||||
|
_OPC_SESSION_DEBUG_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
OPC_SESSION_REVISED_TIMEOUT_MS = Gauge(
|
||||||
|
'opc_session_revised_timeout_milliseconds',
|
||||||
|
'Server-revised OPC UA session timeout (RevisedSessionTimeout) in ms after connect',
|
||||||
|
_OPC_SESSION_DEBUG_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
OPC_WRITE_ATTEMPT_LABELS = [*_OPC_SESSION_DEBUG_LABELS, 'model_id', 'model_name', 'result']
|
||||||
|
|
||||||
|
OPC_WRITE_ATTEMPTS_TOTAL = Counter(
|
||||||
|
'opc_write_attempts_total',
|
||||||
|
'OPC UA write attempts with session and outcome (result=OK or exception class name)',
|
||||||
|
OPC_WRITE_ATTEMPT_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL = Counter(
|
||||||
|
'opc_write_inter_arrival_over_session_timeout_total',
|
||||||
|
'Successful writes where seconds since the previous successful write exceeded RevisedSessionTimeout (ms)',
|
||||||
|
_OPC_SESSION_DEBUG_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ================== Model metrics ==================
|
||||||
|
|
||||||
|
MODEL_READ_LAG = Histogram(
|
||||||
|
'laborious_model_read_lag',
|
||||||
|
'Lag between the start and read of read operations',
|
||||||
|
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',
|
||||||
|
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',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
MODEL_WRITE_COUNT = Counter(
|
||||||
|
'laborious_model_write_count',
|
||||||
|
'Number of writes to the model',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
MODEL_READ_ERROR_COUNT = Counter(
|
||||||
|
'laborious_model_read_error_count',
|
||||||
|
'Number of errors reading from the model',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
MODEL_WRITE_ERROR_COUNT = Counter(
|
||||||
|
'laborious_model_write_error_count',
|
||||||
|
'Number of errors writing to the model',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
MODEL_ANALYZE_LAG = Histogram(
|
||||||
|
'laborious_model_analyze_lag',
|
||||||
|
'Lag between the start and end of analyze operations',
|
||||||
|
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',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
MODEL_ANALYZE_ERROR_COUNT = Counter(
|
||||||
|
'laborious_model_analyze_error_count',
|
||||||
|
'Number of errors during analyze operations',
|
||||||
|
CORE_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ================== PI Web API metrics ==================
|
||||||
|
|
||||||
|
PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name']
|
||||||
|
|
||||||
|
PI_WEB_API_PREDICTION_WRITTEN_COUNT = Counter(
|
||||||
|
'laborious_pi_web_api_prediction_written_count',
|
||||||
|
'Number of predictions written to the PI Web API',
|
||||||
|
PI_WEB_API_LABELS,
|
||||||
|
)
|
||||||
|
|
||||||
|
PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT = Counter(
|
||||||
|
'laborious_pi_web_api_prediction_written_error_count',
|
||||||
|
'Number of errors writing predictions to the PI Web API',
|
||||||
|
PI_WEB_API_LABELS,
|
||||||
|
)
|
||||||
0
laborious/utils/__init__.py
Normal file
0
laborious/utils/__init__.py
Normal file
92
laborious/utils/connectors_config.py
Normal file
92
laborious/utils/connectors_config.py
Normal file
@@ -0,0 +1,92 @@
|
|||||||
|
import json
|
||||||
|
from os import getenv
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
def build_mlflow_config() -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Build MLFlow server configuration from environment variables.
|
||||||
|
|
||||||
|
This function constructs an MLFlow configuration dictionary from
|
||||||
|
environment variables with sensible defaults for local development.
|
||||||
|
It handles server connection and authentication parameters.
|
||||||
|
|
||||||
|
Environment Variables:
|
||||||
|
MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
|
||||||
|
MLFLOW_PORT: MLFlow server port (default: 5080)
|
||||||
|
MLFLOW_USERNAME: MLFlow username (default: aignosi)
|
||||||
|
MLFLOW_PASSWORD: MLFlow password (default: aignosi)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: MLFlow configuration dictionary with all required parameters
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
'host': getenv('MLFLOW_HOST', 'http://localhost'),
|
||||||
|
'port': int(getenv('MLFLOW_PORT', '5080')),
|
||||||
|
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
|
||||||
|
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def build_opc_config() -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Build OPC server configuration from environment variables.
|
||||||
|
|
||||||
|
This function constructs an OPC server configuration dictionary from
|
||||||
|
environment variables. It supports both single server and multi-server
|
||||||
|
configurations with flexible parameter handling.
|
||||||
|
|
||||||
|
Environment Variables:
|
||||||
|
OPC_CONFIG: JSON string containing multiple OPC server configurations
|
||||||
|
OPC_ID: OPC server ID (fallback, default: 1)
|
||||||
|
OPC_URL: Single OPC server URL (fallback, default: opc.tcp://localhost:4840)
|
||||||
|
OPC_SERVER_URI: Single OPC server URI (fallback, default: opc.tcp://localhost:4840)
|
||||||
|
OPC_CERT_PATH: Client certificate path (fallback, default: None)
|
||||||
|
OPC_PRIVATE_KEY_PATH: Client private key path (fallback, default: None)
|
||||||
|
OPC_SERVER_CERT_PATH: Server certificate path (fallback, default: None)
|
||||||
|
OPC_RECONNECTION_INTERVAL: Reconnection interval in seconds (fallback, default: 120)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: OPC server configuration dictionary
|
||||||
|
"""
|
||||||
|
opc_raw = getenv('OPC_CONFIG', None)
|
||||||
|
|
||||||
|
if opc_raw:
|
||||||
|
return json.loads(opc_raw)
|
||||||
|
|
||||||
|
return {
|
||||||
|
getenv('OPC_ID', '1'): {
|
||||||
|
'id': getenv('OPC_ID', '1'),
|
||||||
|
'server_name': getenv('OPC_SERVER_NAME', 'default_server'),
|
||||||
|
'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'),
|
||||||
|
'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'),
|
||||||
|
'cert_path': getenv('OPC_CERT_PATH', None),
|
||||||
|
'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None),
|
||||||
|
'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None),
|
||||||
|
'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def build_minio_config() -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Build MinIO (S3-compatible) configuration from environment variables.
|
||||||
|
|
||||||
|
Environment Variables:
|
||||||
|
MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000)
|
||||||
|
MINIO_ACCESS_KEY: Access key (default: minioadmin)
|
||||||
|
MINIO_SECRET_KEY: Secret key (default: minioadmin)
|
||||||
|
MINIO_REGION: Region name for S3 client (default: us-east-1)
|
||||||
|
MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious)
|
||||||
|
MINIO_SECURE: Whether to use HTTPS (default: false)
|
||||||
|
Returns:
|
||||||
|
dict: MinIO configuration dictionary
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
|
||||||
|
'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
|
||||||
|
'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
|
||||||
|
'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
34
laborious/utils/dataframe_debug.py
Normal 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/filters/__init__.py
Normal file
0
laborious/utils/filters/__init__.py
Normal file
48
laborious/utils/filters/conditional_filters.py
Normal file
48
laborious/utils/filters/conditional_filters.py
Normal file
@@ -0,0 +1,48 @@
|
|||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
|
||||||
|
def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Filter to check if specific variables contain null values.
|
||||||
|
|
||||||
|
This function examines a DataFrame to determine if any of the specified variables
|
||||||
|
contain null (NaN) values. It returns True if null values are found for any of
|
||||||
|
the specified variables, False otherwise.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (DataFrame): The pandas DataFrame to be examined. Must contain columns
|
||||||
|
named 'variable' and 'value'.
|
||||||
|
config (dict): Configuration dictionary containing the following key:
|
||||||
|
- variables (list): List of variable names to check for null values
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if any of the specified variables contain null values,
|
||||||
|
False if none of the specified variables contain null values.
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
if data.empty:
|
||||||
|
return False
|
||||||
|
|
||||||
|
return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
|
||||||
|
|
||||||
|
|
||||||
|
def filter_empty_data(data: DataFrame, _config: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Filter to check if the DataFrame is empty.
|
||||||
|
|
||||||
|
This function determines whether the provided DataFrame contains any data.
|
||||||
|
It's a simple utility function that can be used in conditional logic to
|
||||||
|
handle cases where no data is available.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (DataFrame): The pandas DataFrame to be checked for emptiness.
|
||||||
|
_config (dict): Configuration dictionary (unused in this function).
|
||||||
|
The underscore prefix indicates this parameter is required for
|
||||||
|
interface consistency but not used in the implementation.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if the DataFrame is empty (has no rows), False if it contains data.
|
||||||
|
|
||||||
|
"""
|
||||||
|
return data.empty
|
||||||
64
laborious/utils/filters/mlflow_filters.py
Normal file
64
laborious/utils/filters/mlflow_filters.py
Normal file
@@ -0,0 +1,64 @@
|
|||||||
|
import numpy as np
|
||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
|
||||||
|
def api_error_filter(response: dict, _config: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Filter MLFlow API responses for error conditions.
|
||||||
|
|
||||||
|
This function analyzes MLFlow API responses to detect error conditions
|
||||||
|
and determine if the response should be filtered out due to quality
|
||||||
|
or reliability issues.
|
||||||
|
|
||||||
|
|
||||||
|
Args:
|
||||||
|
response: MLFlow API response data (dict)
|
||||||
|
_config: Filter configuration dictionary
|
||||||
|
Required keys:
|
||||||
|
- error_codes (list, optional): List of error codes to detect
|
||||||
|
- error_keywords (list, optional): List of error keywords to detect
|
||||||
|
- check_structure (bool, optional): Whether to validate response structure
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if data should be filtered (contains errors), False otherwise
|
||||||
|
|
||||||
|
"""
|
||||||
|
if not response:
|
||||||
|
return True
|
||||||
|
|
||||||
|
if not response['success']:
|
||||||
|
return True
|
||||||
|
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Filter data for NaN (Not a Number) values.
|
||||||
|
|
||||||
|
This function detects NaN values in MLFlow prediction results and
|
||||||
|
determines if the data quality is sufficient for further processing
|
||||||
|
or export operations.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
predictions: DataFrame containing prediction data to check for NaN values
|
||||||
|
_config: Filter configuration dictionary
|
||||||
|
Required keys:
|
||||||
|
- max_nan_ratio (float, optional): Maximum allowed NaN value ratio (0.0 to 1.0)
|
||||||
|
- max_nan_count (int, optional): Maximum allowed NaN value count
|
||||||
|
- check_nested (bool, optional): Whether to check nested data structures
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if data should be filtered (too many NaN values), False otherwise
|
||||||
|
|
||||||
|
"""
|
||||||
|
data = (
|
||||||
|
predictions.replace({None: np.nan})
|
||||||
|
.drop(columns=['timestamp'], errors='ignore')
|
||||||
|
.infer_objects()
|
||||||
|
)
|
||||||
|
|
||||||
|
if data.isna().all().all():
|
||||||
|
return True
|
||||||
|
|
||||||
|
return False
|
||||||
0
laborious/utils/models/__init__.py
Normal file
0
laborious/utils/models/__init__.py
Normal file
348
laborious/utils/models/minio_dataframe_payload.py
Normal file
348
laborious/utils/models/minio_dataframe_payload.py
Normal 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
32
laborious/utils/repository/minio_manager.py
Normal 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)
|
||||||
1703
laborious/utils/repository/model_repository.py
Normal file
1703
laborious/utils/repository/model_repository.py
Normal file
File diff suppressed because it is too large
Load Diff
866
laborious/utils/repository/opc_repository.py
Normal file
866
laborious/utils/repository/opc_repository.py
Normal file
@@ -0,0 +1,866 @@
|
|||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import traceback
|
||||||
|
from datetime import datetime
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from asyncua import Client
|
||||||
|
from asyncua.crypto.security_policies import SecurityPolicyBasic256
|
||||||
|
from asyncua.ua import DataValue, Variant, VariantType
|
||||||
|
from asyncua.ua.uaerrors import UaStatusCodeError
|
||||||
|
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 laborious import metrics
|
||||||
|
|
||||||
|
# Requested session and secure channel lifetime (ms) before server revision; 10 minutes.
|
||||||
|
OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS = 10 * 60 * 1000
|
||||||
|
|
||||||
|
|
||||||
|
class OpcClientAlreadyExistsError(RuntimeError):
|
||||||
|
"""Raised when _create_client is called while self.client is already set."""
|
||||||
|
|
||||||
|
|
||||||
|
class OpcSessionAlreadyConnectedError(RuntimeError):
|
||||||
|
"""Raised when _open_session is called while a UA session is already open."""
|
||||||
|
|
||||||
|
|
||||||
|
class OpcClientNotInitializedError(RuntimeError):
|
||||||
|
"""Raised when _open_session is called before _create_client."""
|
||||||
|
|
||||||
|
|
||||||
|
RECONNECTABLE_OPC_BAD_NAMES: frozenset[str] = frozenset(
|
||||||
|
{
|
||||||
|
'BadSessionIdInvalid',
|
||||||
|
'BadSessionClosed',
|
||||||
|
'BadSessionNotActivated',
|
||||||
|
'BadSecureChannelIdInvalid',
|
||||||
|
'BadSecureChannelClosed',
|
||||||
|
'BadSecureChannelTokenUnknown',
|
||||||
|
'BadTcpSecureChannelUnknown',
|
||||||
|
'BadServerNotConnected',
|
||||||
|
'BadConnectionClosed',
|
||||||
|
'BadDisconnect',
|
||||||
|
'BadConnectionRejected',
|
||||||
|
'BadCommunicationError',
|
||||||
|
'BadRequestInterrupted',
|
||||||
|
'BadUnknownResponse',
|
||||||
|
'BadTimeout',
|
||||||
|
'BadRequestTimeout',
|
||||||
|
'BadSequenceNumberInvalid',
|
||||||
|
'BadSequenceNumberUnknown',
|
||||||
|
'BadSecurityModeInsufficient',
|
||||||
|
'BadRequestHeaderInvalid',
|
||||||
|
'BadInvalidState',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _opc_authentication_token_str(client: Client | None) -> str:
|
||||||
|
"""
|
||||||
|
Serialize the current OPC UA authentication token (session handle) for logging and metrics.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
str: Token string, or "unknown" if unavailable.
|
||||||
|
"""
|
||||||
|
if client is None:
|
||||||
|
return 'unknown'
|
||||||
|
try:
|
||||||
|
proto = client.uaclient.protocol
|
||||||
|
if proto is None:
|
||||||
|
return 'unknown'
|
||||||
|
tok = getattr(proto, 'authentication_token', None)
|
||||||
|
if tok is None:
|
||||||
|
return 'unknown'
|
||||||
|
return str(tok)
|
||||||
|
except Exception:
|
||||||
|
return 'unknown'
|
||||||
|
|
||||||
|
|
||||||
|
def _opc_status_from_exception(exc: BaseException) -> str:
|
||||||
|
"""
|
||||||
|
Resolve OPC UA status name from an exception, including chained UaStatusCodeError causes.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
exc (BaseException): Raised error from asyncua.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
str: Status class name or generic Python exception name.
|
||||||
|
"""
|
||||||
|
current: BaseException | None = exc
|
||||||
|
while current is not None:
|
||||||
|
if isinstance(current, UaStatusCodeError):
|
||||||
|
return type(current).__name__
|
||||||
|
current = current.__cause__
|
||||||
|
return type(exc).__name__
|
||||||
|
|
||||||
|
|
||||||
|
def is_reconnectable_opcua_bad(exc: BaseException) -> bool:
|
||||||
|
"""
|
||||||
|
Return whether the exception is a Tier-1 OPC UA Bad* that should trigger reconnect.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
exc (BaseException): Raised error from get_node or write_value.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
bool: True if reconnect should be scheduled.
|
||||||
|
"""
|
||||||
|
return _opc_status_from_exception(exc) in RECONNECTABLE_OPC_BAD_NAMES
|
||||||
|
|
||||||
|
|
||||||
|
def _model_labels_from_write_metadata(metadata: dict[str, Any] | None) -> dict[str, str]:
|
||||||
|
"""
|
||||||
|
Extract model_id and model_name from write metadata for Prometheus labels.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
metadata (dict[str, Any] | None): Context passed into write_data; may omit keys.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict[str, str]: Labels model_id and model_name, defaulting to "unknown".
|
||||||
|
"""
|
||||||
|
if not metadata:
|
||||||
|
return {'model_id': 'unknown', 'model_name': 'unknown'}
|
||||||
|
return {
|
||||||
|
'model_id': str(metadata.get('model_id', 'unknown')),
|
||||||
|
'model_name': str(metadata.get('model_name', 'unknown')),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
data_type_map = {
|
||||||
|
'float': {
|
||||||
|
'converter': float,
|
||||||
|
'opc_type': VariantType.Float,
|
||||||
|
},
|
||||||
|
'double': {
|
||||||
|
'converter': float,
|
||||||
|
'opc_type': VariantType.Double,
|
||||||
|
},
|
||||||
|
'int': {
|
||||||
|
'converter': int,
|
||||||
|
'opc_type': VariantType.Int32,
|
||||||
|
},
|
||||||
|
'bool': {
|
||||||
|
'converter': bool,
|
||||||
|
'opc_type': VariantType.Boolean,
|
||||||
|
},
|
||||||
|
'str': {
|
||||||
|
'converter': str,
|
||||||
|
'opc_type': VariantType.String,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class OpcRepository(SientiaMonitoring):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
opc_id: str,
|
||||||
|
url: str,
|
||||||
|
server_name: str,
|
||||||
|
logger: Logger,
|
||||||
|
notification_handler: NotificationHandler,
|
||||||
|
metrics_controller: MetricsController,
|
||||||
|
reconnection_interval: int = 60,
|
||||||
|
server_uri: str | None = None,
|
||||||
|
cert_path: str | None = None,
|
||||||
|
private_key_path: str | None = None,
|
||||||
|
server_cert_path: str | None = None,
|
||||||
|
):
|
||||||
|
self.url = url
|
||||||
|
self.id = opc_id
|
||||||
|
self.server_name = server_name
|
||||||
|
self.server_uri = server_uri
|
||||||
|
self.cert_path = cert_path
|
||||||
|
self.private_key_path = private_key_path
|
||||||
|
self.server_cert_path = server_cert_path
|
||||||
|
self.reconnection_interval = reconnection_interval
|
||||||
|
self.last_reconnection_time: None | datetime = None
|
||||||
|
self.disconnection_interval = 10.0
|
||||||
|
self.notification_handler = notification_handler
|
||||||
|
self.client: None | Client = None
|
||||||
|
|
||||||
|
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||||
|
|
||||||
|
self.metadata = {
|
||||||
|
'model_name': '-',
|
||||||
|
'model_id': '-',
|
||||||
|
'workflow_name': 'opc_repository',
|
||||||
|
'schedule_name': '-',
|
||||||
|
}
|
||||||
|
self._last_write_mono: float | None = None
|
||||||
|
self._connection_lock = asyncio.Lock()
|
||||||
|
self._session_ready = asyncio.Event()
|
||||||
|
self._reconnect_task: asyncio.Task[None] | None = None
|
||||||
|
self._allow_reconnect = True
|
||||||
|
|
||||||
|
def _opc_debug_tags(self, session_id: str) -> dict[str, str]:
|
||||||
|
"""
|
||||||
|
Build Prometheus/log label tags for OPC session-scoped metrics.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id (str): OPC UA session token string.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict[str, str]: Labels pod_id, server_name, runtime, opc_server_id, session_id.
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
'pod_id': str(getattr(self, 'pod_id', 'unknown')),
|
||||||
|
'server_name': self.server_name,
|
||||||
|
'runtime': str(getattr(self, 'runtime', 'unknown')),
|
||||||
|
'opc_server_id': self.id,
|
||||||
|
'session_id': session_id,
|
||||||
|
}
|
||||||
|
|
||||||
|
def _is_session_open(self) -> bool:
|
||||||
|
"""
|
||||||
|
Return whether the asyncua client has an open transport session.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
bool: True when protocol exists and is not closed.
|
||||||
|
"""
|
||||||
|
if self.client is None:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
proto = self.client.uaclient.protocol
|
||||||
|
return proto is not None and proto.state != 'closed'
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _reconnection_window_elapsed(self) -> bool:
|
||||||
|
"""
|
||||||
|
Return whether enough time has passed since the last reconnect attempt.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
bool: True if a new reconnect is allowed.
|
||||||
|
"""
|
||||||
|
if self.last_reconnection_time is None:
|
||||||
|
return True
|
||||||
|
return (
|
||||||
|
datetime.now() - self.last_reconnection_time
|
||||||
|
).total_seconds() > self.reconnection_interval
|
||||||
|
|
||||||
|
def _not_connected_error(self) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Build the standard error payload when validate_connection finds no open protocol.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict[str, Any]: Notification fields for OPC_CONNECTION_NOT_READY.
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
'notification_id': f'OPC_CONNECTION_NOT_READY_{self.id}',
|
||||||
|
'message': f'OPC server {self.id} is not connected',
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': NotificationLevel.WARNING,
|
||||||
|
}
|
||||||
|
|
||||||
|
async def set_security(self) -> None:
|
||||||
|
"""
|
||||||
|
Configure certificates and timeouts on the asyncua client.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If cert paths or client are missing.
|
||||||
|
"""
|
||||||
|
if self.cert_path is None or self.private_key_path is None:
|
||||||
|
raise ValueError(
|
||||||
|
'Certificate and private key paths must be provided for secure connection.'
|
||||||
|
)
|
||||||
|
|
||||||
|
cert = Path(self.cert_path)
|
||||||
|
private_key = Path(self.private_key_path)
|
||||||
|
server_cert = Path(self.server_cert_path) if self.server_cert_path else None
|
||||||
|
|
||||||
|
if self.client is None:
|
||||||
|
raise ValueError('Client must be initialized before setting security')
|
||||||
|
|
||||||
|
self.client.application_uri = self.server_uri
|
||||||
|
self.info('Setting security...', self.metadata)
|
||||||
|
await self.client.set_security(
|
||||||
|
SecurityPolicyBasic256,
|
||||||
|
certificate=str(cert),
|
||||||
|
private_key=str(private_key),
|
||||||
|
server_certificate=str(server_cert) if server_cert else None,
|
||||||
|
)
|
||||||
|
self.client.secure_channel_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
|
||||||
|
self.client.session_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
|
||||||
|
|
||||||
|
async def _create_client(self) -> None:
|
||||||
|
"""
|
||||||
|
Instantiate the asyncua Client and apply security when configured.
|
||||||
|
|
||||||
|
Caller must hold _connection_lock. Does not open a UA session.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
OpcClientAlreadyExistsError: If self.client is already set.
|
||||||
|
"""
|
||||||
|
if self.client is not None:
|
||||||
|
raise OpcClientAlreadyExistsError(
|
||||||
|
f'OPC client already exists for server {self.id}; '
|
||||||
|
'call disconnect() before creating a new client'
|
||||||
|
)
|
||||||
|
|
||||||
|
self.client = Client(self.url, timeout=10, watchdog_intervall=50) # type: ignore[attr-defined]
|
||||||
|
self.client.name = self.pod_id
|
||||||
|
self.client.application_name = self.pod_id
|
||||||
|
pod_uri = self.pod_id.replace('-', ':')
|
||||||
|
self.client.application_uri = pod_uri
|
||||||
|
self.client.product_uri = pod_uri
|
||||||
|
if self.cert_path:
|
||||||
|
await self.set_security()
|
||||||
|
|
||||||
|
async def _open_session(self) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Open the OPC UA session on the existing client.
|
||||||
|
|
||||||
|
Caller must hold _connection_lock.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
OpcClientNotInitializedError: If self.client is None.
|
||||||
|
OpcSessionAlreadyConnectedError: If a session is already open.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: Success flag and error payload on connect failure.
|
||||||
|
"""
|
||||||
|
if self.client is None:
|
||||||
|
raise OpcClientNotInitializedError(
|
||||||
|
f'OPC client is not initialized for server {self.id}; '
|
||||||
|
'call _create_client() before opening a session'
|
||||||
|
)
|
||||||
|
if self._is_session_open():
|
||||||
|
raise OpcSessionAlreadyConnectedError(
|
||||||
|
f'OPC session already connected for server {self.id}; '
|
||||||
|
'call disconnect() before connecting again'
|
||||||
|
)
|
||||||
|
|
||||||
|
tags = {
|
||||||
|
'pod_id': self.pod_id,
|
||||||
|
'server_name': self.server_name,
|
||||||
|
}
|
||||||
|
await self.emit_metric(metrics.OPC_CONNECTIONS_TOTAL, tags)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await self.client.connect()
|
||||||
|
|
||||||
|
session_id = _opc_authentication_token_str(self.client)
|
||||||
|
revised_session_timeout_ms = int(self.client.session_timeout)
|
||||||
|
revised_secure_channel_timeout_ms = int(self.client.secure_channel_timeout)
|
||||||
|
self.info(
|
||||||
|
f'OPC new session connected opc_server_id={self.id} session_id={session_id} '
|
||||||
|
f'revised_session_timeout_ms={revised_session_timeout_ms} '
|
||||||
|
f'revised_secure_channel_timeout_ms={revised_secure_channel_timeout_ms}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
await self.emit_metric(
|
||||||
|
metrics.OPC_SESSION_CREATED_TOTAL, self._opc_debug_tags(session_id)
|
||||||
|
)
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.OPC_SESSION_REVISED_TIMEOUT_MS,
|
||||||
|
method='set',
|
||||||
|
tags=self._opc_debug_tags(session_id),
|
||||||
|
value=revised_session_timeout_ms,
|
||||||
|
)
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.OPC_CONNECTION_STATUS,
|
||||||
|
method='set',
|
||||||
|
tags={**tags, 'server_url': self.url},
|
||||||
|
value=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
self._last_write_mono = None
|
||||||
|
self._session_ready.set()
|
||||||
|
return True, {}
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
await self._disconnect_locked()
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.error(trace, self.metadata)
|
||||||
|
await self.emit_metric(metrics.OPC_CONNECTIONS_FAILED, tags)
|
||||||
|
return False, {
|
||||||
|
'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
|
||||||
|
'message': f'Failed to connect to OPC server: {e}',
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': NotificationLevel.ERROR,
|
||||||
|
'attachment_content': trace,
|
||||||
|
}
|
||||||
|
|
||||||
|
async def _connect_locked(self) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Create the client when absent, then open a UA session.
|
||||||
|
|
||||||
|
Caller must hold _connection_lock.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
OpcSessionAlreadyConnectedError: If a session is already open.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: Result from _open_session on connect failure.
|
||||||
|
"""
|
||||||
|
if self._is_session_open():
|
||||||
|
raise OpcSessionAlreadyConnectedError(
|
||||||
|
f'OPC session already connected for server {self.id}; '
|
||||||
|
'call disconnect() before connecting again'
|
||||||
|
)
|
||||||
|
if self.client is None:
|
||||||
|
await self._create_client()
|
||||||
|
return await self._open_session()
|
||||||
|
|
||||||
|
async def _disconnection_fallback(self) -> list[dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Try up to five times to disconnect from the OPC UA server.
|
||||||
|
"""
|
||||||
|
assert self.client is not None
|
||||||
|
error_stack: list[dict[str, Any]] = []
|
||||||
|
for i in range(5):
|
||||||
|
try:
|
||||||
|
self.info(
|
||||||
|
f'Disconnecting from OPC UA server, attempt {i + 1} of 5',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
await self.client.disconnect()
|
||||||
|
return []
|
||||||
|
except Exception as e:
|
||||||
|
self.error(
|
||||||
|
f'Failed to disconnect from OPC UA server in attempt {i + 1} of 5: {e}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
error_stack.append(
|
||||||
|
{
|
||||||
|
'attempt': i + 1,
|
||||||
|
'error': str(e),
|
||||||
|
'traceback': traceback.format_exc(),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
await asyncio.sleep(self.disconnection_interval * i)
|
||||||
|
return error_stack
|
||||||
|
|
||||||
|
async def _disconnect_locked(self) -> None:
|
||||||
|
"""
|
||||||
|
Tear down the current session and client.
|
||||||
|
|
||||||
|
Caller must hold _connection_lock.
|
||||||
|
"""
|
||||||
|
self._last_write_mono = None
|
||||||
|
self._session_ready.clear()
|
||||||
|
|
||||||
|
if self.client is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
session_id = _opc_authentication_token_str(self.client)
|
||||||
|
self.info(
|
||||||
|
f'OPC disconnecting opc_server_id={self.id} session_id={session_id}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
await self.emit_metric(metrics.OPC_SESSION_CLOSED_TOTAL, self._opc_debug_tags(session_id))
|
||||||
|
|
||||||
|
errors = await self._disconnection_fallback()
|
||||||
|
if errors:
|
||||||
|
await self.send_notification_async(
|
||||||
|
metadata=self.metadata,
|
||||||
|
notification_id=f'OPC_DISCONNECTION_ERROR_{self.id}',
|
||||||
|
message='Failed to disconnect from OPC server in 5 attempts.',
|
||||||
|
block='opc_repository',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=json.dumps(errors, indent=4),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.warning(f'Disconnected from OPC server {self.id} successfully', self.metadata)
|
||||||
|
|
||||||
|
await self.emit_metric(
|
||||||
|
metric_object=metrics.OPC_CONNECTION_STATUS,
|
||||||
|
method='set',
|
||||||
|
tags={
|
||||||
|
'pod_id': self.pod_id,
|
||||||
|
'server_name': self.server_name,
|
||||||
|
'server_url': self.url,
|
||||||
|
},
|
||||||
|
value=0,
|
||||||
|
)
|
||||||
|
self.client = None
|
||||||
|
|
||||||
|
async def _reconnect_locked(self) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Close the current session and open a new one.
|
||||||
|
|
||||||
|
Caller must hold _connection_lock. Records last_reconnection_time for interval gating.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: Result from _connect_locked after teardown.
|
||||||
|
"""
|
||||||
|
self.last_reconnection_time = datetime.now()
|
||||||
|
await self._disconnect_locked()
|
||||||
|
return await self._connect_locked()
|
||||||
|
|
||||||
|
async def connect(self) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Open an OPC UA session under the connection lock (worker initialization).
|
||||||
|
"""
|
||||||
|
async with self._connection_lock:
|
||||||
|
self.info(
|
||||||
|
f'Starting connection to OPC server {self.id}:{self.server_name}...',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
return await self._connect_locked()
|
||||||
|
|
||||||
|
async def disconnect(self) -> None:
|
||||||
|
"""
|
||||||
|
Gracefully disconnect from the OPC server under the connection lock.
|
||||||
|
|
||||||
|
Disables background reconnect so late writes during worker shutdown do not
|
||||||
|
respawn sessions.
|
||||||
|
"""
|
||||||
|
async with self._connection_lock:
|
||||||
|
self._allow_reconnect = False
|
||||||
|
await self._disconnect_locked()
|
||||||
|
|
||||||
|
async def validate_connection(self) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Read-only check that the asyncua protocol is open.
|
||||||
|
|
||||||
|
Caller must ensure _session_ready before writing. Does not connect or reconnect.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: (True, {}) when open, otherwise (False, error).
|
||||||
|
"""
|
||||||
|
if self._is_session_open():
|
||||||
|
return True, {}
|
||||||
|
self.error(f'OPC server {self.id} is not connected', self.metadata)
|
||||||
|
return False, self._not_connected_error()
|
||||||
|
|
||||||
|
def _reconnect_task_in_progress(self) -> bool:
|
||||||
|
"""
|
||||||
|
Return whether a background reconnect task is currently running.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
bool: True when a reconnect task exists and has not finished.
|
||||||
|
"""
|
||||||
|
return self._reconnect_task is not None and not self._reconnect_task.done()
|
||||||
|
|
||||||
|
async def _start_reconnect(self, reason: str, session_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Schedule a background reconnect when allowed by interval and task state.
|
||||||
|
|
||||||
|
Clears _session_ready before starting the task. No-op when _allow_reconnect is
|
||||||
|
False, the reconnection window has not elapsed, or a reconnect is already running.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
reason (str): Trigger for reconnect (OPC status name or synthetic reason).
|
||||||
|
session_id (str): Session token before failure.
|
||||||
|
"""
|
||||||
|
if not self._allow_reconnect:
|
||||||
|
return
|
||||||
|
if not self._reconnection_window_elapsed():
|
||||||
|
self.warning(
|
||||||
|
f'OPC reconnect skipped reason=reconnection_window opc_server_id={self.id} '
|
||||||
|
f'reconnect_reason={reason}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
if self._reconnect_task_in_progress():
|
||||||
|
self.warning(
|
||||||
|
f'OPC reconnect skipped reason=in_progress opc_server_id={self.id} '
|
||||||
|
f'reconnect_reason={reason}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
self._session_ready.clear()
|
||||||
|
self.info(
|
||||||
|
f'OPC reconnect scheduled reconnect_reason={reason} opc_server_id={self.id} '
|
||||||
|
f'old_session_id={session_id}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
self._reconnect_task = asyncio.create_task(self._run_reconnect(reason, session_id))
|
||||||
|
|
||||||
|
async def _run_reconnect(self, reason: str, session_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Background task that tears down and re-establishes the OPC UA session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
reason (str): Trigger for reconnect (OPC status or ProtocolClosed).
|
||||||
|
session_id (str): Previous session token string for logging.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
async with self._connection_lock:
|
||||||
|
self.info(
|
||||||
|
f'OPC reconnect started reconnect_reason={reason} opc_server_id={self.id} '
|
||||||
|
f'old_session_id={session_id}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
success, error = await self._reconnect_locked()
|
||||||
|
if not success:
|
||||||
|
self.error(
|
||||||
|
f'OPC reconnect failed reconnect_reason={reason} opc_server_id={self.id}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
if error:
|
||||||
|
self.error(error.get('message', ''), self.metadata)
|
||||||
|
except Exception:
|
||||||
|
self.error(
|
||||||
|
f'OPC reconnect task failed opc_server_id={self.id} reconnect_reason={reason}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
self.error(traceback.format_exc(), self.metadata)
|
||||||
|
|
||||||
|
async def _log_write_inter_arrival(self, session_id: str, node: str) -> None:
|
||||||
|
"""
|
||||||
|
Log elapsed wall time since the previous successful OPC write on this repository.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id (str): Current OPC UA session token string.
|
||||||
|
node (str): Node id written in this operation.
|
||||||
|
"""
|
||||||
|
now = time.monotonic()
|
||||||
|
if self._last_write_mono is not None:
|
||||||
|
delta_s = now - self._last_write_mono
|
||||||
|
self.info(
|
||||||
|
f'OPC write inter-arrival_s={delta_s:.6f} opc_server_id={self.id} '
|
||||||
|
f'session_id={session_id} node={node}',
|
||||||
|
self.metadata,
|
||||||
|
)
|
||||||
|
if self.client is not None:
|
||||||
|
session_timeout_ms = float(self.client.session_timeout)
|
||||||
|
if session_timeout_ms > 0 and delta_s > (session_timeout_ms / 1000.0):
|
||||||
|
await self.emit_metric(
|
||||||
|
metrics.OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL,
|
||||||
|
self._opc_debug_tags(session_id),
|
||||||
|
)
|
||||||
|
self._last_write_mono = now
|
||||||
|
|
||||||
|
async def _emit_opc_write_metric(
|
||||||
|
self, session_id: str, result: str, metadata: dict[str, Any] | None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Emit opc_write_attempts_total for a single write attempt outcome.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id (str): OPC UA session token string, or "unknown".
|
||||||
|
result (str): Outcome label (OK, OPC status name, ProtocolClosed, etc.).
|
||||||
|
metadata (dict[str, Any] | None): Write context for model_id/model_name labels.
|
||||||
|
"""
|
||||||
|
await self.emit_metric(
|
||||||
|
metrics.OPC_WRITE_ATTEMPTS_TOTAL,
|
||||||
|
{
|
||||||
|
**self._opc_debug_tags(session_id),
|
||||||
|
**_model_labels_from_write_metadata(metadata),
|
||||||
|
'result': result,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
def _write_failure_payload(
|
||||||
|
self,
|
||||||
|
notification_id: str,
|
||||||
|
message: str,
|
||||||
|
level: NotificationLevel = NotificationLevel.ERROR,
|
||||||
|
attachment_content: str | None = None,
|
||||||
|
opc_error_kind: str | None = None,
|
||||||
|
opc_status: str | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Build a structured error dict returned from failed write_data paths.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
notification_id (str): Stable notification identifier.
|
||||||
|
message (str): Human-readable failure message.
|
||||||
|
level (NotificationLevel): Severity for downstream notifications.
|
||||||
|
attachment_content (str | None): Optional traceback or diagnostic text.
|
||||||
|
opc_error_kind (str | None): Classifier (session_bad, connection_lost, etc.).
|
||||||
|
opc_status (str | None): OPC UA status name or synthetic reason.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
dict[str, Any]: Error payload consumed by the OPC activity layer.
|
||||||
|
"""
|
||||||
|
payload: dict[str, Any] = {
|
||||||
|
'notification_id': notification_id,
|
||||||
|
'message': message,
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': level,
|
||||||
|
}
|
||||||
|
if attachment_content is not None:
|
||||||
|
payload['attachment_content'] = attachment_content
|
||||||
|
if opc_error_kind is not None:
|
||||||
|
payload['opc_error_kind'] = opc_error_kind
|
||||||
|
if opc_status is not None:
|
||||||
|
payload['opc_status'] = opc_status
|
||||||
|
return payload
|
||||||
|
|
||||||
|
async def _handle_tier1_bad(
|
||||||
|
self,
|
||||||
|
exc: BaseException,
|
||||||
|
session_id: str,
|
||||||
|
node: str,
|
||||||
|
metadata: dict[str, Any],
|
||||||
|
phase: str,
|
||||||
|
) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Record metrics/logs and schedule reconnect after a Tier-1 Bad* error.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
exc (BaseException): Tier-1 OPC UA error.
|
||||||
|
session_id (str): Session token at failure time.
|
||||||
|
node (str): Node id being written.
|
||||||
|
metadata (dict[str, Any]): Write context.
|
||||||
|
phase (str): get_node or write_value.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: Always (False, error payload).
|
||||||
|
"""
|
||||||
|
opc_status = _opc_status_from_exception(exc)
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.error(trace, metadata)
|
||||||
|
await self._emit_opc_write_metric(session_id, opc_status, metadata)
|
||||||
|
self.error(
|
||||||
|
f'OPC write failed opc_status={opc_status} opc_server_id={self.id} '
|
||||||
|
f'session_id={session_id} model_id={metadata.get("model_id", "unknown")} '
|
||||||
|
f'model_name={metadata.get("model_name", "unknown")} node={node} phase={phase}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
await self._start_reconnect(opc_status, session_id)
|
||||||
|
return False, self._write_failure_payload(
|
||||||
|
notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}',
|
||||||
|
message=f'Failed to {phase} on OPC server: {exc} | metadata: {metadata}',
|
||||||
|
attachment_content=trace,
|
||||||
|
opc_error_kind='session_bad',
|
||||||
|
opc_status=opc_status,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _write_reconnect_in_progress(
|
||||||
|
self, metadata: dict[str, Any]
|
||||||
|
) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Fail a write because a background reconnect task is already running.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
metadata (dict[str, Any]): Write context passed through to the activity.
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind reconnect_in_progress).
|
||||||
|
"""
|
||||||
|
await self._emit_opc_write_metric('unknown', 'ReconnectInProgress', metadata)
|
||||||
|
self.warning(
|
||||||
|
f'OPC write rejected reconnect_in_progress opc_server_id={self.id} '
|
||||||
|
f'model_id={metadata.get("model_id", "unknown")} '
|
||||||
|
f'model_name={metadata.get("model_name", "unknown")}',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
return False, {
|
||||||
|
'notification_id': f'OPC_WRITE_RECONNECT_IN_PROGRESS_{self.id}',
|
||||||
|
'message': f'OPC write skipped: reconnect in progress | metadata: {metadata}',
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': NotificationLevel.WARNING,
|
||||||
|
'opc_error_kind': 'reconnect_in_progress',
|
||||||
|
}
|
||||||
|
|
||||||
|
async def _write_connection_lost(
|
||||||
|
self, metadata: dict[str, Any], opc_status: str
|
||||||
|
) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Fail a write after scheduling reconnect for a closed or stale session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
metadata (dict[str, Any]): Write context passed through to the activity.
|
||||||
|
opc_status (str): Synthetic reason (ProtocolClosed, SessionNotReady).
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind connection_lost).
|
||||||
|
"""
|
||||||
|
await self._emit_opc_write_metric('unknown', opc_status, metadata)
|
||||||
|
return False, self._write_failure_payload(
|
||||||
|
notification_id=f'OPC_WRITE_CONNECTION_LOST_{self.id}',
|
||||||
|
message=f'OPC write skipped: connection lost ({opc_status}) | metadata: {metadata}',
|
||||||
|
level=NotificationLevel.WARNING,
|
||||||
|
opc_error_kind='connection_lost',
|
||||||
|
opc_status=opc_status,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def write_data(
|
||||||
|
self, node: str, value: Any, data_type: str, metadata: dict[str, Any]
|
||||||
|
) -> tuple[bool, dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Write data to OPC server with a single attempt and background reconnect scheduling.
|
||||||
|
|
||||||
|
Reconnect is scheduled on Tier-1 Bad*, closed protocol, or stale session readiness.
|
||||||
|
There is no retry within the same call.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
node (str): OPC UA node id to write.
|
||||||
|
value (Any): Value to convert and send.
|
||||||
|
data_type (str): Logical type key (float, int, bool, str, double).
|
||||||
|
metadata (dict[str, Any]): Activity context (model_id, model_name, etc.).
|
||||||
|
|
||||||
|
Return:
|
||||||
|
tuple[bool, dict[str, Any]]: (True, {response_time}) on success, or
|
||||||
|
(False, structured error info) on failure.
|
||||||
|
"""
|
||||||
|
if self._reconnect_task_in_progress():
|
||||||
|
return await self._write_reconnect_in_progress(metadata)
|
||||||
|
|
||||||
|
if not self._session_ready.is_set():
|
||||||
|
session_id = _opc_authentication_token_str(self.client)
|
||||||
|
await self._start_reconnect('SessionNotReady', session_id)
|
||||||
|
if self._reconnect_task_in_progress():
|
||||||
|
return await self._write_reconnect_in_progress(metadata)
|
||||||
|
return await self._write_connection_lost(metadata, 'SessionNotReady')
|
||||||
|
|
||||||
|
is_connected, _error = await self.validate_connection()
|
||||||
|
if not is_connected:
|
||||||
|
session_id = _opc_authentication_token_str(self.client)
|
||||||
|
await self._start_reconnect('ProtocolClosed', session_id)
|
||||||
|
return await self._write_connection_lost(metadata, 'ProtocolClosed')
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
session_id = _opc_authentication_token_str(self.client)
|
||||||
|
|
||||||
|
try:
|
||||||
|
node_obj = self.client.get_node(node) # type: ignore[union-attr]
|
||||||
|
except Exception as e:
|
||||||
|
if is_reconnectable_opcua_bad(e):
|
||||||
|
return await self._handle_tier1_bad(e, session_id, node, metadata, 'get_node')
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.error(trace, metadata)
|
||||||
|
await self._emit_opc_write_metric(
|
||||||
|
session_id, f'GetNodeError:{type(e).__name__}', metadata
|
||||||
|
)
|
||||||
|
return False, self._write_failure_payload(
|
||||||
|
notification_id=f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
|
||||||
|
message=f'Failed to get node from OPC server: {e} | metadata: {metadata}',
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
if data_type not in data_type_map:
|
||||||
|
await self._emit_opc_write_metric(session_id, 'UnsupportedDataType', metadata)
|
||||||
|
return False, self._write_failure_payload(
|
||||||
|
notification_id=f'OPC_WRITE_DATA_TYPE_ERROR_{self.id}',
|
||||||
|
message=f'Unsupported data type: {data_type} | metadata: {metadata}',
|
||||||
|
)
|
||||||
|
|
||||||
|
data = data_type_map[data_type]['converter'](value)
|
||||||
|
self.info(f'Writing {data} - {type(data)} to {node}', metadata)
|
||||||
|
ua_data = DataValue(
|
||||||
|
Variant(data, data_type_map[data_type]['opc_type']),
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await node_obj.write_value(ua_data)
|
||||||
|
end_time = time.time()
|
||||||
|
response_time = end_time - start_time
|
||||||
|
except Exception as e:
|
||||||
|
if is_reconnectable_opcua_bad(e):
|
||||||
|
return await self._handle_tier1_bad(e, session_id, node, metadata, 'write_value')
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.error(trace, metadata)
|
||||||
|
await self._emit_opc_write_metric(session_id, type(e).__name__, metadata)
|
||||||
|
return False, self._write_failure_payload(
|
||||||
|
notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}',
|
||||||
|
message=f'Failed to write data to OPC server: {e} | metadata: {metadata}',
|
||||||
|
attachment_content=trace,
|
||||||
|
)
|
||||||
|
|
||||||
|
await self._emit_opc_write_metric(session_id, 'OK', metadata)
|
||||||
|
await self._log_write_inter_arrival(session_id, node)
|
||||||
|
|
||||||
|
return True, {
|
||||||
|
'response_time': response_time,
|
||||||
|
}
|
||||||
0
laborious/worker/__init__.py
Normal file
0
laborious/worker/__init__.py
Normal file
289
laborious/worker/worker.py
Normal file
289
laborious/worker/worker.py
Normal file
@@ -0,0 +1,289 @@
|
|||||||
|
"""
|
||||||
|
Laborious Worker Module
|
||||||
|
|
||||||
|
This module provides the main worker implementation for the Sientia DataOps Laborious system.
|
||||||
|
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
|
||||||
|
prediction and retraining workflows.
|
||||||
|
|
||||||
|
The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``):
|
||||||
|
- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload)
|
||||||
|
- minimal_retrain-{runtime}-queue: Model retraining workflows
|
||||||
|
- drift-{runtime}-queue: Drift detection workflows
|
||||||
|
- simple_metrics-{runtime}-queue: Simple metrics workflows
|
||||||
|
|
||||||
|
``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the
|
||||||
|
same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``).
|
||||||
|
|
||||||
|
Key Features:
|
||||||
|
- Resource-based scaling with WorkerTuner (CPU and memory aware)
|
||||||
|
- Automatic polling scaling with PollerBehaviorAutoscaling
|
||||||
|
- Prometheus metrics integration
|
||||||
|
- Comprehensive error handling and logging
|
||||||
|
- Graceful shutdown with cleanup
|
||||||
|
- Multiple worker instances for different workflow types
|
||||||
|
|
||||||
|
Environment Variables:
|
||||||
|
- RUNTIME: Required non-empty string; suffix for all task queue names
|
||||||
|
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
|
||||||
|
- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious)
|
||||||
|
- POD_ID: Kubernetes pod identifier for metrics
|
||||||
|
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
|
||||||
|
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
|
||||||
|
- PROJECT_NAME: Project name for notifications (default: laborious)
|
||||||
|
"""
|
||||||
|
|
||||||
|
from temporalio import client, workflow
|
||||||
|
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
from prometheus_client import start_http_server
|
||||||
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||||
|
from sientia_do.observability.logger import get_logger
|
||||||
|
from sientia_do.temporal.worker.prepare_worker import prepare_worker
|
||||||
|
from sientia_do.utils.connectors_config import (
|
||||||
|
build_api_config,
|
||||||
|
build_mongodb_config,
|
||||||
|
build_postgres_config,
|
||||||
|
)
|
||||||
|
|
||||||
|
from laborious import metrics
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.utils.connectors_config import (
|
||||||
|
build_minio_config,
|
||||||
|
build_mlflow_config,
|
||||||
|
build_opc_config,
|
||||||
|
)
|
||||||
|
from laborious.workflows.drift import Drift
|
||||||
|
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||||
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||||
|
from laborious.workflows.simple_metrics import SimpleMetrics
|
||||||
|
from laborious.workflows.sub_workflows.format_and_export_prediction import (
|
||||||
|
FormatAndExportPrediction,
|
||||||
|
)
|
||||||
|
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||||
|
|
||||||
|
POD_ID = os.getenv('HOSTNAME')
|
||||||
|
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
|
||||||
|
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
"""
|
||||||
|
Main entry point for the Laborious worker application.
|
||||||
|
|
||||||
|
This function initializes and starts all components of the worker:
|
||||||
|
1. Sets up logging and metadata
|
||||||
|
2. Starts Prometheus metrics server
|
||||||
|
3. Initializes notification handler
|
||||||
|
4. Creates and configures activities
|
||||||
|
5. Initializes OPC connections
|
||||||
|
6. Starts Temporal client and workers
|
||||||
|
7. Manages worker lifecycle and graceful shutdown
|
||||||
|
|
||||||
|
The function runs indefinitely until interrupted or an error occurs.
|
||||||
|
On error, it performs cleanup and exits with a non-zero status code.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: Any unhandled exception during worker execution
|
||||||
|
SystemExit: On graceful shutdown or error conditions
|
||||||
|
"""
|
||||||
|
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
|
||||||
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'pod_id': POD_ID,
|
||||||
|
'model_name': '-',
|
||||||
|
'model_id': '-',
|
||||||
|
'workflow_name': '-',
|
||||||
|
'schedule_name': '-',
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
|
||||||
|
|
||||||
|
runtime = os.getenv('RUNTIME', '').strip()
|
||||||
|
if not runtime:
|
||||||
|
logger.custom_critical(
|
||||||
|
'RUNTIME environment variable is required and must be non-empty',
|
||||||
|
metadata,
|
||||||
|
)
|
||||||
|
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
metadata_runtime = {**metadata, 'runtime': runtime}
|
||||||
|
|
||||||
|
logger.custom_info('Starting prometheus client...', metadata_runtime)
|
||||||
|
start_prometheus_server()
|
||||||
|
|
||||||
|
logger.custom_info('Starting Notification Handler...', metadata_runtime)
|
||||||
|
|
||||||
|
mongo_config = build_mongodb_config()
|
||||||
|
notification_handler = NotificationHandler(
|
||||||
|
connection_string=mongo_config['connection_string'],
|
||||||
|
database=mongo_config['database_name'],
|
||||||
|
logger=logger,
|
||||||
|
project_name=os.getenv('PROJECT_NAME', 'laborious'),
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.custom_info('Starting Activities...', metadata_runtime)
|
||||||
|
|
||||||
|
activities = Activities(
|
||||||
|
postgres_config=build_postgres_config(),
|
||||||
|
mlflow_config=build_mlflow_config(),
|
||||||
|
minio_config=build_minio_config(),
|
||||||
|
opc_config=build_opc_config(),
|
||||||
|
pi_web_api_config=build_api_config(),
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.custom_info('Initializing OPC...', metadata_runtime)
|
||||||
|
await activities.init_opc()
|
||||||
|
|
||||||
|
logger.custom_info(
|
||||||
|
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...',
|
||||||
|
metadata_runtime,
|
||||||
|
)
|
||||||
|
|
||||||
|
new_runtime = Runtime(
|
||||||
|
telemetry=TelemetryConfig(
|
||||||
|
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.custom_info(f'Starting Temporal Client at {host}...', metadata_runtime)
|
||||||
|
|
||||||
|
temporal_client = await client.Client.connect(
|
||||||
|
target_host=host,
|
||||||
|
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
|
||||||
|
runtime=new_runtime,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.custom_info(f'Starting Workers (runtime={runtime})...', metadata_runtime)
|
||||||
|
|
||||||
|
workers = [
|
||||||
|
prepare_worker(
|
||||||
|
temporal_client=temporal_client,
|
||||||
|
main_workflow=MinimalRetrain,
|
||||||
|
other_workflows=[],
|
||||||
|
activities=[
|
||||||
|
activities.load_query_with_minio_offload,
|
||||||
|
activities.retrain_model,
|
||||||
|
activities.update_production_model,
|
||||||
|
activities.format_retrain_report,
|
||||||
|
activities.export_data_to_postgres,
|
||||||
|
],
|
||||||
|
logger=logger,
|
||||||
|
runtime=runtime,
|
||||||
|
),
|
||||||
|
prepare_worker(
|
||||||
|
temporal_client=temporal_client,
|
||||||
|
main_workflow=SimpleMetrics,
|
||||||
|
other_workflows=[],
|
||||||
|
activities=[
|
||||||
|
activities.load_custom_query,
|
||||||
|
activities.calculate_simple_metrics,
|
||||||
|
activities.export_data_to_postgres,
|
||||||
|
],
|
||||||
|
logger=logger,
|
||||||
|
runtime=runtime,
|
||||||
|
),
|
||||||
|
prepare_worker(
|
||||||
|
temporal_client=temporal_client,
|
||||||
|
main_workflow=Drift,
|
||||||
|
other_workflows=[],
|
||||||
|
activities=[
|
||||||
|
activities.load_custom_query,
|
||||||
|
activities.get_reference_data,
|
||||||
|
activities.calculate_drift,
|
||||||
|
activities.export_data_to_postgres,
|
||||||
|
],
|
||||||
|
logger=logger,
|
||||||
|
runtime=runtime,
|
||||||
|
),
|
||||||
|
prepare_worker(
|
||||||
|
temporal_client=temporal_client,
|
||||||
|
main_workflow=PredictionsBatch,
|
||||||
|
other_workflows=[PredictionProcess, FormatAndExportPrediction],
|
||||||
|
activities=[
|
||||||
|
# MLFlow
|
||||||
|
activities.request_predict,
|
||||||
|
activities.request_transform,
|
||||||
|
# Gates
|
||||||
|
activities.input_gate,
|
||||||
|
activities.mlflow_response_gate,
|
||||||
|
activities.mlflow_content_gate,
|
||||||
|
activities.format_transformed_data,
|
||||||
|
activities.format_prediction,
|
||||||
|
activities.format_default_prediction,
|
||||||
|
# OPC
|
||||||
|
activities.write_opc_data,
|
||||||
|
# Postgres / MinIO offload
|
||||||
|
activities.load_query_with_minio_offload,
|
||||||
|
activities.cleanup_minio_objects_expired,
|
||||||
|
activities.repeat_last_prediction,
|
||||||
|
activities.export_data_to_postgres,
|
||||||
|
activities.export_payload_to_postgres,
|
||||||
|
activities.write_metrics,
|
||||||
|
# Pi Web API
|
||||||
|
activities.write_pi_web_api_data,
|
||||||
|
],
|
||||||
|
logger=logger,
|
||||||
|
runtime=runtime,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
handlers = []
|
||||||
|
for w in workers:
|
||||||
|
handlers.append(w.run())
|
||||||
|
|
||||||
|
logger.custom_info('Workers started successfully', metadata_runtime)
|
||||||
|
|
||||||
|
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()
|
||||||
|
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
|
||||||
|
sys.exit(exit_code)
|
||||||
|
|
||||||
|
|
||||||
|
def start_prometheus_server():
|
||||||
|
"""
|
||||||
|
Starts the Prometheus metrics server for monitoring and observability.
|
||||||
|
|
||||||
|
This function initializes the Prometheus HTTP server on the configured port
|
||||||
|
and sets the application health metric to indicate the service is running.
|
||||||
|
|
||||||
|
The server exposes metrics that can be scraped by Prometheus for monitoring
|
||||||
|
the health and performance of the Laborious worker.
|
||||||
|
|
||||||
|
Environment Variables:
|
||||||
|
HTTP_METRICS_PORT: Port for the metrics server (default: 9090)
|
||||||
|
POD_ID: Pod identifier for metrics labeling
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SystemExit: If the metrics server fails to start
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
|
||||||
|
start_http_server(port)
|
||||||
|
print(f'Prometheus server started on port {port}.')
|
||||||
|
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
|
||||||
|
except Exception as e:
|
||||||
|
print(f'Failed to start Prometheus server: {e}')
|
||||||
|
os._exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
asyncio.run(main())
|
||||||
0
laborious/workflows/__init__.py
Normal file
0
laborious/workflows/__init__.py
Normal file
107
laborious/workflows/drift.py
Normal file
107
laborious/workflows/drift.py
Normal file
@@ -0,0 +1,107 @@
|
|||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from datetime import timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||||
|
from sientia_do.temporal.policies import retry_policy
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
|
||||||
|
|
||||||
|
@workflow.defn(name='drift')
|
||||||
|
class Drift:
|
||||||
|
@workflow.run
|
||||||
|
async def run(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Execute the drift workflow.
|
||||||
|
|
||||||
|
This method orchestrates the complete drift process by:
|
||||||
|
1. Loading data using the provided custom SQL query
|
||||||
|
2. Preparing prediction configuration and filters
|
||||||
|
3. Delegating to the PredictionProcess workflow for ML operations
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'schedule_name': input_data['schedule_name'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'workflow_name': 'drift',
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
print(f'Input data: {input_data}', metadata)
|
||||||
|
|
||||||
|
model_config = input_data['model_config']
|
||||||
|
target_name = model_config['target']
|
||||||
|
|
||||||
|
gathering_query = f"""
|
||||||
|
SELECT *
|
||||||
|
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
|
||||||
|
WHERE
|
||||||
|
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_activity_method(
|
||||||
|
Activities.load_custom_query,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'query': gathering_query,
|
||||||
|
'datetime_columns': ['timestamp', 'created_at'],
|
||||||
|
'orient': 'records',
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
|
|
||||||
|
reference_data_handler = workflow.start_activity_method(
|
||||||
|
Activities.get_reference_data,
|
||||||
|
{**metadata, 'model_name': input_data['model_name']},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
|
|
||||||
|
target_data = await target_data_handler
|
||||||
|
reference_data = await reference_data_handler
|
||||||
|
|
||||||
|
if not target_data:
|
||||||
|
return
|
||||||
|
|
||||||
|
drift_data = await workflow.execute_local_activity_method(
|
||||||
|
Activities.calculate_drift,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'target_data': target_data,
|
||||||
|
'reference_data': reference_data,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'target_name': target_name,
|
||||||
|
'drift_metrics': input_data.get(
|
||||||
|
'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
|
||||||
|
),
|
||||||
|
'chunk_period': input_data.get('chunk_period', 'min'),
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
|
|
||||||
|
if drift_data:
|
||||||
|
await workflow.execute_activity_method(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': drift_data,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['target_table_name'],
|
||||||
|
'timestamp_conversion': {
|
||||||
|
'column': 'timestamp',
|
||||||
|
'format': DATETIME_FORMAT_WITH_TZ,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
137
laborious/workflows/minimal_retrain.py
Normal file
137
laborious/workflows/minimal_retrain.py
Normal file
@@ -0,0 +1,137 @@
|
|||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from datetime import timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
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')
|
||||||
|
class MinimalRetrain:
|
||||||
|
"""
|
||||||
|
Automated model retraining workflow for the Laborious system.
|
||||||
|
|
||||||
|
This workflow implements a complete model retraining pipeline that loads
|
||||||
|
training data, executes model retraining, updates production models,
|
||||||
|
and maintains comprehensive audit trails. It's designed for automated
|
||||||
|
model lifecycle management with minimal manual intervention.
|
||||||
|
|
||||||
|
The workflow provides a robust retraining process with:
|
||||||
|
- Automated data loading from configured data sources
|
||||||
|
- MLFlow model retraining with quality validation
|
||||||
|
- Production model updates with version control
|
||||||
|
- Comprehensive reporting and audit trail maintenance
|
||||||
|
- Error handling and notification integration
|
||||||
|
"""
|
||||||
|
|
||||||
|
@workflow.run
|
||||||
|
async def run(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Execute the automated model retraining workflow.
|
||||||
|
|
||||||
|
This method orchestrates the complete model retraining process by:
|
||||||
|
1. Loading training data using the provided custom SQL query
|
||||||
|
2. Executing MLFlow model retraining with the loaded data
|
||||||
|
3. Updating production models with newly trained versions
|
||||||
|
4. Persisting comprehensive retraining reports to database
|
||||||
|
|
||||||
|
The method implements comprehensive error handling and ensures all
|
||||||
|
required parameters are properly configured before proceeding.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Complete configuration for the retraining workflow
|
||||||
|
Required keys:
|
||||||
|
- schedule_name (str): Schedule identifier for the retraining
|
||||||
|
- model_name (str): Name of the ML model to retrain
|
||||||
|
- model_id (int): Unique identifier for the model version
|
||||||
|
- query (str): SQL query for training data loading
|
||||||
|
- schema (str, optional): Database schema for report storage
|
||||||
|
- table_name (str, optional): Target table for retraining reports
|
||||||
|
- datetime_columns (list[str], optional): Columns to treat as datetime
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None: The workflow completes successfully when all steps finish
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If any required parameters are missing or if the workflow fails
|
||||||
|
during data loading, retraining, or model update operations
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'schedule_name': input_data['schedule_name'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'workflow_name': 'minimal_retrain',
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
model_config = input_data.get('model_config', {})
|
||||||
|
|
||||||
|
storage_result = 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': model_name,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=600),
|
||||||
|
)
|
||||||
|
|
||||||
|
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,
|
||||||
|
'data': storage_result,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(hours=1),
|
||||||
|
)
|
||||||
|
|
||||||
|
if experiment_response['success']:
|
||||||
|
update_report = await workflow.execute_activity_method(
|
||||||
|
Activities.update_production_model,
|
||||||
|
{**metadata, 'model_name': model_name, **experiment_response},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
update_report = {}
|
||||||
|
|
||||||
|
report = await workflow.execute_local_activity_method(
|
||||||
|
Activities.format_retrain_report,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'experiment_response': experiment_response,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'update_report': update_report,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
await workflow.execute_activity_method(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': report,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=600),
|
||||||
|
)
|
||||||
127
laborious/workflows/predictions_batch.py
Normal file
127
laborious/workflows/predictions_batch.py
Normal file
@@ -0,0 +1,127 @@
|
|||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from datetime import timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sientia_do.temporal.policies import retry_policy
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
|
||||||
|
|
||||||
|
@workflow.defn(name='predictions_batch')
|
||||||
|
class PredictionsBatch:
|
||||||
|
"""
|
||||||
|
Main batch prediction workflow for the Laborious system.
|
||||||
|
|
||||||
|
This workflow orchestrates the complete batch prediction process, handling
|
||||||
|
data loading, configuration management, and workflow delegation. It serves
|
||||||
|
as the primary entry point for batch prediction operations and ensures
|
||||||
|
proper data preparation before ML model inference.
|
||||||
|
|
||||||
|
The workflow implements a robust data processing pipeline with:
|
||||||
|
- Custom SQL query execution for data loading
|
||||||
|
- Comprehensive configuration management
|
||||||
|
- Data quality filter application
|
||||||
|
- MLFlow model integration
|
||||||
|
- Workflow delegation to specialized sub-workflows
|
||||||
|
|
||||||
|
Workflow Execution:
|
||||||
|
1. Data Loading: Executes custom SQL query to load prediction data
|
||||||
|
2. Configuration Preparation: Sets up prediction parameters and filters
|
||||||
|
3. Workflow Delegation: Spawns PredictionProcess child workflow
|
||||||
|
4. Error Handling: Implements comprehensive error handling and retry policies
|
||||||
|
"""
|
||||||
|
|
||||||
|
@workflow.run
|
||||||
|
async def run(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Execute the batch prediction workflow.
|
||||||
|
|
||||||
|
This method orchestrates the complete batch prediction process by:
|
||||||
|
1. Loading data using the provided custom SQL query
|
||||||
|
2. Preparing prediction configuration and filters
|
||||||
|
3. Delegating to the PredictionProcess workflow for ML operations
|
||||||
|
|
||||||
|
The method implements comprehensive error handling and ensures all
|
||||||
|
required parameters are properly configured before proceeding.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Complete configuration for the batch prediction
|
||||||
|
Required keys:
|
||||||
|
- schedule_name (str): Schedule identifier for the prediction
|
||||||
|
- model_name (str): Name of the ML model to use
|
||||||
|
- model_id (int): Unique identifier for the model
|
||||||
|
- query (str): SQL query for data loading
|
||||||
|
- schema (dict, optional): Data schema definition
|
||||||
|
- table_name (str, optional): Target table for predictions
|
||||||
|
- input_filters (dict, optional): Data quality filters
|
||||||
|
- mlflow_transform_filters (dict, optional): MLFlow transform filters
|
||||||
|
- mlflow_predict_filters (dict, optional): MLFlow prediction filters
|
||||||
|
- model_retention (int, optional): Model retention period in minutes
|
||||||
|
- path_priority (list[str]): Decision path priority configuration
|
||||||
|
- opc_output_config (dict, optional): OPC server export configuration
|
||||||
|
- pi_web_api_output_config (dict, optional): PI Web API export configuration
|
||||||
|
- datetime_columns (list[str], optional): Columns to treat as datetime
|
||||||
|
- save_transform (bool, optional): Whether to save transformed data (default: True)
|
||||||
|
- prediction_store_policy (str, optional): Data retention policy (default: 'lts:1')
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None: The workflow completes successfully when the child workflow finishes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If any required parameters are missing or if the workflow fails
|
||||||
|
during data loading or workflow delegation
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'schedule_name': input_data['schedule_name'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'workflow_name': 'predictions_batch',
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
# Load data using custom query 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),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Prepare input for prediction_process workflow
|
||||||
|
prediction_input = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'data': data,
|
||||||
|
'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', 'CONFIG': {}}}
|
||||||
|
),
|
||||||
|
'mlflow_transform_filters': input_data.get(
|
||||||
|
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
|
||||||
|
),
|
||||||
|
'mlflow_predict_filters': input_data.get(
|
||||||
|
'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),
|
||||||
|
}
|
||||||
|
|
||||||
|
# Execute prediction process workflow
|
||||||
|
await workflow.execute_child_workflow('subworkflow.prediction_process', prediction_input)
|
||||||
95
laborious/workflows/simple_metrics.py
Normal file
95
laborious/workflows/simple_metrics.py
Normal file
@@ -0,0 +1,95 @@
|
|||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from datetime import timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||||
|
from sientia_do.temporal.policies import retry_policy
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
|
||||||
|
|
||||||
|
@workflow.defn(name='simple_metrics')
|
||||||
|
class SimpleMetrics:
|
||||||
|
@workflow.run
|
||||||
|
async def run(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Execute the simple metrics workflow.
|
||||||
|
"""
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'workflow_name': 'simple_metrics',
|
||||||
|
'schedule_name': input_data['schedule_name'],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
model_id = input_data['model_id']
|
||||||
|
interval_minutes = input_data['interval_minutes']
|
||||||
|
|
||||||
|
model_config = input_data['model_config']
|
||||||
|
target_name = model_config['target']
|
||||||
|
|
||||||
|
query = f"""
|
||||||
|
select p."timestamp", p.prediction, ld.value as "target"
|
||||||
|
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
|
||||||
|
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
|
||||||
|
on p."timestamp" = ld."timestamp"
|
||||||
|
where
|
||||||
|
p.model_id = '{model_id}' and
|
||||||
|
p.prediction is not null and
|
||||||
|
ld.variable = '{target_name}' and
|
||||||
|
ld.value is not null and
|
||||||
|
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_activity_method(
|
||||||
|
Activities.load_custom_query,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'query': query,
|
||||||
|
'datetime_columns': ['timestamp'],
|
||||||
|
'orient': 'records',
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
|
|
||||||
|
if not target_data:
|
||||||
|
return
|
||||||
|
|
||||||
|
simple_metrics = await workflow.execute_local_activity_method(
|
||||||
|
Activities.calculate_simple_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'model_id': model_id,
|
||||||
|
'target_data': target_data,
|
||||||
|
'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
|
||||||
|
'interval_minutes': interval_minutes,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
|
|
||||||
|
if not simple_metrics:
|
||||||
|
return
|
||||||
|
|
||||||
|
await workflow.execute_activity_method(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': simple_metrics,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['target_table_name'],
|
||||||
|
'timestamp_conversion': {
|
||||||
|
'column': 'timestamp',
|
||||||
|
'format': DATETIME_FORMAT_WITH_TZ,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=300),
|
||||||
|
)
|
||||||
0
laborious/workflows/sub_workflows/__init__.py
Normal file
0
laborious/workflows/sub_workflows/__init__.py
Normal file
@@ -0,0 +1,220 @@
|
|||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from datetime import timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||||
|
from sientia_do.temporal.policies import retry_policy
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
|
||||||
|
|
||||||
|
@workflow.defn(name='subworkflow.format_and_export_prediction')
|
||||||
|
class FormatAndExportPrediction:
|
||||||
|
"""
|
||||||
|
Data formatting and export workflow for prediction results.
|
||||||
|
|
||||||
|
This workflow handles the final stages of the prediction pipeline, including
|
||||||
|
data formatting, database persistence, OPC server export, and metrics recording.
|
||||||
|
It implements flexible formatting based on prediction quality and provides
|
||||||
|
comprehensive export capabilities to multiple destinations.
|
||||||
|
|
||||||
|
The workflow supports two main prediction paths:
|
||||||
|
1. Normal Prediction: Formats and exports successful prediction results
|
||||||
|
2. Default Prediction: Creates fallback predictions for error conditions
|
||||||
|
|
||||||
|
Export Destinations:
|
||||||
|
- PostgreSQL Database: Persistent storage with timestamp conversion
|
||||||
|
- PI Web API: Real-time industrial system integration for prediction and confidence values
|
||||||
|
- OPC Servers: Real-time industrial system integration
|
||||||
|
- Prometheus Metrics: Performance monitoring and operational visibility
|
||||||
|
"""
|
||||||
|
|
||||||
|
@workflow.run
|
||||||
|
async def run(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Execute the prediction formatting and export workflow.
|
||||||
|
|
||||||
|
This method orchestrates the complete data export process by:
|
||||||
|
1. Determining the appropriate formatting strategy based on path_flag
|
||||||
|
2. Formatting prediction data according to quality and requirements
|
||||||
|
3. Exporting data to PI Web API for real-time industrial access (if configured)
|
||||||
|
4. Exporting data to OPC servers for real-time industrial access (if configured)
|
||||||
|
5. Persisting data to PostgreSQL database with comprehensive metadata
|
||||||
|
6. Recording performance metrics for operational monitoring
|
||||||
|
|
||||||
|
The method implements flexible formatting strategies:
|
||||||
|
- Normal predictions: Full data formatting with confidence scores
|
||||||
|
- Error predictions: Default formatting with error indicators
|
||||||
|
- Comprehensive export: Multi-destination data distribution
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Complete configuration for the export workflow
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- path_flag (str | None): Decision path flag for formatting strategy
|
||||||
|
- None: Normal prediction path with full formatting
|
||||||
|
- Any other value: Default prediction path for error conditions
|
||||||
|
- data (dict[str, Any]): Prediction data to format and export
|
||||||
|
- prediction_confidence (float): Confidence score for the prediction
|
||||||
|
- timestamp (str): ISO-formatted timestamp for the prediction
|
||||||
|
- model_id (int): Unique identifier for the ML model
|
||||||
|
- model_name (str): Name of the ML model
|
||||||
|
- schema (str): Database schema for data storage
|
||||||
|
- table_name (str): Target table for data persistence
|
||||||
|
Optional keys:
|
||||||
|
- opc_output_config (dict[str, Any]): OPC server export configuration
|
||||||
|
- pi_web_api_output_config (dict[str, Any]): PI Web API export configuration
|
||||||
|
Contains endpoint, prediction_tags, and confidence_tags mappings
|
||||||
|
- transformed_data (dict[str, Any]): Transformed data to export separately
|
||||||
|
Only processed when path_flag is None
|
||||||
|
- transform_table_name (str): Target table for transformed data export
|
||||||
|
Required if transformed_data is provided
|
||||||
|
- prediction_store_policy (str): Data retention policy (e.g., 'lts:1', 'erl:2')
|
||||||
|
Required when path_flag is None
|
||||||
|
- comment (str): Operational comment or error description
|
||||||
|
Required when path_flag is not None
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None: The workflow completes successfully when all export operations finish
|
||||||
|
|
||||||
|
Note:
|
||||||
|
When transformed_data is provided and path_flag is None, the workflow will:
|
||||||
|
1. Format the transformed data using format_transformed_data
|
||||||
|
2. Export it to a separate table (transform_table_name) asynchronously
|
||||||
|
3. Wait for both prediction and transformed data exports to complete
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
path_flag = input_data['path_flag']
|
||||||
|
data = input_data['data']
|
||||||
|
transformed_data = input_data.get('transformed_data', None)
|
||||||
|
prediction_confidence = input_data['prediction_confidence']
|
||||||
|
|
||||||
|
opc_output_config = input_data.get('opc_output_config', None)
|
||||||
|
pi_web_api_output_config = input_data.get('pi_web_api_output_config', None)
|
||||||
|
|
||||||
|
if path_flag is None:
|
||||||
|
# Normal prediction path: format prediction data with full metadata
|
||||||
|
prediction = await workflow.execute_local_activity_method(
|
||||||
|
Activities.format_prediction,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': data,
|
||||||
|
'timestamp': input_data['timestamp'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'prediction_confidence': prediction_confidence,
|
||||||
|
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Optionally format and export transformed data to separate table
|
||||||
|
if transformed_data is not None:
|
||||||
|
transformed = await workflow.execute_local_activity_method(
|
||||||
|
Activities.format_transformed_data,
|
||||||
|
{
|
||||||
|
**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_payload_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['transform_table_name'],
|
||||||
|
'data': transformed,
|
||||||
|
'timestamp_conversion': {
|
||||||
|
'column': 'timestamp',
|
||||||
|
'format': DATETIME_FORMAT_WITH_TZ,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
else:
|
||||||
|
write_transformed_handler = None
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Error path: create default prediction with error indicators
|
||||||
|
prediction = await workflow.execute_local_activity_method(
|
||||||
|
Activities.format_default_prediction,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'timestamp': input_data['timestamp'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'prediction_confidence': prediction_confidence,
|
||||||
|
'comment': input_data['comment'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
write_transformed_handler = None
|
||||||
|
|
||||||
|
opc_metrics = {}
|
||||||
|
|
||||||
|
# write to pi web api
|
||||||
|
if pi_web_api_output_config:
|
||||||
|
prediction = await workflow.execute_activity_method(
|
||||||
|
Activities.write_pi_web_api_data,
|
||||||
|
{
|
||||||
|
'pi_web_api_output_config': pi_web_api_output_config,
|
||||||
|
'data': prediction,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
# write to opc
|
||||||
|
if opc_output_config:
|
||||||
|
prediction, opc_metrics = await workflow.execute_activity_method(
|
||||||
|
Activities.write_opc_data,
|
||||||
|
{
|
||||||
|
'opc_output_config': opc_output_config,
|
||||||
|
'data': prediction,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
# write to postgres
|
||||||
|
await workflow.execute_activity_method(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': prediction,
|
||||||
|
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
|
||||||
|
'on_conflict': input_data.get('on_conflict', 'error'),
|
||||||
|
'unique_columns': ['model_id', 'timestamp'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=180),
|
||||||
|
)
|
||||||
|
|
||||||
|
if write_transformed_handler is not None:
|
||||||
|
await write_transformed_handler
|
||||||
|
|
||||||
|
await workflow.execute_activity_method(
|
||||||
|
Activities.write_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'prediction': prediction,
|
||||||
|
'opc_metrics': opc_metrics,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(seconds=60),
|
||||||
|
)
|
||||||
346
laborious/workflows/sub_workflows/prediction_process.py
Normal file
346
laborious/workflows/sub_workflows/prediction_process.py
Normal file
@@ -0,0 +1,346 @@
|
|||||||
|
from temporalio import workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from datetime import timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sientia_do.temporal.policies import retry_policy
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
|
||||||
|
|
||||||
|
@workflow.defn(name='subworkflow.prediction_process')
|
||||||
|
class PredictionProcess:
|
||||||
|
"""
|
||||||
|
Core prediction processing workflow for the Laborious system.
|
||||||
|
|
||||||
|
This workflow implements the complete ML model inference pipeline, handling
|
||||||
|
data quality validation, MLFlow model interactions, and prediction processing.
|
||||||
|
It serves as the central orchestrator for all prediction operations and ensures
|
||||||
|
data quality throughout the entire process.
|
||||||
|
|
||||||
|
The workflow implements a robust data processing pipeline with:
|
||||||
|
- Data quality validation using configurable filters
|
||||||
|
- MLFlow model transformation and prediction
|
||||||
|
- Response validation and quality assurance
|
||||||
|
- Flexible decision path handling
|
||||||
|
- Comprehensive error handling and retry policies
|
||||||
|
|
||||||
|
Workflow Execution:
|
||||||
|
1. Timestamp Retrieval: Gets last processed timestamp for incremental processing
|
||||||
|
2. Input Data Gate: Applies data quality filters
|
||||||
|
3. Path Decision: Determines processing path based on filter results
|
||||||
|
4. MLFlow Transform: Requests data transformation using MLFlow models
|
||||||
|
5. Response Validation: Filters transform responses for quality assurance
|
||||||
|
6. MLFlow Prediction: Executes prediction using transformed data
|
||||||
|
7. Content Validation: Filters prediction responses for final quality check
|
||||||
|
8. Export Delegation: Delegates to FormatAndExportPrediction workflow
|
||||||
|
"""
|
||||||
|
|
||||||
|
@workflow.run
|
||||||
|
async def run(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Execute the prediction process workflow.
|
||||||
|
|
||||||
|
This method orchestrates the complete prediction processing pipeline by:
|
||||||
|
1. Retrieving the last processed timestamp for incremental processing
|
||||||
|
2. Applying data quality filters to validate input data
|
||||||
|
3. Executing MLFlow model transformation and prediction
|
||||||
|
4. Validating all responses for quality assurance
|
||||||
|
5. Delegating to export workflow for data persistence
|
||||||
|
|
||||||
|
The method implements comprehensive error handling and ensures all
|
||||||
|
data quality requirements are met before proceeding with ML operations.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data: Complete configuration for the prediction process
|
||||||
|
Required keys:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- data (dict): Input data for prediction processing
|
||||||
|
- schema (dict): Data schema definition
|
||||||
|
- table_name (str): Target table for predictions
|
||||||
|
- model_id (str): ML model identifier
|
||||||
|
- model_name (str): ML model name
|
||||||
|
- input_filters (dict): Data quality filters
|
||||||
|
- mlflow_transform_filters (dict): MLFlow transform filters
|
||||||
|
- mlflow_predict_filters (dict): MLFlow prediction filters
|
||||||
|
- model_retention (int): Model retention period in minutes
|
||||||
|
- path_priority (list[str]): Decision path priority configuration
|
||||||
|
- opc_output_config (dict, optional): OPC server export configuration
|
||||||
|
- pi_web_api_output_config (dict, optional): PI Web API export configuration
|
||||||
|
- save_transform (bool, optional): Whether to save transformed data (default: True)
|
||||||
|
- prediction_store_policy (str, optional): Data retention policy (default: 'lts:1')
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None: The workflow completes successfully when export workflow finishes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Exception: If any required parameters are missing or if the workflow fails
|
||||||
|
during data processing, MLFlow operations, or workflow delegation
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
data = input_data['data']
|
||||||
|
model_id = input_data['model_id']
|
||||||
|
model_name = input_data['model_name']
|
||||||
|
model_config = input_data.get('model_config', {})
|
||||||
|
save_transform = input_data.get('save_transform', True)
|
||||||
|
|
||||||
|
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 = {
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['input_filters'],
|
||||||
|
'data': data,
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
}
|
||||||
|
|
||||||
|
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||||
|
Activities.input_gate,
|
||||||
|
gate_input,
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(minutes=1),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handle path decision based on filter results
|
||||||
|
if await self.path_flag_handler(
|
||||||
|
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||||
|
):
|
||||||
|
return
|
||||||
|
|
||||||
|
# Request MLFlow model transformation
|
||||||
|
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,
|
||||||
|
start_to_close_timeout=timedelta(minutes=5),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate MLFlow transform response
|
||||||
|
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': transformed_data,
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(minutes=1),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handle path decision based on transform validation
|
||||||
|
if await self.path_flag_handler(
|
||||||
|
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||||
|
):
|
||||||
|
return
|
||||||
|
|
||||||
|
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||||
|
Activities.mlflow_content_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': transformed_data,
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(minutes=1),
|
||||||
|
)
|
||||||
|
|
||||||
|
if await self.path_flag_handler(
|
||||||
|
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||||
|
):
|
||||||
|
return
|
||||||
|
|
||||||
|
predicted_data = await workflow.execute_activity_method(
|
||||||
|
Activities.request_predict,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': transformed_data,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(minutes=5),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate MLFlow prediction response
|
||||||
|
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['mlflow_predict_filters'],
|
||||||
|
'data': predicted_data,
|
||||||
|
'type': 'predict',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(minutes=1),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handle path decision based on prediction validation
|
||||||
|
if await self.path_flag_handler(
|
||||||
|
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||||
|
):
|
||||||
|
return
|
||||||
|
|
||||||
|
# Delegate to export workflow for data persistence
|
||||||
|
await workflow.execute_child_workflow(
|
||||||
|
'subworkflow.format_and_export_prediction',
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'on_conflict': input_data.get('on_conflict', 'error'),
|
||||||
|
'path_flag': path_flag,
|
||||||
|
'data': predicted_data,
|
||||||
|
'transformed_data': transformed_data if save_transform else None,
|
||||||
|
'prediction_confidence': confidence,
|
||||||
|
'timestamp': last_timestamp,
|
||||||
|
'model_id': model_id,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
'opc_output_config': input_data['opc_output_config'],
|
||||||
|
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'transform_table_name': input_data['transform_table_name'],
|
||||||
|
'comment': comment,
|
||||||
|
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
async def path_flag_handler(
|
||||||
|
self,
|
||||||
|
data: dict[str, Any],
|
||||||
|
path_flag: str,
|
||||||
|
input_data: dict,
|
||||||
|
confidence: int,
|
||||||
|
last_timestamp: str,
|
||||||
|
comment: str,
|
||||||
|
) -> bool:
|
||||||
|
"""
|
||||||
|
Handle path decisions based on filter results and confidence levels.
|
||||||
|
|
||||||
|
This method determines the appropriate action based on the path flag
|
||||||
|
returned by data quality filters. It can stop processing, continue,
|
||||||
|
or repeat operations based on the configured path priority.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data: Input data for processing
|
||||||
|
path_flag: Path decision from filter (STOP, CONTINUE, REPEAT)
|
||||||
|
input_data: Complete workflow input configuration including:
|
||||||
|
- metadata (dict): Workflow execution metadata
|
||||||
|
- schema (str): Database schema
|
||||||
|
- table_name (str): Target table for predictions
|
||||||
|
- transform_table_name (str): Target table for transformed data
|
||||||
|
- model_id (str): ML model identifier
|
||||||
|
- model_name (str): ML model name
|
||||||
|
- model_config (dict, optional): Model configuration
|
||||||
|
- opc_output_config (dict, optional): OPC server export configuration
|
||||||
|
- pi_web_api_output_config (dict, optional): PI Web API export configuration
|
||||||
|
- prediction_store_policy (str, optional): Data retention policy
|
||||||
|
confidence: Confidence level from filter validation
|
||||||
|
last_timestamp: Last processed timestamp
|
||||||
|
comment: Additional information about the filter result
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if processing should stop, False to continue
|
||||||
|
|
||||||
|
Path Handling:
|
||||||
|
- STOP: Terminates workflow execution
|
||||||
|
- CONTINUE: Delegates to FormatAndExportPrediction workflow with current data
|
||||||
|
- REPEAT: Repeats last prediction if available
|
||||||
|
"""
|
||||||
|
metadata = input_data['metadata']
|
||||||
|
|
||||||
|
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']
|
||||||
|
model_config = input_data.get('model_config', {})
|
||||||
|
|
||||||
|
path_flag = path_flag.upper() if path_flag else ''
|
||||||
|
|
||||||
|
if path_flag == 'STOP':
|
||||||
|
# Stop processing and exit workflow
|
||||||
|
return True
|
||||||
|
elif path_flag == 'REPEAT':
|
||||||
|
# Repeat last prediction if available
|
||||||
|
await workflow.execute_activity_method(
|
||||||
|
Activities.repeat_last_prediction,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'model': model_id,
|
||||||
|
'last_timestamp': last_timestamp,
|
||||||
|
},
|
||||||
|
retry_policy=retry_policy,
|
||||||
|
start_to_close_timeout=timedelta(minutes=1),
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
elif path_flag == 'CONTINUE':
|
||||||
|
# call write workflow
|
||||||
|
await workflow.execute_child_workflow(
|
||||||
|
'subworkflow.format_and_export_prediction',
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': path_flag,
|
||||||
|
'data': data,
|
||||||
|
'prediction_confidence': confidence,
|
||||||
|
'timestamp': last_timestamp,
|
||||||
|
'model_id': model_id,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'transform_table_name': transform_table_name,
|
||||||
|
'comment': comment,
|
||||||
|
'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
|
||||||
|
|
||||||
|
return False
|
||||||
106
model_convert.ipynb
Normal file
106
model_convert.ipynb
Normal file
@@ -0,0 +1,106 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 23,
|
||||||
|
"id": "e838ff21",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import csv\n",
|
||||||
|
"\n",
|
||||||
|
"def csv_to_tag_lists(csv_path: str) -> dict:\n",
|
||||||
|
" read_tags = []\n",
|
||||||
|
" write_tags = []\n",
|
||||||
|
"\n",
|
||||||
|
" def to_float(val):\n",
|
||||||
|
" try:\n",
|
||||||
|
" return float(str(val).strip())\n",
|
||||||
|
" except Exception:\n",
|
||||||
|
" return None\n",
|
||||||
|
"\n",
|
||||||
|
" with open(csv_path, newline=\"\", encoding=\"utf-8\") as f:\n",
|
||||||
|
" reader = csv.DictReader(f)\n",
|
||||||
|
" for row in reader:\n",
|
||||||
|
" # Basic normalization\n",
|
||||||
|
" op = (row.get(\"operation\") or \"\").strip()\n",
|
||||||
|
"\n",
|
||||||
|
" if op == \"READ\":\n",
|
||||||
|
" # Build common tag payload with required mappings\n",
|
||||||
|
" tag = {\n",
|
||||||
|
" \"server_id\": \"1\",\n",
|
||||||
|
" \"tag_address\": row.get(\"opc_tag\"),\n",
|
||||||
|
" \"tag_name\": row.get(\"name\"),\n",
|
||||||
|
" \"data_range\": [to_float(row.get(\"min_value\")), to_float(row.get(\"max_value\"))],\n",
|
||||||
|
" \"aggr_func\": row.get(\"aggregation_func\").lower(),\n",
|
||||||
|
" # keep other fields with their original names\n",
|
||||||
|
" \"frequency\": row.get(\"frequency\"),\n",
|
||||||
|
" \"local\": row.get(\"local\"),\n",
|
||||||
|
" \"area\": row.get(\"area\"),\n",
|
||||||
|
" \"description\": row.get(\"description\"),\n",
|
||||||
|
" }\n",
|
||||||
|
"\n",
|
||||||
|
" read_tags.append(tag)\n",
|
||||||
|
"\n",
|
||||||
|
" else:\n",
|
||||||
|
" tag = {\n",
|
||||||
|
" \"server_id\": \"1\",\n",
|
||||||
|
" \"addr\": row.get(\"opc_tag\"),\n",
|
||||||
|
" \"tag_name\": row.get(\"name\"),\n",
|
||||||
|
" \"local\": row.get(\"local\"),\n",
|
||||||
|
" \"area\": row.get(\"area\"),\n",
|
||||||
|
" \"description\": row.get(\"description\"),\n",
|
||||||
|
" }\n",
|
||||||
|
" \n",
|
||||||
|
" if op == \"WRITE_PREDICTION\":\n",
|
||||||
|
" tag[\"type\"] = \"prediction\"\n",
|
||||||
|
" write_tags.append(tag)\n",
|
||||||
|
" elif op == \"WRITE_CONFIDENCE\":\n",
|
||||||
|
" tag[\"type\"] = \"confidence\"\n",
|
||||||
|
" write_tags.append(tag)\n",
|
||||||
|
" # ignore any other operation values silently\n",
|
||||||
|
"\n",
|
||||||
|
" return {\"read_tags\": read_tags, \"write_tags\": write_tags}"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "4621cd43",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import json\n",
|
||||||
|
"\n",
|
||||||
|
"file_names = [\"Courier - Página1.csv\"]\n",
|
||||||
|
"\n",
|
||||||
|
"for file_name in file_names:\n",
|
||||||
|
" write_file = file_name.replace(\".csv\", \".json\")\n",
|
||||||
|
"\n",
|
||||||
|
" with open(write_file, \"w\", encoding=\"utf-8\") as f:\n",
|
||||||
|
" json.dump(csv_to_tag_lists(file_name), f, indent=2, ensure_ascii=False)"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "venv",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"codemirror_mode": {
|
||||||
|
"name": "ipython",
|
||||||
|
"version": 3
|
||||||
|
},
|
||||||
|
"file_extension": ".py",
|
||||||
|
"mimetype": "text/x-python",
|
||||||
|
"name": "python",
|
||||||
|
"nbconvert_exporter": "python",
|
||||||
|
"pygments_lexer": "ipython3",
|
||||||
|
"version": "3.11.13"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 5
|
||||||
|
}
|
||||||
169
pyproject.toml
Normal file
169
pyproject.toml
Normal file
@@ -0,0 +1,169 @@
|
|||||||
|
[build-system]
|
||||||
|
requires = ["setuptools>=61.0"]
|
||||||
|
build-backend = "setuptools.build_meta"
|
||||||
|
|
||||||
|
[project]
|
||||||
|
name = "laborious"
|
||||||
|
version = "0.0.0"
|
||||||
|
description = "Sientia DataOps Laborious - ML Model Orchestration System"
|
||||||
|
readme = "README.md"
|
||||||
|
requires-python = ">=3.11"
|
||||||
|
authors = [
|
||||||
|
{name = "Aignosi", email = "dev@aignosi.com"}
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.ruff]
|
||||||
|
line-length = 100
|
||||||
|
target-version = "py311"
|
||||||
|
exclude = [
|
||||||
|
".git",
|
||||||
|
".venv",
|
||||||
|
"venv",
|
||||||
|
"__pycache__",
|
||||||
|
"*.pyc",
|
||||||
|
".pytest_cache",
|
||||||
|
"htmlcov",
|
||||||
|
"tests/laborious/workflows/subworkflows/test_prediction_process.py",
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.ruff.lint]
|
||||||
|
select = [
|
||||||
|
"E", # pycodestyle errors
|
||||||
|
"W", # pycodestyle warnings
|
||||||
|
"F", # pyflakes
|
||||||
|
"I", # isort
|
||||||
|
"B", # flake8-bugbear
|
||||||
|
"C4", # flake8-comprehensions
|
||||||
|
"UP", # pyupgrade
|
||||||
|
"N", # pep8-naming
|
||||||
|
"YTT", # flake8-2020
|
||||||
|
"S", # flake8-bandit
|
||||||
|
"BLE", # flake8-blind-except
|
||||||
|
"A", # flake8-builtins
|
||||||
|
"C90", # mccabe complexity
|
||||||
|
]
|
||||||
|
|
||||||
|
ignore = [
|
||||||
|
"BLE001", # ignore blind except, we need to send notifications with any error
|
||||||
|
"E501", # line too long (handled by formatter)
|
||||||
|
"S101", # use of assert (needed for tests)
|
||||||
|
"S105", # possible hardcoded password (false positives)
|
||||||
|
"S106", # possible hardcoded password (false positives)
|
||||||
|
"S608", # potential sql injection (false positives)
|
||||||
|
"N802", # function name should be lowercase (temporal decorators)
|
||||||
|
"N806", # variable in function should be lowercase
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.ruff.lint.per-file-ignores]
|
||||||
|
"tests/**/*.py" = [
|
||||||
|
"S101", # assert allowed in tests
|
||||||
|
"S105", # hardcoded passwords ok in tests
|
||||||
|
"S106", # hardcoded passwords ok in tests
|
||||||
|
]
|
||||||
|
"e2e/conftest.py" = [
|
||||||
|
# Imports below the sientia/ModelAnalysis stub are load-order critical: the stub
|
||||||
|
# must be in sys.modules before Activities is imported, so they cannot move up.
|
||||||
|
"E402",
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.ruff.lint.mccabe]
|
||||||
|
max-complexity = 15
|
||||||
|
|
||||||
|
[tool.ruff.format]
|
||||||
|
quote-style = "single"
|
||||||
|
indent-style = "space"
|
||||||
|
line-ending = "auto"
|
||||||
|
|
||||||
|
[tool.mypy]
|
||||||
|
python_version = "3.11"
|
||||||
|
warn_return_any = false
|
||||||
|
warn_unused_configs = true
|
||||||
|
disallow_untyped_defs = false
|
||||||
|
disallow_incomplete_defs = false
|
||||||
|
check_untyped_defs = true
|
||||||
|
no_implicit_optional = true
|
||||||
|
warn_redundant_casts = true
|
||||||
|
warn_unused_ignores = false
|
||||||
|
warn_no_return = true
|
||||||
|
strict_equality = true
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
# Ignore missing imports for external packages
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "temporalio.*"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "sientia_do.*"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "mlflow.*"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "prometheus_client.*"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "sientia.*"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "pandas.*"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[tool.pytest.ini_options]
|
||||||
|
testpaths = ["tests"]
|
||||||
|
python_files = ["test_*.py"]
|
||||||
|
python_classes = ["Test*"]
|
||||||
|
python_functions = ["test_*"]
|
||||||
|
addopts = [
|
||||||
|
"-v",
|
||||||
|
"--strict-markers",
|
||||||
|
# pytest>=9.1 has a known bug where its unraisableexception plugin crashes
|
||||||
|
# (tracemalloc partially-initialized AttributeError) when 2+ unraisable
|
||||||
|
# exceptions land close together — e.g. "coroutine was never awaited" from
|
||||||
|
# AsyncMock-mocked sync methods (metrics_controller, minio_repository) being
|
||||||
|
# GC'd. Harmless mock artifacts turned into a hard ERROR by the plugin itself.
|
||||||
|
"-p", "no:unraisableexception",
|
||||||
|
]
|
||||||
|
markers = [
|
||||||
|
"asyncio: marks tests as async",
|
||||||
|
"integration: marks tests as integration tests",
|
||||||
|
"unit: marks tests as unit tests",
|
||||||
|
"opc: marks tests that use the in-process OPC UA server (OpcRepository E2E)",
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.coverage.run]
|
||||||
|
source = ["laborious"]
|
||||||
|
omit = [
|
||||||
|
"*/tests/*",
|
||||||
|
"*/venv/*",
|
||||||
|
"*/__pycache__/*",
|
||||||
|
"*/site-packages/*",
|
||||||
|
]
|
||||||
|
branch = true
|
||||||
|
|
||||||
|
[tool.coverage.report]
|
||||||
|
precision = 2
|
||||||
|
show_missing = true
|
||||||
|
skip_covered = false
|
||||||
|
exclude_lines = [
|
||||||
|
"pragma: no cover",
|
||||||
|
"def __repr__",
|
||||||
|
"def __str__",
|
||||||
|
"raise AssertionError",
|
||||||
|
"raise NotImplementedError",
|
||||||
|
"if __name__ == .__main__.:",
|
||||||
|
"if TYPE_CHECKING:",
|
||||||
|
"class .*\\bProtocol\\):",
|
||||||
|
"@(abc\\.)?abstractmethod",
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.coverage.html]
|
||||||
|
directory = "htmlcov"
|
||||||
|
|
||||||
|
[tool.bandit]
|
||||||
|
exclude_dirs = ["tests", "venv", ".venv"]
|
||||||
|
skips = ["B101", "B601", "B608"] # Skip assert, shell injection, and SQL injection (false positives)
|
||||||
21
requirements-dev.txt
Normal file
21
requirements-dev.txt
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
# Development and Testing Dependencies
|
||||||
|
# These packages are only needed for development, testing, and code quality checks
|
||||||
|
# Install with: pip install -r requirements-dev.txt
|
||||||
|
|
||||||
|
# Code Quality & Linting
|
||||||
|
ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort)
|
||||||
|
mypy>=1.7.0 # Static type checker
|
||||||
|
bandit>=1.7.5 # Security vulnerability scanner
|
||||||
|
pandas-stubs>=2.0.0 # Type stubs for pandas
|
||||||
|
types-requests>=2.31.0 # Type stubs for requests
|
||||||
|
|
||||||
|
# Testing
|
||||||
|
pytest>=7.4.0 # Testing framework
|
||||||
|
pytest-cov>=4.1.0 # Coverage plugin for pytest
|
||||||
|
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
|
||||||
|
testcontainers[postgres,minio] # PostgreSQL and MinIO containers for E2E tests
|
||||||
|
|
||||||
|
# Development Tools
|
||||||
|
ipython>=8.12.0 # Enhanced Python shell
|
||||||
|
ipdb>=0.13.13 # IPython debugger
|
||||||
|
ipykernel==6.30.1 # IPython kernel for Jupyter notebooks
|
||||||
19
requirements-light.txt
Normal file
19
requirements-light.txt
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
temporalio
|
||||||
|
psycopg2-binary
|
||||||
|
sqlalchemy
|
||||||
|
asyncua==1.0.6
|
||||||
|
redis
|
||||||
|
sientia_do>=1.12.1
|
||||||
|
mlflow
|
||||||
|
prometheus-client
|
||||||
|
botocore
|
||||||
|
boto3
|
||||||
|
s3fs
|
||||||
|
pyarrow
|
||||||
|
kaleido
|
||||||
|
hyperopt
|
||||||
|
shap
|
||||||
|
pycurl
|
||||||
|
scipy<1.14.0
|
||||||
|
scikit-learn==1.5.2
|
||||||
|
joblib
|
||||||
20
requirements-local.txt
Normal file
20
requirements-local.txt
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
temporalio
|
||||||
|
psycopg2-binary
|
||||||
|
sqlalchemy
|
||||||
|
asyncua==1.0.6
|
||||||
|
redis
|
||||||
|
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.1
|
||||||
|
git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.10.0
|
||||||
|
prometheus-client
|
||||||
|
botocore
|
||||||
|
boto3
|
||||||
|
s3fs
|
||||||
|
pyarrow
|
||||||
|
kaleido
|
||||||
|
hyperopt
|
||||||
|
shap
|
||||||
|
pycurl
|
||||||
|
scipy<1.14.0
|
||||||
|
scikit-learn==1.4.2
|
||||||
|
pycaret==3.3.2
|
||||||
|
joblib
|
||||||
19
requirements.txt
Normal file
19
requirements.txt
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
temporalio
|
||||||
|
psycopg2-binary
|
||||||
|
sqlalchemy
|
||||||
|
asyncua==1.0.6
|
||||||
|
redis
|
||||||
|
sientia_do>=1.12.1
|
||||||
|
sientia>0.40.0
|
||||||
|
prometheus-client
|
||||||
|
botocore
|
||||||
|
boto3
|
||||||
|
s3fs
|
||||||
|
pyarrow
|
||||||
|
kaleido
|
||||||
|
hyperopt
|
||||||
|
shap
|
||||||
|
pycurl
|
||||||
|
scipy<1.14.0
|
||||||
|
scikit-learn==1.5.2
|
||||||
|
joblib
|
||||||
11
run_coverage.sh
Executable file
11
run_coverage.sh
Executable file
@@ -0,0 +1,11 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
|
||||||
|
# Exit on any error
|
||||||
|
set -e
|
||||||
|
|
||||||
|
echo "Activating virtual environment..."
|
||||||
|
source ./venv/bin/activate
|
||||||
|
|
||||||
|
pytest --cov=laborious --cov-report=html
|
||||||
|
|
||||||
|
xdg-open htmlcov/index.html
|
||||||
18
run_local.sh
Executable file
18
run_local.sh
Executable file
@@ -0,0 +1,18 @@
|
|||||||
|
#!/bin/bash
|
||||||
|
|
||||||
|
# Exit on any error
|
||||||
|
set -e
|
||||||
|
|
||||||
|
echo "Activating virtual environment..."
|
||||||
|
source ./venv/bin/activate
|
||||||
|
|
||||||
|
echo "Loading environment variables from .env..."
|
||||||
|
if [ -f .env ]; then
|
||||||
|
export $(cat .env | grep -v '^#' | xargs)
|
||||||
|
echo "Environment variables loaded from .env"
|
||||||
|
else
|
||||||
|
echo "Warning: .env file not found. Continuing without environment variables."
|
||||||
|
fi
|
||||||
|
|
||||||
|
echo "Starting ingestor application..."
|
||||||
|
python -m laborious.worker.worker
|
||||||
11
sonar-project.properties
Normal file
11
sonar-project.properties
Normal file
@@ -0,0 +1,11 @@
|
|||||||
|
sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_ca1a7039-6db9-49e5-be78-54d29bc93e4f
|
||||||
|
sonar.projectName=sientia-dataops-laborious_temporal
|
||||||
|
sonar.sources=laborious
|
||||||
|
sonar.tests=tests
|
||||||
|
sonar.projectVersion=1.0.0
|
||||||
|
sonar.coverage.exclusions=laborious/worker/*
|
||||||
|
sonar.qualitygate.wait=true
|
||||||
|
sonar.qualitygate.timeout=300
|
||||||
|
sonar.python.coverage.reportPaths=coverage.xml
|
||||||
|
sonar.python.xunit.reportPath=pytest.xml
|
||||||
|
sonar.python.version=3.11
|
||||||
938
tests.ipynb
Normal file
938
tests.ipynb
Normal file
File diff suppressed because one or more lines are too long
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
59
tests/conftest.py
Normal file
59
tests/conftest.py
Normal file
@@ -0,0 +1,59 @@
|
|||||||
|
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.
|
||||||
|
|
||||||
|
This module mocks the 'sientia' module to avoid requiring its installation
|
||||||
|
during unit tests. The mock is registered in sys.modules before any test
|
||||||
|
imports are executed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Mock sientia module
|
||||||
|
sientia_mock = MagicMock()
|
||||||
|
sientia_mock.ModelAnalysis = MagicMock
|
||||||
|
sys.modules['sientia'] = sientia_mock
|
||||||
|
sys.modules['sientia.ModelAnalysis'] = MagicMock()
|
||||||
0
tests/laborious/__init__.py
Normal file
0
tests/laborious/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
229
tests/laborious/activities/test_activities.py
Normal file
229
tests/laborious/activities/test_activities.py
Normal file
@@ -0,0 +1,229 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
|
from pytest import mark
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.activities.api import API
|
||||||
|
from laborious.activities.gates import Gates
|
||||||
|
from laborious.activities.mlflow import MLFlow
|
||||||
|
from laborious.activities.model_metrics import ModelMetrics
|
||||||
|
from laborious.activities.opc import OPC
|
||||||
|
from laborious.activities.storage import Storage
|
||||||
|
|
||||||
|
|
||||||
|
@patch('laborious.activities.activities.Storage.__init__')
|
||||||
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||||
|
@patch('laborious.activities.activities.OPC.__init__')
|
||||||
|
@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,
|
||||||
|
mock_opc_init,
|
||||||
|
mock_mlflow_init,
|
||||||
|
mock_storage_init,
|
||||||
|
):
|
||||||
|
postgres_config = {
|
||||||
|
'host': 'localhost',
|
||||||
|
'port': 5432,
|
||||||
|
'user': 'postgres',
|
||||||
|
'password': 'postgres',
|
||||||
|
'dbname': 'postgres',
|
||||||
|
'min_connections': 1,
|
||||||
|
'max_connections': 10,
|
||||||
|
}
|
||||||
|
|
||||||
|
minio_config = {
|
||||||
|
'endpoint_url': 'localhost:9000',
|
||||||
|
'access_key': 'minio',
|
||||||
|
'secret_key': 'minio123',
|
||||||
|
'default_bucket': 'test',
|
||||||
|
'retention_hours': 24,
|
||||||
|
'secure': False,
|
||||||
|
}
|
||||||
|
|
||||||
|
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||||
|
|
||||||
|
opc_config = {
|
||||||
|
'bootstrap_servers': 'localhost:9092',
|
||||||
|
'polling_time': 1000,
|
||||||
|
'group_id': 'test-group',
|
||||||
|
}
|
||||||
|
|
||||||
|
pi_web_api_config = {
|
||||||
|
'base_url': 'https://test-pi-server.com',
|
||||||
|
'auth_type': 'bearer',
|
||||||
|
'auth_token': 'test_token',
|
||||||
|
}
|
||||||
|
|
||||||
|
logger = MagicMock()
|
||||||
|
notification_handler = MagicMock()
|
||||||
|
|
||||||
|
activities = Activities(
|
||||||
|
postgres_config=postgres_config,
|
||||||
|
mlflow_config=mlflow_config,
|
||||||
|
minio_config=minio_config,
|
||||||
|
opc_config=opc_config,
|
||||||
|
pi_web_api_config=pi_web_api_config,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(activities, Activities)
|
||||||
|
assert isinstance(activities, Storage)
|
||||||
|
assert isinstance(activities, MLFlow)
|
||||||
|
assert isinstance(activities, OPC)
|
||||||
|
assert isinstance(activities, Gates)
|
||||||
|
assert isinstance(activities, ModelMetrics)
|
||||||
|
assert isinstance(activities, API)
|
||||||
|
|
||||||
|
mock_storage_init.assert_called_once_with(
|
||||||
|
ANY,
|
||||||
|
host=postgres_config['host'],
|
||||||
|
port=postgres_config['port'],
|
||||||
|
user=postgres_config['user'],
|
||||||
|
password=postgres_config['password'],
|
||||||
|
dbname=postgres_config['dbname'],
|
||||||
|
min_connections=postgres_config['min_connections'],
|
||||||
|
max_connections=postgres_config['max_connections'],
|
||||||
|
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,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_mlflow_init.assert_called_once_with(
|
||||||
|
ANY,
|
||||||
|
mlflow_host=mlflow_config['host'],
|
||||||
|
mlflow_port=mlflow_config['port'],
|
||||||
|
mlflow_username=mlflow_config['username'],
|
||||||
|
mlflow_password=mlflow_config['password'],
|
||||||
|
minio_repository=mock_minio_repository.return_value,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=mock_metrics_controller.return_value,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_opc_init.assert_called_once_with(
|
||||||
|
ANY,
|
||||||
|
opc_servers=opc_config,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=mock_metrics_controller.return_value,
|
||||||
|
)
|
||||||
|
|
||||||
|
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,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_model_metrics_init.assert_called_once_with(
|
||||||
|
ANY,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=mock_metrics_controller.return_value,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_api_init.assert_called_once_with(
|
||||||
|
ANY,
|
||||||
|
base_url=pi_web_api_config['base_url'],
|
||||||
|
auth_type=pi_web_api_config['auth_type'],
|
||||||
|
auth_token=pi_web_api_config['auth_token'],
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
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')
|
||||||
|
@patch('laborious.activities.activities.MLFlow')
|
||||||
|
@patch('laborious.activities.activities.OPC')
|
||||||
|
@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,
|
||||||
|
mock_opc_init,
|
||||||
|
mock_mlflow_init,
|
||||||
|
mock_storage_init,
|
||||||
|
):
|
||||||
|
mock_opc_init.aclose = AsyncMock()
|
||||||
|
postgres_config = {
|
||||||
|
'host': 'localhost',
|
||||||
|
'port': 5432,
|
||||||
|
'user': 'postgres',
|
||||||
|
'password': 'postgres',
|
||||||
|
'dbname': 'postgres',
|
||||||
|
'min_connections': 1,
|
||||||
|
'max_connections': 10,
|
||||||
|
}
|
||||||
|
|
||||||
|
minio_config = {
|
||||||
|
'endpoint_url': 'localhost:9000',
|
||||||
|
'access_key': 'minio',
|
||||||
|
'secret_key': 'minio123',
|
||||||
|
'default_bucket': 'test',
|
||||||
|
'retention_hours': 24,
|
||||||
|
'secure': False,
|
||||||
|
}
|
||||||
|
|
||||||
|
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
||||||
|
|
||||||
|
opc_config = {
|
||||||
|
'bootstrap_servers': 'localhost:9092',
|
||||||
|
'polling_time': 1000,
|
||||||
|
'group_id': 'test-group',
|
||||||
|
}
|
||||||
|
|
||||||
|
pi_web_api_config = {
|
||||||
|
'base_url': 'https://test-pi-server.com',
|
||||||
|
'auth_type': 'bearer',
|
||||||
|
'auth_token': 'test_token',
|
||||||
|
}
|
||||||
|
|
||||||
|
logger = MagicMock()
|
||||||
|
notification_handler = MagicMock()
|
||||||
|
|
||||||
|
activities = Activities(
|
||||||
|
postgres_config=postgres_config,
|
||||||
|
mlflow_config=mlflow_config,
|
||||||
|
minio_config=minio_config,
|
||||||
|
opc_config=opc_config,
|
||||||
|
pi_web_api_config=pi_web_api_config,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
)
|
||||||
|
|
||||||
|
await activities.shutdown()
|
||||||
|
mock_opc_init.aclose.assert_called_once()
|
||||||
|
mock_storage_init.close.assert_called_once()
|
||||||
|
mock_mlflow_init.close.assert_called_once()
|
||||||
|
mock_gates_init.close.assert_called_once()
|
||||||
|
mock_model_metrics_init.close.assert_called_once()
|
||||||
|
mock_api_init.close.assert_called_once()
|
||||||
492
tests/laborious/activities/test_api.py
Normal file
492
tests/laborious/activities/test_api.py
Normal file
@@ -0,0 +1,492 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||||
|
|
||||||
|
import pytest_asyncio
|
||||||
|
from pytest import fixture, mark
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
|
||||||
|
from laborious.activities.api import API, PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schema_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _create_mock_dataframe(to_dict_return=None):
|
||||||
|
"""Helper function to create a mocked DataFrame for testing."""
|
||||||
|
mock_df = MagicMock()
|
||||||
|
mock_head = MagicMock()
|
||||||
|
|
||||||
|
def get_column_values(key):
|
||||||
|
if key == 'prediction':
|
||||||
|
return MagicMock(values=[0.75])
|
||||||
|
elif key == 'prediction_confidence':
|
||||||
|
return MagicMock(values=[0.95])
|
||||||
|
else:
|
||||||
|
return MagicMock(values=['2024-01-01T00:00:00+00:00'])
|
||||||
|
|
||||||
|
mock_head.__getitem__.side_effect = get_column_values
|
||||||
|
mock_df.head.return_value = mock_head
|
||||||
|
|
||||||
|
if to_dict_return is None:
|
||||||
|
to_dict_return = {
|
||||||
|
'prediction': [0.75],
|
||||||
|
'prediction_confidence': [0.95],
|
||||||
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
||||||
|
}
|
||||||
|
mock_df.to_dict.return_value = to_dict_return
|
||||||
|
|
||||||
|
return mock_df
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def base_input_data():
|
||||||
|
"""Base input data for PI Web API tests."""
|
||||||
|
return {
|
||||||
|
**metadata,
|
||||||
|
'data': {
|
||||||
|
'prediction': [0.75],
|
||||||
|
'prediction_confidence': [0.95],
|
||||||
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
||||||
|
},
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
||||||
|
'prediction_tags': {'tag1': 'web_id_1'},
|
||||||
|
'confidence_tags': {'tag2': 'web_id_2'},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@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',
|
||||||
|
auth_type='bearer',
|
||||||
|
auth_token='test_token',
|
||||||
|
logger=MagicMock(),
|
||||||
|
notification_handler=MagicMock(),
|
||||||
|
metrics_controller=AsyncMock(),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert api.pi_web_api_client is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
@patch('laborious.activities.api.PIWebAPIClient')
|
||||||
|
def api(mock_pi_web_api_client):
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client.write_value = AsyncMock()
|
||||||
|
mock_client.close = MagicMock()
|
||||||
|
mock_client.base_url = 'https://test-pi-server.com'
|
||||||
|
mock_pi_web_api_client.return_value = mock_client
|
||||||
|
|
||||||
|
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(),
|
||||||
|
)
|
||||||
|
api_instance.send_notification_async = AsyncMock()
|
||||||
|
api_instance.info = MagicMock()
|
||||||
|
api_instance.error = MagicMock()
|
||||||
|
api_instance.emit_metric = AsyncMock()
|
||||||
|
api_instance.get_core_labels = MagicMock(
|
||||||
|
return_value={
|
||||||
|
'pod_id': 'test_pod',
|
||||||
|
'runtime': 'local',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return api_instance
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.api.DataFrame')
|
||||||
|
async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_data):
|
||||||
|
input_data = {
|
||||||
|
**base_input_data,
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
||||||
|
'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
|
||||||
|
'confidence_tags': {'tag3': 'web_id_3', 'tag4': 'web_id_4'},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
mock_dataframe.return_value = _create_mock_dataframe()
|
||||||
|
|
||||||
|
# Mock successful responses
|
||||||
|
api.pi_web_api_client.write_value.side_effect = [
|
||||||
|
[{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}],
|
||||||
|
[{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}],
|
||||||
|
]
|
||||||
|
|
||||||
|
result = await api.write_pi_web_api_data(input_data)
|
||||||
|
|
||||||
|
api.pi_web_api_client.write_value.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_1', 'web_id_2'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
||||||
|
'Value': 0.75,
|
||||||
|
},
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
web_ids=['web_id_3', 'web_id_4'],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
||||||
|
'Value': 0.95,
|
||||||
|
},
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result == {
|
||||||
|
'prediction': [0.75],
|
||||||
|
'prediction_confidence': [0.95],
|
||||||
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.api.DataFrame')
|
||||||
|
async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
|
||||||
|
mock_dataframe.return_value = _create_mock_dataframe(
|
||||||
|
{
|
||||||
|
'prediction': [0.75],
|
||||||
|
'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
|
||||||
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')
|
||||||
|
|
||||||
|
result = await api.write_pi_web_api_data(base_input_data)
|
||||||
|
|
||||||
|
api.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||||
|
message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
|
||||||
|
block='write_pi_web_api_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result['prediction_confidence'][0] == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
assert api.pi_web_api_client.write_value.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.api.DataFrame')
|
||||||
|
async def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_input_data):
|
||||||
|
mock_dataframe.return_value = _create_mock_dataframe()
|
||||||
|
|
||||||
|
# First call succeeds, second fails
|
||||||
|
api.pi_web_api_client.write_value.side_effect = [
|
||||||
|
[{'WebId': 'web_id_1', 'Errors': []}],
|
||||||
|
Exception('Confidence write failed'),
|
||||||
|
]
|
||||||
|
|
||||||
|
result = await api.write_pi_web_api_data(base_input_data)
|
||||||
|
|
||||||
|
api.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
|
||||||
|
message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
|
||||||
|
block='write_pi_web_api_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result == {
|
||||||
|
'prediction': [0.75],
|
||||||
|
'prediction_confidence': [0.95],
|
||||||
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
||||||
|
}
|
||||||
|
assert api.pi_web_api_client.write_value.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.api.DataFrame')
|
||||||
|
async def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
|
||||||
|
input_data = {
|
||||||
|
**base_input_data,
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
mock_dataframe.return_value = _create_mock_dataframe()
|
||||||
|
|
||||||
|
# Mock empty responses
|
||||||
|
api.pi_web_api_client.write_value.side_effect = [
|
||||||
|
[],
|
||||||
|
[],
|
||||||
|
]
|
||||||
|
|
||||||
|
result = await api.write_pi_web_api_data(input_data)
|
||||||
|
|
||||||
|
api.pi_web_api_client.write_value.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
web_ids=[],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
||||||
|
'Value': 0.75,
|
||||||
|
},
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
web_ids=[],
|
||||||
|
value={
|
||||||
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
||||||
|
'Value': 0.95,
|
||||||
|
},
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result == {
|
||||||
|
'prediction': [0.75],
|
||||||
|
'prediction_confidence': [0.95],
|
||||||
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_close(api):
|
||||||
|
api.close()
|
||||||
|
|
||||||
|
api.pi_web_api_client.close.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_process_pi_web_api_response_success(api):
|
||||||
|
"""Test successful processing of PI Web API response with all tags written."""
|
||||||
|
response_data = [
|
||||||
|
{'WebId': 'web_id_1', 'Errors': []},
|
||||||
|
{'WebId': 'web_id_2', 'Errors': []},
|
||||||
|
]
|
||||||
|
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',
|
||||||
|
}
|
||||||
|
|
||||||
|
confidence, message = await api.process_pi_web_api_response(
|
||||||
|
response_data=response_data,
|
||||||
|
tags=tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert confidence == 0
|
||||||
|
assert message == ''
|
||||||
|
assert api.emit_metric.call_count == 2
|
||||||
|
# Verify that emit_metric was called with correct tags structure
|
||||||
|
call_args_list = api.emit_metric.call_args_list
|
||||||
|
assert len(call_args_list) == 2
|
||||||
|
# Check that all calls include core_labels and tag_name
|
||||||
|
for call_args in call_args_list:
|
||||||
|
assert 'tag_name' in call_args.kwargs['tags']
|
||||||
|
assert call_args.kwargs['tags']['tag_name'] in ['tag1', 'tag2']
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_process_pi_web_api_response_with_errors(api):
|
||||||
|
"""Test processing response with errors in some tags."""
|
||||||
|
response_data = [
|
||||||
|
{'WebId': 'web_id_1', 'Errors': ['Error writing tag']},
|
||||||
|
{'WebId': 'web_id_2', 'Errors': []},
|
||||||
|
]
|
||||||
|
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',
|
||||||
|
}
|
||||||
|
|
||||||
|
confidence, message = await api.process_pi_web_api_response(
|
||||||
|
response_data=response_data,
|
||||||
|
tags=tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
assert (
|
||||||
|
message
|
||||||
|
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
|
||||||
|
)
|
||||||
|
assert api.emit_metric.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_process_pi_web_api_response_missing_tags(api):
|
||||||
|
"""Test processing response when number of written tags doesn't match expected."""
|
||||||
|
response_data = [
|
||||||
|
{'WebId': 'web_id_1', 'Errors': []},
|
||||||
|
]
|
||||||
|
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',
|
||||||
|
}
|
||||||
|
|
||||||
|
confidence, message = await api.process_pi_web_api_response(
|
||||||
|
response_data=response_data,
|
||||||
|
tags=tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
assert (
|
||||||
|
message
|
||||||
|
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
|
||||||
|
)
|
||||||
|
api.send_notification_async.assert_called_once()
|
||||||
|
call_args = api.send_notification_async.call_args
|
||||||
|
assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
|
||||||
|
assert call_args.kwargs['level'] == NotificationLevel.ERROR
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_process_pi_web_api_response_missing_webid(api):
|
||||||
|
"""Test processing response when WebId is missing in response item."""
|
||||||
|
response_data = [
|
||||||
|
{'Errors': []},
|
||||||
|
{'WebId': 'web_id_2', 'Errors': []},
|
||||||
|
]
|
||||||
|
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',
|
||||||
|
}
|
||||||
|
|
||||||
|
confidence, message = await api.process_pi_web_api_response(
|
||||||
|
response_data=response_data,
|
||||||
|
tags=tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
assert (
|
||||||
|
message
|
||||||
|
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
|
||||||
|
)
|
||||||
|
api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_process_pi_web_api_response_missing_tag_name(api):
|
||||||
|
"""Test processing response when tag name is not found for WebId."""
|
||||||
|
response_data = [
|
||||||
|
{'WebId': 'unknown_web_id', 'Errors': []},
|
||||||
|
]
|
||||||
|
tags = {'tag1': 'web_id_1'}
|
||||||
|
core_labels = {
|
||||||
|
'pod_id': 'test_pod',
|
||||||
|
'runtime': 'local',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
}
|
||||||
|
|
||||||
|
confidence, message = await api.process_pi_web_api_response(
|
||||||
|
response_data=response_data,
|
||||||
|
tags=tags,
|
||||||
|
core_labels=core_labels,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||||
|
assert (
|
||||||
|
message
|
||||||
|
== "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
|
||||||
|
)
|
||||||
|
api.error.assert_any_call(
|
||||||
|
'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
|
||||||
|
)
|
||||||
1008
tests/laborious/activities/test_gates.py
Normal file
1008
tests/laborious/activities/test_gates.py
Normal file
File diff suppressed because it is too large
Load Diff
629
tests/laborious/activities/test_mlflow.py
Normal file
629
tests/laborious/activities/test_mlflow.py
Normal file
@@ -0,0 +1,629 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from pytest import fixture, mark, raises
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
|
||||||
|
|
||||||
|
from laborious.activities.mlflow import MLFlow
|
||||||
|
|
||||||
|
|
||||||
|
@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_repository=minio_repo,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert mlflow.mlflow_host == 'http://localhost'
|
||||||
|
assert mlflow.mlflow_port == 5000
|
||||||
|
assert mlflow.mlflow_username == 'admin'
|
||||||
|
assert mlflow.mlflow_password == 'admin'
|
||||||
|
|
||||||
|
mock_mlflow_repository.assert_called_once_with(
|
||||||
|
'http://localhost:5000', 'admin', 'admin', ANY, ANY, ANY
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_minio_repository.assert_called_once_with(
|
||||||
|
endpoint='localhost:9000',
|
||||||
|
access_key='minio',
|
||||||
|
secret_key='minio123',
|
||||||
|
logger=ANY,
|
||||||
|
notification_handler=ANY,
|
||||||
|
metrics_controller=ANY,
|
||||||
|
bucket='test',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
@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_repository=minio_repo,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository = AsyncMock()
|
||||||
|
mlflow.minio_repository = AsyncMock()
|
||||||
|
|
||||||
|
mlflow.send_notification = MagicMock()
|
||||||
|
mlflow.emit_metric = AsyncMock()
|
||||||
|
mlflow.send_notification_async = AsyncMock()
|
||||||
|
mlflow.error = MagicMock()
|
||||||
|
mlflow.debug = MagicMock()
|
||||||
|
mlflow.info = MagicMock()
|
||||||
|
mlflow.warning = MagicMock()
|
||||||
|
mlflow.critical = MagicMock()
|
||||||
|
|
||||||
|
return mlflow
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schema_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.activities.mlflow.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': payload,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_config': {},
|
||||||
|
}
|
||||||
|
|
||||||
|
transform_response = {'success': True, 'content': MagicMock()}
|
||||||
|
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)
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.transform.assert_called_once_with(
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
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': payload,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_config': {},
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
response_data = await mlflow.request_predict(input_data)
|
||||||
|
|
||||||
|
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
||||||
|
mock_to_datetime.assert_called_once_with(
|
||||||
|
data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
||||||
|
)
|
||||||
|
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
||||||
|
'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
|
||||||
|
@patch('laborious.activities.mlflow.to_datetime')
|
||||||
|
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
|
||||||
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
||||||
|
'success': True,
|
||||||
|
'experiment': 'test_experiment',
|
||||||
|
'message': 'Model retrained successfully.',
|
||||||
|
}
|
||||||
|
|
||||||
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
|
||||||
|
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
|
||||||
|
|
||||||
|
raw_data.sort_values.assert_not_called()
|
||||||
|
raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
|
||||||
|
raw_data = raw_data.drop_duplicates.return_value
|
||||||
|
|
||||||
|
raw_data.drop.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(columns=['model_id'], inplace=True, errors='ignore'),
|
||||||
|
call(columns=['created_at'], inplace=True, errors='ignore'),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
|
||||||
|
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
|
||||||
|
|
||||||
|
raw_data = raw_data.pivot.return_value
|
||||||
|
|
||||||
|
raw_data.__setitem__.assert_has_calls(
|
||||||
|
[
|
||||||
|
call('timestamp', raw_data.index),
|
||||||
|
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
|
||||||
|
call('timestamp', mock_to_datetime.return_value),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
mock_to_datetime.assert_has_calls(
|
||||||
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
|
||||||
|
)
|
||||||
|
mock_to_datetime.assert_has_calls(
|
||||||
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
|
||||||
|
)
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
|
||||||
|
data=raw_data,
|
||||||
|
model_name='test_model',
|
||||||
|
model_config={
|
||||||
|
'target': 'target',
|
||||||
|
'transform_flavor': 'sklearn',
|
||||||
|
'predict_flavor': 'pyfunc',
|
||||||
|
},
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response == {
|
||||||
|
'success': True,
|
||||||
|
'experiment': 'test_experiment',
|
||||||
|
'message': 'Model retrained successfully.',
|
||||||
|
'timestamp': timestamp,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.mlflow.to_datetime')
|
||||||
|
async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow):
|
||||||
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
||||||
|
'success': True,
|
||||||
|
'experiment': 'test_experiment',
|
||||||
|
'message': 'Model retrained successfully.',
|
||||||
|
}
|
||||||
|
|
||||||
|
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',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
|
||||||
|
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
|
||||||
|
subset=['variable', 'timestamp'], keep='first'
|
||||||
|
)
|
||||||
|
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
|
||||||
|
|
||||||
|
raw_data.drop.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(columns=['model_id'], inplace=True, errors='ignore'),
|
||||||
|
call(columns=['created_at'], inplace=True, errors='ignore'),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
|
||||||
|
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
|
||||||
|
|
||||||
|
raw_data = raw_data.pivot.return_value
|
||||||
|
|
||||||
|
raw_data.__setitem__.assert_has_calls(
|
||||||
|
[
|
||||||
|
call('timestamp', raw_data.index),
|
||||||
|
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
|
||||||
|
call('timestamp', mock_to_datetime.return_value),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
mock_to_datetime.assert_has_calls(
|
||||||
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
|
||||||
|
)
|
||||||
|
mock_to_datetime.assert_has_calls(
|
||||||
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
|
||||||
|
)
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
|
||||||
|
data=raw_data,
|
||||||
|
model_name='test_model',
|
||||||
|
model_config={
|
||||||
|
'target': 'target',
|
||||||
|
'transform_flavor': 'sklearn',
|
||||||
|
'predict_flavor': 'pyfunc',
|
||||||
|
},
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
mlflow.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
notification_id='RETRAIN_MODEL_ERROR',
|
||||||
|
message='Error retraining model test_model: Model retrained failed.',
|
||||||
|
block='retrain_model',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response == {
|
||||||
|
'success': False,
|
||||||
|
'traceback': 'test_traceback',
|
||||||
|
'message': 'Model retrained failed.',
|
||||||
|
'timestamp': timestamp,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_retrain_model_data_error(mlflow):
|
||||||
|
response = await mlflow.retrain_model(
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_config': {
|
||||||
|
'target': 'target',
|
||||||
|
'transform_flavor': 'sklearn',
|
||||||
|
'predict_flavor': 'pyfunc',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response == {
|
||||||
|
'success': False,
|
||||||
|
'message': "Error loading retrain data: 'data'",
|
||||||
|
'traceback': ANY,
|
||||||
|
'timestamp': ANY,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_retrain_model_data_error_no_minio_repository(mlflow):
|
||||||
|
mlflow.minio_repository = None
|
||||||
|
|
||||||
|
with raises(ValueError) as e:
|
||||||
|
await mlflow.retrain_model(
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'object_key': 'test_object_key',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_config': {
|
||||||
|
'target': 'target',
|
||||||
|
'transform_flavor': 'sklearn',
|
||||||
|
'predict_flavor': 'pyfunc',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert str(e.value) == 'Minio repository not initialized'
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_update_production_model(mlflow):
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 1,
|
||||||
|
'experiment': 'test',
|
||||||
|
'timestamp': 2,
|
||||||
|
'status': 'success',
|
||||||
|
}
|
||||||
|
|
||||||
|
response = await mlflow.update_production_model(input_data)
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
|
||||||
|
experiment='test', model_name='test_model', metadata=metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response == mlflow.model_monitoring_repository.update_production_model.return_value
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_update_production_model_error(mlflow):
|
||||||
|
mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
|
||||||
|
'Error updating production model'
|
||||||
|
)
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 1,
|
||||||
|
'experiment': 'test',
|
||||||
|
'timestamp': 2,
|
||||||
|
'status': 'success',
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
await mlflow.update_production_model(input_data)
|
||||||
|
except Exception as e:
|
||||||
|
assert str(e) == 'Error updating production model'
|
||||||
|
mlflow.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
||||||
|
message='Error updating production model test_model: Error updating production model',
|
||||||
|
block='update_production_model',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise AssertionError('No exception raised')
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.mlflow.to_datetime')
|
||||||
|
async def test_get_reference_data_success(mock_to_datetime, mlflow):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
}
|
||||||
|
|
||||||
|
# Mock reference data DataFrame
|
||||||
|
mock_reference_data = MagicMock()
|
||||||
|
mock_reference_data.__getitem__.return_value = MagicMock()
|
||||||
|
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
|
||||||
|
mock_reference_data.to_dict.return_value = [
|
||||||
|
{'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
|
||||||
|
{'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
|
||||||
|
]
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
|
||||||
|
|
||||||
|
# Act
|
||||||
|
result = await mlflow.get_reference_data(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
||||||
|
model_name='test_model',
|
||||||
|
artifact_path='evaluation_data.csv',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
|
||||||
|
|
||||||
|
mock_reference_data.to_dict.assert_called_once_with(orient='records')
|
||||||
|
assert result == mock_reference_data.to_dict.return_value
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_get_reference_data_not_found(mlflow):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
}
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
|
||||||
|
|
||||||
|
# Act
|
||||||
|
result = await mlflow.get_reference_data(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
||||||
|
model_name='test_model',
|
||||||
|
artifact_path='evaluation_data.csv',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
mlflow.warning.assert_called_once_with(
|
||||||
|
'Reference data not found for model test_model', metadata['metadata']
|
||||||
|
)
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_get_reference_data_exception(mlflow):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'model_name': 'test_model',
|
||||||
|
}
|
||||||
|
|
||||||
|
mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
|
||||||
|
'Error loading artifact'
|
||||||
|
)
|
||||||
|
|
||||||
|
# Act & Assert
|
||||||
|
with raises(Exception) as e:
|
||||||
|
await mlflow.get_reference_data(input_data)
|
||||||
|
|
||||||
|
assert str(e.value) == 'Error loading artifact'
|
||||||
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
||||||
|
model_name='test_model',
|
||||||
|
artifact_path='evaluation_data.csv',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
1108
tests/laborious/activities/test_model_metrics.py
Normal file
1108
tests/laborious/activities/test_model_metrics.py
Normal file
File diff suppressed because it is too large
Load Diff
738
tests/laborious/activities/test_opc.py
Normal file
738
tests/laborious/activities/test_opc.py
Normal file
@@ -0,0 +1,738 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||||
|
|
||||||
|
import pytest_asyncio
|
||||||
|
from pandas import DataFrame
|
||||||
|
from pytest import mark
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
|
||||||
|
from laborious.activities.opc import (
|
||||||
|
OPC,
|
||||||
|
OPC_COMMENT_SEPARATOR,
|
||||||
|
OPC_RECONNECT_IN_PROGRESS_COMMENT,
|
||||||
|
OPC_SESSION_BAD_COMMENT_PREFIX,
|
||||||
|
OPC_SESSION_BAD_CONFIDENCE,
|
||||||
|
OPC_WRITTING_ERROR_CONFIDENCE,
|
||||||
|
OPC_WRITTING_ERROR_MESSAGE,
|
||||||
|
_apply_opc_write_error,
|
||||||
|
)
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schema_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test__init__():
|
||||||
|
servers = {'server1': {'id': 'server1'}}
|
||||||
|
opc = OPC(
|
||||||
|
opc_servers=servers,
|
||||||
|
logger=MagicMock(),
|
||||||
|
notification_handler=MagicMock(),
|
||||||
|
metrics_controller=AsyncMock(),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert opc.opc_servers == servers
|
||||||
|
assert opc.opc_repository == {}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.opc.OpcRepository')
|
||||||
|
@patch('laborious.activities.opc.OPC.send_notification_async')
|
||||||
|
async def test_init_opc(mock_send_notification, mock_opc_repository):
|
||||||
|
mock_logger = MagicMock()
|
||||||
|
mock_metrics_controller = AsyncMock()
|
||||||
|
server1 = MagicMock(
|
||||||
|
connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
|
||||||
|
)
|
||||||
|
server2 = MagicMock(
|
||||||
|
connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
|
||||||
|
)
|
||||||
|
server3 = MagicMock(
|
||||||
|
connect=AsyncMock(
|
||||||
|
return_value=(
|
||||||
|
False,
|
||||||
|
{
|
||||||
|
'notification_id': 'OPC_CONNECTION_ERROR_server3',
|
||||||
|
'message': 'Failed to connect to OPC server: Test error',
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': NotificationLevel.ERROR,
|
||||||
|
'attachment_content': 'Test error',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
),
|
||||||
|
write_data=AsyncMock(return_value=(True, {})),
|
||||||
|
)
|
||||||
|
mock_opc_repository.side_effect = [server1, server2, server3]
|
||||||
|
mock_notification_handler = MagicMock()
|
||||||
|
servers = {
|
||||||
|
'server1': {
|
||||||
|
'server_name': 'server1',
|
||||||
|
'id': 'server1',
|
||||||
|
'url': 'http://localhost:8080',
|
||||||
|
'server_uri': 'opc.tcp://localhost:4840',
|
||||||
|
'cert_path': '',
|
||||||
|
'private_key_path': '',
|
||||||
|
'server_cert_path': '',
|
||||||
|
'reconnection_interval': 60,
|
||||||
|
},
|
||||||
|
'server2': {
|
||||||
|
'server_name': 'server2',
|
||||||
|
'id': 'server2',
|
||||||
|
'url': 'http://localhost:8080',
|
||||||
|
'server_uri': 'opc.tcp://localhost:4840',
|
||||||
|
'cert_path': '',
|
||||||
|
'private_key_path': '',
|
||||||
|
'server_cert_path': '',
|
||||||
|
'reconnection_interval': 60,
|
||||||
|
},
|
||||||
|
'server3': {
|
||||||
|
'server_name': 'server3',
|
||||||
|
'id': 'server3',
|
||||||
|
'url': 'http://localhost:8080',
|
||||||
|
'server_uri': 'opc.tcp://localhost:4840',
|
||||||
|
'cert_path': '',
|
||||||
|
'private_key_path': '',
|
||||||
|
'server_cert_path': '',
|
||||||
|
'reconnection_interval': 60,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
opc = OPC(
|
||||||
|
opc_servers=servers,
|
||||||
|
logger=mock_logger,
|
||||||
|
notification_handler=mock_notification_handler,
|
||||||
|
metrics_controller=mock_metrics_controller,
|
||||||
|
)
|
||||||
|
await opc.init_opc()
|
||||||
|
|
||||||
|
assert opc.opc_servers == servers
|
||||||
|
assert opc.logger == mock_logger
|
||||||
|
assert opc.notification_handler == mock_notification_handler
|
||||||
|
assert opc.opc_repository['server1'] == server1
|
||||||
|
assert opc.opc_repository['server2'] == server2
|
||||||
|
|
||||||
|
mock_opc_repository.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
opc_id='server1',
|
||||||
|
server_name='server1',
|
||||||
|
url='http://localhost:8080',
|
||||||
|
logger=mock_logger,
|
||||||
|
server_uri='opc.tcp://localhost:4840',
|
||||||
|
cert_path='',
|
||||||
|
private_key_path='',
|
||||||
|
server_cert_path='',
|
||||||
|
notification_handler=mock_notification_handler,
|
||||||
|
reconnection_interval=60,
|
||||||
|
metrics_controller=mock_metrics_controller,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
mock_opc_repository.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
opc_id='server2',
|
||||||
|
server_name='server2',
|
||||||
|
url='http://localhost:8080',
|
||||||
|
logger=mock_logger,
|
||||||
|
server_uri='opc.tcp://localhost:4840',
|
||||||
|
cert_path='',
|
||||||
|
private_key_path='',
|
||||||
|
server_cert_path='',
|
||||||
|
notification_handler=mock_notification_handler,
|
||||||
|
reconnection_interval=60,
|
||||||
|
metrics_controller=mock_metrics_controller,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
server1.connect.assert_called_once()
|
||||||
|
server2.connect.assert_called_once()
|
||||||
|
|
||||||
|
mock_send_notification.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
metadata={
|
||||||
|
'model_id': '-',
|
||||||
|
'model_name': '-',
|
||||||
|
'workflow_name': '-',
|
||||||
|
'schedule_name': 'INITIALIZATION',
|
||||||
|
},
|
||||||
|
notification_id='OPC_CONNECTION_ERROR_server3',
|
||||||
|
message='Failed to connect to OPC server: Test error',
|
||||||
|
block='opc_repository',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
@patch('laborious.activities.opc.OpcRepository')
|
||||||
|
async def opc(mock_opc_repository):
|
||||||
|
servers = {
|
||||||
|
'server1': {
|
||||||
|
'id': 'server1',
|
||||||
|
'server_name': 'server1',
|
||||||
|
'url': 'http://localhost:8080',
|
||||||
|
'server_uri': 'opc.tcp://localhost:4840',
|
||||||
|
'cert_path': '',
|
||||||
|
'private_key_path': '',
|
||||||
|
'server_cert_path': '',
|
||||||
|
'reconnection_interval': 60,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
|
||||||
|
mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
|
||||||
|
opc = OPC(
|
||||||
|
opc_servers=servers,
|
||||||
|
logger=MagicMock(),
|
||||||
|
notification_handler=MagicMock(),
|
||||||
|
metrics_controller=AsyncMock(),
|
||||||
|
)
|
||||||
|
await opc.init_opc()
|
||||||
|
opc.send_notification = MagicMock()
|
||||||
|
opc.send_notification_async = AsyncMock()
|
||||||
|
opc.emit_metric = AsyncMock()
|
||||||
|
return opc
|
||||||
|
|
||||||
|
|
||||||
|
WRITE_DATA_CASES = [
|
||||||
|
('tag1', 'int', 50),
|
||||||
|
('tag2', 'float', 50.5),
|
||||||
|
('tag3', 'bool', True),
|
||||||
|
('tag4', 'string', 'test'),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_data_success(opc, tag, data_type, data):
|
||||||
|
opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})
|
||||||
|
|
||||||
|
response_time, error_info = await opc.write_data(
|
||||||
|
server_id='server1',
|
||||||
|
tag=tag,
|
||||||
|
data=data,
|
||||||
|
data_type=data_type,
|
||||||
|
tag_type='prediction',
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
assert response_time == 0.1
|
||||||
|
assert error_info is None
|
||||||
|
opc.opc_repository['server1'].write_data.assert_called_once_with(tag, data, data_type, metadata)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_data_failed(opc):
|
||||||
|
opc.opc_repository['server1'].write_data.return_value = (
|
||||||
|
False,
|
||||||
|
{
|
||||||
|
'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
|
||||||
|
'message': 'Failed to write data to OPC server: Test error',
|
||||||
|
'block': 'opc_repository',
|
||||||
|
'level': NotificationLevel.ERROR,
|
||||||
|
'attachment_content': 'Test error',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
response_time, error_info = await opc.write_data(
|
||||||
|
server_id='server1',
|
||||||
|
tag='tag1',
|
||||||
|
data=50,
|
||||||
|
data_type='int',
|
||||||
|
tag_type='prediction',
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
assert response_time is None
|
||||||
|
assert error_info is not None
|
||||||
|
|
||||||
|
opc.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='OPC_WRITE_DATA_ERROR_server1',
|
||||||
|
message='Failed to write data to OPC server: Test error',
|
||||||
|
block='opc_repository',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_data_exception(opc):
|
||||||
|
opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
|
||||||
|
|
||||||
|
try:
|
||||||
|
await opc.write_data(
|
||||||
|
server_id='server1',
|
||||||
|
tag='tag1',
|
||||||
|
data=50,
|
||||||
|
data_type='int',
|
||||||
|
tag_type='prediction',
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
opc.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=metadata,
|
||||||
|
notification_id='WRITE_OPC_PREDICTION_ERROR',
|
||||||
|
message='Error writing data to OPC server: Test error',
|
||||||
|
block='write_opc_data',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
else:
|
||||||
|
raise AssertionError('Expected an exception to be raised')
|
||||||
|
|
||||||
|
|
||||||
|
@mark.parametrize(
|
||||||
|
'error_info,initial_seen,initial_status,initial_reconnect,expected',
|
||||||
|
[
|
||||||
|
(None, False, None, False, (False, None, False)),
|
||||||
|
({}, False, None, False, (False, None, False)),
|
||||||
|
(
|
||||||
|
{'opc_error_kind': 'session_bad', 'opc_status': 'BadSessionIdInvalid'},
|
||||||
|
False,
|
||||||
|
None,
|
||||||
|
False,
|
||||||
|
(True, 'BadSessionIdInvalid', False),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{'opc_error_kind': 'session_bad', 'opc_status': 'NewStatus'},
|
||||||
|
True,
|
||||||
|
'OldStatus',
|
||||||
|
False,
|
||||||
|
(True, 'NewStatus', False),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{'opc_error_kind': 'session_bad'},
|
||||||
|
True,
|
||||||
|
'KeptStatus',
|
||||||
|
False,
|
||||||
|
(True, 'KeptStatus', False),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{'opc_error_kind': 'reconnect_in_progress'},
|
||||||
|
False,
|
||||||
|
None,
|
||||||
|
False,
|
||||||
|
(False, None, True),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
{'opc_error_kind': 'other'},
|
||||||
|
True,
|
||||||
|
'Status',
|
||||||
|
True,
|
||||||
|
(True, 'Status', True),
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_apply_opc_write_error(
|
||||||
|
error_info, initial_seen, initial_status, initial_reconnect, expected
|
||||||
|
):
|
||||||
|
result = _apply_opc_write_error(
|
||||||
|
error_info,
|
||||||
|
initial_seen,
|
||||||
|
initial_status,
|
||||||
|
initial_reconnect,
|
||||||
|
)
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_tags_from_config_prediction_success(opc):
|
||||||
|
opc.write_data = AsyncMock(return_value=(0.1, None))
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
tags_config = {'tag1': {'data_type': 'float'}}
|
||||||
|
|
||||||
|
response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
|
||||||
|
server_id='server1',
|
||||||
|
tags_config=tags_config,
|
||||||
|
data=data,
|
||||||
|
data_column='prediction',
|
||||||
|
tag_type='prediction',
|
||||||
|
log_label='Prediction data',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response_times == {'tag1': 0.1}
|
||||||
|
assert session_bad is False
|
||||||
|
assert opc_status is None
|
||||||
|
assert reconnect is False
|
||||||
|
opc.write_data.assert_called_once_with(
|
||||||
|
server_id='server1',
|
||||||
|
tag='tag1',
|
||||||
|
data=0.75,
|
||||||
|
data_type='float',
|
||||||
|
tag_type='prediction',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_tags_from_config_confidence_success(opc):
|
||||||
|
opc.write_data = AsyncMock(return_value=(0.2, None))
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
tags_config = {'tag2': {'data_type': 'float'}}
|
||||||
|
|
||||||
|
response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
|
||||||
|
server_id='server1',
|
||||||
|
tags_config=tags_config,
|
||||||
|
data=data,
|
||||||
|
data_column='prediction_confidence',
|
||||||
|
tag_type='confidence',
|
||||||
|
log_label='Confidence data',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response_times == {'tag2': 0.2}
|
||||||
|
assert session_bad is False
|
||||||
|
assert opc_status is None
|
||||||
|
assert reconnect is False
|
||||||
|
opc.write_data.assert_called_once_with(
|
||||||
|
server_id='server1',
|
||||||
|
tag='tag2',
|
||||||
|
data=0.95,
|
||||||
|
data_type='float',
|
||||||
|
tag_type='confidence',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_tags_from_config_write_failure(opc):
|
||||||
|
opc.write_data = AsyncMock(return_value=(None, {}))
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
|
||||||
|
response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
|
||||||
|
server_id='server1',
|
||||||
|
tags_config={'tag1': {'data_type': 'float'}},
|
||||||
|
data=data,
|
||||||
|
data_column='prediction',
|
||||||
|
tag_type='prediction',
|
||||||
|
log_label='Prediction data',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response_times == {'tag1': None}
|
||||||
|
assert session_bad is False
|
||||||
|
assert opc_status is None
|
||||||
|
assert reconnect is False
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_tags_from_config_session_bad(opc):
|
||||||
|
opc.write_data = AsyncMock(
|
||||||
|
return_value=(
|
||||||
|
None,
|
||||||
|
{
|
||||||
|
'opc_error_kind': 'session_bad',
|
||||||
|
'opc_status': 'BadSessionIdInvalid',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
|
||||||
|
response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
|
||||||
|
server_id='server1',
|
||||||
|
tags_config={'tag1': {'data_type': 'float'}},
|
||||||
|
data=data,
|
||||||
|
data_column='prediction',
|
||||||
|
tag_type='prediction',
|
||||||
|
log_label='Prediction data',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response_times == {'tag1': None}
|
||||||
|
assert session_bad is True
|
||||||
|
assert opc_status == 'BadSessionIdInvalid'
|
||||||
|
assert reconnect is False
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_tags_from_config_reconnect_in_progress(opc):
|
||||||
|
opc.write_data = AsyncMock(
|
||||||
|
return_value=(
|
||||||
|
None,
|
||||||
|
{'opc_error_kind': 'reconnect_in_progress'},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
|
||||||
|
response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
|
||||||
|
server_id='server1',
|
||||||
|
tags_config={'tag1': {'data_type': 'float'}},
|
||||||
|
data=data,
|
||||||
|
data_column='prediction',
|
||||||
|
tag_type='prediction',
|
||||||
|
log_label='Prediction data',
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response_times == {'tag1': None}
|
||||||
|
assert session_bad is False
|
||||||
|
assert opc_status is None
|
||||||
|
assert reconnect is True
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_manage_output_tags_success(opc):
|
||||||
|
opc._write_tags_from_config = AsyncMock(
|
||||||
|
side_effect=[
|
||||||
|
({'tag1': 0.1}, False, None, False),
|
||||||
|
({'tag2': 0.1}, False, None, False),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
config = {
|
||||||
|
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||||
|
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||||
|
}
|
||||||
|
|
||||||
|
output_data, opc_metrics, session_bad, opc_status, reconnect = await opc.manage_output_tags(
|
||||||
|
server_id='server1',
|
||||||
|
config=config,
|
||||||
|
data=data,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert output_data is True
|
||||||
|
assert opc_metrics == {'tag1': 0.1, 'tag2': 0.1}
|
||||||
|
assert session_bad is False
|
||||||
|
assert opc_status is None
|
||||||
|
assert reconnect is False
|
||||||
|
assert opc._write_tags_from_config.await_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_manage_output_tags_failed(opc):
|
||||||
|
opc._write_tags_from_config = AsyncMock(
|
||||||
|
side_effect=[
|
||||||
|
({'tag1': 0.1}, False, None, False),
|
||||||
|
({'tag2': None}, False, None, False),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
config = {
|
||||||
|
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||||
|
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||||
|
}
|
||||||
|
|
||||||
|
output_data, opc_metrics, _, _, _ = await opc.manage_output_tags(
|
||||||
|
server_id='server1',
|
||||||
|
config=config,
|
||||||
|
data=data,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert output_data is False
|
||||||
|
assert opc_metrics == {'tag1': 0.1, 'tag2': None}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_manage_output_tags_do_nothing(opc):
|
||||||
|
opc._write_tags_from_config = AsyncMock()
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
config = {'_invalid_key': {'tag1': {'data_type': 'float'}}}
|
||||||
|
|
||||||
|
output_data, opc_metrics, _, _, _ = await opc.manage_output_tags(
|
||||||
|
server_id='server1',
|
||||||
|
config=config,
|
||||||
|
data=data,
|
||||||
|
metadata=metadata['metadata'],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert output_data is True
|
||||||
|
assert opc_metrics == {}
|
||||||
|
opc._write_tags_from_config.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.opc.DataFrame')
|
||||||
|
async def test_write_opc_data_success(mock_dataframe, opc):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
|
||||||
|
'opc_output_config': {
|
||||||
|
'server1': {
|
||||||
|
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||||
|
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
# Act
|
||||||
|
opc.manage_output_tags = AsyncMock(
|
||||||
|
return_value=(True, {'tag1': 0.1, 'tag2': 0.2}, False, None, False)
|
||||||
|
)
|
||||||
|
|
||||||
|
opc.process_confidence = MagicMock(return_value={'data': 'data'})
|
||||||
|
output_data, opc_metrics = await opc.write_opc_data(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert output_data == {'data': 'data'}
|
||||||
|
assert opc_metrics == {'server1': {'tag1': 0.1, 'tag2': 0.2}}
|
||||||
|
opc.manage_output_tags.assert_called_once_with(
|
||||||
|
'server1',
|
||||||
|
input_data['opc_output_config']['server1'],
|
||||||
|
mock_dataframe.return_value,
|
||||||
|
metadata['metadata'],
|
||||||
|
)
|
||||||
|
opc.process_confidence.assert_called_once_with(
|
||||||
|
mock_dataframe.return_value,
|
||||||
|
True,
|
||||||
|
metadata['metadata'],
|
||||||
|
session_bad=False,
|
||||||
|
opc_status=None,
|
||||||
|
reconnect_in_progress=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_opc_data_empty_config(opc):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
|
||||||
|
'opc_servers': ['server1'],
|
||||||
|
'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
|
||||||
|
}
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await opc.write_opc_data(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
opc.opc_repository['server1'].write_data.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_write_opc_data_no_validate_server(opc):
|
||||||
|
opc.validate_server = AsyncMock(return_value=False)
|
||||||
|
input_data = {
|
||||||
|
**metadata,
|
||||||
|
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
|
||||||
|
'opc_output_config': {
|
||||||
|
'server1': {
|
||||||
|
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||||
|
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await opc.write_opc_data(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
opc.opc_repository['server1'].write_data.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.parametrize(
|
||||||
|
'data,success,expected',
|
||||||
|
[
|
||||||
|
(DataFrame({'prediction_confidence': [0]}), True, 0),
|
||||||
|
(DataFrame({'prediction_confidence': [0]}), False, 12),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_process_confidence(opc, data, success, expected):
|
||||||
|
result = opc.process_confidence(data, success, metadata['metadata'])
|
||||||
|
assert result['prediction_confidence'][0] == expected
|
||||||
|
|
||||||
|
|
||||||
|
def test_process_confidence_session_bad(opc):
|
||||||
|
data = DataFrame({'prediction_confidence': [0.9]})
|
||||||
|
result = opc.process_confidence(
|
||||||
|
data,
|
||||||
|
False,
|
||||||
|
metadata['metadata'],
|
||||||
|
session_bad=True,
|
||||||
|
opc_status='BadSessionIdInvalid',
|
||||||
|
)
|
||||||
|
assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
|
||||||
|
assert result['comments'][0].startswith(OPC_SESSION_BAD_COMMENT_PREFIX)
|
||||||
|
assert 'BadSessionIdInvalid' in result['comments'][0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_process_confidence_generic_failure(opc):
|
||||||
|
data = DataFrame({'prediction_confidence': [0.9]})
|
||||||
|
result = opc.process_confidence(data, False, metadata['metadata'])
|
||||||
|
assert result['prediction_confidence'][0] == OPC_WRITTING_ERROR_CONFIDENCE
|
||||||
|
assert result['comments'][0] == OPC_WRITTING_ERROR_MESSAGE
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_manage_output_tags_merges_error_flags(opc):
|
||||||
|
opc._write_tags_from_config = AsyncMock(
|
||||||
|
side_effect=[
|
||||||
|
({'tag1': None}, True, 'BadSessionIdInvalid', False),
|
||||||
|
({'tag2': 0.2}, False, None, True),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
|
||||||
|
config = {
|
||||||
|
'prediction_tags': {'tag1': {'data_type': 'float'}},
|
||||||
|
'confidence_tags': {'tag2': {'data_type': 'float'}},
|
||||||
|
}
|
||||||
|
|
||||||
|
(
|
||||||
|
success,
|
||||||
|
metrics,
|
||||||
|
session_bad_seen,
|
||||||
|
opc_status,
|
||||||
|
reconnect_in_progress,
|
||||||
|
) = await opc.manage_output_tags('server1', config, data, metadata['metadata'])
|
||||||
|
|
||||||
|
assert success is False
|
||||||
|
assert session_bad_seen is True
|
||||||
|
assert reconnect_in_progress is True
|
||||||
|
assert opc_status == 'BadSessionIdInvalid'
|
||||||
|
assert metrics == {'tag1': None, 'tag2': 0.2}
|
||||||
|
|
||||||
|
|
||||||
|
def test_process_confidence_reconnect_in_progress(opc):
|
||||||
|
data = DataFrame({'prediction_confidence': [0.9]})
|
||||||
|
result = opc.process_confidence(
|
||||||
|
data,
|
||||||
|
False,
|
||||||
|
metadata['metadata'],
|
||||||
|
reconnect_in_progress=True,
|
||||||
|
)
|
||||||
|
assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
|
||||||
|
assert result['comments'][0] == OPC_RECONNECT_IN_PROGRESS_COMMENT
|
||||||
|
|
||||||
|
|
||||||
|
def test_process_confidence_concatenates_multiple_comments(opc):
|
||||||
|
data = DataFrame({'prediction_confidence': [0.9]})
|
||||||
|
session_comment = f'{OPC_SESSION_BAD_COMMENT_PREFIX} BadSessionIdInvalid'
|
||||||
|
|
||||||
|
result = opc.process_confidence(
|
||||||
|
data,
|
||||||
|
False,
|
||||||
|
metadata['metadata'],
|
||||||
|
session_bad=True,
|
||||||
|
opc_status='BadSessionIdInvalid',
|
||||||
|
reconnect_in_progress=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
|
||||||
|
assert result['comments'][0] == OPC_COMMENT_SEPARATOR.join(
|
||||||
|
[session_comment, OPC_RECONNECT_IN_PROGRESS_COMMENT]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_validate_server(opc):
|
||||||
|
assert await opc.validate_server('server1', metadata) is True
|
||||||
|
assert await opc.validate_server('server2', metadata) is False
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
async def test_close(opc):
|
||||||
|
opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
|
||||||
|
await opc.aclose()
|
||||||
|
opc.opc_repository['server1'].disconnect.assert_called_once()
|
||||||
323
tests/laborious/activities/test_storage.py
Normal file
323
tests/laborious/activities/test_storage.py
Normal file
@@ -0,0 +1,323 @@
|
|||||||
|
import datetime
|
||||||
|
import os
|
||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark, raises
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
from sientia_do.temporal.activities.postgres import Postgres
|
||||||
|
|
||||||
|
from laborious.activities.storage import Storage
|
||||||
|
|
||||||
|
|
||||||
|
@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',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
@patch('laborious.activities.storage.MinioRepository')
|
||||||
|
def storage(mock_minio_repository):
|
||||||
|
return Storage(
|
||||||
|
host='localhost',
|
||||||
|
port=5432,
|
||||||
|
user='postgres',
|
||||||
|
password='postgres',
|
||||||
|
dbname='postgres',
|
||||||
|
min_connections=1,
|
||||||
|
max_connections=10,
|
||||||
|
retention_hours=24,
|
||||||
|
minio_repository=mock_minio_repository.return_value,
|
||||||
|
logger=MagicMock(),
|
||||||
|
notification_handler=MagicMock(),
|
||||||
|
metrics_controller=AsyncMock(),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@patch('laborious.activities.storage.MinioRepository')
|
||||||
|
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,
|
||||||
|
user='postgres',
|
||||||
|
password='postgres',
|
||||||
|
dbname='postgres',
|
||||||
|
min_connections=1,
|
||||||
|
max_connections=10,
|
||||||
|
retention_hours=24,
|
||||||
|
minio_repository=minio_repo,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
assert isinstance(storage, Postgres)
|
||||||
|
|
||||||
|
assert storage.minio_repository is minio_repo
|
||||||
|
mock_minio_repository.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@patch('laborious.activities.storage.MinioRepository')
|
||||||
|
def test___init___none_minio_repository(mock_minio_repository, storage):
|
||||||
|
storage.minio_repository = None
|
||||||
|
logger = MagicMock()
|
||||||
|
notification_handler = MagicMock()
|
||||||
|
metrics_controller = AsyncMock()
|
||||||
|
storage.__init__(
|
||||||
|
host='localhost',
|
||||||
|
port=5432,
|
||||||
|
user='postgres',
|
||||||
|
password='postgres',
|
||||||
|
dbname='postgres',
|
||||||
|
min_connections=1,
|
||||||
|
max_connections=10,
|
||||||
|
retention_hours=24,
|
||||||
|
minio_repository=None,
|
||||||
|
logger=logger,
|
||||||
|
notification_handler=notification_handler,
|
||||||
|
metrics_controller=metrics_controller,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert storage.minio_repository is None
|
||||||
|
mock_minio_repository.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@patch('laborious.activities.storage.MinioRepository')
|
||||||
|
def test___init___done_repository(mock_minio_repository, storage):
|
||||||
|
storage.__init__(
|
||||||
|
host='localhost',
|
||||||
|
port=5432,
|
||||||
|
user='postgres',
|
||||||
|
password='postgres',
|
||||||
|
dbname='postgres',
|
||||||
|
min_connections=1,
|
||||||
|
max_connections=10,
|
||||||
|
retention_hours=24,
|
||||||
|
minio_repository=mock_minio_repository.return_value,
|
||||||
|
logger=MagicMock(),
|
||||||
|
notification_handler=MagicMock(),
|
||||||
|
metrics_controller=AsyncMock(),
|
||||||
|
)
|
||||||
|
mock_minio_repository.assert_not_called()
|
||||||
|
assert storage.minio_repository is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_close(storage):
|
||||||
|
storage.minio_repository = MagicMock()
|
||||||
|
|
||||||
|
storage.close()
|
||||||
|
|
||||||
|
assert storage.minio_repository is None
|
||||||
|
|
||||||
|
|
||||||
|
def test___del__(storage):
|
||||||
|
storage.close = MagicMock()
|
||||||
|
|
||||||
|
storage.__del__()
|
||||||
|
|
||||||
|
storage.close.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@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()
|
||||||
0
tests/laborious/utils/__init__.py
Normal file
0
tests/laborious/utils/__init__.py
Normal file
0
tests/laborious/utils/filters/__init__.py
Normal file
0
tests/laborious/utils/filters/__init__.py
Normal file
44
tests/laborious/utils/filters/test_conditional_filters.py
Normal file
44
tests/laborious/utils/filters/test_conditional_filters.py
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
from laborious.utils.filters.conditional_filters import (
|
||||||
|
filter_empty_data,
|
||||||
|
filter_specific_variables_null_values,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_specific_variables_null_values():
|
||||||
|
assert (
|
||||||
|
filter_specific_variables_null_values(
|
||||||
|
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
|
||||||
|
config={'variables': ['variable2']},
|
||||||
|
)
|
||||||
|
is False
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_specific_variables_null_values_with_empty_data():
|
||||||
|
assert (
|
||||||
|
filter_specific_variables_null_values(DataFrame(), config={'variables': ['variable2']})
|
||||||
|
is False
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_specific_variables_null_values_with_null_values():
|
||||||
|
assert (
|
||||||
|
filter_specific_variables_null_values(
|
||||||
|
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
|
||||||
|
config={'variables': ['variable2']},
|
||||||
|
)
|
||||||
|
is True
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_empty_data():
|
||||||
|
assert filter_empty_data(DataFrame(), {}) is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_empty_data_with_data():
|
||||||
|
assert (
|
||||||
|
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
|
||||||
|
is False
|
||||||
|
)
|
||||||
23
tests/laborious/utils/filters/test_mlflow_filters.py
Normal file
23
tests/laborious/utils/filters/test_mlflow_filters.py
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
|
||||||
|
|
||||||
|
|
||||||
|
def test_api_error_filter_invalid_response():
|
||||||
|
assert api_error_filter(None, {}) is True # NOSONAR
|
||||||
|
|
||||||
|
|
||||||
|
def test_api_error_filter_valid_response_fail():
|
||||||
|
assert api_error_filter({'success': False}, {}) is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_api_error_filter_valid_response_success():
|
||||||
|
assert api_error_filter({'success': True}, {}) is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_nan_values_filter_all_nan_values():
|
||||||
|
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_nan_values_filter_no_nan_values():
|
||||||
|
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False
|
||||||
266
tests/laborious/utils/models/test_minio_dataframe_payload.py
Normal file
266
tests/laborious/utils/models/test_minio_dataframe_payload.py
Normal 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
|
||||||
2121
tests/laborious/utils/repository/test_model_repository.py
Normal file
2121
tests/laborious/utils/repository/test_model_repository.py
Normal file
File diff suppressed because it is too large
Load Diff
610
tests/laborious/utils/repository/test_opc_repository.py
Normal file
610
tests/laborious/utils/repository/test_opc_repository.py
Normal file
@@ -0,0 +1,610 @@
|
|||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from datetime import datetime
|
||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from asyncua.crypto.security_policies import SecurityPolicyBasic256
|
||||||
|
from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
|
||||||
|
from laborious.utils.repository.opc_repository import (
|
||||||
|
OpcClientAlreadyExistsError,
|
||||||
|
OpcClientNotInitializedError,
|
||||||
|
OpcRepository,
|
||||||
|
OpcSessionAlreadyConnectedError,
|
||||||
|
is_reconnectable_opcua_bad,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_logger():
|
||||||
|
return Mock()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def opc_repository(mock_logger):
|
||||||
|
repository = OpcRepository(
|
||||||
|
opc_id='test_repo',
|
||||||
|
server_name='test_server',
|
||||||
|
url='opc.tcp://localhost:4840',
|
||||||
|
logger=mock_logger,
|
||||||
|
notification_handler=Mock(),
|
||||||
|
reconnection_interval=60,
|
||||||
|
server_uri='urn:test:server',
|
||||||
|
cert_path='/path/to/cert.pem',
|
||||||
|
private_key_path='/path/to/key.pem',
|
||||||
|
server_cert_path='/path/to/server_cert.pem',
|
||||||
|
metrics_controller=AsyncMock(),
|
||||||
|
)
|
||||||
|
repository.disconnection_interval = 0.1
|
||||||
|
repository.send_notification = MagicMock()
|
||||||
|
repository.send_notification_async = AsyncMock()
|
||||||
|
repository.emit_metric = AsyncMock()
|
||||||
|
repository.info = MagicMock()
|
||||||
|
repository.error = MagicMock()
|
||||||
|
repository.warning = MagicMock()
|
||||||
|
repository.debug = MagicMock()
|
||||||
|
repository._session_ready.set()
|
||||||
|
return repository
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_client():
|
||||||
|
with patch('laborious.utils.repository.opc_repository.Client') as mock:
|
||||||
|
client_instance = AsyncMock()
|
||||||
|
mock.return_value = client_instance
|
||||||
|
yield client_instance
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schema_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_init(opc_repository):
|
||||||
|
assert opc_repository.id == 'test_repo'
|
||||||
|
assert opc_repository.server_name == 'test_server'
|
||||||
|
assert opc_repository.url == 'opc.tcp://localhost:4840'
|
||||||
|
assert opc_repository.server_uri == 'urn:test:server'
|
||||||
|
assert opc_repository.cert_path == '/path/to/cert.pem'
|
||||||
|
assert opc_repository.private_key_path == '/path/to/key.pem'
|
||||||
|
assert opc_repository.server_cert_path == '/path/to/server_cert.pem'
|
||||||
|
assert opc_repository.reconnection_interval == 60
|
||||||
|
assert opc_repository.client is None
|
||||||
|
assert opc_repository.last_reconnection_time is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_set_security(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
await opc_repository.set_security()
|
||||||
|
|
||||||
|
mock_client.application_uri = 'urn:test:server'
|
||||||
|
mock_client.set_security.assert_called_once_with(
|
||||||
|
SecurityPolicyBasic256,
|
||||||
|
certificate='/path/to/cert.pem',
|
||||||
|
private_key='/path/to/key.pem',
|
||||||
|
server_certificate='/path/to/server_cert.pem',
|
||||||
|
)
|
||||||
|
assert mock_client.secure_channel_timeout == 600_000
|
||||||
|
assert mock_client.session_timeout == 600_000
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_set_security_missing_certificates(opc_repository):
|
||||||
|
opc_repository.cert_path = None
|
||||||
|
opc_repository.private_key_path = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
await opc_repository.set_security()
|
||||||
|
except ValueError as e:
|
||||||
|
assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_set_security_missing_client(opc_repository):
|
||||||
|
opc_repository.client = None
|
||||||
|
try:
|
||||||
|
await opc_repository.set_security()
|
||||||
|
except ValueError as e:
|
||||||
|
assert str(e) == 'Client must be initialized before setting security'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_connect_with_security(opc_repository, mock_client):
|
||||||
|
opc_repository._create_client = AsyncMock()
|
||||||
|
opc_repository._open_session = AsyncMock(return_value=(True, {}))
|
||||||
|
result = await opc_repository.connect()
|
||||||
|
|
||||||
|
opc_repository._create_client.assert_called_once()
|
||||||
|
opc_repository._open_session.assert_called_once()
|
||||||
|
assert result == (True, {})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_connect_without_security(opc_repository, mock_client):
|
||||||
|
opc_repository.cert_path = None
|
||||||
|
opc_repository._create_client = AsyncMock()
|
||||||
|
opc_repository._open_session = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.set_security = AsyncMock()
|
||||||
|
result = await opc_repository.connect()
|
||||||
|
|
||||||
|
opc_repository._create_client.assert_called_once()
|
||||||
|
opc_repository._open_session.assert_called_once()
|
||||||
|
opc_repository.set_security.assert_not_called()
|
||||||
|
assert result == (True, {})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_connect_raises_when_session_already_open(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
proto = MagicMock()
|
||||||
|
proto.state = 'open'
|
||||||
|
mock_client.uaclient = MagicMock(protocol=proto)
|
||||||
|
|
||||||
|
with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
|
||||||
|
await opc_repository.connect()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_create_client_raises_when_client_exists(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
|
||||||
|
with pytest.raises(OpcClientAlreadyExistsError, match='already exists'):
|
||||||
|
await opc_repository._create_client()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_open_session_success(opc_repository):
|
||||||
|
closed_proto = MagicMock()
|
||||||
|
closed_proto.state = 'closed'
|
||||||
|
opc_repository.client = AsyncMock()
|
||||||
|
opc_repository.client.uaclient = MagicMock(protocol=closed_proto)
|
||||||
|
opc_repository.client.session_timeout = 600_000
|
||||||
|
opc_repository.client.secure_channel_timeout = 600_000
|
||||||
|
|
||||||
|
open_proto = MagicMock()
|
||||||
|
open_proto.state = 'open'
|
||||||
|
open_proto.authentication_token = 'tok'
|
||||||
|
|
||||||
|
async def connect_side_effect():
|
||||||
|
opc_repository.client.uaclient.protocol = open_proto
|
||||||
|
|
||||||
|
opc_repository.client.connect = AsyncMock(side_effect=connect_side_effect)
|
||||||
|
|
||||||
|
result = await opc_repository._open_session()
|
||||||
|
|
||||||
|
opc_repository.client.connect.assert_called_once()
|
||||||
|
assert opc_repository.last_reconnection_time is None
|
||||||
|
assert result == (True, {})
|
||||||
|
assert opc_repository._session_ready.is_set()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_open_session_raises_when_already_connected(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
proto = MagicMock()
|
||||||
|
proto.state = 'open'
|
||||||
|
mock_client.uaclient = MagicMock(protocol=proto)
|
||||||
|
|
||||||
|
with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
|
||||||
|
await opc_repository._open_session()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_open_session_fail(opc_repository):
|
||||||
|
opc_repository._disconnect_locked = AsyncMock()
|
||||||
|
opc_repository.client = MagicMock()
|
||||||
|
opc_repository.client.uaclient = MagicMock(protocol=MagicMock(state='closed'))
|
||||||
|
opc_repository.client.connect = AsyncMock(side_effect=Exception('Test error'))
|
||||||
|
|
||||||
|
is_connected, error_data = await opc_repository._open_session()
|
||||||
|
|
||||||
|
opc_repository._disconnect_locked.assert_called_once()
|
||||||
|
opc_repository.client.connect.assert_called_once()
|
||||||
|
assert is_connected is False
|
||||||
|
assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'
|
||||||
|
assert error_data['message'] == 'Failed to connect to OPC server: Test error'
|
||||||
|
assert error_data['block'] == 'opc_repository'
|
||||||
|
assert error_data['level'] == NotificationLevel.ERROR
|
||||||
|
assert error_data['attachment_content'] is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_open_session_raises_when_no_client(opc_repository):
|
||||||
|
opc_repository.client = None
|
||||||
|
|
||||||
|
with pytest.raises(OpcClientNotInitializedError, match='not initialized'):
|
||||||
|
await opc_repository._open_session()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_disconnection_fallback_success(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
mock_client.disconnect.return_value = True
|
||||||
|
result = await opc_repository._disconnection_fallback()
|
||||||
|
|
||||||
|
mock_client.disconnect.assert_called_once()
|
||||||
|
assert result == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_disconnection_fallback_fail(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
mock_client.disconnect.side_effect = Exception('Test error')
|
||||||
|
result = await opc_repository._disconnection_fallback()
|
||||||
|
assert result == [
|
||||||
|
{'attempt': 1, 'error': 'Test error', 'traceback': ANY},
|
||||||
|
{'attempt': 2, 'error': 'Test error', 'traceback': ANY},
|
||||||
|
{'attempt': 3, 'error': 'Test error', 'traceback': ANY},
|
||||||
|
{'attempt': 4, 'error': 'Test error', 'traceback': ANY},
|
||||||
|
{'attempt': 5, 'error': 'Test error', 'traceback': ANY},
|
||||||
|
]
|
||||||
|
assert mock_client.disconnect.call_count == 5
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_disconnect(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
opc_repository._disconnection_fallback = AsyncMock(return_value=[])
|
||||||
|
await opc_repository.disconnect()
|
||||||
|
|
||||||
|
opc_repository._disconnection_fallback.assert_called_once()
|
||||||
|
assert opc_repository.client is None
|
||||||
|
assert opc_repository._allow_reconnect is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_disconnect_no_client(opc_repository):
|
||||||
|
opc_repository.client = None
|
||||||
|
assert await opc_repository.disconnect() is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_disconnect_error(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
opc_repository._disconnection_fallback = AsyncMock(
|
||||||
|
return_value=[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}]
|
||||||
|
)
|
||||||
|
await opc_repository.disconnect()
|
||||||
|
|
||||||
|
opc_repository._disconnection_fallback.assert_called_once()
|
||||||
|
opc_repository.send_notification_async.assert_called_once_with(
|
||||||
|
metadata=opc_repository.metadata,
|
||||||
|
notification_id=f'OPC_DISCONNECTION_ERROR_{opc_repository.id}',
|
||||||
|
message='Failed to disconnect from OPC server in 5 attempts.',
|
||||||
|
block='opc_repository',
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=json.dumps(
|
||||||
|
[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}], indent=4
|
||||||
|
),
|
||||||
|
)
|
||||||
|
assert opc_repository.client is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_validate_connection_none_client(opc_repository):
|
||||||
|
opc_repository.client = None
|
||||||
|
response = await opc_repository.validate_connection()
|
||||||
|
assert response == (False, opc_repository._not_connected_error())
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_validate_connection_session_not_open(opc_repository):
|
||||||
|
opc_repository.client = MagicMock()
|
||||||
|
opc_repository.client.uaclient.protocol = None
|
||||||
|
|
||||||
|
response = await opc_repository.validate_connection()
|
||||||
|
|
||||||
|
assert response == (False, opc_repository._not_connected_error())
|
||||||
|
opc_repository.error.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_validate_connection_success(opc_repository):
|
||||||
|
opc_repository.client = MagicMock()
|
||||||
|
opc_repository.client.uaclient.protocol = MagicMock()
|
||||||
|
opc_repository.client.uaclient.protocol.state = 'open'
|
||||||
|
|
||||||
|
output = await opc_repository.validate_connection()
|
||||||
|
assert output == (True, {})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_validate_connection_do_nothing(opc_repository):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = AsyncMock(get_node=MagicMock())
|
||||||
|
mock_node = AsyncMock()
|
||||||
|
opc_repository.client.get_node.return_value = mock_node
|
||||||
|
|
||||||
|
result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||||
|
|
||||||
|
opc_repository.validate_connection.assert_called_once()
|
||||||
|
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
|
||||||
|
assert result == (True, {'response_time': ANY})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_validate_connection_failed(opc_repository):
|
||||||
|
opc_repository.client = MagicMock()
|
||||||
|
opc_repository.client.uaclient.protocol = MagicMock(state='closed')
|
||||||
|
opc_repository._start_reconnect = AsyncMock()
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_repository._start_reconnect.assert_called_once()
|
||||||
|
assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||||
|
assert error_data['opc_status'] == 'ProtocolClosed'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_get_node_failed(opc_repository):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = AsyncMock()
|
||||||
|
opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_repository.validate_connection.assert_called_once()
|
||||||
|
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['notification_id'] == f'OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}'
|
||||||
|
assert (
|
||||||
|
error_data['message']
|
||||||
|
== "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
|
||||||
|
)
|
||||||
|
assert error_data['block'] == 'opc_repository'
|
||||||
|
assert error_data['level'] == NotificationLevel.ERROR
|
||||||
|
assert error_data['attachment_content'] is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_invalid_data_type(opc_repository, mock_client):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
mock_node = AsyncMock()
|
||||||
|
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'invalid_type', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_repository.validate_connection.assert_called_once()
|
||||||
|
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
|
||||||
|
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['notification_id'] == f'OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}'
|
||||||
|
assert (
|
||||||
|
error_data['message']
|
||||||
|
== "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
|
||||||
|
)
|
||||||
|
assert error_data['block'] == 'opc_repository'
|
||||||
|
assert error_data['level'] == NotificationLevel.ERROR
|
||||||
|
assert error_data.get('attachment_content') is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data(opc_repository, mock_client):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
mock_node = AsyncMock()
|
||||||
|
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||||
|
|
||||||
|
result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||||
|
|
||||||
|
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
|
||||||
|
mock_node.write_value.assert_called_once()
|
||||||
|
assert result == (True, {'response_time': ANY})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_write_value_failed(opc_repository, mock_client):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
mock_node = AsyncMock()
|
||||||
|
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||||
|
mock_node.write_value.side_effect = Exception('Test error')
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_repository.validate_connection.assert_called_once()
|
||||||
|
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
|
||||||
|
mock_node.write_value.assert_called_once()
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['notification_id'] == f'OPC_WRITE_DATA_ERROR_{opc_repository.id}'
|
||||||
|
assert (
|
||||||
|
error_data['message']
|
||||||
|
== "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
|
||||||
|
)
|
||||||
|
assert error_data['block'] == 'opc_repository'
|
||||||
|
assert error_data['level'] == NotificationLevel.ERROR
|
||||||
|
assert error_data['attachment_content'] is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_is_reconnectable_opcua_bad():
|
||||||
|
assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
|
||||||
|
assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
|
||||||
|
assert is_reconnectable_opcua_bad(Exception('other')) is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_bad_session_id_invalid_schedules_reconnect(opc_repository, mock_client):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
opc_repository._start_reconnect = AsyncMock()
|
||||||
|
mock_node = AsyncMock()
|
||||||
|
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||||
|
mock_node.write_value.side_effect = BadSessionIdInvalid()
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_node.write_value.assert_called_once()
|
||||||
|
opc_repository._start_reconnect.assert_called_once()
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['opc_error_kind'] == 'session_bad'
|
||||||
|
assert error_data['opc_status'] == 'BadSessionIdInvalid'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_reconnect_in_progress_immediate(opc_repository):
|
||||||
|
opc_repository._session_ready.clear()
|
||||||
|
opc_repository._reconnect_task = asyncio.create_task(asyncio.sleep(60))
|
||||||
|
opc_repository.validate_connection = AsyncMock()
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_repository._reconnect_task.cancel()
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await opc_repository._reconnect_task
|
||||||
|
opc_repository._reconnect_task = None
|
||||||
|
|
||||||
|
opc_repository.validate_connection.assert_not_called()
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['opc_error_kind'] == 'reconnect_in_progress'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_start_reconnect_skips_within_interval(opc_repository):
|
||||||
|
opc_repository.last_reconnection_time = datetime.now()
|
||||||
|
opc_repository.reconnection_interval = 3600
|
||||||
|
|
||||||
|
await opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
|
||||||
|
|
||||||
|
assert opc_repository._reconnect_task is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
|
||||||
|
opc_repository.client = MagicMock()
|
||||||
|
opc_repository.client.uaclient.protocol = MagicMock(state='closed')
|
||||||
|
opc_repository._start_reconnect = AsyncMock()
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
opc_repository._start_reconnect.assert_called_once()
|
||||||
|
assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||||
|
assert error_data['opc_status'] == 'ProtocolClosed'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_protocol_closed_skips_reconnect_within_interval(opc_repository):
|
||||||
|
opc_repository.client = MagicMock()
|
||||||
|
opc_repository.client.uaclient.protocol = MagicMock(state='closed')
|
||||||
|
opc_repository.last_reconnection_time = datetime.now()
|
||||||
|
opc_repository.reconnection_interval = 3600
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
assert opc_repository._reconnect_task is None
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_after_failed_reconnect_schedules_again(opc_repository):
|
||||||
|
opc_repository._session_ready.clear()
|
||||||
|
opc_repository.reconnection_interval = 0
|
||||||
|
opc_repository.last_reconnection_time = None
|
||||||
|
opc_repository._reconnect_locked = AsyncMock(
|
||||||
|
return_value=(False, {'message': 'connect failed'})
|
||||||
|
)
|
||||||
|
|
||||||
|
await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||||
|
await asyncio.sleep(0.1)
|
||||||
|
assert opc_repository._reconnect_locked.call_count == 1
|
||||||
|
assert not opc_repository._reconnect_task_in_progress()
|
||||||
|
|
||||||
|
await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||||
|
await asyncio.sleep(0.1)
|
||||||
|
assert opc_repository._reconnect_locked.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_write_data_after_disconnect_does_not_schedule_reconnect(opc_repository, mock_client):
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
proto = MagicMock()
|
||||||
|
proto.state = 'closed'
|
||||||
|
mock_client.uaclient = MagicMock(protocol=proto)
|
||||||
|
opc_repository._disconnection_fallback = AsyncMock(return_value=[])
|
||||||
|
await opc_repository.disconnect()
|
||||||
|
|
||||||
|
is_success, error_data = await opc_repository.write_data(
|
||||||
|
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||||
|
)
|
||||||
|
|
||||||
|
assert opc_repository._reconnect_task is None
|
||||||
|
assert is_success is False
|
||||||
|
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parallel_bad_writes_single_reconnect_task(opc_repository, mock_client):
|
||||||
|
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||||
|
opc_repository.client = mock_client
|
||||||
|
opc_repository.reconnection_interval = 0
|
||||||
|
opc_repository.last_reconnection_time = None
|
||||||
|
mock_node = AsyncMock()
|
||||||
|
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||||
|
mock_node.write_value.side_effect = BadSessionIdInvalid()
|
||||||
|
|
||||||
|
connect_count = 0
|
||||||
|
|
||||||
|
async def slow_reconnect():
|
||||||
|
nonlocal connect_count
|
||||||
|
connect_count += 1
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
opc_repository._session_ready.set()
|
||||||
|
return True, {}
|
||||||
|
|
||||||
|
opc_repository._reconnect_locked = slow_reconnect
|
||||||
|
|
||||||
|
results = await asyncio.gather(
|
||||||
|
opc_repository.write_data('ns=2;s=TestNode', 1.0, 'float', metadata['metadata']),
|
||||||
|
opc_repository.write_data('ns=2;s=TestNode2', 2.0, 'float', metadata['metadata']),
|
||||||
|
)
|
||||||
|
await asyncio.sleep(0.15)
|
||||||
|
|
||||||
|
assert connect_count <= 1
|
||||||
|
assert 1 <= mock_node.write_value.call_count <= 2
|
||||||
|
error_kinds = [r[1].get('opc_error_kind') for r in results]
|
||||||
|
assert error_kinds.count('session_bad') >= 1
|
||||||
|
assert all(k in ('session_bad', 'reconnect_in_progress') for k in error_kinds)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@patch('laborious.utils.repository.opc_repository.datetime')
|
||||||
|
async def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_repository):
|
||||||
|
mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 12, 0, 0))
|
||||||
|
opc_repository._disconnect_locked = AsyncMock()
|
||||||
|
opc_repository._connect_locked = AsyncMock(return_value=(True, {}))
|
||||||
|
|
||||||
|
result = await opc_repository._reconnect_locked()
|
||||||
|
|
||||||
|
opc_repository._disconnect_locked.assert_called_once()
|
||||||
|
opc_repository._connect_locked.assert_called_once()
|
||||||
|
assert result == (True, {})
|
||||||
|
assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)
|
||||||
122
tests/laborious/utils/test_connectors_config.py
Normal file
122
tests/laborious/utils/test_connectors_config.py
Normal file
@@ -0,0 +1,122 @@
|
|||||||
|
from os import environ
|
||||||
|
|
||||||
|
from laborious.utils.connectors_config import (
|
||||||
|
build_minio_config,
|
||||||
|
build_mlflow_config,
|
||||||
|
build_opc_config,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_mlflow_config_with_env_vars():
|
||||||
|
# Arrange
|
||||||
|
environ['MLFLOW_HOST'] = 'http://test-host'
|
||||||
|
environ['MLFLOW_PORT'] = '8080'
|
||||||
|
environ['MLFLOW_USERNAME'] = 'test-user'
|
||||||
|
environ['MLFLOW_PASSWORD'] = 'test-pass'
|
||||||
|
|
||||||
|
# Act
|
||||||
|
config = build_mlflow_config()
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert config['host'] == 'http://test-host'
|
||||||
|
assert config['port'] == 8080
|
||||||
|
assert config['username'] == 'test-user'
|
||||||
|
assert config['password'] == 'test-pass'
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_mlflow_config_with_defaults():
|
||||||
|
# Arrange
|
||||||
|
# Clear any existing env vars
|
||||||
|
environ.pop('MLFLOW_HOST', None)
|
||||||
|
environ.pop('MLFLOW_PORT', None)
|
||||||
|
environ.pop('MLFLOW_USERNAME', None)
|
||||||
|
environ.pop('MLFLOW_PASSWORD', None)
|
||||||
|
|
||||||
|
# Act
|
||||||
|
config = build_mlflow_config()
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert config['host'] == 'http://localhost'
|
||||||
|
assert config['port'] == 5080
|
||||||
|
assert config['username'] == 'aignosi'
|
||||||
|
assert config['password'] == 'aignosi'
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_opc_config_with_env_vars():
|
||||||
|
# Arrange
|
||||||
|
environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
|
||||||
|
|
||||||
|
# Act
|
||||||
|
config = build_opc_config()
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert config['opc']['name'] == 'test-opc'
|
||||||
|
assert config['opc']['url'] == 'opc.tcp://test:4840'
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_opc_config_with_individual_env_vars():
|
||||||
|
# Arrange
|
||||||
|
environ.pop('OPC_CONFIG', None)
|
||||||
|
environ['OPC_ID'] = '1'
|
||||||
|
environ['OPC_URL'] = 'opc.tcp://test:4840'
|
||||||
|
environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
|
||||||
|
environ['OPC_RECONNECTION_INTERVAL'] = '300'
|
||||||
|
|
||||||
|
# Act
|
||||||
|
config = build_opc_config()
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert config['1']['id'] == '1'
|
||||||
|
assert config['1']['url'] == 'opc.tcp://test:4840'
|
||||||
|
assert config['1']['server_uri'] == 'opc.tcp://test:4840'
|
||||||
|
assert config['1']['reconnection_interval'] == 300
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_opc_config_with_defaults():
|
||||||
|
# Arrange
|
||||||
|
environ.pop('OPC_CONFIG', None)
|
||||||
|
environ.pop('OPC_ID', None)
|
||||||
|
environ.pop('OPC_URL', None)
|
||||||
|
environ.pop('OPC_SERVER_URI', None)
|
||||||
|
environ.pop('OPC_RECONNECTION_INTERVAL', None)
|
||||||
|
|
||||||
|
# Act
|
||||||
|
config = build_opc_config()
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert config['1']['id'] == '1'
|
||||||
|
assert config['1']['url'] == 'opc.tcp://localhost:4840'
|
||||||
|
assert config['1']['server_uri'] == 'opc.tcp://localhost:4840'
|
||||||
|
assert config['1']['reconnection_interval'] == 120
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_minio_config_with_env_vars():
|
||||||
|
environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
|
||||||
|
environ['MINIO_ACCESS_KEY'] = 'test-key'
|
||||||
|
environ['MINIO_SECRET_KEY'] = 'test-secret'
|
||||||
|
environ['MINIO_REGION_NAME'] = 'test-region'
|
||||||
|
environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
|
||||||
|
assert build_minio_config() == {
|
||||||
|
'endpoint_url': 'http://test-host',
|
||||||
|
'access_key': 'test-key',
|
||||||
|
'secret_key': 'test-secret',
|
||||||
|
'default_bucket': 'test-bucket',
|
||||||
|
'retention_hours': 24,
|
||||||
|
'secure': False,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_minio_config_with_defaults():
|
||||||
|
environ.pop('MINIO_ENDPOINT_URL', None)
|
||||||
|
environ.pop('MINIO_ACCESS_KEY', None)
|
||||||
|
environ.pop('MINIO_SECRET_KEY', None)
|
||||||
|
environ.pop('MINIO_REGION_NAME', None)
|
||||||
|
environ.pop('MINIO_DEFAULT_BUCKET', None)
|
||||||
|
assert build_minio_config() == {
|
||||||
|
'endpoint_url': 'http://localhost:9000',
|
||||||
|
'access_key': 'minioadmin',
|
||||||
|
'secret_key': 'minioadmin',
|
||||||
|
'default_bucket': 'laborious',
|
||||||
|
'retention_hours': 24,
|
||||||
|
'secure': False,
|
||||||
|
}
|
||||||
14
tests/laborious/worker/test_runtime_task_queues.py
Normal file
14
tests/laborious/worker/test_runtime_task_queues.py
Normal file
@@ -0,0 +1,14 @@
|
|||||||
|
from sientia_do.temporal.worker.prepare_worker import build_queue_name
|
||||||
|
|
||||||
|
from laborious.workflows.drift import Drift
|
||||||
|
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||||
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||||
|
from laborious.workflows.simple_metrics import SimpleMetrics
|
||||||
|
|
||||||
|
|
||||||
|
def test_runtime_scoped_queue_names():
|
||||||
|
runtime = 'prod-a'
|
||||||
|
assert build_queue_name(PredictionsBatch.__name__, runtime) == 'predictions_batch-prod-a-queue'
|
||||||
|
assert build_queue_name(MinimalRetrain.__name__, runtime) == 'minimal_retrain-prod-a-queue'
|
||||||
|
assert build_queue_name(Drift.__name__, runtime) == 'drift-prod-a-queue'
|
||||||
|
assert build_queue_name(SimpleMetrics.__name__, runtime) == 'simple_metrics-prod-a-queue'
|
||||||
@@ -0,0 +1,679 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def format_and_export_prediction():
|
||||||
|
return FormatAndExportPrediction()
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schema_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
)
|
||||||
|
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
|
||||||
|
input_data = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': None,
|
||||||
|
'data': {'test': 'data'},
|
||||||
|
'timestamp': '2021-01-01',
|
||||||
|
'model_id': 1,
|
||||||
|
'model_name': metadata['metadata']['model_name'],
|
||||||
|
'prediction_confidence': 0,
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'table_name': 'test_table',
|
||||||
|
'opc_servers': ['test_server'],
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'prediction_store_policy': 'erl:1',
|
||||||
|
}
|
||||||
|
|
||||||
|
prediction_data = MagicMock()
|
||||||
|
opc_metrics = MagicMock()
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.side_effect = [
|
||||||
|
(prediction_data, opc_metrics),
|
||||||
|
MagicMock(),
|
||||||
|
MagicMock(),
|
||||||
|
]
|
||||||
|
|
||||||
|
await format_and_export_prediction.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
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'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_opc_data,
|
||||||
|
{
|
||||||
|
'opc_output_config': input_data['opc_output_config'],
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': prediction_data,
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'prediction': prediction_data,
|
||||||
|
'opc_metrics': opc_metrics,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 3
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
)
|
||||||
|
async def test_run_none_path_flag_with_transformed_data(
|
||||||
|
workflow_mock, format_and_export_prediction
|
||||||
|
):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': None,
|
||||||
|
'data': {'test': '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',
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'opc_servers': ['test_server'],
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'prediction_store_policy': 'lts:1',
|
||||||
|
}
|
||||||
|
|
||||||
|
prediction_data = MagicMock()
|
||||||
|
opc_metrics = MagicMock()
|
||||||
|
transformed_data = MagicMock()
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.side_effect = [
|
||||||
|
prediction_data, # format_prediction
|
||||||
|
transformed_data, # format_transformed_data
|
||||||
|
]
|
||||||
|
|
||||||
|
write_transformed_handler = AsyncMock()
|
||||||
|
workflow_mock.start_activity_method.return_value = write_transformed_handler
|
||||||
|
workflow_mock.execute_activity_method.side_effect = [
|
||||||
|
(prediction_data, opc_metrics), # write_opc_data
|
||||||
|
MagicMock(), # export_data_to_postgres (prediction)
|
||||||
|
MagicMock(), # write_metrics
|
||||||
|
]
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await format_and_export_prediction.run(input_data)
|
||||||
|
|
||||||
|
# Assert - format_prediction call
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
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'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
Activities.format_transformed_data,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': input_data['transformed_data'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert - start_activity_method for transformed data export
|
||||||
|
workflow_mock.start_activity_method.assert_called_once_with(
|
||||||
|
Activities.export_payload_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['transform_table_name'],
|
||||||
|
'data': transformed_data,
|
||||||
|
'timestamp_conversion': {
|
||||||
|
'column': 'timestamp',
|
||||||
|
'format': DATETIME_FORMAT_WITH_TZ,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert - write_opc_data call
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_opc_data,
|
||||||
|
{
|
||||||
|
'opc_output_config': input_data['opc_output_config'],
|
||||||
|
'data': prediction_data,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert - export_data_to_postgres for prediction call
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': prediction_data,
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert - write_metrics call
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'prediction': prediction_data,
|
||||||
|
'opc_metrics': opc_metrics,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert - verify counts
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 3
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 2
|
||||||
|
assert workflow_mock.start_activity_method.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
)
|
||||||
|
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
|
||||||
|
input_data = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': 'default',
|
||||||
|
'data': {'test': 'data'},
|
||||||
|
'timestamp': '2021-01-01',
|
||||||
|
'model_id': 1,
|
||||||
|
'model_name': metadata['metadata']['model_name'],
|
||||||
|
'prediction_confidence': 0,
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'table_name': 'test_table',
|
||||||
|
'opc_servers': ['test_server'],
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'comment': 'test_comment',
|
||||||
|
}
|
||||||
|
|
||||||
|
prediction_data = MagicMock()
|
||||||
|
opc_metrics = MagicMock()
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.side_effect = [
|
||||||
|
(prediction_data, opc_metrics),
|
||||||
|
MagicMock(),
|
||||||
|
MagicMock(),
|
||||||
|
]
|
||||||
|
|
||||||
|
await format_and_export_prediction.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
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'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_opc_data,
|
||||||
|
{
|
||||||
|
'opc_output_config': input_data['opc_output_config'],
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': prediction_data,
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'prediction': prediction_data,
|
||||||
|
'opc_metrics': opc_metrics,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 3
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
)
|
||||||
|
async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
|
||||||
|
input_data = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': None,
|
||||||
|
'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',
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
'prediction_store_policy': 'erl:1',
|
||||||
|
}
|
||||||
|
|
||||||
|
pi_web_api_data = MagicMock()
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.side_effect = [
|
||||||
|
pi_web_api_data, # write_pi_web_api_data
|
||||||
|
MagicMock(), # export_data_to_postgres
|
||||||
|
MagicMock(), # write_metrics
|
||||||
|
]
|
||||||
|
|
||||||
|
await format_and_export_prediction.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
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'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_pi_web_api_data,
|
||||||
|
{
|
||||||
|
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': pi_web_api_data,
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'prediction': pi_web_api_data,
|
||||||
|
'opc_metrics': {},
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 3
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
)
|
||||||
|
async def test_run_none_path_flag_with_pi_web_api_and_opc(
|
||||||
|
workflow_mock, format_and_export_prediction
|
||||||
|
):
|
||||||
|
input_data = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': None,
|
||||||
|
'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',
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
'prediction_store_policy': 'erl:1',
|
||||||
|
}
|
||||||
|
|
||||||
|
prediction_data = MagicMock()
|
||||||
|
pi_web_api_data = MagicMock()
|
||||||
|
opc_metrics = MagicMock()
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.side_effect = [
|
||||||
|
pi_web_api_data, # write_pi_web_api_data
|
||||||
|
(prediction_data, opc_metrics), # write_opc_data
|
||||||
|
MagicMock(), # export_data_to_postgres
|
||||||
|
MagicMock(), # write_metrics
|
||||||
|
]
|
||||||
|
|
||||||
|
await format_and_export_prediction.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_pi_web_api_data,
|
||||||
|
{
|
||||||
|
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
Activities.write_opc_data,
|
||||||
|
{
|
||||||
|
'opc_output_config': input_data['opc_output_config'],
|
||||||
|
'data': pi_web_api_data,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': prediction_data,
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_metrics,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'prediction': prediction_data,
|
||||||
|
'opc_metrics': opc_metrics,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 4
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch(
|
||||||
|
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
)
|
||||||
|
async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
|
||||||
|
input_data = {
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': 'default',
|
||||||
|
'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',
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
'comment': 'test_comment',
|
||||||
|
}
|
||||||
|
|
||||||
|
pi_web_api_data = MagicMock()
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.side_effect = [
|
||||||
|
pi_web_api_data, # write_pi_web_api_data
|
||||||
|
MagicMock(), # export_data_to_postgres
|
||||||
|
MagicMock(), # write_metrics
|
||||||
|
]
|
||||||
|
|
||||||
|
await format_and_export_prediction.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.write_pi_web_api_data,
|
||||||
|
{
|
||||||
|
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'data': pi_web_api_data,
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 3
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
@@ -0,0 +1,842 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def prediction_process():
|
||||||
|
return PredictionProcess()
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'test_workflow',
|
||||||
|
'schema_name': 'test_schedule',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run(workflow_mock, prediction_process):
|
||||||
|
prediction_process.path_flag_handler = AsyncMock(return_value=False)
|
||||||
|
# Arrange
|
||||||
|
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',
|
||||||
|
}
|
||||||
|
|
||||||
|
# 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 = [
|
||||||
|
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||||
|
# mlflow_response_gate (transform)
|
||||||
|
('continue', 0.95, 'Error'),
|
||||||
|
# mlflow_content_gate (transform)
|
||||||
|
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||||
|
# mlflow_response_gate (predict)
|
||||||
|
('continue', 0.95, 'Error'),
|
||||||
|
]
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await prediction_process.run(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 4
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.input_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['input_filters'],
|
||||||
|
'data': input_data['data'],
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.request_transform,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': input_data['data'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_content_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.request_predict,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'filters': input_data['mlflow_predict_filters'],
|
||||||
|
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'predict',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||||
|
'subworkflow.format_and_export_prediction',
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'on_conflict': 'error',
|
||||||
|
'path_flag': 'continue',
|
||||||
|
'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,
|
||||||
|
'model_name': 'test_model_name',
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
'opc_output_config': input_data['opc_output_config'],
|
||||||
|
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
'transform_table_name': input_data['transform_table_name'],
|
||||||
|
'comment': 'Error',
|
||||||
|
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
|
||||||
|
prediction_process.path_flag_handler = AsyncMock(return_value=True)
|
||||||
|
# Arrange
|
||||||
|
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'},
|
||||||
|
}
|
||||||
|
|
||||||
|
# Mock the activity responses
|
||||||
|
workflow_mock.execute_local_activity_method.side_effect = [
|
||||||
|
('stop', 0.95, 'Input data with bad quality'), # input_gate
|
||||||
|
]
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await prediction_process.run(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.input_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['input_filters'],
|
||||||
|
'data': input_data['data'],
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
|
||||||
|
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
|
||||||
|
# Arrange
|
||||||
|
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'},
|
||||||
|
}
|
||||||
|
|
||||||
|
# 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 = [
|
||||||
|
('repeat', 0.95, 'Input data with bad quality'), # input_gate
|
||||||
|
('continue', 0.95, 'Error'), # mlflow_response_gate (transform)
|
||||||
|
]
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await prediction_process.run(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 2
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.input_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['input_filters'],
|
||||||
|
'data': input_data['data'],
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.request_transform,
|
||||||
|
{
|
||||||
|
'data': input_data['data'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
|
||||||
|
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
|
||||||
|
# Arrange
|
||||||
|
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'},
|
||||||
|
}
|
||||||
|
|
||||||
|
# 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 = [
|
||||||
|
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||||
|
# mlflow_response_gate (transform)
|
||||||
|
('continue', 0.95, 'Error'),
|
||||||
|
# mlflow_content_gate (transform)
|
||||||
|
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||||
|
]
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await prediction_process.run(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 3
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.input_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['input_filters'],
|
||||||
|
'data': input_data['data'],
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.request_transform,
|
||||||
|
{
|
||||||
|
'data': input_data['data'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_content_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_stop_at_mlflow_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': 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'},
|
||||||
|
}
|
||||||
|
|
||||||
|
# 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 = [
|
||||||
|
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||||
|
# mlflow_response_gate (transform)
|
||||||
|
('continue', 0.95, 'Error'),
|
||||||
|
# mlflow_content_gate (transform)
|
||||||
|
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||||
|
('continue', 0.95, 'Error'), # mlflow_response_gate (predict)
|
||||||
|
]
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await prediction_process.run(input_data)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 4
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.input_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['input_filters'],
|
||||||
|
'data': input_data['data'],
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.request_transform,
|
||||||
|
{
|
||||||
|
'data': input_data['data'],
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_content_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['mlflow_transform_filters'],
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'transform',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.request_predict,
|
||||||
|
{
|
||||||
|
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.mlflow_response_gate,
|
||||||
|
{
|
||||||
|
'filters': input_data['mlflow_predict_filters'],
|
||||||
|
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||||
|
'type': 'predict',
|
||||||
|
'path_priority': input_data['path_priority'],
|
||||||
|
**metadata,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
|
||||||
|
# Arrange
|
||||||
|
data = {'test': 'data'}
|
||||||
|
path_flag = 'STOP'
|
||||||
|
confidence = 0.95
|
||||||
|
schema = 'test_schema'
|
||||||
|
table_name = 'test_table'
|
||||||
|
model = 'test_model'
|
||||||
|
last_timestamp = '2024-01-01'
|
||||||
|
model_name = 'test_model_name'
|
||||||
|
model_config = {'retention': '30'}
|
||||||
|
|
||||||
|
# Act
|
||||||
|
result = await prediction_process.path_flag_handler(
|
||||||
|
data,
|
||||||
|
path_flag,
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'model_id': model,
|
||||||
|
'last_timestamp': last_timestamp,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
},
|
||||||
|
confidence,
|
||||||
|
last_timestamp,
|
||||||
|
'',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert result is True
|
||||||
|
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||||
|
# Arrange
|
||||||
|
data = {'test': 'data'}
|
||||||
|
path_flag = 'repeat'
|
||||||
|
confidence = 0.95
|
||||||
|
schema = 'test_schema'
|
||||||
|
table_name = 'test_table'
|
||||||
|
model = 'test_model'
|
||||||
|
last_timestamp = '2024-01-01'
|
||||||
|
model_name = 'test_model_name'
|
||||||
|
model_config = {'retention': '30'}
|
||||||
|
|
||||||
|
# Act
|
||||||
|
result = await prediction_process.path_flag_handler(
|
||||||
|
data,
|
||||||
|
path_flag,
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'model_id': model,
|
||||||
|
'last_timestamp': last_timestamp,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
},
|
||||||
|
confidence,
|
||||||
|
last_timestamp,
|
||||||
|
'',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert result is True
|
||||||
|
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||||
|
Activities.repeat_last_prediction,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'model': model,
|
||||||
|
'last_timestamp': last_timestamp,
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||||
|
# Arrange
|
||||||
|
data = {'test': 'data'}
|
||||||
|
path_flag = 'CONTINUE'
|
||||||
|
confidence = 0.95
|
||||||
|
schema = 'test_schema'
|
||||||
|
table_name = 'test_table'
|
||||||
|
model = 'test_model'
|
||||||
|
last_timestamp = '2024-01-01'
|
||||||
|
model_name = 'test_model_name'
|
||||||
|
model_config = {'retention': '30'}
|
||||||
|
prediction_store_policy = 'erl:1'
|
||||||
|
|
||||||
|
# Act
|
||||||
|
result = await prediction_process.path_flag_handler(
|
||||||
|
data,
|
||||||
|
path_flag,
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'model_id': model,
|
||||||
|
'last_timestamp': last_timestamp,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
'prediction_store_policy': prediction_store_policy,
|
||||||
|
},
|
||||||
|
confidence,
|
||||||
|
last_timestamp,
|
||||||
|
'Prediction Process',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert result is True
|
||||||
|
workflow_mock.execute_activity_method.assert_not_called()
|
||||||
|
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||||
|
'subworkflow.format_and_export_prediction',
|
||||||
|
{
|
||||||
|
'metadata': metadata,
|
||||||
|
'path_flag': path_flag,
|
||||||
|
'data': data,
|
||||||
|
'prediction_confidence': confidence,
|
||||||
|
'timestamp': last_timestamp,
|
||||||
|
'model_id': model,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'comment': 'Prediction Process',
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
'prediction_store_policy': prediction_store_policy,
|
||||||
|
'on_conflict': 'error',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
|
||||||
|
# Arrange
|
||||||
|
data = {'test': 'data'}
|
||||||
|
path_flag = 'unknown'
|
||||||
|
confidence = 0.95
|
||||||
|
schema = 'test_schema'
|
||||||
|
table_name = 'test_table'
|
||||||
|
model = 'test_model'
|
||||||
|
last_timestamp = '2024-01-01'
|
||||||
|
model_name = 'test_model_name'
|
||||||
|
model_config = {'retention': '30'}
|
||||||
|
prediction_store_policy = 'erl:1'
|
||||||
|
# Act
|
||||||
|
result = await prediction_process.path_flag_handler(
|
||||||
|
data,
|
||||||
|
path_flag,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'schema': schema,
|
||||||
|
'table_name': table_name,
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'model_id': model,
|
||||||
|
'last_timestamp': last_timestamp,
|
||||||
|
'model_name': model_name,
|
||||||
|
'model_config': model_config,
|
||||||
|
'opc_output_config': {'test': 'config'},
|
||||||
|
'pi_web_api_output_config': {
|
||||||
|
'endpoint': 'https://test-pi-server.com',
|
||||||
|
'prediction_tags': {},
|
||||||
|
'confidence_tags': {},
|
||||||
|
},
|
||||||
|
'prediction_store_policy': prediction_store_policy,
|
||||||
|
},
|
||||||
|
confidence,
|
||||||
|
last_timestamp,
|
||||||
|
'',
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert
|
||||||
|
assert result is False
|
||||||
|
workflow_mock.execute_activity_method.assert_not_called()
|
||||||
|
workflow_mock.execute_child_workflow.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@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,
|
||||||
|
)
|
||||||
248
tests/laborious/workflows/test_drift.py
Normal file
248
tests/laborious/workflows/test_drift.py
Normal file
@@ -0,0 +1,248 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, call, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.drift import Drift
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def drift() -> Drift:
|
||||||
|
return Drift()
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'drift',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run(workflow_mock: AsyncMock, drift: Drift):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'source_table_name': 'test_source_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'interval': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'drift_metrics': ['psi', 'ks'],
|
||||||
|
'chunk_period': 'hour',
|
||||||
|
}
|
||||||
|
|
||||||
|
target_name = input_data['model_config']['target']
|
||||||
|
|
||||||
|
target_data = {'data': 'test_target_data'}
|
||||||
|
reference_data = {'data': 'test_reference_data'}
|
||||||
|
drift_data = {'drift': 'test_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(return_value=None)
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await drift.run(input_data)
|
||||||
|
|
||||||
|
# Assert - Check start_local_activity_method calls
|
||||||
|
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
|
||||||
|
expected_gathering_query = f"""
|
||||||
|
SELECT *
|
||||||
|
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
|
||||||
|
WHERE
|
||||||
|
model_id = '{input_data['model_id']}' AND
|
||||||
|
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
|
||||||
|
ORDER BY timestamp ASC
|
||||||
|
"""
|
||||||
|
|
||||||
|
workflow_mock.start_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.load_custom_query,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'query': expected_gathering_query,
|
||||||
|
'datetime_columns': ['timestamp', 'created_at'],
|
||||||
|
'orient': 'records',
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
Activities.get_reference_data,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Assert - Check calculate_drift call
|
||||||
|
workflow_mock.execute_local_activity_method.assert_called_once_with(
|
||||||
|
Activities.calculate_drift,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'target_data': target_data,
|
||||||
|
'reference_data': reference_data,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'target_name': target_name,
|
||||||
|
'drift_metrics': input_data['drift_metrics'],
|
||||||
|
'chunk_period': input_data['chunk_period'],
|
||||||
|
},
|
||||||
|
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,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': drift_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,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'source_table_name': 'test_source_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'interval': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'drift_metrics': ['psi', 'ks'],
|
||||||
|
}
|
||||||
|
|
||||||
|
target_data = None
|
||||||
|
reference_data = {'data': 'test_reference_data'}
|
||||||
|
|
||||||
|
workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
|
||||||
|
|
||||||
|
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_activity_method.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'source_table_name': 'test_source_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'interval': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'drift_metrics': ['psi', 'ks'],
|
||||||
|
}
|
||||||
|
|
||||||
|
target_name = input_data['model_config']['target']
|
||||||
|
|
||||||
|
target_data = {'data': 'test_target_data'}
|
||||||
|
reference_data = {'data': 'test_reference_data'}
|
||||||
|
drift_data = None
|
||||||
|
|
||||||
|
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
|
||||||
|
await drift.run(input_data)
|
||||||
|
|
||||||
|
# Assert - Should call calculate_drift but not export
|
||||||
|
workflow_mock.execute_local_activity_method.assert_called_once_with(
|
||||||
|
Activities.calculate_drift,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'target_data': target_data,
|
||||||
|
'reference_data': reference_data,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'target_name': target_name,
|
||||||
|
'drift_metrics': input_data['drift_metrics'],
|
||||||
|
'chunk_period': input_data.get('chunk_period', 'min'),
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'source_table_name': 'test_source_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'interval': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'drift_metrics': ['psi', 'ks'],
|
||||||
|
# chunk_period not provided, should default to 'min'
|
||||||
|
}
|
||||||
|
|
||||||
|
target_name = input_data['model_config']['target']
|
||||||
|
|
||||||
|
target_data = {'data': 'test_target_data'}
|
||||||
|
reference_data = {'data': 'test_reference_data'}
|
||||||
|
drift_data = {'drift': 'test_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(return_value=None)
|
||||||
|
|
||||||
|
# Act
|
||||||
|
await drift.run(input_data)
|
||||||
|
|
||||||
|
# Assert - Check calculate_drift call with default chunk_period
|
||||||
|
workflow_mock.execute_local_activity_method.assert_called_once_with(
|
||||||
|
Activities.calculate_drift,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'target_data': target_data,
|
||||||
|
'reference_data': reference_data,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'target_name': target_name,
|
||||||
|
'drift_metrics': input_data['drift_metrics'],
|
||||||
|
'chunk_period': 'min', # Default value
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
327
tests/laborious/workflows/test_minimal_retrain.py
Normal file
327
tests/laborious/workflows/test_minimal_retrain.py
Normal file
@@ -0,0 +1,327 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, call, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def minimal_retrain() -> MinimalRetrain:
|
||||||
|
return MinimalRetrain()
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'minimal_retrain',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
|
||||||
|
input_data = {
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'minimal_retrain',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'query': 'test_query',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'table_name': 'test_table',
|
||||||
|
'model_config': {
|
||||||
|
'target': 'test_target',
|
||||||
|
'transform_flavor': 'test_transform_flavor',
|
||||||
|
'predict_flavor': 'test_predict_flavor',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
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=[
|
||||||
|
storage_result,
|
||||||
|
{'success': True, 'experiment': 'test_experiment'},
|
||||||
|
{
|
||||||
|
'success': True,
|
||||||
|
'version': 'test_version',
|
||||||
|
'mlflow_run_id': 'test_mlflow_run_id',
|
||||||
|
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
||||||
|
},
|
||||||
|
{'report': 'test_report'},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
await minimal_retrain.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
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=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.retrain_model,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': storage_result,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.update_production_model,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'success': True,
|
||||||
|
'experiment': 'test_experiment',
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.format_retrain_report,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'experiment_response': {'success': True, 'experiment': 'test_experiment'},
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'update_report': {
|
||||||
|
'success': True,
|
||||||
|
'version': 'test_version',
|
||||||
|
'mlflow_run_id': 'test_mlflow_run_id',
|
||||||
|
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
|
||||||
|
input_data = {
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'minimal_retrain',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'query': 'test_query',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'table_name': 'test_table',
|
||||||
|
'model_config': {
|
||||||
|
'target': 'test_target',
|
||||||
|
'transform_flavor': 'test_transform_flavor',
|
||||||
|
'predict_flavor': 'test_predict_flavor',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
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=[
|
||||||
|
storage_result,
|
||||||
|
{'success': True, 'experiment': 'test_experiment'},
|
||||||
|
{
|
||||||
|
'success': True,
|
||||||
|
'version': 'test_version',
|
||||||
|
'mlflow_run_id': 'test_mlflow_run_id',
|
||||||
|
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
||||||
|
},
|
||||||
|
{'report': 'test_report'},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
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.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=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
|
||||||
|
input_data = {
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'minimal_retrain',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'query': 'test_query',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'table_name': 'test_table',
|
||||||
|
'model_config': {
|
||||||
|
'target': 'test_target',
|
||||||
|
'transform_flavor': 'test_transform_flavor',
|
||||||
|
'predict_flavor': 'test_predict_flavor',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
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=[
|
||||||
|
storage_result,
|
||||||
|
{'success': False, 'experiment': 'test_experiment'},
|
||||||
|
{
|
||||||
|
'success': True,
|
||||||
|
'version': 'test_version',
|
||||||
|
'mlflow_run_id': 'test_mlflow_run_id',
|
||||||
|
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
||||||
|
},
|
||||||
|
{'report': 'test_report'},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
await minimal_retrain.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.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=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.retrain_model,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': storage_result,
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_config': input_data['model_config'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.format_retrain_report,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'experiment_response': {'success': False, 'experiment': 'test_experiment'},
|
||||||
|
'model_name': input_data['model_name'],
|
||||||
|
'model_id': input_data['model_id'],
|
||||||
|
'update_report': {},
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||||
|
'schema': input_data['schema'],
|
||||||
|
'table_name': input_data['table_name'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert workflow_mock.execute_activity_method.call_count == 3
|
||||||
|
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||||
106
tests/laborious/workflows/test_predictions_batch.py
Normal file
106
tests/laborious/workflows/test_predictions_batch.py
Normal file
@@ -0,0 +1,106 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def predictions_batch() -> PredictionsBatch:
|
||||||
|
return PredictionsBatch()
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'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):
|
||||||
|
activity_return = MagicMock()
|
||||||
|
workflow_mock.execute_activity_method.return_value = activity_return
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'query': 'SELECT * FROM test',
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'table_name': 'test_table',
|
||||||
|
'transform_table_name': 'test_transform_table',
|
||||||
|
'opc_output_config': 'test_opc_output_config',
|
||||||
|
'pi_web_api_output_config': 'test_pi_web_api_output_config',
|
||||||
|
'datetime_columns': ['timestamp', 'created_at'],
|
||||||
|
'prediction_store_policy': 'erl:1',
|
||||||
|
'model_config': {'retention': '30'},
|
||||||
|
}
|
||||||
|
|
||||||
|
await predictions_batch.run(input_data)
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.load_query_with_minio_offload,
|
||||||
|
{
|
||||||
|
'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,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
prediction_input = {
|
||||||
|
'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',
|
||||||
|
'CONFIG': {},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
),
|
||||||
|
'mlflow_transform_filters': input_data.get(
|
||||||
|
'mlflow_transform_filters',
|
||||||
|
{
|
||||||
|
'API_ERROR': {
|
||||||
|
'POLICY': 'STOP',
|
||||||
|
'CONFIG': {},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
),
|
||||||
|
'mlflow_predict_filters': input_data.get(
|
||||||
|
'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),
|
||||||
|
}
|
||||||
|
|
||||||
|
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||||
|
[call('subworkflow.prediction_process', prediction_input)]
|
||||||
|
)
|
||||||
215
tests/laborious/workflows/test_simple_metrics.py
Normal file
215
tests/laborious/workflows/test_simple_metrics.py
Normal file
@@ -0,0 +1,215 @@
|
|||||||
|
from unittest.mock import ANY, AsyncMock, call, patch
|
||||||
|
|
||||||
|
from pytest import fixture, mark
|
||||||
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||||
|
|
||||||
|
from laborious.activities.activities import Activities
|
||||||
|
from laborious.workflows.simple_metrics import SimpleMetrics
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def simple_metrics() -> SimpleMetrics:
|
||||||
|
return SimpleMetrics()
|
||||||
|
|
||||||
|
|
||||||
|
metadata = {
|
||||||
|
'metadata': {
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'workflow_name': 'simple_metrics',
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'interval_minutes': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'predictions_table_name': 'test_predictions_table',
|
||||||
|
'data_table_name': 'test_data_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'metrics': ['rmse', 'mse', 'mae', 'r2'],
|
||||||
|
}
|
||||||
|
|
||||||
|
target_data = {'data': 'test_target_data'}
|
||||||
|
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
|
||||||
|
|
||||||
|
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 load_custom_query call
|
||||||
|
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
|
||||||
|
expected_query = f"""
|
||||||
|
select p."timestamp", p.prediction, ld.value as "target"
|
||||||
|
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
|
||||||
|
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
|
||||||
|
on p."timestamp" = ld."timestamp"
|
||||||
|
where
|
||||||
|
p.model_id = '{input_data['model_id']}' and
|
||||||
|
p.prediction is not null and
|
||||||
|
ld.variable = '{input_data['model_config']['target']}' and
|
||||||
|
ld.value is not null and
|
||||||
|
p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
|
||||||
|
order by
|
||||||
|
p."timestamp" desc;
|
||||||
|
"""
|
||||||
|
|
||||||
|
workflow_mock.execute_activity_method.assert_has_calls(
|
||||||
|
[
|
||||||
|
call(
|
||||||
|
Activities.load_custom_query,
|
||||||
|
{
|
||||||
|
**metadata,
|
||||||
|
'query': expected_query,
|
||||||
|
'datetime_columns': ['timestamp'],
|
||||||
|
'orient': 'records',
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
),
|
||||||
|
call(
|
||||||
|
Activities.export_data_to_postgres,
|
||||||
|
{
|
||||||
|
**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,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
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': input_data['metrics'],
|
||||||
|
'interval_minutes': input_data['interval_minutes'],
|
||||||
|
},
|
||||||
|
retry_policy=ANY,
|
||||||
|
start_to_close_timeout=ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'interval_minutes': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'predictions_table_name': 'test_predictions_table',
|
||||||
|
'data_table_name': 'test_data_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'metrics': ['rmse', 'mse'],
|
||||||
|
}
|
||||||
|
|
||||||
|
target_data = None
|
||||||
|
|
||||||
|
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_activity_method.call_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'interval_minutes': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'predictions_table_name': 'test_predictions_table',
|
||||||
|
'data_table_name': 'test_data_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
'metrics': ['rmse', 'mse'],
|
||||||
|
}
|
||||||
|
|
||||||
|
target_data = {'data': 'test_target_data'}
|
||||||
|
simple_metrics_data = None
|
||||||
|
|
||||||
|
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
|
||||||
|
workflow_mock.execute_activity_method.assert_called_once()
|
||||||
|
workflow_mock.execute_local_activity_method.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
|
||||||
|
async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
|
||||||
|
# Arrange
|
||||||
|
input_data = {
|
||||||
|
'schedule_name': 'test_schedule',
|
||||||
|
'model_name': 'test_model',
|
||||||
|
'model_id': 'test_model_id',
|
||||||
|
'interval_minutes': 60,
|
||||||
|
'model_config': {'target': 'test_target'},
|
||||||
|
'schema': 'test_schema',
|
||||||
|
'predictions_table_name': 'test_predictions_table',
|
||||||
|
'data_table_name': 'test_data_table',
|
||||||
|
'target_table_name': 'test_target_table',
|
||||||
|
# metrics not provided, should default to ['rmse', 'mse', 'mae', 'r2']
|
||||||
|
}
|
||||||
|
|
||||||
|
target_data = {'data': 'test_target_data'}
|
||||||
|
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
|
||||||
|
|
||||||
|
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_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,
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user