Code import - branch 0.6.0

This commit is contained in:
2026-08-05 13:53:39 +00:00
commit d775f6b264
89 changed files with 23672 additions and 0 deletions

35
.env.example Normal file
View 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"

17
.github/workflows/quality-gate.yml vendored Normal file
View File

@@ -0,0 +1,17 @@
name: Quality gate
on:
pull_request:
branches:
- main
types: [ opened, synchronize, reopened ]
jobs:
quality-gate:
uses: Aignosi/github_workflow_templates/.github/workflows/python-quality-gate.yml@main
permissions: write-all
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
View 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

58
.gitignore vendored Normal file
View File

@@ -0,0 +1,58 @@
# 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/

7
Makefile Normal file
View 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)

1223
README.md Normal file

File diff suppressed because it is too large Load Diff

170
docs/opc-communication.md Normal file
View 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
View File

@@ -0,0 +1,3 @@
"""
End-to-end tests for laborious temporal workflows.
"""

681
e2e/conftest.py Normal file
View File

@@ -0,0 +1,681 @@
"""
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 pandas as pd
import pytest
import pytest_asyncio
from sqlalchemy import create_engine, text
from testcontainers.minio import MinioContainer
from testcontainers.postgres import PostgresContainer
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
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.prediction_process import PredictionProcess
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
# Test constants
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'
@pytest_asyncio.fixture(scope='session')
def minio_container():
"""
MinIO S3-compatible storage for E2E tests that exercise real offload uploads.
"""
minio = MinioContainer()
minio.start()
yield minio
minio.stop()
@pytest_asyncio.fixture(scope='session')
def postgres_container():
"""
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={}):
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,
}
@pytest_asyncio.fixture
def patch_mlflow(mock_mlflow_models):
"""Patch mlflow module in repository with load_model mocks."""
mock_mlflow = MagicMock()
# Mock sklearn.load_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 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 '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.tracking.MlflowClient = MagicMock(return_value=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'
)])
))
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
yield mock_mlflow
@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

181
e2e/helpers.py Normal file
View File

@@ -0,0 +1,181 @@
"""
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).
"""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
f'ORDER BY created_at ASC'
)
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
row = prediction_rows[0]
assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
f'Expected prediction={prediction}, got {row[1]}'
)
assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
str(prediction_confidence)
), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
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."""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
)
count = result_query.scalar()
assert count == 0, f'Expected no predictions, but found {count} records'
def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None:
"""
Assert two prediction rows for model_id both match last_prediction.
Rows are compared in created_at order for stability.
Args:
postgres_engine: SQLAlchemy engine.
model_id: Model id.
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
"""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
text(
f'SELECT model_id, prediction, prediction_confidence, prediction_status '
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
f'ORDER BY created_at ASC'
)
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 2, 'Expected two prediction records'
assert prediction_rows[0] == last_prediction, (
f'Expected first row {last_prediction}, got {prediction_rows[0]}'
)
assert prediction_rows[1] == last_prediction, (
f'Expected second row {last_prediction}, got {prediction_rows[1]}'
)
def make_workflow_id(prefix: str) -> str:
"""Build a unique workflow id using a prefix and current timestamp."""
return f'{prefix}-{datetime.now().timestamp()}'

189
e2e/opc_test_server.py Normal file
View File

@@ -0,0 +1,189 @@
"""
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:
if self._session_bad_on_write:
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
View 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
---

View File

