refactor(activities): move Mongo activities into MongoActivities (SIENTIAPDE-2191)
Rename PredictionMetadata (prediction_metadata.py) to MongoActivities (mongo.py) and move update_simple_metrics_adaptive_state there from ModelMetrics, so both Mongo-backed activities share one class and repository (PR #59 review, agreed with the class author). write_prediction_metadata is unchanged; the pure hysteresis helpers stay in model_metrics.py. References in activities.py, docstrings and tests updated.
This commit is contained in:
PedroHMCosme
@@ -16,22 +16,22 @@ with workflow.unsafe.imports_passed_through():
|
||||
from laborious.activities.gates import Gates
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
from laborious.activities.mongo import MongoActivities
|
||||
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, PredictionMetadata):
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, MongoActivities):
|
||||
"""
|
||||
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, 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.
|
||||
drift/simple metrics, PI Web API writes, and MongoActivities (the single write point for
|
||||
gate metadata and the simple metrics adaptive state — see `laborious.activities.mongo`).
|
||||
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
|
||||
@@ -73,7 +73,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, PredictionMetad
|
||||
- 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`` (prediction metadata, best-effort) for
|
||||
- mongo_repository: Optional ``MongoDBRepository`` (prediction metadata, adaptive state) 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.
|
||||
|
||||
@@ -153,7 +153,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, PredictionMetad
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
PredictionMetadata.__init__(
|
||||
MongoActivities.__init__(
|
||||
self,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
@@ -202,4 +202,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, PredictionMetad
|
||||
OPC.close(self)
|
||||
ModelMetrics.close(self)
|
||||
API.close(self)
|
||||
PredictionMetadata.close(self)
|
||||
MongoActivities.close(self)
|
||||
|
||||
@@ -219,7 +219,7 @@ class Gates(SientiaMonitoring):
|
||||
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)
|
||||
(see `MongoActivities.write_prediction_metadata` for the actual write)
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for input validation
|
||||
|
||||
@@ -14,8 +14,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
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
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
from sientia_model.analytics.drift_analysis import DriftAnalysis, DriftInsufficientDataError
|
||||
from sientia_model.metrics.regression import RegressionMetrics
|
||||
|
||||
@@ -104,12 +103,6 @@ class ModelMetrics(SientiaMonitoring):
|
||||
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
|
||||
# ponytail: same declaration as PredictionMetadata; the Activities composite sets the
|
||||
# instance value in PredictionMetadata.__init__. A standalone ModelMetrics (unit tests) keeps
|
||||
# None and the adaptive-state activity becomes a no-op.
|
||||
mongo_repository: MongoDBRepository | None = None
|
||||
_PIPELINES_COLLECTION = 'pipelines'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logger: Logger,
|
||||
@@ -520,83 +513,3 @@ class ModelMetrics(SientiaMonitoring):
|
||||
)
|
||||
|
||||
return data.to_dict(orient='records')
|
||||
|
||||
@activity.defn(name='update_simple_metrics_adaptive_state')
|
||||
def update_simple_metrics_adaptive_state(self, input_data: dict[str, Any]) -> dict | None:
|
||||
"""
|
||||
Update the adaptive frequency state on the schedule's `pipelines` doc.
|
||||
|
||||
The orchestrator derives the schedule frequency from `adaptive_state.escalated` and
|
||||
`frequency_alert`. State is read from Mongo, not from the workflow input, because the
|
||||
input only refreshes on the orchestrator tick and would be stale while escalated.
|
||||
`updated_at` is bumped only when `escalated` flips: that is what makes the orchestrator
|
||||
rewrite the schedule, and a streak change alone does not change the frequency.
|
||||
|
||||
Args:
|
||||
input_data: Dict with:
|
||||
- metadata (dict): Workflow metadata; `schedule_name` identifies the doc.
|
||||
- simple_metrics (list[dict]): Rows from `calculate_simple_metrics`; empty for an
|
||||
empty or under-sampled window, which counts as a non-breach run.
|
||||
- thresholds (dict): `simple_metrics_thresholds` from `model_config` (may be None: no breach).
|
||||
- interval_minutes (int): Window size, used to compute the recovery runs.
|
||||
|
||||
Return:
|
||||
The new state, or `None` when adaptive mode is off, Mongo is not configured, the doc
|
||||
is missing or `frequency_alert` is invalid.
|
||||
|
||||
Raises:
|
||||
Exception: Mongo read/write failures, so the workflow run fails explicitly.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
if self.mongo_repository is None:
|
||||
return None
|
||||
|
||||
# Mongo errors propagate on purpose: the run must fail explicitly so the problem shows
|
||||
# up in the Temporal metrics (PR #59 review). Metrics were already exported by then.
|
||||
query = {'schedule_name': metadata['schedule_name']}
|
||||
pipeline = self.mongo_repository.find_one(self._PIPELINES_COLLECTION, query, metadata)
|
||||
if not pipeline:
|
||||
# Likely misconfig: MONGODB_DATABASE_NAME differs from the orchestrator's, or the
|
||||
# schedule_name does not match its doc.
|
||||
self.warning(
|
||||
f'No pipelines doc for schedule {metadata["schedule_name"]}: '
|
||||
'adaptive state skipped',
|
||||
metadata,
|
||||
)
|
||||
return None
|
||||
frequency_alert = pipeline.get('frequency_alert')
|
||||
if not frequency_alert:
|
||||
return None
|
||||
|
||||
try:
|
||||
runs = recovery_runs(input_data['interval_minutes'], str(frequency_alert))
|
||||
except ValueError:
|
||||
# Config error, not a failure: the orchestrator ignores this value too (keeps
|
||||
# `frequency`) and already sends SIMPLE_METRICS_FREQUENCY_ALERT_BELOW_FLOOR for it.
|
||||
self.warning(
|
||||
f'Invalid frequency_alert={frequency_alert!r}: adaptive state skipped', metadata
|
||||
)
|
||||
return None
|
||||
|
||||
breached = bool(
|
||||
threshold_breaches(input_data['simple_metrics'], input_data.get('thresholds'))
|
||||
)
|
||||
state = pipeline.get('adaptive_state')
|
||||
if not isinstance(state, dict):
|
||||
# Missing or hand-edited garbage: not escalated, same as the orchestrator reads it.
|
||||
state = {'escalated': False, 'healthy_streak': 0}
|
||||
new_state = next_adaptive_state(state, breached, runs)
|
||||
if new_state == state:
|
||||
return new_state
|
||||
|
||||
update: dict[str, Any] = {'adaptive_state': new_state}
|
||||
if new_state['escalated'] != bool(state.get('escalated')):
|
||||
# BSON datetime, never a string: the orchestrator normalizes `updated_at` with
|
||||
# `.replace(tzinfo=UTC).strftime(...)` and a string would fail its whole tick.
|
||||
update['updated_at'] = now()
|
||||
self.mongo_repository.update_one(
|
||||
self._PIPELINES_COLLECTION, query, {'$set': update}, metadata
|
||||
)
|
||||
|
||||
self.info(f'Adaptive state of {query["schedule_name"]}: {state} -> {new_state}', metadata)
|
||||
return new_state
|
||||
|
||||
224
laborious/activities/mongo.py
Normal file
224
laborious/activities/mongo.py
Normal file
@@ -0,0 +1,224 @@
|
||||
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
|
||||
from sientia_do.temporal.constants import now
|
||||
|
||||
from laborious.activities.model_metrics import (
|
||||
next_adaptive_state,
|
||||
recovery_runs,
|
||||
threshold_breaches,
|
||||
)
|
||||
|
||||
|
||||
class MongoActivities(SientiaMonitoring):
|
||||
"""
|
||||
Activities that read from or write to MongoDB, sharing one `MongoDBRepository`.
|
||||
|
||||
- `write_prediction_metadata`: centralized, best-effort writer for prediction gate metadata
|
||||
(triggered filter, policy, counter). 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.
|
||||
- `update_simple_metrics_adaptive_state`: adaptive frequency state on the `pipelines` doc of
|
||||
a simple metrics schedule (SIENTIAPDE-1995).
|
||||
"""
|
||||
|
||||
mongo_repository: MongoDBRepository | None = None
|
||||
_PREDICTION_METADATA_COLLECTION = 'prediction_metadata'
|
||||
_PIPELINES_COLLECTION = 'pipelines'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the Mongo activities.
|
||||
|
||||
Args:
|
||||
mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
|
||||
written best-effort and the simple metrics adaptive state is kept. `None` is a
|
||||
valid state: both activities become no-ops.
|
||||
|
||||
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 Mongo repository 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
|
||||
|
||||
@activity.defn(name='update_simple_metrics_adaptive_state')
|
||||
def update_simple_metrics_adaptive_state(self, input_data: dict[str, Any]) -> dict | None:
|
||||
"""
|
||||
Update the adaptive frequency state on the schedule's `pipelines` doc.
|
||||
|
||||
The orchestrator derives the schedule frequency from `adaptive_state.escalated` and
|
||||
`frequency_alert`. State is read from Mongo, not from the workflow input, because the
|
||||
input only refreshes on the orchestrator tick and would be stale while escalated.
|
||||
`updated_at` is bumped only when `escalated` flips: that is what makes the orchestrator
|
||||
rewrite the schedule, and a streak change alone does not change the frequency.
|
||||
|
||||
Args:
|
||||
input_data: Dict with:
|
||||
- metadata (dict): Workflow metadata; `schedule_name` identifies the doc.
|
||||
- simple_metrics (list[dict]): Rows from `calculate_simple_metrics`; empty for an
|
||||
empty or under-sampled window, which counts as a non-breach run.
|
||||
- thresholds (dict): `simple_metrics_thresholds` from `model_config` (may be None: no breach).
|
||||
- interval_minutes (int): Window size, used to compute the recovery runs.
|
||||
|
||||
Return:
|
||||
The new state, or `None` when adaptive mode is off, Mongo is not configured, the doc
|
||||
is missing or `frequency_alert` is invalid.
|
||||
|
||||
Raises:
|
||||
Exception: Mongo read/write failures, so the workflow run fails explicitly.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
if self.mongo_repository is None:
|
||||
return None
|
||||
|
||||
# Mongo errors propagate on purpose: the run must fail explicitly so the problem shows
|
||||
# up in the Temporal metrics (PR #59 review). Metrics were already exported by then.
|
||||
query = {'schedule_name': metadata['schedule_name']}
|
||||
pipeline = self.mongo_repository.find_one(self._PIPELINES_COLLECTION, query, metadata)
|
||||
if not pipeline:
|
||||
# Likely misconfig: MONGODB_DATABASE_NAME differs from the orchestrator's, or the
|
||||
# schedule_name does not match its doc.
|
||||
self.warning(
|
||||
f'No pipelines doc for schedule {metadata["schedule_name"]}: '
|
||||
'adaptive state skipped',
|
||||
metadata,
|
||||
)
|
||||
return None
|
||||
frequency_alert = pipeline.get('frequency_alert')
|
||||
if not frequency_alert:
|
||||
return None
|
||||
|
||||
try:
|
||||
runs = recovery_runs(input_data['interval_minutes'], str(frequency_alert))
|
||||
except ValueError:
|
||||
# Config error, not a failure: the orchestrator ignores this value too (keeps
|
||||
# `frequency`) and already sends SIMPLE_METRICS_FREQUENCY_ALERT_BELOW_FLOOR for it.
|
||||
self.warning(
|
||||
f'Invalid frequency_alert={frequency_alert!r}: adaptive state skipped', metadata
|
||||
)
|
||||
return None
|
||||
|
||||
breached = bool(
|
||||
threshold_breaches(input_data['simple_metrics'], input_data.get('thresholds'))
|
||||
)
|
||||
state = pipeline.get('adaptive_state')
|
||||
if not isinstance(state, dict):
|
||||
# Missing or hand-edited garbage: not escalated, same as the orchestrator reads it.
|
||||
state = {'escalated': False, 'healthy_streak': 0}
|
||||
new_state = next_adaptive_state(state, breached, runs)
|
||||
if new_state == state:
|
||||
return new_state
|
||||
|
||||
update: dict[str, Any] = {'adaptive_state': new_state}
|
||||
if new_state['escalated'] != bool(state.get('escalated')):
|
||||
# BSON datetime, never a string: the orchestrator normalizes `updated_at` with
|
||||
# `.replace(tzinfo=UTC).strftime(...)` and a string would fail its whole tick.
|
||||
update['updated_at'] = now()
|
||||
self.mongo_repository.update_one(
|
||||
self._PIPELINES_COLLECTION, query, {'$set': update}, metadata
|
||||
)
|
||||
|
||||
self.info(f'Adaptive state of {query["schedule_name"]}: {state} -> {new_state}', metadata)
|
||||
return new_state
|
||||
@@ -1,132 +0,0 @@
|
||||
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
|
||||
@@ -184,7 +184,7 @@ class Storage(Postgres, SientiaMonitoring):
|
||||
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`);
|
||||
document for this cycle, if any (see `MongoActivities.write_prediction_metadata`);
|
||||
`None` is a valid value — the column is nullable.
|
||||
|
||||
Args (query_items):
|
||||
|
||||
@@ -267,7 +267,7 @@ class PredictionProcess:
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, including when the
|
||||
activity call itself fails (see `PredictionMetadata.write_prediction_metadata`).
|
||||
activity call itself fails (see `MongoActivities.write_prediction_metadata`).
|
||||
"""
|
||||
if not document:
|
||||
# Skip the remote activity call entirely rather than dispatching it just to have
|
||||
|
||||
@@ -5,8 +5,8 @@ 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.mongo import MongoActivities
|
||||
from laborious.activities.opc import OPC
|
||||
from laborious.activities.prediction_metadata import PredictionMetadata
|
||||
from laborious.activities.storage import Storage
|
||||
|
||||
|
||||
@@ -14,7 +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.MongoActivities.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -24,7 +24,7 @@ def test___init__(
|
||||
mock_minio_repository,
|
||||
mock_api_init,
|
||||
mock_model_metrics_init,
|
||||
mock_prediction_metadata_init,
|
||||
mock_mongo_activities_init,
|
||||
mock_gates_init,
|
||||
mock_opc_init,
|
||||
mock_mlflow_init,
|
||||
@@ -87,7 +87,7 @@ def test___init__(
|
||||
assert isinstance(activities, Gates)
|
||||
assert isinstance(activities, ModelMetrics)
|
||||
assert isinstance(activities, API)
|
||||
assert isinstance(activities, PredictionMetadata)
|
||||
assert isinstance(activities, MongoActivities)
|
||||
|
||||
mock_storage_init.assert_called_once_with(
|
||||
ANY,
|
||||
@@ -131,7 +131,7 @@ def test___init__(
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_prediction_metadata_init.assert_called_once_with(
|
||||
mock_mongo_activities_init.assert_called_once_with(
|
||||
ANY,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
@@ -172,7 +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.MongoActivities')
|
||||
@patch('laborious.activities.activities.ModelMetrics')
|
||||
@patch('laborious.activities.activities.API')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -180,7 +180,7 @@ def test_shutdown(
|
||||
_mock_minio_repository,
|
||||
mock_api_init,
|
||||
mock_model_metrics_init,
|
||||
mock_prediction_metadata_init,
|
||||
mock_mongo_activities_init,
|
||||
mock_gates_init,
|
||||
mock_opc_init,
|
||||
mock_mlflow_init,
|
||||
@@ -243,7 +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()
|
||||
mock_mongo_activities_init.close.assert_called_once()
|
||||
|
||||
|
||||
@patch('laborious.activities.activities.SientiaMLflowRepository')
|
||||
@@ -252,7 +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.MongoActivities.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -262,7 +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_mongo_activities_init,
|
||||
_mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
@@ -326,7 +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.MongoActivities.__init__')
|
||||
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
||||
@patch('laborious.activities.activities.API.__init__')
|
||||
@patch('laborious.activities.activities.MinioRepository')
|
||||
@@ -336,7 +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_mongo_activities_init,
|
||||
mock_gates_init,
|
||||
_mock_opc_init,
|
||||
_mock_mlflow_init,
|
||||
@@ -398,9 +398,9 @@ 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 PredictionMetadata, not a second
|
||||
# The same repository built here is the one handed to MongoActivities, not a second
|
||||
# instance.
|
||||
mock_prediction_metadata_init.assert_called_once_with(
|
||||
mock_mongo_activities_init.assert_called_once_with(
|
||||
ANY,
|
||||
mongo_repository=mock_mongo_repository_cls.return_value,
|
||||
logger=logger,
|
||||
|
||||
@@ -147,8 +147,8 @@ def test_input_gate_with_filter(gates_activity):
|
||||
|
||||
# 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.
|
||||
# `MongoActivities.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()
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
from pandas import DataFrame, Series, Timestamp
|
||||
@@ -1301,156 +1300,3 @@ def test_frequency_to_seconds_rejects_invalid(frequency):
|
||||
)
|
||||
def test_recovery_runs(interval_minutes, frequency_alert, expected):
|
||||
assert recovery_runs(interval_minutes, frequency_alert) == expected
|
||||
|
||||
|
||||
_SM_METADATA = {**metadata['metadata'], 'schedule_name': 'm1-simple_metrics-60m'}
|
||||
_BREACH_INPUT = {
|
||||
'metadata': _SM_METADATA,
|
||||
'simple_metrics': [{'metric': 'rmse', 'value': 2.0}],
|
||||
'thresholds': {'rmse_max': 1.0},
|
||||
'interval_minutes': 60,
|
||||
}
|
||||
_HEALTHY_INPUT = {**_BREACH_INPUT, 'simple_metrics': [{'metric': 'rmse', 'value': 0.5}]}
|
||||
_EMPTY_INPUT = {**_BREACH_INPUT, 'simple_metrics': []}
|
||||
|
||||
|
||||
@fixture
|
||||
def mongo(model_metrics_activity):
|
||||
repository = MagicMock()
|
||||
model_metrics_activity.mongo_repository = repository
|
||||
return repository
|
||||
|
||||
|
||||
def _pipeline(**fields) -> dict:
|
||||
return {'schedule_name': 'm1-simple_metrics-60m', 'frequency': '1h', **fields}
|
||||
|
||||
|
||||
def _set_of(mongo: MagicMock) -> dict:
|
||||
collection, query, update, _metadata = mongo.update_one.call_args.args
|
||||
assert collection == 'pipelines'
|
||||
assert query == {'schedule_name': 'm1-simple_metrics-60m'}
|
||||
return update['$set']
|
||||
|
||||
|
||||
def test_adaptive_state_no_mongo_is_noop(model_metrics_activity):
|
||||
assert model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
|
||||
|
||||
def test_adaptive_state_missing_pipeline_doc_is_noop(model_metrics_activity, mongo):
|
||||
mongo.find_one.return_value = None
|
||||
assert model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
mongo.update_one.assert_not_called()
|
||||
# Likely misconfig (other MONGODB_DATABASE_NAME or schedule_name mismatch): not silent.
|
||||
model_metrics_activity.warning.assert_called_once()
|
||||
assert 'm1-simple_metrics-60m' in model_metrics_activity.warning.call_args.args[0]
|
||||
|
||||
|
||||
def test_adaptive_state_without_frequency_alert_is_noop(model_metrics_activity, mongo):
|
||||
mongo.find_one.return_value = _pipeline()
|
||||
assert model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
mongo.update_one.assert_not_called()
|
||||
model_metrics_activity.warning.assert_not_called()
|
||||
|
||||
|
||||
def test_adaptive_state_breach_escalates_and_bumps_updated_at(model_metrics_activity, mongo):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m')
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 0}
|
||||
update = _set_of(mongo)
|
||||
assert update['adaptive_state'] == {'escalated': True, 'healthy_streak': 0}
|
||||
# BSON datetime, never a string: the orchestrator calls .replace(tzinfo=...) on it.
|
||||
assert isinstance(update['updated_at'], datetime)
|
||||
assert update['updated_at'].tzinfo is not None
|
||||
|
||||
|
||||
def test_adaptive_state_healthy_run_increments_streak_without_updated_at(
|
||||
model_metrics_activity, mongo
|
||||
):
|
||||
# 60 min window / 5m = N 12, so one healthy run only moves the streak.
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='5m', adaptive_state={'escalated': True, 'healthy_streak': 0}
|
||||
)
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_HEALTHY_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 1}
|
||||
assert _set_of(mongo) == {'adaptive_state': {'escalated': True, 'healthy_streak': 1}}
|
||||
|
||||
|
||||
def test_adaptive_state_recovers_after_window_renewed(model_metrics_activity, mongo):
|
||||
# 60 min window / 30m = N 2: second consecutive healthy run recovers.
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='30m', adaptive_state={'escalated': True, 'healthy_streak': 1}
|
||||
)
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_HEALTHY_INPUT)
|
||||
|
||||
assert result == {'escalated': False, 'healthy_streak': 0}
|
||||
update = _set_of(mongo)
|
||||
assert update['adaptive_state'] == {'escalated': False, 'healthy_streak': 0}
|
||||
assert isinstance(update['updated_at'], datetime)
|
||||
|
||||
|
||||
def test_adaptive_state_empty_window_counts_as_healthy_run(model_metrics_activity, mongo):
|
||||
# Spec decision 8: an empty/under-sampled window never escalates but counts toward recovery.
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='5m', adaptive_state={'escalated': True, 'healthy_streak': 3}
|
||||
)
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_EMPTY_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 4}
|
||||
|
||||
|
||||
def test_adaptive_state_empty_window_never_escalates(model_metrics_activity, mongo):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m')
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_EMPTY_INPUT)
|
||||
|
||||
assert result == {'escalated': False, 'healthy_streak': 0}
|
||||
mongo.update_one.assert_not_called()
|
||||
|
||||
|
||||
def test_adaptive_state_unchanged_does_not_write(model_metrics_activity, mongo):
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='5m', adaptive_state={'escalated': False, 'healthy_streak': 0}
|
||||
)
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_HEALTHY_INPUT)
|
||||
|
||||
assert result == {'escalated': False, 'healthy_streak': 0}
|
||||
mongo.update_one.assert_not_called()
|
||||
|
||||
|
||||
@mark.parametrize('frequency_alert', ['5min', '0m', 5])
|
||||
def test_adaptive_state_invalid_frequency_alert_is_skipped(
|
||||
model_metrics_activity, mongo, frequency_alert
|
||||
):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert=frequency_alert)
|
||||
|
||||
assert model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
mongo.update_one.assert_not_called()
|
||||
model_metrics_activity.warning.assert_called_once()
|
||||
|
||||
|
||||
@mark.parametrize('adaptive_state', ['garbage', ['escalated']])
|
||||
def test_adaptive_state_malformed_state_is_not_escalated(
|
||||
model_metrics_activity, mongo, adaptive_state
|
||||
):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m', adaptive_state=adaptive_state)
|
||||
|
||||
result = model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 0}
|
||||
mongo.update_one.assert_called_once()
|
||||
|
||||
|
||||
@mark.parametrize('failing_call', ['find_one', 'update_one'])
|
||||
def test_adaptive_state_mongo_error_propagates(model_metrics_activity, mongo, failing_call):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m')
|
||||
getattr(mongo, failing_call).side_effect = RuntimeError('mongo down')
|
||||
|
||||
with raises(RuntimeError, match='mongo down'):
|
||||
model_metrics_activity.update_simple_metrics_adaptive_state(_BREACH_INPUT)
|
||||
|
||||
341
tests/laborious/activities/test_mongo.py
Normal file
341
tests/laborious/activities/test_mongo.py
Normal file
@@ -0,0 +1,341 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
from pytest import fixture, mark, raises
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from laborious.activities.mongo import MongoActivities
|
||||
|
||||
|
||||
@fixture
|
||||
def mongo() -> MagicMock:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@fixture
|
||||
def prediction_metadata(mongo: MagicMock):
|
||||
activity = MongoActivities(
|
||||
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 = MongoActivities(
|
||||
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 == MongoActivities._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):
|
||||
"""``MongoActivities.close`` should close the Mongo client and clear the reference."""
|
||||
mongo = MagicMock()
|
||||
activity = MongoActivities(
|
||||
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 = MongoActivities(
|
||||
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
|
||||
|
||||
|
||||
_SM_METADATA = {**metadata['metadata'], 'schedule_name': 'm1-simple_metrics-60m'}
|
||||
_BREACH_INPUT = {
|
||||
'metadata': _SM_METADATA,
|
||||
'simple_metrics': [{'metric': 'rmse', 'value': 2.0}],
|
||||
'thresholds': {'rmse_max': 1.0},
|
||||
'interval_minutes': 60,
|
||||
}
|
||||
_HEALTHY_INPUT = {**_BREACH_INPUT, 'simple_metrics': [{'metric': 'rmse', 'value': 0.5}]}
|
||||
_EMPTY_INPUT = {**_BREACH_INPUT, 'simple_metrics': []}
|
||||
|
||||
|
||||
@fixture
|
||||
def mongo_activities(mongo: MagicMock):
|
||||
activity = MongoActivities(
|
||||
mongo_repository=mongo,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
activity.info = MagicMock()
|
||||
activity.warning = MagicMock()
|
||||
return activity
|
||||
|
||||
|
||||
def _pipeline(**fields) -> dict:
|
||||
return {'schedule_name': 'm1-simple_metrics-60m', 'frequency': '1h', **fields}
|
||||
|
||||
|
||||
def _set_of(mongo: MagicMock) -> dict:
|
||||
collection, query, update, _metadata = mongo.update_one.call_args.args
|
||||
assert collection == 'pipelines'
|
||||
assert query == {'schedule_name': 'm1-simple_metrics-60m'}
|
||||
return update['$set']
|
||||
|
||||
|
||||
def test_adaptive_state_no_mongo_is_noop():
|
||||
activity = MongoActivities(
|
||||
mongo_repository=None,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=MagicMock(spec=MetricsController),
|
||||
)
|
||||
assert activity.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
|
||||
|
||||
def test_adaptive_state_missing_pipeline_doc_is_noop(mongo_activities, mongo):
|
||||
mongo.find_one.return_value = None
|
||||
assert mongo_activities.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
mongo.update_one.assert_not_called()
|
||||
# Likely misconfig (other MONGODB_DATABASE_NAME or schedule_name mismatch): not silent.
|
||||
mongo_activities.warning.assert_called_once()
|
||||
assert 'm1-simple_metrics-60m' in mongo_activities.warning.call_args.args[0]
|
||||
|
||||
|
||||
def test_adaptive_state_without_frequency_alert_is_noop(mongo_activities, mongo):
|
||||
mongo.find_one.return_value = _pipeline()
|
||||
assert mongo_activities.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
mongo.update_one.assert_not_called()
|
||||
mongo_activities.warning.assert_not_called()
|
||||
|
||||
|
||||
def test_adaptive_state_breach_escalates_and_bumps_updated_at(mongo_activities, mongo):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m')
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_BREACH_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 0}
|
||||
update = _set_of(mongo)
|
||||
assert update['adaptive_state'] == {'escalated': True, 'healthy_streak': 0}
|
||||
# BSON datetime, never a string: the orchestrator calls .replace(tzinfo=...) on it.
|
||||
assert isinstance(update['updated_at'], datetime)
|
||||
assert update['updated_at'].tzinfo is not None
|
||||
|
||||
|
||||
def test_adaptive_state_healthy_run_increments_streak_without_updated_at(mongo_activities, mongo):
|
||||
# 60 min window / 5m = N 12, so one healthy run only moves the streak.
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='5m', adaptive_state={'escalated': True, 'healthy_streak': 0}
|
||||
)
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_HEALTHY_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 1}
|
||||
assert _set_of(mongo) == {'adaptive_state': {'escalated': True, 'healthy_streak': 1}}
|
||||
|
||||
|
||||
def test_adaptive_state_recovers_after_window_renewed(mongo_activities, mongo):
|
||||
# 60 min window / 30m = N 2: second consecutive healthy run recovers.
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='30m', adaptive_state={'escalated': True, 'healthy_streak': 1}
|
||||
)
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_HEALTHY_INPUT)
|
||||
|
||||
assert result == {'escalated': False, 'healthy_streak': 0}
|
||||
update = _set_of(mongo)
|
||||
assert update['adaptive_state'] == {'escalated': False, 'healthy_streak': 0}
|
||||
assert isinstance(update['updated_at'], datetime)
|
||||
|
||||
|
||||
def test_adaptive_state_empty_window_counts_as_healthy_run(mongo_activities, mongo):
|
||||
# Spec decision 8: an empty/under-sampled window never escalates but counts toward recovery.
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='5m', adaptive_state={'escalated': True, 'healthy_streak': 3}
|
||||
)
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_EMPTY_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 4}
|
||||
|
||||
|
||||
def test_adaptive_state_empty_window_never_escalates(mongo_activities, mongo):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m')
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_EMPTY_INPUT)
|
||||
|
||||
assert result == {'escalated': False, 'healthy_streak': 0}
|
||||
mongo.update_one.assert_not_called()
|
||||
|
||||
|
||||
def test_adaptive_state_unchanged_does_not_write(mongo_activities, mongo):
|
||||
mongo.find_one.return_value = _pipeline(
|
||||
frequency_alert='5m', adaptive_state={'escalated': False, 'healthy_streak': 0}
|
||||
)
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_HEALTHY_INPUT)
|
||||
|
||||
assert result == {'escalated': False, 'healthy_streak': 0}
|
||||
mongo.update_one.assert_not_called()
|
||||
|
||||
|
||||
@mark.parametrize('frequency_alert', ['5min', '0m', 5])
|
||||
def test_adaptive_state_invalid_frequency_alert_is_skipped(
|
||||
mongo_activities, mongo, frequency_alert
|
||||
):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert=frequency_alert)
|
||||
|
||||
assert mongo_activities.update_simple_metrics_adaptive_state(_BREACH_INPUT) is None
|
||||
mongo.update_one.assert_not_called()
|
||||
mongo_activities.warning.assert_called_once()
|
||||
|
||||
|
||||
@mark.parametrize('adaptive_state', ['garbage', ['escalated']])
|
||||
def test_adaptive_state_malformed_state_is_not_escalated(mongo_activities, mongo, adaptive_state):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m', adaptive_state=adaptive_state)
|
||||
|
||||
result = mongo_activities.update_simple_metrics_adaptive_state(_BREACH_INPUT)
|
||||
|
||||
assert result == {'escalated': True, 'healthy_streak': 0}
|
||||
mongo.update_one.assert_called_once()
|
||||
|
||||
|
||||
@mark.parametrize('failing_call', ['find_one', 'update_one'])
|
||||
def test_adaptive_state_mongo_error_propagates(mongo_activities, mongo, failing_call):
|
||||
mongo.find_one.return_value = _pipeline(frequency_alert='5m')
|
||||
getattr(mongo, failing_call).side_effect = RuntimeError('mongo down')
|
||||
|
||||
with raises(RuntimeError, match='mongo down'):
|
||||
mongo_activities.update_simple_metrics_adaptive_state(_BREACH_INPUT)
|
||||
@@ -1,179 +0,0 @@
|
||||
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
|
||||
Reference in New Issue
Block a user