Compare commits
14 Commits
release/SI
...
79605db260
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
79605db260 | ||
| 4ae76cfd83 | |||
|
|
14b27e4c30 | ||
|
|
09349b6274 | ||
|
|
2ace55c241 | ||
|
|
6981d09ac6 | ||
| 860bce1ba3 | |||
|
|
c190267292 | ||
|
|
206b47dfae | ||
|
|
9c628e70a6 | ||
|
|
db32450c68 | ||
|
|
a01763b69c | ||
|
|
e5ebf243bd | ||
|
|
cff81e3594 |
@@ -40,6 +40,7 @@ CREATE TABLE IF NOT EXISTS sientia_data.predictions (
|
||||
"timestamp" timestamptz NOT NULL,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||
"comments" text NULL,
|
||||
metadata_id varchar(24) NULL,
|
||||
CONSTRAINT unique_model_id_timestamp
|
||||
UNIQUE (model_id, "timestamp")
|
||||
);
|
||||
|
||||
@@ -17,18 +17,21 @@ with workflow.unsafe.imports_passed_through():
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
from laborious.activities.opc import OPC
|
||||
from laborious.activities.prediction_metadata import PredictionMetadata
|
||||
from laborious.activities.storage import Storage
|
||||
from laborious.utils.connectors_config import build_mlflow_config
|
||||
|
||||
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, PredictionMetadata):
|
||||
"""
|
||||
Central orchestrator for all Temporal activities used by Laborious workflows.
|
||||
|
||||
Composes Storage (Postgres + MinIO offload), MLFlow (wrapper-based inference and retrain
|
||||
via ``SientiaMLflowRepository``), Gates (data quality and ML response filters), OPC exports,
|
||||
drift/simple metrics, and PI Web API writes. The worker constructs one ``Activities`` instance
|
||||
per process and registers its callables on multiple workers bound to different task queues.
|
||||
drift/simple metrics, PI Web API writes, and PredictionMetadata (the single write point for
|
||||
gate metadata — see `laborious.activities.prediction_metadata`). The worker constructs one
|
||||
``Activities`` instance per process and registers its callables on multiple workers bound to
|
||||
different task queues.
|
||||
|
||||
MLflow connectivity: unless ``mlflow_repository`` is injected (tests only), this class builds
|
||||
``SientiaMLflowRepository`` from ``build_mlflow_config()`` so tracking credentials and URL
|
||||
@@ -70,7 +73,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
- metrics_controller: Optional shared controller; if ``None``, a new one is created.
|
||||
- mlflow_repository: Optional ``SientiaMLflowRepository`` for unit/e2e tests; in production
|
||||
leave unset so the repository is built from environment via ``build_mlflow_config()``.
|
||||
- mongo_repository: Optional ``MongoDBRepository`` (gate metadata, best-effort) for
|
||||
- mongo_repository: Optional ``MongoDBRepository`` (prediction metadata, best-effort) for
|
||||
unit/e2e tests; in production leave unset so the repository is built from environment
|
||||
via ``build_mongodb_config()``, same pattern as ``mlflow_repository`` above.
|
||||
|
||||
@@ -145,6 +148,13 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
Gates.__init__(
|
||||
self,
|
||||
minio_repository=minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
PredictionMetadata.__init__(
|
||||
self,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
@@ -192,3 +202,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
OPC.close(self)
|
||||
ModelMetrics.close(self)
|
||||
API.close(self)
|
||||
PredictionMetadata.close(self)
|
||||
|
||||
@@ -5,7 +5,6 @@ with workflow.unsafe.imports_passed_through():
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any
|
||||
|
||||
from bson import ObjectId
|
||||
from pandas import DataFrame
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
@@ -13,7 +12,6 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
|
||||
from laborious import metrics
|
||||
@@ -29,6 +27,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||
ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
|
||||
ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||
InputCounterFunc = Callable[[DataFrame, dict[str, Any]], int]
|
||||
|
||||
# Input filter function mappings
|
||||
input_filter_functions: dict[str, InputFilterFunc] = {
|
||||
@@ -36,6 +35,42 @@ input_filter_functions: dict[str, InputFilterFunc] = {
|
||||
'EMPTY_DATA': filter_empty_data,
|
||||
}
|
||||
|
||||
|
||||
def _count_specific_variables_null_values(data: DataFrame, config: dict[str, Any]) -> int:
|
||||
"""
|
||||
Count rows with a null value in one of the configured variables.
|
||||
|
||||
Mirrors `filter_specific_variables_null_values`'s own trigger condition, so the count
|
||||
reflects exactly what made that filter fire.
|
||||
|
||||
Return:
|
||||
int: Number of rows with a null value in `config['variables']`.
|
||||
"""
|
||||
if data.empty:
|
||||
return 0
|
||||
return int((data['variable'].isin(config['variables']) & data['value'].isna()).sum())
|
||||
|
||||
|
||||
def _count_empty_data(_data: DataFrame, _config: dict[str, Any]) -> int:
|
||||
"""
|
||||
Count rows for the EMPTY_DATA filter.
|
||||
|
||||
`EMPTY_DATA` only triggers when there are no rows at all, so the count is always 0.
|
||||
|
||||
Return:
|
||||
int: Always 0.
|
||||
"""
|
||||
return 0
|
||||
|
||||
|
||||
# Per-filter row counts for the `counter` field in prediction metadata (gate-local: mirrors each
|
||||
# filter's own trigger condition instead of changing `conditional_filters.py`'s bool contract,
|
||||
# which `mlflow_content_gate` also depends on for EMPTY_DATA).
|
||||
input_filter_counters: dict[str, InputCounterFunc] = {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': _count_specific_variables_null_values,
|
||||
'EMPTY_DATA': _count_empty_data,
|
||||
}
|
||||
|
||||
# Confidence mappings kept separate from function maps to avoid Union types
|
||||
input_path_confidence: Mapping[str, int] = {
|
||||
'STOP': -1,
|
||||
@@ -88,14 +123,11 @@ class Gates(SientiaMonitoring):
|
||||
"""
|
||||
|
||||
minio_repository: MinioRepository | None = None
|
||||
mongo_repository: MongoDBRepository | None = None
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
_PREDICTION_METADATA_COLLECTION = 'prediction_metadata'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
@@ -106,14 +138,11 @@ class Gates(SientiaMonitoring):
|
||||
Args:
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
|
||||
written best-effort. `None` is a valid state (see `_write_prediction_metadata`).
|
||||
|
||||
Raises:
|
||||
Exception: If BaseActivity initialization fails
|
||||
"""
|
||||
self.minio_repository = minio_repository
|
||||
self.mongo_repository = mongo_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
@@ -131,11 +160,6 @@ class Gates(SientiaMonitoring):
|
||||
self.minio_repository.close()
|
||||
finally:
|
||||
self.minio_repository = None
|
||||
if self.mongo_repository is not None:
|
||||
try:
|
||||
self.mongo_repository.close()
|
||||
finally:
|
||||
self.mongo_repository = None
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
@@ -159,56 +183,6 @@ class Gates(SientiaMonitoring):
|
||||
metadata,
|
||||
)
|
||||
|
||||
def _write_prediction_metadata(
|
||||
self, document: dict[str, Any] | None, metadata: dict[str, Any]
|
||||
) -> str | None:
|
||||
"""
|
||||
Best-effort write of gate metadata (triggered filter, policy, counter) to Mongo.
|
||||
|
||||
Implements both business rules for this feature in one place:
|
||||
- The prediction write must never be blocked or failed by a metadata write. A failed
|
||||
`insert(...)` is swallowed, never re-raised, logged here with gate-level context (which
|
||||
`mongo_repository.insert(...)` itself can't know) — matching how `storage.py`'s
|
||||
`cleanup_minio_objects_expired` also logs again at the call site even though the
|
||||
repository already logs internally and emits its own Prometheus metric
|
||||
(`MONGODB_WRITE_ERROR_COUNT`); the generic library log is for the metric/trace, this
|
||||
one is for whoever is reading logs for this specific gate.
|
||||
- Nothing is written when there's nothing relevant to report. A falsy `document` (e.g.
|
||||
no filter triggered) skips Mongo entirely, so callers can call this unconditionally
|
||||
and get `None` back for "no metadata this time" instead of deciding upfront whether
|
||||
there's something worth saving.
|
||||
|
||||
`insert(...)` doesn't hand back the inserted document's `_id`, so it is generated here,
|
||||
client-side, before the insert, as a plain string (not a BSON ObjectId) — matching this
|
||||
codebase's convention of string `_id`s (e.g. the `models` collection) and the 24-hex-char
|
||||
size of the `predictions.metadata_id` Postgres column that will reference it.
|
||||
|
||||
Args:
|
||||
document: Metadata payload (filter, policy, counter, etc.); `None`/empty skips the
|
||||
write entirely. No `_id` key expected.
|
||||
metadata: Workflow metadata for contextual logging.
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, Mongo isn't configured, or
|
||||
the insert failed.
|
||||
"""
|
||||
if not document or self.mongo_repository is None:
|
||||
return None
|
||||
|
||||
generated_id = str(ObjectId())
|
||||
|
||||
try:
|
||||
self.mongo_repository.insert(
|
||||
self._PREDICTION_METADATA_COLLECTION,
|
||||
{**document, '_id': generated_id},
|
||||
metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Failed to write prediction metadata: {e}', metadata)
|
||||
return None
|
||||
|
||||
return generated_id
|
||||
|
||||
@staticmethod
|
||||
def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]:
|
||||
"""
|
||||
@@ -226,7 +200,9 @@ class Gates(SientiaMonitoring):
|
||||
return policy, filter_config
|
||||
|
||||
@activity.defn(name='input_gate')
|
||||
def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
def input_gate(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> tuple[str | None, int, str, dict[str, Any] | None]:
|
||||
"""
|
||||
Apply input data quality filters and validation.
|
||||
|
||||
@@ -241,6 +217,9 @@ class Gates(SientiaMonitoring):
|
||||
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
|
||||
4. Provides confidence scores and detailed comments
|
||||
5. Handles errors gracefully with notification integration
|
||||
6. Builds a metadata document describing which filter(s) triggered, with their policy
|
||||
and a counter, whenever one does — this activity only computes it, it never writes it
|
||||
(see `PredictionMetadata.write_prediction_metadata` for the actual write)
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for input validation
|
||||
@@ -251,10 +230,14 @@ class Gates(SientiaMonitoring):
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
tuple: (path_flag, confidence, comment, document)
|
||||
- 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
|
||||
- document (dict | None): Metadata describing the filter(s) that produced
|
||||
`path_flag`, shaped `{'filters': [{'name', 'policy', 'counter'}, ...]}`;
|
||||
`None` when nothing was filtered. Not written anywhere by this activity —
|
||||
the caller decides whether and how to persist it.
|
||||
|
||||
Raises:
|
||||
Exception: If filter execution fails or configuration is invalid
|
||||
@@ -268,7 +251,8 @@ class Gates(SientiaMonitoring):
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
# (filter_name, policy, filter_config) per filter that triggered, in configured order
|
||||
triggered_filters: list[tuple[str, str, dict[str, Any]]] = []
|
||||
|
||||
self._debug_dataframe('Input data:', data, metadata)
|
||||
self.debug(f'Filters: {filters}', metadata)
|
||||
@@ -282,7 +266,7 @@ class Gates(SientiaMonitoring):
|
||||
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)
|
||||
triggered_filters.append((fil, policy, filter_config))
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
@@ -294,16 +278,46 @@ class Gates(SientiaMonitoring):
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
triggered_policies = [policy for _, policy, _ in triggered_filters]
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
if path_flag in triggered_policies:
|
||||
self.info(f'Input gate result: {path_flag}', metadata)
|
||||
return path_flag, input_path_confidence[path_flag], 'Input data with bad quality'
|
||||
|
||||
matched = [
|
||||
(fil, filter_config)
|
||||
for fil, policy, filter_config in triggered_filters
|
||||
if policy == path_flag
|
||||
]
|
||||
document = {
|
||||
'filters': [
|
||||
{
|
||||
'name': fil,
|
||||
'policy': path_flag,
|
||||
# `.get(..., lambda *_: 0)`, not `[fil]`: a filter added to
|
||||
# `input_filter_functions` without a matching counter here degrades to
|
||||
# counter=0 instead of a raw KeyError crashing this prediction (review
|
||||
# finding, 2026-09-30 — see test enforcing key parity between the two).
|
||||
'counter': input_filter_counters.get(fil, lambda *_: 0)(
|
||||
data, filter_config
|
||||
),
|
||||
}
|
||||
for fil, filter_config in matched
|
||||
],
|
||||
}
|
||||
|
||||
return (
|
||||
path_flag,
|
||||
input_path_confidence[path_flag],
|
||||
'Input data with bad quality',
|
||||
document,
|
||||
)
|
||||
|
||||
self.info('Nothing was filtered by the input gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
return None, 0, '', None
|
||||
|
||||
@activity.defn(name='mlflow_response_gate')
|
||||
def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
@@ -699,6 +713,9 @@ class Gates(SientiaMonitoring):
|
||||
- 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
|
||||
- metadata_id (str | None, optional): Id of the `prediction_metadata` Mongo
|
||||
document to reference for this row; `None`/absent when there's nothing to
|
||||
reference
|
||||
|
||||
Returns:
|
||||
dict: Formatted default prediction data with error indicators
|
||||
@@ -716,6 +733,7 @@ class Gates(SientiaMonitoring):
|
||||
'prediction_confidence': [input_data['prediction_confidence']],
|
||||
'prediction_status': ['Bad'],
|
||||
'comments': [input_data['comment']],
|
||||
'metadata_id': [input_data.get('metadata_id')],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
132
laborious/activities/prediction_metadata.py
Normal file
132
laborious/activities/prediction_metadata.py
Normal file
@@ -0,0 +1,132 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from typing import Any
|
||||
|
||||
from bson import ObjectId
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
|
||||
|
||||
class PredictionMetadata(SientiaMonitoring):
|
||||
"""
|
||||
Centralized, best-effort writer for prediction gate metadata (triggered filter, policy,
|
||||
counter) to MongoDB.
|
||||
|
||||
Gates only compute and return a metadata document; they never write it themselves. Keeping
|
||||
the write here, instead of in each gate, puts the one real write point in the subworkflow
|
||||
that's about to persist a `predictions` row — the only place that reliably knows whether one
|
||||
is being written at all — and lets other metadata sources reuse the same activity instead of
|
||||
duplicating write logic wherever they originate.
|
||||
"""
|
||||
|
||||
mongo_repository: MongoDBRepository | None = None
|
||||
_PREDICTION_METADATA_COLLECTION = 'prediction_metadata'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the prediction metadata writer.
|
||||
|
||||
Args:
|
||||
mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
|
||||
written best-effort. `None` is a valid state (see `write_prediction_metadata`).
|
||||
|
||||
Raises:
|
||||
Exception: If BaseActivity initialization fails
|
||||
"""
|
||||
self.mongo_repository = mongo_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the prediction metadata writer and clean up resources.
|
||||
"""
|
||||
if self.mongo_repository is not None:
|
||||
try:
|
||||
self.mongo_repository.close()
|
||||
finally:
|
||||
self.mongo_repository = None
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
@activity.defn(name='write_prediction_metadata')
|
||||
def write_prediction_metadata(self, input_data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Best-effort write of one or more gate metadata documents to Mongo, combined into a single
|
||||
`prediction_metadata` document.
|
||||
|
||||
Implements both business rules for this feature in one place:
|
||||
- The prediction write must never be blocked or failed by a metadata write. A failed
|
||||
`insert(...)` is swallowed, never re-raised, logged here with call-site context (which
|
||||
`mongo_repository.insert(...)` itself can't know) — matching how `storage.py`'s
|
||||
`cleanup_minio_objects_expired` also logs again at the call site even though the
|
||||
repository already logs internally and emits its own Prometheus metric
|
||||
(`MONGODB_WRITE_ERROR_COUNT`); the generic library log is for the metric/trace, this
|
||||
one is for whoever is reading logs for this specific call.
|
||||
- Nothing is written when there's nothing relevant to report. An empty/missing
|
||||
`documents` list (or one full of falsy/empty entries) skips Mongo entirely, so callers
|
||||
can call this unconditionally and get `None` back for "no metadata this time" instead
|
||||
of deciding upfront whether there's something worth saving.
|
||||
|
||||
Takes a *list* of documents, not a single one: each caller contributes its own document
|
||||
(shaped like `{'filters': [...]}`), and this flattens every `filters` entry across all of
|
||||
them into one `prediction_metadata` document before the single insert. Today only one
|
||||
gate ever contributes a document per prediction cycle (the gate chain short-circuits on
|
||||
the first non-`None` `path_flag`), but the list shape means this activity won't need to
|
||||
change again once more than one source contributes metadata for the same cycle.
|
||||
|
||||
`insert(...)` doesn't hand back the inserted document's `_id`, so it is generated here,
|
||||
client-side, before the insert, as a plain string (not a BSON ObjectId) — matching this
|
||||
codebase's convention of string `_id`s (e.g. the `models` collection) and the 24-hex-char
|
||||
size of the `predictions.metadata_id` Postgres column that will reference it.
|
||||
|
||||
Args:
|
||||
input_data: Dict with:
|
||||
- documents (list[dict[str, Any]]): Metadata payloads to combine, each shaped
|
||||
like `{'filters': [...]}`. Missing/empty, or entries with no `filters`, skip
|
||||
the write entirely.
|
||||
- metadata (dict): Workflow metadata for contextual logging.
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, Mongo isn't configured, or
|
||||
the insert failed.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
documents = input_data.get('documents') or []
|
||||
|
||||
filters = [
|
||||
entry for document in documents if document for entry in document.get('filters', [])
|
||||
]
|
||||
|
||||
if not filters or self.mongo_repository is None:
|
||||
return None
|
||||
|
||||
generated_id = str(ObjectId())
|
||||
|
||||
try:
|
||||
self.mongo_repository.insert(
|
||||
self._PREDICTION_METADATA_COLLECTION,
|
||||
{'filters': filters, '_id': generated_id},
|
||||
metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Failed to write prediction metadata: {e}', metadata)
|
||||
return None
|
||||
|
||||
return generated_id
|
||||
@@ -2,6 +2,7 @@ from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
||||
import time
|
||||
import traceback
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
@@ -9,12 +10,14 @@ with workflow.unsafe.imports_passed_through():
|
||||
import pandas as pd
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability import metrics
|
||||
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_sync import MinioRepository
|
||||
from sientia_do.temporal.activities.postgres_sync import Postgres
|
||||
from sientia_do.temporal.constants import now
|
||||
from sqlalchemy import text
|
||||
|
||||
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
@@ -26,7 +29,9 @@ _LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_thr
|
||||
class Storage(Postgres, SientiaMonitoring):
|
||||
"""
|
||||
Extensions for Postgres activities with a helper to export query results
|
||||
directly to MinIO as Parquet and return the object name.
|
||||
directly to MinIO as Parquet and return the object name, and with
|
||||
``repeat_last_prediction``, a query specific to this repo's REPEAT path
|
||||
(moved here from the shared library — nothing else consumes it).
|
||||
"""
|
||||
|
||||
minio_repository: MinioRepository | None = None
|
||||
@@ -141,6 +146,110 @@ class Storage(Postgres, SientiaMonitoring):
|
||||
}
|
||||
)
|
||||
|
||||
def _build_repeat_last_prediction_query(self, schema: str, table_name: str):
|
||||
table_identifier = f'"{schema}".{table_name}'
|
||||
return text( # noqa: S608
|
||||
f"""
|
||||
INSERT INTO {table_identifier} (
|
||||
model_id,
|
||||
prediction,
|
||||
timestamp,
|
||||
response_time,
|
||||
prediction_status,
|
||||
prediction_confidence,
|
||||
metadata_id,
|
||||
created_at
|
||||
)
|
||||
SELECT
|
||||
model_id,
|
||||
prediction,
|
||||
:last_timestamp,
|
||||
response_time,
|
||||
prediction_status,
|
||||
prediction_confidence,
|
||||
:metadata_id,
|
||||
NOW()
|
||||
FROM {table_identifier}
|
||||
WHERE model_id = :model
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT 1;
|
||||
""" # nosec B608
|
||||
)
|
||||
|
||||
@activity.defn(name='repeat_last_prediction')
|
||||
def repeat_last_prediction(self, query_items: dict[str, Any]):
|
||||
"""
|
||||
Insert a new `predictions` row that copies the last one for a model forward in time.
|
||||
|
||||
REPEAT persists a new row (unlike STOP, which persists nothing), reusing the previous
|
||||
row's `prediction`/`response_time`/`prediction_status`/`prediction_confidence` values with
|
||||
a fresh `timestamp`/`created_at`. `metadata_id` references the `prediction_metadata`
|
||||
document for this cycle, if any (see `PredictionMetadata.write_prediction_metadata`);
|
||||
`None` is a valid value — the column is nullable.
|
||||
|
||||
Args (query_items):
|
||||
metadata (dict): Workflow metadata for logging and metrics.
|
||||
schema (str): Postgres schema of the `predictions` table.
|
||||
table_name (str): `predictions` table name.
|
||||
model (str): `model_id` to repeat the last prediction for.
|
||||
last_timestamp: Timestamp to stamp the new row with.
|
||||
metadata_id (str | None): `prediction_metadata` document id, or `None`.
|
||||
|
||||
Raises:
|
||||
Exception: Re-raised after logging/metrics if the insert fails.
|
||||
"""
|
||||
metadata = query_items['metadata']
|
||||
schema = query_items['schema']
|
||||
table_name = query_items['table_name']
|
||||
model = query_items['model']
|
||||
last_timestamp = query_items['last_timestamp']
|
||||
metadata_id = query_items.get('metadata_id')
|
||||
|
||||
repeat_query = self._build_repeat_last_prediction_query(schema, table_name)
|
||||
self.info(f'Repeating last prediction for model {model}', metadata)
|
||||
self.debug(
|
||||
f'Query parameters: schema={schema}, model={model}, last_timestamp={last_timestamp}',
|
||||
metadata,
|
||||
)
|
||||
core_labels = {
|
||||
**self.get_core_labels(metadata, 'execute_sql'),
|
||||
'conflict_mode': 'error',
|
||||
}
|
||||
|
||||
with self.session_factory.begin() as session:
|
||||
try:
|
||||
start_time = time.time()
|
||||
session.execute(
|
||||
repeat_query,
|
||||
{
|
||||
'last_timestamp': last_timestamp,
|
||||
'model': model,
|
||||
'metadata_id': metadata_id,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='ERROR_REPEATING_LAST_PREDICTION',
|
||||
message=f'Error repeating last prediction: {e}',
|
||||
block='repeat_last_prediction',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata)
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.POSTGRES_WRITE_ERROR_COUNT,
|
||||
tags=core_labels,
|
||||
)
|
||||
raise
|
||||
else:
|
||||
self.observe_lag_sync(start_time, metrics.POSTGRES_WRITE_LAG, core_labels)
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.POSTGRES_WRITE_COUNT,
|
||||
tags=core_labels,
|
||||
)
|
||||
|
||||
@activity.defn(name='cleanup_minio_objects_expired')
|
||||
def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
|
||||
@@ -249,6 +249,8 @@ async def main():
|
||||
activities.format_transformed_data,
|
||||
activities.format_prediction,
|
||||
activities.format_default_prediction,
|
||||
# Prediction metadata
|
||||
activities.write_prediction_metadata,
|
||||
# OPC
|
||||
activities.write_opc_data,
|
||||
# Postgres / MinIO offload
|
||||
|
||||
@@ -75,6 +75,11 @@ class FormatAndExportPrediction:
|
||||
Required when path_flag is None
|
||||
- comment (str): Operational comment or error description
|
||||
Required when path_flag is not None
|
||||
- metadata_id (str | None): Id of the `prediction_metadata` Mongo document to
|
||||
reference for this row, already written by the caller; `None` when there's
|
||||
nothing to reference. Only meaningful when path_flag is not None — the
|
||||
normal (path_flag is None) path never has metadata to reference, so this
|
||||
key is irrelevant there.
|
||||
|
||||
Returns:
|
||||
None: The workflow completes successfully when all export operations finish
|
||||
@@ -154,6 +159,7 @@ class FormatAndExportPrediction:
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': prediction_confidence,
|
||||
'comment': input_data['comment'],
|
||||
'metadata_id': input_data.get('metadata_id'),
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
|
||||
@@ -132,7 +132,9 @@ class PredictionProcess:
|
||||
'path_priority': input_data['path_priority'],
|
||||
}
|
||||
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
# `document` describes which filter(s) produced `path_flag` (`None` when nothing was
|
||||
# filtered); `path_flag_handler` is the one that decides whether and how to persist it.
|
||||
path_flag, confidence, comment, document = await workflow.execute_local_activity_method(
|
||||
Activities.input_gate,
|
||||
gate_input,
|
||||
retry_policy=retry_policy,
|
||||
@@ -141,7 +143,7 @@ class PredictionProcess:
|
||||
|
||||
# Handle path decision based on filter results
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment, document
|
||||
):
|
||||
return
|
||||
|
||||
@@ -247,6 +249,51 @@ class PredictionProcess:
|
||||
},
|
||||
)
|
||||
|
||||
async def _write_prediction_metadata(
|
||||
self, metadata: dict[str, Any], document: dict[str, Any] | None
|
||||
) -> str | None:
|
||||
"""
|
||||
Write a gate's metadata document via the centralized write activity.
|
||||
|
||||
Shared by the CONTINUE and REPEAT branches of `path_flag_handler` — both persist a new
|
||||
`predictions` row (`REPEAT` reuses the previous row's values but still inserts a new
|
||||
one), so both need a `metadata_id` to reference it.
|
||||
|
||||
Args:
|
||||
metadata: Workflow metadata for contextual logging.
|
||||
document: Metadata produced by the gate that resolved this `path_flag`; `None` when
|
||||
nothing was filtered, or when the gate that produced this `path_flag` doesn't
|
||||
build one yet.
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, including when the
|
||||
activity call itself fails (see `PredictionMetadata.write_prediction_metadata`).
|
||||
"""
|
||||
if not document:
|
||||
# Skip the remote activity call entirely rather than dispatching it just to have
|
||||
# `write_prediction_metadata` decide there's nothing to write — `mlflow_response_gate`
|
||||
# and `mlflow_content_gate` don't build a `document` yet, so this is the common case
|
||||
# today, not a rare edge case (review finding, 2026-10-01).
|
||||
return None
|
||||
|
||||
try:
|
||||
return await workflow.execute_activity_method(
|
||||
Activities.write_prediction_metadata,
|
||||
{**metadata, 'documents': [document]},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
except Exception as e:
|
||||
# Business rule 1: writing metadata must never block or fail the prediction write.
|
||||
# `write_prediction_metadata` already swallows Mongo-specific failures internally and
|
||||
# returns `None`, but that protection doesn't cover the activity call itself failing
|
||||
# to run at all (its `retry_policy` has `maximum_attempts=1` — no retry — so a
|
||||
# timeout, a momentarily unavailable worker, or any other dispatch failure raises
|
||||
# straight into the workflow). Catching it here, instead of just at the Mongo client
|
||||
# level, is what actually makes the rule hold (review finding, 2026-10-01).
|
||||
workflow.logger.warning(f'Failed to write prediction metadata, continuing anyway: {e}')
|
||||
return None
|
||||
|
||||
async def path_flag_handler(
|
||||
self,
|
||||
data: dict[str, Any],
|
||||
@@ -255,6 +302,7 @@ class PredictionProcess:
|
||||
confidence: int,
|
||||
last_timestamp: str,
|
||||
comment: str,
|
||||
document: dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Handle path decisions based on filter results and confidence levels.
|
||||
@@ -280,6 +328,13 @@ class PredictionProcess:
|
||||
confidence: Confidence level from filter validation
|
||||
last_timestamp: Last processed timestamp
|
||||
comment: Additional information about the filter result
|
||||
document: Metadata describing the filter(s) that produced `path_flag`, as returned by
|
||||
the gate (`None` when nothing was filtered, or when the gate that produced this
|
||||
`path_flag` doesn't build one yet). Written via `write_prediction_metadata` for
|
||||
`CONTINUE` and `REPEAT` alike — both persist a `predictions` row (`REPEAT` reuses
|
||||
the previous row's values but still inserts a new one), so both can reference the
|
||||
resulting `metadata_id`; `STOP` never reaches a write at all, so its `document` is
|
||||
simply never used here.
|
||||
|
||||
Returns:
|
||||
bool: True if processing should stop, False to continue
|
||||
@@ -304,7 +359,10 @@ class PredictionProcess:
|
||||
# Stop processing and exit workflow
|
||||
return True
|
||||
elif path_flag == 'REPEAT':
|
||||
# Repeat last prediction if available
|
||||
metadata_id = await self._write_prediction_metadata(metadata, document)
|
||||
# Repeat last prediction, copying the previous row's values forward with a new
|
||||
# timestamp. `repeat_last_prediction` now lives in `Storage` (laborious), not the
|
||||
# shared library — its INSERT...SELECT references `metadata_id` directly.
|
||||
await workflow.execute_activity_method(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
@@ -313,12 +371,14 @@ class PredictionProcess:
|
||||
'table_name': table_name,
|
||||
'model': model_id,
|
||||
'last_timestamp': last_timestamp,
|
||||
'metadata_id': metadata_id,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
return True
|
||||
elif path_flag == 'CONTINUE':
|
||||
metadata_id = await self._write_prediction_metadata(metadata, document)
|
||||
# call write workflow
|
||||
await workflow.execute_child_workflow(
|
||||
'subworkflow.format_and_export_prediction',
|
||||
@@ -335,6 +395,7 @@ class PredictionProcess:
|
||||
'table_name': table_name,
|
||||
'transform_table_name': transform_table_name,
|
||||
'comment': comment,
|
||||
'metadata_id': metadata_id,
|
||||
'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'],
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "laborious"
|
||||
version = "0.9.0"
|
||||
version = "0.9.0-rc2"
|
||||
description = "Sientia DataOps Laborious - ML Model Orchestration System"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
@@ -6,6 +6,7 @@ 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.prediction_metadata import PredictionMetadata
|
||||
from laborious.activities.storage import Storage
|
||||
|
||||
|
||||
@@ -13,6 +14,7 @@ from laborious.activities.storage import Storage
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.PredictionMetadata.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -22,6 +24,7 @@ def test___init__(
|
||||
mock_minio_repository,
|
||||
mock_api_init,
|
||||
mock_model_metrics_init,
|
||||
mock_prediction_metadata_init,
|
||||
mock_gates_init,
|
||||
mock_opc_init,
|
||||
mock_mlflow_init,
|
||||
@@ -84,6 +87,7 @@ def test___init__(
|
||||
assert isinstance(activities, Gates)
|
||||
assert isinstance(activities, ModelMetrics)
|
||||
assert isinstance(activities, API)
|
||||
assert isinstance(activities, PredictionMetadata)
|
||||
|
||||
mock_storage_init.assert_called_once_with(
|
||||
ANY,
|
||||
@@ -122,6 +126,13 @@ def test___init__(
|
||||
mock_gates_init.assert_called_once_with(
|
||||
ANY,
|
||||
minio_repository=mock_minio_repository.return_value,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_prediction_metadata_init.assert_called_once_with(
|
||||
ANY,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
@@ -161,6 +172,7 @@ def test___init__(
|
||||
@patch('laborious.activities.activities.MLFlow')
|
||||
@patch('laborious.activities.activities.OPC')
|
||||
@patch('laborious.activities.activities.Gates')
|
||||
@patch('laborious.activities.activities.PredictionMetadata')
|
||||
@patch('laborious.activities.activities.ModelMetrics')
|
||||
@patch('laborious.activities.activities.API')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -168,6 +180,7 @@ def test_shutdown(
|
||||
_mock_minio_repository,
|
||||
mock_api_init,
|
||||
mock_model_metrics_init,
|
||||
mock_prediction_metadata_init,
|
||||
mock_gates_init,
|
||||
mock_opc_init,
|
||||
mock_mlflow_init,
|
||||
@@ -230,6 +243,7 @@ def test_shutdown(
|
||||
mock_gates_init.close.assert_called_once()
|
||||
mock_model_metrics_init.close.assert_called_once()
|
||||
mock_api_init.close.assert_called_once()
|
||||
mock_prediction_metadata_init.close.assert_called_once()
|
||||
|
||||
|
||||
@patch('laborious.activities.activities.SientiaMLflowRepository')
|
||||
@@ -238,6 +252,7 @@ def test_shutdown(
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.PredictionMetadata.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -247,6 +262,7 @@ def test___init___builds_mlflow_repository_when_not_provided(
|
||||
mock_minio_repository,
|
||||
_mock_api_init,
|
||||
_mock_model_metrics_init,
|
||||
_mock_prediction_metadata_init,
|
||||
_mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
@@ -310,6 +326,7 @@ def test___init___builds_mlflow_repository_when_not_provided(
|
||||
@patch('laborious.activities.activities.MLFlow.__init__')
|
||||
@patch('laborious.activities.activities.OPC.__init__')
|
||||
@patch('laborious.activities.activities.Gates.__init__')
|
||||
@patch('laborious.activities.activities.PredictionMetadata.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -319,6 +336,7 @@ def test___init___builds_mongo_repository_when_not_provided(
|
||||
mock_minio_repository,
|
||||
_mock_api_init,
|
||||
_mock_model_metrics_init,
|
||||
mock_prediction_metadata_init,
|
||||
mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
@@ -373,10 +391,17 @@ def test___init___builds_mongo_repository_when_not_provided(
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
# The same repository built here is the one handed to Gates, not a second instance.
|
||||
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,
|
||||
)
|
||||
# The same repository built here is the one handed to PredictionMetadata, not a second
|
||||
# instance.
|
||||
mock_prediction_metadata_init.assert_called_once_with(
|
||||
ANY,
|
||||
mongo_repository=mock_mongo_repository_cls.return_value,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
|
||||
@@ -6,7 +6,7 @@ from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from laborious.activities.gates import Gates
|
||||
from laborious.activities.gates import Gates, input_filter_counters, input_filter_functions
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
@@ -35,14 +35,8 @@ def _minio_payload(retrieve_return, status=None):
|
||||
|
||||
|
||||
@fixture
|
||||
def mongo() -> MagicMock:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@fixture
|
||||
def gates_activity(mongo: MagicMock):
|
||||
def gates_activity():
|
||||
gates = Gates(
|
||||
mongo_repository=mongo,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
@@ -67,73 +61,11 @@ metadata = {
|
||||
}
|
||||
|
||||
|
||||
def test_write_prediction_metadata_no_document_skips_mongo(gates_activity, mongo: MagicMock):
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata(None, metadata)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_empty_document_skips_mongo(gates_activity, mongo: MagicMock):
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata({}, metadata)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_not_configured_returns_none():
|
||||
# Arrange
|
||||
gates = Gates(
|
||||
mongo_repository=None,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
gates.error = MagicMock()
|
||||
|
||||
# Act
|
||||
result = gates._write_prediction_metadata({'filter_triggered': 'EMPTY_DATA'}, metadata)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
gates.error.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_success_generates_id_and_inserts(
|
||||
gates_activity, mongo: MagicMock
|
||||
):
|
||||
# Arrange
|
||||
document = {'filter_triggered': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 3}
|
||||
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata(document, metadata)
|
||||
|
||||
# Assert: a plain 24-hex-char string id (matches predictions.metadata_id's varchar(24)),
|
||||
# not a BSON ObjectId — `insert(...)` never hands one back, so it's generated client-side.
|
||||
assert result is not None
|
||||
assert len(result) == 24
|
||||
int(result, 16) # raises ValueError if not hex
|
||||
|
||||
collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
|
||||
assert collection == Gates._PREDICTION_METADATA_COLLECTION
|
||||
assert inserted_document == {**document, '_id': result}
|
||||
assert inserted_metadata == metadata
|
||||
|
||||
|
||||
def test_write_prediction_metadata_insert_failure_swallowed(gates_activity, mongo: MagicMock):
|
||||
# Arrange
|
||||
mongo.insert.side_effect = RuntimeError('connection reset')
|
||||
|
||||
# Act
|
||||
result = gates_activity._write_prediction_metadata({'filter_triggered': 'EMPTY_DATA'}, metadata)
|
||||
|
||||
# Assert: the prediction write must never be blocked or failed by this
|
||||
assert result is None
|
||||
gates_activity.error.assert_called_once()
|
||||
def test_input_filter_counters_covers_every_input_filter_function():
|
||||
# A filter added to input_filter_functions without a matching entry here would degrade to
|
||||
# counter=0 (see the `.get(..., lambda *_: 0)` fallback in input_gate) instead of crashing —
|
||||
# but this test should still catch the gap before it ships (review finding, 2026-09-30).
|
||||
assert set(input_filter_counters) == set(input_filter_functions)
|
||||
|
||||
|
||||
def test_input_gate_invalid_filter(gates_activity):
|
||||
@@ -149,7 +81,7 @@ def test_input_gate_invalid_filter(gates_activity):
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == (None, 0, '')
|
||||
assert result == (None, 0, '', None)
|
||||
gates_activity.error.assert_called_once_with(
|
||||
'Filter INVALID_FILTER not found', metadata['metadata']
|
||||
)
|
||||
@@ -173,7 +105,7 @@ def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == (None, 0, '')
|
||||
assert result == (None, 0, '', None)
|
||||
gates_activity.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
|
||||
@@ -197,7 +129,7 @@ def test_input_gate_no_filters(gates_activity):
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == (None, 0, '')
|
||||
assert result == (None, 0, '', None)
|
||||
gates_activity.debug.assert_called()
|
||||
|
||||
|
||||
@@ -211,10 +143,14 @@ def test_input_gate_with_filter(gates_activity):
|
||||
}
|
||||
|
||||
# Act
|
||||
result = gates_activity.input_gate(input_data)
|
||||
path_flag, confidence, comment, document = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == ('STOP', -1, 'Input data with bad quality')
|
||||
# Assert: the gate only computes the document now — it never writes anywhere, and never
|
||||
# checks path_flag to decide whether to build one (that decision moved to the caller, see
|
||||
# `PredictionMetadata.write_prediction_metadata` and `prediction_process.py`), so even a STOP
|
||||
# resolution still gets its document built.
|
||||
assert (path_flag, confidence, comment) == ('STOP', -1, 'Input data with bad quality')
|
||||
assert document == {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
|
||||
gates_activity.debug.assert_called()
|
||||
|
||||
|
||||
@@ -231,7 +167,12 @@ def test_input_gate_with_filter_lowercase_keys(gates_activity):
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == ('STOP', -1, 'Input data with bad quality')
|
||||
assert result == (
|
||||
'STOP',
|
||||
-1,
|
||||
'Input data with bad quality',
|
||||
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]},
|
||||
)
|
||||
|
||||
|
||||
def test_input_gate_with_filter_capitalized_keys(gates_activity):
|
||||
@@ -247,7 +188,12 @@ def test_input_gate_with_filter_capitalized_keys(gates_activity):
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == ('STOP', -1, 'Input data with bad quality')
|
||||
assert result == (
|
||||
'STOP',
|
||||
-1,
|
||||
'Input data with bad quality',
|
||||
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]},
|
||||
)
|
||||
|
||||
|
||||
def test_input_gate_with_filter_not_caught(gates_activity):
|
||||
@@ -263,10 +209,94 @@ def test_input_gate_with_filter_not_caught(gates_activity):
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == (None, 0, '')
|
||||
assert result == (None, 0, '', None)
|
||||
gates_activity.debug.assert_called()
|
||||
|
||||
|
||||
def test_input_gate_repeat_also_builds_document(gates_activity):
|
||||
# Arrange: REPEAT used to be treated like STOP (no document at all) — it isn't anymore, since
|
||||
# REPEAT does insert a new `predictions` row (reusing the previous row's values), so it needs
|
||||
# a document/metadata_id just like CONTINUE (2026-10-01 redesign).
|
||||
input_data = {
|
||||
**metadata,
|
||||
'filters': {'EMPTY_DATA': {'POLICY': 'REPEAT', 'CONFIG': {}}},
|
||||
'data': _minio_payload(DataFrame({'value': []})),
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
}
|
||||
|
||||
# Act
|
||||
result = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == (
|
||||
'REPEAT',
|
||||
-1,
|
||||
'Input data with bad quality',
|
||||
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'REPEAT', 'counter': 0}]},
|
||||
)
|
||||
|
||||
|
||||
def test_input_gate_continue_builds_document_with_counter(gates_activity):
|
||||
# Arrange: two rows with a null in one of the configured variables ('x'); a third row is
|
||||
# clean, to confirm the counter reflects only the rows that actually matched the condition.
|
||||
input_data = {
|
||||
**metadata,
|
||||
'filters': {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||
'POLICY': 'CONTINUE',
|
||||
'CONFIG': {'variables': ['x']},
|
||||
}
|
||||
},
|
||||
'data': _minio_payload(
|
||||
DataFrame({'variable': ['x', 'x', 'y'], 'value': [None, None, None]})
|
||||
),
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
}
|
||||
|
||||
# Act
|
||||
path_flag, confidence, comment, document = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert (path_flag, confidence, comment) == ('CONTINUE', 2, 'Input data with bad quality')
|
||||
assert document == {
|
||||
'filters': [{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}]
|
||||
}
|
||||
|
||||
|
||||
def test_input_gate_multiple_filters_same_policy_lists_each(gates_activity):
|
||||
# Arrange: force both input filters to trigger (not reachable with today's real filter
|
||||
# conditions — EMPTY_DATA and SPECIFIC_VARIABLES_NULL_VALUES are mutually exclusive) to lock
|
||||
# in the list-building behavior approved by Vitor (2026-09-30): one entry per filter sharing
|
||||
# the winning policy, instead of a single joined/aggregated value.
|
||||
with patch('laborious.activities.gates.input_filter_functions') as mock_filters:
|
||||
mock_filters.__contains__.return_value = True
|
||||
mock_filters.__getitem__.return_value = MagicMock(return_value=True)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'filters': {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||
'POLICY': 'CONTINUE',
|
||||
'CONFIG': {'variables': ['x']},
|
||||
},
|
||||
'EMPTY_DATA': {'POLICY': 'CONTINUE', 'CONFIG': {}},
|
||||
},
|
||||
'data': _minio_payload(DataFrame({'variable': ['x'], 'value': [None]})),
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
}
|
||||
|
||||
# Act
|
||||
path_flag, _, _, document = gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert path_flag == 'CONTINUE'
|
||||
assert {f['name'] for f in document['filters']} == {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES',
|
||||
'EMPTY_DATA',
|
||||
}
|
||||
assert all(f['policy'] == 'CONTINUE' for f in document['filters'])
|
||||
|
||||
|
||||
def test_mlflow_response_gate_invalid_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
|
||||
179
tests/laborious/activities/test_prediction_metadata.py
Normal file
179
tests/laborious/activities/test_prediction_metadata.py
Normal file
@@ -0,0 +1,179 @@
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
from pytest import fixture
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from laborious.activities.prediction_metadata import PredictionMetadata
|
||||
|
||||
|
||||
@fixture
|
||||
def mongo() -> MagicMock:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@fixture
|
||||
def prediction_metadata(mongo: MagicMock):
|
||||
activity = PredictionMetadata(
|
||||
mongo_repository=mongo,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
activity.error = MagicMock()
|
||||
return activity
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_write_prediction_metadata_no_documents_skips_mongo(prediction_metadata, mongo: MagicMock):
|
||||
# Act
|
||||
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': []})
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_missing_documents_key_skips_mongo(
|
||||
prediction_metadata, mongo: MagicMock
|
||||
):
|
||||
# Act
|
||||
result = prediction_metadata.write_prediction_metadata({**metadata})
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_documents_with_no_filters_skips_mongo(
|
||||
prediction_metadata, mongo: MagicMock
|
||||
):
|
||||
# Arrange: a falsy document (None) and one with an empty `filters` list — neither contributes
|
||||
# anything to write.
|
||||
documents = [None, {'filters': []}]
|
||||
|
||||
# Act
|
||||
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': documents})
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
mongo.insert.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_not_configured_returns_none():
|
||||
# Arrange
|
||||
activity = PredictionMetadata(
|
||||
mongo_repository=None,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
activity.error = MagicMock()
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
|
||||
|
||||
# Act
|
||||
result = activity.write_prediction_metadata({**metadata, 'documents': [document]})
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
activity.error.assert_not_called()
|
||||
|
||||
|
||||
def test_write_prediction_metadata_success_generates_id_and_inserts(
|
||||
prediction_metadata, mongo: MagicMock
|
||||
):
|
||||
# Arrange
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 3}]}
|
||||
|
||||
# Act
|
||||
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': [document]})
|
||||
|
||||
# Assert: a plain 24-hex-char string id (matches predictions.metadata_id's varchar(24)),
|
||||
# not a BSON ObjectId — `insert(...)` never hands one back, so it's generated client-side.
|
||||
assert result is not None
|
||||
assert len(result) == 24
|
||||
int(result, 16) # raises ValueError if not hex
|
||||
|
||||
collection, inserted_document, inserted_metadata = mongo.insert.call_args.args
|
||||
assert collection == PredictionMetadata._PREDICTION_METADATA_COLLECTION
|
||||
assert inserted_document == {'filters': document['filters'], '_id': result}
|
||||
assert inserted_metadata == metadata['metadata']
|
||||
|
||||
|
||||
def test_write_prediction_metadata_combines_multiple_documents(
|
||||
prediction_metadata, mongo: MagicMock
|
||||
):
|
||||
# Arrange: two documents (e.g. from two different sources contributing metadata for the same
|
||||
# prediction cycle) — their `filters` entries are flattened into one combined document.
|
||||
documents = [
|
||||
{'filters': [{'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0}]},
|
||||
{
|
||||
'filters': [
|
||||
{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}
|
||||
]
|
||||
},
|
||||
]
|
||||
|
||||
# Act
|
||||
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': documents})
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
_, inserted_document, _ = mongo.insert.call_args.args
|
||||
assert inserted_document['filters'] == [
|
||||
{'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0},
|
||||
{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2},
|
||||
]
|
||||
|
||||
|
||||
def test_write_prediction_metadata_insert_failure_swallowed(prediction_metadata, mongo: MagicMock):
|
||||
# Arrange
|
||||
mongo.insert.side_effect = RuntimeError('connection reset')
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
|
||||
|
||||
# Act
|
||||
result = prediction_metadata.write_prediction_metadata({**metadata, 'documents': [document]})
|
||||
|
||||
# Assert: the prediction write must never be blocked or failed by this
|
||||
assert result is None
|
||||
prediction_metadata.error.assert_called_once()
|
||||
|
||||
|
||||
@patch.object(SientiaMonitoring, 'shutdown')
|
||||
def test_close_disposes_mongo_repository(mock_shutdown):
|
||||
"""``PredictionMetadata.close`` should close the Mongo client and clear the reference."""
|
||||
mongo = MagicMock()
|
||||
activity = PredictionMetadata(
|
||||
mongo_repository=mongo,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
activity.close()
|
||||
mongo.close.assert_called_once()
|
||||
assert activity.mongo_repository is None
|
||||
# Other instances collected by GC during the test also hit the patched shutdown.
|
||||
assert mock_shutdown.call_args_list.count(call(activity)) == 1
|
||||
|
||||
|
||||
@patch.object(SientiaMonitoring, 'shutdown')
|
||||
def test_close_without_mongo_repository(mock_shutdown):
|
||||
"""When no Mongo repository is configured, ``close`` only shuts down monitoring."""
|
||||
activity = PredictionMetadata(
|
||||
mongo_repository=None,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
activity.close()
|
||||
# Other instances collected by GC during the test also hit the patched shutdown.
|
||||
assert mock_shutdown.call_args_list.count(call(activity)) == 1
|
||||
@@ -257,6 +257,86 @@ def test_export_payload_to_postgres(storage):
|
||||
assert result == {'success': True}
|
||||
|
||||
|
||||
def test_build_repeat_last_prediction_query(storage):
|
||||
query = storage._build_repeat_last_prediction_query('public', 'predictions')
|
||||
query_str = str(query)
|
||||
|
||||
assert 'INSERT INTO "public".predictions' in query_str
|
||||
assert 'metadata_id' in query_str
|
||||
assert ':last_timestamp' in query_str
|
||||
assert ':metadata_id' in query_str
|
||||
assert ':model' in query_str
|
||||
|
||||
|
||||
@patch('laborious.activities.storage.text')
|
||||
def test_repeat_last_prediction_success(text_mock, storage):
|
||||
storage.emit_metric_sync = MagicMock()
|
||||
storage.observe_lag_sync = MagicMock()
|
||||
query_items = {
|
||||
**metadata,
|
||||
'schema': 'public',
|
||||
'table_name': 'predictions',
|
||||
'model': 1,
|
||||
'last_timestamp': '2022-01-01',
|
||||
'metadata_id': 'abc123',
|
||||
}
|
||||
|
||||
with patch('sqlalchemy.orm.session.Session.execute') as mock_execute:
|
||||
storage.repeat_last_prediction(query_items)
|
||||
|
||||
text_mock.assert_called_once()
|
||||
mock_execute.assert_called_once_with(
|
||||
text_mock.return_value,
|
||||
{'last_timestamp': '2022-01-01', 'model': 1, 'metadata_id': 'abc123'},
|
||||
)
|
||||
storage.observe_lag_sync.assert_called_once()
|
||||
storage.emit_metric_sync.assert_called_once()
|
||||
|
||||
|
||||
@patch('laborious.activities.storage.text')
|
||||
def test_repeat_last_prediction_without_metadata_id(text_mock, storage):
|
||||
storage.emit_metric_sync = MagicMock()
|
||||
storage.observe_lag_sync = MagicMock()
|
||||
query_items = {
|
||||
**metadata,
|
||||
'schema': 'public',
|
||||
'table_name': 'predictions',
|
||||
'model': 1,
|
||||
'last_timestamp': '2022-01-01',
|
||||
}
|
||||
|
||||
with patch('sqlalchemy.orm.session.Session.execute') as mock_execute:
|
||||
storage.repeat_last_prediction(query_items)
|
||||
|
||||
mock_execute.assert_called_once_with(
|
||||
text_mock.return_value,
|
||||
{'last_timestamp': '2022-01-01', 'model': 1, 'metadata_id': None},
|
||||
)
|
||||
|
||||
|
||||
@patch('laborious.activities.storage.text')
|
||||
def test_repeat_last_prediction_error(text_mock, storage):
|
||||
storage.emit_metric_sync = MagicMock()
|
||||
storage.send_notification = MagicMock()
|
||||
query_items = {
|
||||
**metadata,
|
||||
'schema': 'public',
|
||||
'table_name': 'predictions',
|
||||
'model': 1,
|
||||
'last_timestamp': '2022-01-01',
|
||||
'metadata_id': 'abc123',
|
||||
}
|
||||
|
||||
with (
|
||||
patch('sqlalchemy.orm.session.Session.execute', side_effect=ValueError('db')),
|
||||
raises(ValueError),
|
||||
):
|
||||
storage.repeat_last_prediction(query_items)
|
||||
|
||||
storage.send_notification.assert_called_once()
|
||||
storage.emit_metric_sync.assert_called()
|
||||
|
||||
|
||||
def test_cleanup_minio_objects_expired_minio_not_initialized(storage):
|
||||
storage.minio_repository = None
|
||||
|
||||
|
||||
@@ -326,6 +326,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': input_data['prediction_confidence'],
|
||||
'comment': input_data['comment'],
|
||||
'metadata_id': None,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
@@ -389,6 +390,56 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
|
||||
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_passes_metadata_id(
|
||||
workflow_mock, format_and_export_prediction
|
||||
):
|
||||
# Arrange: metadata_id reaches format_default_prediction, which is the row that ends up in
|
||||
# Postgres via export_data_to_postgres (CONTINUE/REPEAT path).
|
||||
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_output_config': {'test': 'config'},
|
||||
'comment': 'test_comment',
|
||||
'metadata_id': 'a' * 24,
|
||||
}
|
||||
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
(MagicMock(), MagicMock()),
|
||||
MagicMock(),
|
||||
MagicMock(),
|
||||
]
|
||||
|
||||
# Act
|
||||
await format_and_export_prediction.run(input_data)
|
||||
|
||||
# Assert
|
||||
workflow_mock.execute_local_activity_method.assert_any_call(
|
||||
Activities.format_default_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': input_data['prediction_confidence'],
|
||||
'comment': input_data['comment'],
|
||||
'metadata_id': 'a' * 24,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
|
||||
@@ -638,6 +689,20 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e
|
||||
|
||||
await format_and_export_prediction.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_any_call(
|
||||
Activities.format_default_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': input_data['prediction_confidence'],
|
||||
'comment': input_data['comment'],
|
||||
'metadata_id': None,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
|
||||
@@ -57,7 +57,7 @@ async def test_run(workflow_mock, prediction_process):
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
('continue', 0.95, 'Input data with bad quality', None), # input_gate
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
@@ -217,7 +217,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('stop', 0.95, 'Input data with bad quality'), # input_gate
|
||||
('stop', 0.95, 'Input data with bad quality', None), # input_gate
|
||||
]
|
||||
|
||||
# Act
|
||||
@@ -275,7 +275,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('repeat', 0.95, 'Input data with bad quality'), # input_gate
|
||||
('repeat', 0.95, 'Input data with bad quality', None), # input_gate
|
||||
('continue', 0.95, 'Error'), # mlflow_response_gate (transform)
|
||||
]
|
||||
|
||||
@@ -365,7 +365,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
('continue', 0.95, 'Input data with bad quality', None), # input_gate
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
@@ -476,7 +476,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
('continue', 0.95, 'Input data with bad quality', None), # input_gate
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
@@ -598,6 +598,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'STOP', 'counter': 0}]}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
@@ -616,18 +617,22 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
document,
|
||||
)
|
||||
|
||||
# Assert
|
||||
# Assert: STOP never reaches a write, even though it now gets a `document` same as CONTINUE
|
||||
# and REPEAT do — it just never has anywhere to use it.
|
||||
assert result is True
|
||||
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||
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_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
# Arrange: REPEAT now writes metadata just like CONTINUE — it does insert a new `predictions`
|
||||
# row (reusing the previous row's values), so it needs a `metadata_id` too.
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'repeat'
|
||||
confidence = 0.95
|
||||
@@ -637,6 +642,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'REPEAT', 'counter': 0}]}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
@@ -655,11 +661,20 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
document,
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
metadata_id = workflow_mock.execute_activity_method.return_value
|
||||
assert workflow_mock.execute_activity_method.call_count == 2
|
||||
workflow_mock.execute_activity_method.assert_any_call(
|
||||
Activities.write_prediction_metadata,
|
||||
{**metadata, 'documents': [document]},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_any_call(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
@@ -667,6 +682,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
'table_name': table_name,
|
||||
'model': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'metadata_id': metadata_id,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
@@ -674,6 +690,56 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
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_no_document_skips_write_activity_call(
|
||||
workflow_mock, prediction_process
|
||||
):
|
||||
# Arrange: no document (e.g. REPEAT triggered by a gate that doesn't build one yet — today,
|
||||
# `mlflow_response_gate`/`mlflow_content_gate`) must not even dispatch the remote
|
||||
# write_prediction_metadata activity — that would be a wasted Temporal round-trip on every
|
||||
# REPEAT caused by those two gates, not a rare edge case (review finding, 2026-10-01).
|
||||
data = {'test': 'data'}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
'repeat',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 'test_model',
|
||||
'last_timestamp': '2024-01-01',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
},
|
||||
0.95,
|
||||
'2024-01-01',
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
called_activities = [c.args[0] for c in workflow_mock.execute_activity_method.call_args_list]
|
||||
assert Activities.write_prediction_metadata not in called_activities
|
||||
# repeat_last_prediction still happens, just with metadata_id=None.
|
||||
workflow_mock.execute_activity_method.assert_any_call(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model': 'test_model',
|
||||
'last_timestamp': '2024-01-01',
|
||||
'metadata_id': None,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
@@ -688,6 +754,9 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
prediction_store_policy = 'erl:1'
|
||||
document = {
|
||||
'filters': [{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 2}]
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
@@ -713,11 +782,18 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'Prediction Process',
|
||||
document,
|
||||
)
|
||||
|
||||
# Assert
|
||||
metadata_id = workflow_mock.execute_activity_method.return_value
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
Activities.write_prediction_metadata,
|
||||
{**metadata, 'documents': [document]},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'subworkflow.format_and_export_prediction',
|
||||
{
|
||||
@@ -733,6 +809,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
'table_name': table_name,
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'comment': 'Prediction Process',
|
||||
'metadata_id': metadata_id,
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'https://test-pi-server.com',
|
||||
@@ -745,6 +822,153 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_continue_no_document_skips_write_activity_call(
|
||||
workflow_mock, prediction_process
|
||||
):
|
||||
# Arrange: same gap as the REPEAT case above (review finding, 2026-10-01) — no document must
|
||||
# not dispatch the remote write_prediction_metadata activity.
|
||||
data = {'test': 'data'}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
'continue',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 'test_model',
|
||||
'last_timestamp': '2024-01-01',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'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',
|
||||
},
|
||||
0.95,
|
||||
'2024-01-01',
|
||||
'',
|
||||
)
|
||||
|
||||
# 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',
|
||||
ANY,
|
||||
)
|
||||
_, export_input = workflow_mock.execute_child_workflow.call_args.args
|
||||
assert export_input['metadata_id'] is None
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_continue_write_metadata_failure_does_not_block_prediction(
|
||||
workflow_mock, prediction_process
|
||||
):
|
||||
# Arrange: write_prediction_metadata's retry_policy has maximum_attempts=1 — a timeout, an
|
||||
# unavailable worker, or any dispatch failure raises straight into the workflow with no
|
||||
# retry. Business rule 1 says this must never block the prediction write (review finding,
|
||||
# 2026-10-01) — confirm the CONTINUE branch still proceeds instead of letting it propagate.
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'CONTINUE', 'counter': 0}]}
|
||||
workflow_mock.execute_activity_method.side_effect = RuntimeError('activity dispatch failed')
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
{'test': 'data'},
|
||||
'continue',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 'test_model',
|
||||
'last_timestamp': '2024-01-01',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'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',
|
||||
},
|
||||
0.95,
|
||||
'2024-01-01',
|
||||
'',
|
||||
document,
|
||||
)
|
||||
|
||||
# Assert: the prediction write still happens, with metadata_id=None, instead of raising.
|
||||
assert result is True
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'subworkflow.format_and_export_prediction',
|
||||
ANY,
|
||||
)
|
||||
_, export_input = workflow_mock.execute_child_workflow.call_args.args
|
||||
assert export_input['metadata_id'] is None
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_repeat_write_metadata_failure_does_not_block_prediction(
|
||||
workflow_mock, prediction_process
|
||||
):
|
||||
# Arrange: only the write_prediction_metadata call (first) fails — repeat_last_prediction
|
||||
# (second) must still run normally; its own failure is a separate, intentional
|
||||
# non-best-effort path, not what this test is about.
|
||||
document = {'filters': [{'name': 'EMPTY_DATA', 'policy': 'REPEAT', 'counter': 0}]}
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
RuntimeError('activity dispatch failed'),
|
||||
MagicMock(),
|
||||
]
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
{'test': 'data'},
|
||||
'repeat',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 'test_model',
|
||||
'last_timestamp': '2024-01-01',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
},
|
||||
0.95,
|
||||
'2024-01-01',
|
||||
'',
|
||||
document,
|
||||
)
|
||||
|
||||
# Assert: repeat_last_prediction still runs, with metadata_id=None, instead of raising.
|
||||
assert result is True
|
||||
assert workflow_mock.execute_activity_method.call_count == 2
|
||||
workflow_mock.execute_activity_method.assert_any_call(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model': 'test_model',
|
||||
'last_timestamp': '2024-01-01',
|
||||
'metadata_id': None,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
|
||||
@@ -850,7 +1074,7 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process):
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'ok'),
|
||||
('continue', 0.95, 'ok', None), # input_gate
|
||||
('continue', 0.95, ''),
|
||||
('continue', 0.95, ''),
|
||||
('continue', 0.95, ''),
|
||||
|
||||
Reference in New Issue
Block a user