@@ -0,0 +1,74 @@
"""
Direct E2E execution of child workflows (smaller surface than PredictionsBatch).
"""
from decimal import Decimal
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.helpers import make_workflow_id, start_and_await_workflow
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
@pytest.mark.asyncio
@pytest.mark.integration
async def test_format_and_export_prediction_default_path_e2e(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
postgres_engine,
):
"""
Run FormatAndExportPrediction with path_flag set (format_default_prediction path).
"""
client = temporal_test_env.client
model_id = 401
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
metadata = {
'metadata': {
'model_id': model_id,
'model_name': 'test_model',
'schedule_name': 'test-schedule',
'workflow_name': 'subworkflow.format_and_export_prediction',
}
}
input_data = {
'metadata': metadata,
'path_flag': 'CONTINUE',
'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'},
'prediction_confidence': 2,
'timestamp': '2024-01-01 12:00:00+00:00',
'model_id': model_id,
'model_name': 'test_model',
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
'comment': 'e2e child workflow default path',
'opc_output_config': {},
'pi_web_api_output_config': {},
'prediction_store_policy': 'lts:1',
}
await start_and_await_workflow(
client,
FormatAndExportPrediction.run,
input_data,
make_workflow_id('e2e-format-export-child'),
)
with postgres_engine.connect() as conn:
row = conn.execute(
text(
f'SELECT prediction, prediction_confidence, prediction_status, comments '
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
)
).fetchone()
assert row is not None
assert row[0] == 0
assert row[1] == Decimal(2)
assert row[2] == 'Bad'
assert row[3] == 'e2e child workflow default path'

124
e2e/test_minio_offload.py Normal file
View File

@@ -0,0 +1,124 @@
"""
E2E-style tests for MinIO offload using a real MinIO testcontainer.
"""
from unittest.mock import patch
import pytest
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow
from laborious.activities.activities import Activities
from laborious.utils.models import minio_dataframe_payload as mdp
from laborious.workflows.predictions_batch import PredictionsBatch
@pytest.mark.asyncio
@pytest.mark.integration
async def test_load_query_with_minio_offload_writes_object_to_bucket(
postgres_engine,
minio_container,
test_activities_real_minio: Activities,
):
"""
With a tiny offload threshold, query results are uploaded as Parquet to MinIO.
Uses real MinioRepository against testcontainers MinIO (no MinIO mock).
"""
model_id = 501
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
metadata = {
'metadata': {
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': model_id,
'workflow_name': 'predictions_batch',
}
}
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
payload = await test_activities_real_minio.load_query_with_minio_offload(
{
**metadata,
'query': (
'SELECT timestamp, variable, value, created_at '
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
),
'model_name': 'test_model',
'datetime_columns': ['timestamp', 'created_at'],
}
)
assert payload.object_key, 'offloaded payload must reference a MinIO object'
assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'
df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
assert len(df) >= 1
client = minio_container.get_client()
listed = list(client.list_objects('test-bucket', recursive=True))
names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'
@pytest.mark.asyncio
@pytest.mark.integration
async def test_predictions_batch_with_minio_offload_path(
temporal_test_env: WorkflowEnvironment,
temporal_worker_real_minio: Worker,
postgres_engine,
test_activities_real_minio: Activities,
):
"""
Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes.
"""
model_id = 502
with postgres_engine.begin() as conn:
conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}'))
insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
input_data = {
'schedule_name': 'test-schedule',
'model_name': 'test_model',
'model_id': model_id,
'query': (
'SELECT timestamp, variable, value, created_at '
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
),
'schema': 'predictions_schema',
'table_name': 'predictions',
'transform_table_name': 'transformed_data',
'input_filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
'mlflow_transform_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
'mlflow_predict_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
'opc_output_config': {},
'pi_web_api_output_config': {},
'save_transform': True,
'prediction_store_policy': 'lts:1',
'model_config': {
'retention_minutes': 0,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
},
'datetime_columns': ['timestamp', 'created_at'],
}
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
await start_and_await_workflow(
temporal_test_env.client,
PredictionsBatch.run,
input_data,
make_workflow_id('test-batch-minio-offload'),
)
with postgres_engine.connect() as conn:
count = conn.execute(
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
).scalar()
assert count == 1

194
e2e/test_opc_real_server.py Normal file
View File

@@ -0,0 +1,194 @@
"""
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 _slow_reconnect_under_lock(repo: OpcRepository, hold_seconds: float = 0.75) -> None:
"""
Hold the connection lock briefly so concurrent writes see reconnect_in_progress.
Args:
repo (OpcRepository): Connected repository.
hold_seconds (float): Time to keep the lock before reconnecting.
"""
async with repo._connection_lock:
await asyncio.sleep(hold_seconds)
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()
reconnect_task = asyncio.create_task(_slow_reconnect_under_lock(repo))
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:
await reconnect_task
assert_prediction(
postgres_engine,
model_id,
prediction_confidence=14,
comments_contains=OPC_RECONNECT_IN_PROGRESS_COMMENT,
)

View File

@@ -0,0 +1,786 @@
"""
End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
"""
from decimal import Decimal
from typing import Any, cast
from unittest.mock import ANY, 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.",
)

View File

@@ -0,0 +1,303 @@
"""
End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
"""
import asyncio
from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
import pytest
from e2e.helpers import make_workflow_id, start_and_await_workflow
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@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)')

View File

@@ -0,0 +1,379 @@
"""
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
"""
from decimal import Decimal
from unittest.mock import MagicMock, patch
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
View 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
View 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()

View 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
View 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
View File

View File

View 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.close(self)
ModelMetrics.close(self)
API.close(self)

305
laborious/activities/api.py Normal file
View 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()

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

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

View 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
View 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 close(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()

View File

@@ -0,0 +1,209 @@
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)."""
Postgres.close(self)
MinioManager.close(self)
def __del__(self):
self.close()

200
laborious/metrics.py Normal file
View 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,
)

View File

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

View File

@@ -0,0 +1,34 @@
from typing import Any
from pandas import DataFrame
DEFAULT_MAX_DEBUG_DATAFRAME_ROWS = 100
def build_dataframe_debug_message(
message: str,
data: Any,
max_rows: int = DEFAULT_MAX_DEBUG_DATAFRAME_ROWS,
) -> str:
"""
Build a safe debug message for dataframe payloads
Args:
- message (str): Base message to identify the logged payload
- data (Any): Payload to evaluate for dataframe-aware logging
- max_rows (int): Maximum dataframe row count allowed for full payload logging
Return:
Formatted debug message with full dataframe content or compact summary
"""
if not isinstance(data, DataFrame):
return f'{message} {data}'
rows = data.shape[0]
if rows <= max_rows:
return f'{message}\n{data.to_csv()}'
return (
f'{message} skipped because dataframe has {rows} rows '
f'(max: {max_rows}). Shape: {data.shape}'
)

View File

View 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

View 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

View File

View File

@@ -0,0 +1,348 @@
"""
MinIO-backed DataFrame payload for Temporal workflows.
Data is never stored as a pandas ``DataFrame`` field on the dataclass.
Instead, the DataFrame is only provided as an input to:
`from_dataframe` / `from_dataframe_to_dict`.
At build time, the DataFrame is evaluated for its serialized size; if it exceeds
the configured threshold, it is serialized to parquet bytes and uploaded to MinIO.
Otherwise, it is inlined as a Temporal-friendly ``dict``.
"""
import pickle
import re
from collections.abc import Hashable
from dataclasses import dataclass
from datetime import datetime
from io import BytesIO
from os import getenv
from typing import Any, Literal
from pandas import DataFrame, read_parquet
from sientia_do.observability.logger import Logger
from sientia_do.repository.minio_repository import MinioRepository
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now
# Keys that are part of the serialized wire format (not arbitrary metadata).
_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'})
_OBJECT_TIMESTAMP_PATTERN = re.compile(
r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$'
)
OFFLOAD_THRESHOLD_BYTES = int(
float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024
)
# Relative prefix used for storing offloaded prediction datasets in MinIO.
# It is also the root directory for retention cleanup listing.
PREDICTION_DATASETS_PREFIX = 'prediction_datasets'
OperationKind = Literal['initial', 'transform', 'predict']
def _build_object_key(
model_name: str, operation: OperationKind, timestamp: str
) -> tuple[str, str | None]:
"""
Build the MinIO object key and the directory prefix used for retention listing.
Args:
model_name: Registered model name used in the pipeline.
operation: Either initial (pre-transform load) or transform (post-MLFlow transform).
timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME.
Return:
tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root).
"""
# Naming convention:
# - Directory is always `prediction_datasets/<model_name>`
# - Filename follows the retention-parsing pattern
basename = f'{model_name}-{operation}-{timestamp}.parquet'
model_dir = model_name.strip().strip('/')
prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}'
return f'{prefix}/{basename}', prefix
@dataclass
class MinioDataFramePayload:
"""
Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys.
Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`.
Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class.
"""
last_timestamp: str
status: dict[str, Any] | None = None
data: dict[Hashable, Any] | None = None
bucket: str | None = None
object_key: str | None = None
object_prefix: str | None = None
uri: str | None = None
@staticmethod
def _debug(
logger: Logger | None,
message: str,
metadata: dict[str, Any] | None = None,
) -> None:
"""
Emit debug logs only when logger is provided
Args:
- logger (Logger | None): Logger instance used for debug messages
- message (str): Message to be logged
- metadata (dict[str, Any] | None): Optional workflow metadata context
"""
if logger is None:
return
logger.custom_debug(message, metadata)
@classmethod
def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload':
"""
Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization.
Temporal converts dataclass return values into plain dicts when crossing
workflow/activity boundaries. This method rebuilds the typed instance so
that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are
available on the receiving side.
If the argument is already a MinioDataFramePayload, it is returned as-is.
Args:
raw: Dict with keys matching the dataclass fields
(last_timestamp, status, data, bucket, object_key, object_prefix, uri),
or an existing MinioDataFramePayload instance.
Return:
MinioDataFramePayload: Reconstructed (or original) instance.
"""
if isinstance(raw, MinioDataFramePayload):
return raw
return cls(
last_timestamp=raw['last_timestamp'],
status=raw.get('status'),
data=raw.get('data'),
bucket=raw.get('bucket'),
object_key=raw.get('object_key'),
object_prefix=raw.get('object_prefix'),
uri=raw.get('uri'),
)
@staticmethod
def estimate_size_bytes(
df: DataFrame,
metadata: dict[str, Any] | None = None,
logger: Logger | None = None,
) -> int:
"""
Approximate serialized size of the DataFrame as the default-orient dict.
Args:
df: DataFrame whose tabular content size is estimated.
Return:
int: Estimated size in bytes (pickle of dict representation).
"""
try:
size = len(pickle.dumps(df.to_dict()))
except Exception:
size = len(pickle.dumps(df))
MinioDataFramePayload._debug(
logger,
f'DataFrame size: {size} bytes',
metadata,
)
return size
@staticmethod
def parse_object_timestamp(object_key: str) -> datetime | None:
"""
Parse the timestamp embedded in the object key basename (before .parquet).
Args:
object_key: S3/MinIO object key whose basename follows
``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``.
Return:
datetime | None: Parsed UTC-naive datetime from the key, or None if not matched.
"""
basename = object_key.rsplit('/', 1)[-1]
match = _OBJECT_TIMESTAMP_PATTERN.search(basename)
if not match:
return None
try:
return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME)
except ValueError:
return None
def cleanup_prefix(self) -> str | None:
"""
Return True if cleanup is enabled for this payload.
"""
if self.object_key is not None and self.data is None:
return self.object_prefix
return None
def has_data(self) -> bool:
"""
Return True if the payload has some data internally or in MinIO.
"""
return (self.data is not None and self.data != {}) or self.object_key is not None
@classmethod
async def from_dataframe(
cls,
dataframe: DataFrame | None,
minio_repo: MinioRepository,
model_name: str,
operation: OperationKind,
status: dict[str, Any] | None = None,
workflow_metadata: dict | None = None,
last_timestamp: str | None = None,
logger: Logger | None = None,
) -> 'MinioDataFramePayload':
"""
Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO.
The DataFrame is not stored on the returned instance.
Args:
dataframe: Tabular data to evaluate and persist (inline or MinIO).
metadata: Small metadata dict merged into the payload (e.g. success, message).
minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`.
workflow_metadata: Metadata passed to MinIO store for logging/metrics.
model_name: Registered model name used in the object basename.
operation: Either ``initial`` (query load) or ``transform`` (post-transform).
key_prefix: Backward-compatible parameter (currently ignored for object naming).
size_threshold_bytes: Byte limit before offload. When None, the module-level
environment-derived default is used.
Return:
MinioDataFramePayload: Instance with data and/or MinIO fields set.
"""
if dataframe is None or dataframe.empty:
cls._debug(
logger,
'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload',
workflow_metadata,
)
return cls(
data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
)
if last_timestamp is None:
last_timestamp = max(dataframe['timestamp'].values.tolist())
dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger)
cls._debug(
logger,
(
f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes '
f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)'
),
workflow_metadata,
)
if dataframe_size <= OFFLOAD_THRESHOLD_BYTES:
cls._debug(
logger,
'MinioDataFramePayload.from_dataframe using inline payload',
workflow_metadata,
)
return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status)
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
object_key, object_prefix = _build_object_key(model_name, operation, timestamp)
cls._debug(
logger,
(
'MinioDataFramePayload.from_dataframe offloading payload to MinIO '
f'with key {object_key}'
),
workflow_metadata,
)
# Upload using the relative object key. The upstream repository will
# prefix it internally under its MinIO namespace.
parquet_buffer = BytesIO()
dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
file_bytes = parquet_buffer.getvalue()
upload_result = await minio_repo.upload_file(
file_bytes=file_bytes,
relative_key=object_key,
metadata=workflow_metadata,
)
bucket = minio_repo.bucket
object_key_full = upload_result.get('minio_object_name', object_key)
uri = f's3://{bucket}/{object_key_full}' if bucket else None
cls._debug(
logger,
f'MinioDataFramePayload.from_dataframe upload completed: {uri}',
workflow_metadata,
)
return cls(
data=None,
bucket=bucket,
object_key=object_key_full,
object_prefix=object_prefix,
uri=uri,
last_timestamp=last_timestamp,
status=status,
)
async def retrieve(
self,
minio_repo: MinioRepository,
workflow_metadata: dict[str, Any] | None = None,
logger: Logger | None = None,
) -> DataFrame:
"""
Load parquet from MinIO when object_key is set and populate inline data.
Args:
minio_repo: sientia_do MinioRepository (or compatible) with download_file().
workflow_metadata: Metadata passed to MinIO read for logging/metrics.
Return:
dict[str, Any]: Flat dict with data filled (same keys as to_dict after load).
"""
if self.data is not None:
self._debug(
logger,
'MinioDataFramePayload.retrieve using inline payload data',
workflow_metadata,
)
return DataFrame(self.data)
if not self.has_data():
self._debug(
logger,
'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe',
workflow_metadata,
)
return DataFrame()
self._debug(
logger,
f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}',
workflow_metadata,
)
file_bytes = await minio_repo.download_file(
object_name=self.object_key, metadata=workflow_metadata
)
df = read_parquet(BytesIO(file_bytes))
self._debug(
logger,
f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}',
workflow_metadata,
)
return df

View File

@@ -0,0 +1,32 @@
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.minio_repository import MinioRepository
class MinioManager(SientiaMonitoring):
minio_repository: MinioRepository | None = None
def __init__(
self,
minio_repository: MinioRepository | None = None,
logger: Logger | None = None,
notification_handler: NotificationHandler | None = None,
metrics_controller: MetricsController | None = None,
):
if self.minio_repository is None:
self.minio_repository = minio_repository
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
def close(self) -> None:
"""
Close the MinioManager and clean up resources.
"""
if self.minio_repository is not None:
try:
self.minio_repository.close()
finally:
self.minio_repository = None
SientiaMonitoring.shutdown(self)

File diff suppressed because it is too large Load Diff

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

View File

289
laborious/worker/worker.py Normal file
View 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())

View File

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

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

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

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

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

View 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
View 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
}

158
pyproject.toml Normal file
View File

@@ -0,0 +1,158 @@
[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
]
[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",
]
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
View 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

18
requirements-light.txt Normal file
View File

@@ -0,0 +1,18 @@
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

18
requirements-local.txt Normal file
View File

@@ -0,0 +1,18 @@
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.5.2

18
requirements.txt Normal file
View File

@@ -0,0 +1,18 @@
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

11
run_coverage.sh Executable file
View 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
View 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
View 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

File diff suppressed because one or more lines are too long

0
tests/__init__.py Normal file
View File

59
tests/conftest.py Normal file
View 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()

View File

View File

View 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.close = 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.close.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()

View 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']
)

File diff suppressed because it is too large Load Diff

View 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'],
)

File diff suppressed because it is too large Load Diff

View 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.close()
opc.opc_repository['server1'].disconnect.assert_called_once()

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

View File

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

View 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

View File

@@ -0,0 +1,266 @@
from datetime import datetime
from io import BytesIO
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pandas import DataFrame
from laborious.utils.models.minio_dataframe_payload import (
MinioDataFramePayload,
_build_object_key,
)
def test_parse_object_timestamp_hyphenated_model():
key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
ts = MinioDataFramePayload.parse_object_timestamp(key)
assert ts == datetime(2024, 6, 15, 10, 30, 45)
def test_parse_object_timestamp_transform():
key = 'p/m-transform-2024-01-02_03-04-05.parquet'
ts = MinioDataFramePayload.parse_object_timestamp(key)
assert ts == datetime(2024, 1, 2, 3, 4, 5)
def test_parse_object_timestamp_invalid():
assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
def test_estimate_size_bytes_returns_positive_for_nonempty_frame():
df = DataFrame({'a': [1, 2]})
size = MinioDataFramePayload.estimate_size_bytes(df)
assert isinstance(size, int)
assert size > 0
def test_cleanup_prefix_when_offloaded_returns_object_prefix():
payload = MinioDataFramePayload(
last_timestamp='t',
data=None,
object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
object_prefix='training_datasets/m',
)
assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m'
def test_cleanup_prefix_when_inline_returns_none():
payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
assert MinioDataFramePayload.cleanup_prefix(payload) is None
def test_has_data_true_when_object_key_set():
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
assert payload.has_data() is True
@pytest.mark.asyncio
async def test_retrieve_inline_dict_as_dataframe():
payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
minio = AsyncMock()
out = await payload.retrieve(minio, {'metadata': {}})
assert list(out.columns) == ['a']
minio.download_file.assert_not_called()
@pytest.mark.asyncio
async def test_retrieve_downloads_parquet_when_offloaded():
source = DataFrame({'a': [1, 2]})
buf = BytesIO()
source.to_parquet(buf, engine='pyarrow', index=True)
file_bytes = buf.getvalue()
payload = MinioDataFramePayload(
last_timestamp='t',
data=None,
object_key='training_datasets/m/f.parquet',
object_prefix='training_datasets/m',
)
minio = AsyncMock()
minio.download_file = AsyncMock(return_value=file_bytes)
out = await payload.retrieve(minio, {'metadata': {}})
minio.download_file.assert_awaited_once_with(
object_name='training_datasets/m/f.parquet',
metadata={'metadata': {}},
)
assert list(out.columns) == ['a']
def test_build_object_key():
key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
assert prefix == 'prediction_datasets/my-model'
def test_build_object_key_strips_slashes():
key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
assert prefix == 'prediction_datasets/my-model'
assert key.startswith('prediction_datasets/my-model/')
def test_estimate_size_bytes_fallback():
df = DataFrame({'a': [1, 2]})
with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')):
size = MinioDataFramePayload.estimate_size_bytes(df)
assert isinstance(size, int)
assert size > 0
def test_parse_object_timestamp_bad_datetime():
key = 'p/m-initial-9999-99-99_99-99-99.parquet'
assert MinioDataFramePayload.parse_object_timestamp(key) is None
@pytest.mark.asyncio
async def test_retrieve_empty_when_no_data():
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
minio = AsyncMock()
out = await payload.retrieve(minio, {})
assert out.empty
minio.download_file.assert_not_called()
@pytest.mark.asyncio
@patch('laborious.utils.models.minio_dataframe_payload.now')
async def test_from_dataframe_none(mock_now):
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
minio = AsyncMock()
result = await MinioDataFramePayload.from_dataframe(
dataframe=None,
minio_repo=minio,
model_name='m',
operation='initial',
status={'success': False, 'message': 'no data'},
)
assert result.data is None
assert result.status == {'success': False, 'message': 'no data'}
assert result.object_key is None
@pytest.mark.asyncio
@patch('laborious.utils.models.minio_dataframe_payload.now')
async def test_from_dataframe_empty(mock_now):
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
minio = AsyncMock()
mock_df = MagicMock()
mock_df.__bool__ = MagicMock(return_value=True)
mock_df.empty = True
result = await MinioDataFramePayload.from_dataframe(
dataframe=mock_df,
minio_repo=minio,
model_name='m',
operation='initial',
)
assert result.data is None
assert result.object_key is None
def _mock_dataframe(data_dict, timestamp_values=None):
"""Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
mock_df = MagicMock()
mock_df.__bool__ = MagicMock(return_value=True)
mock_df.empty = False
if timestamp_values is None:
timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
ts_col = MagicMock()
ts_col.values.tolist.return_value = timestamp_values
mock_df.__getitem__ = MagicMock(return_value=ts_col)
mock_df.to_dict.return_value = data_dict
buf = BytesIO()
DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
return mock_df
@pytest.mark.asyncio
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
async def test_from_dataframe_inline():
minio = AsyncMock()
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
result = await MinioDataFramePayload.from_dataframe(
dataframe=df,
minio_repo=minio,
model_name='m',
operation='initial',
)
assert result.data is not None
assert result.object_key is None
assert result.last_timestamp == '2024-01-01'
@pytest.mark.asyncio
@patch('laborious.utils.models.minio_dataframe_payload.now')
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
async def test_from_dataframe_offloaded(mock_now):
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
minio = AsyncMock()
minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'})
minio.bucket = 'test-bucket'
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
result = await MinioDataFramePayload.from_dataframe(
dataframe=df,
minio_repo=minio,
model_name='m',
operation='initial',
workflow_metadata={'wf': 'data'},
)
assert result.data is None
assert result.object_key == 'full/key.parquet'
assert result.bucket == 'test-bucket'
assert result.uri == 's3://test-bucket/full/key.parquet'
minio.upload_file.assert_awaited_once()
def test_from_dict_inline():
raw = {
'last_timestamp': '2024-01-01T00:00:00+00:00',
'status': None,
'data': {'col1': {0: 'val1'}},
'bucket': None,
'object_key': None,
'object_prefix': None,
'uri': None,
}
payload = MinioDataFramePayload.from_dict(raw)
assert isinstance(payload, MinioDataFramePayload)
assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
assert payload.data == {'col1': {0: 'val1'}}
assert payload.object_key is None
def test_from_dict_offloaded():
raw = {
'last_timestamp': '2024-06-15T10:30:45+00:00',
'status': {'success': True},
'data': None,
'bucket': 'my-bucket',
'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
'object_prefix': 'training_datasets/model',
'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
}
payload = MinioDataFramePayload.from_dict(raw)
assert isinstance(payload, MinioDataFramePayload)
assert payload.data is None
assert payload.bucket == 'my-bucket'
assert payload.object_key == raw['object_key']
assert payload.object_prefix == 'training_datasets/model'
assert payload.uri == raw['uri']
assert payload.status == {'success': True}
def test_from_dict_minimal_keys():
raw = {'last_timestamp': '2024-01-01'}
payload = MinioDataFramePayload.from_dict(raw)
assert payload.last_timestamp == '2024-01-01'
assert payload.data is None
assert payload.bucket is None
assert payload.object_key is None
def test_from_dict_passthrough_existing_instance():
original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
result = MinioDataFramePayload.from_dict(original)
assert result is original

File diff suppressed because it is too large Load Diff

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

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

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

View File

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

View File

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

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

View 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

View 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)]
)

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