Refactor ModelMetrics to utilize DriftAnalysis for drift detection - Replaced ModelAnalysis with DriftAnalysis in the ModelMetrics class to enhance drift detection capabilities. - Updated method signatures and documentation to reflect the changes in target_name and return values. - Adjusted data handling to ensure compatibility with the new analysis methods and improved clarity in the drift metrics dataframe preparation.
500 lines
18 KiB
Python
500 lines
18 KiB
Python
"""
|
|
End-to-end tests for the Drift workflow.
|
|
|
|
The drift suite drives the **real** ``sientia_model.analytics.drift_analysis.DriftAnalysis``
|
|
analyzer (no mocking). Each scenario exercises the full pipeline:
|
|
|
|
laborious_data (Postgres)
|
|
-> load_custom_query
|
|
-> calculate_drift (DriftAnalysis univariate + multivariate)
|
|
-> export_data_to_postgres (sientia_data.drift_metrics)
|
|
|
|
Coverage focus:
|
|
|
|
- Happy path persists every column required by ``sientia_data.drift_metrics``
|
|
with a valid reference dataset downloaded from MLflow.
|
|
- 30% fallback path activates when the MLflow reference is unavailable and
|
|
emits the ``MODEL_METRICS_REFERENCE_DATA_WARNING`` notification.
|
|
- Empty target data short-circuits the workflow without persisting anything.
|
|
- Invalid ``chunk_period`` is rejected by ``calculate_drift``.
|
|
- ``chunk_period='s'`` preserves second-level precision in
|
|
``chunk_start_date``.
|
|
|
|
Tests assert behavioral / structural properties (column presence, NOT NULL
|
|
constraints, business-key invariants) rather than exact numeric values, since
|
|
those depend on the real analyzer implementation and synthetic data.
|
|
"""
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import (
|
|
insert_target_data_for_drift,
|
|
load_scenario_input,
|
|
make_workflow_id,
|
|
start_and_await_workflow,
|
|
)
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.drift import Drift
|
|
|
|
# Drift columns persisted on every row in ``sientia_data.drift_metrics`` —
|
|
# mirrors the production DDL.
|
|
EXPECTED_DRIFT_COLUMNS = [
|
|
'id',
|
|
'model_id',
|
|
'feature',
|
|
'method',
|
|
'value',
|
|
'alert',
|
|
'chunk_index',
|
|
'chunk_start_date',
|
|
'chunk_end_date',
|
|
'accurate',
|
|
'timestamp',
|
|
'created_at',
|
|
]
|
|
|
|
# Columns the DDL marks as NOT NULL. ``feature`` and ``timestamp`` are
|
|
# nullable in the production schema (multivariate rows do not bind to a
|
|
# single feature; ``timestamp`` is allowed to be empty when upstream data has
|
|
# no usable instant).
|
|
NON_NULL_DRIFT_COLUMNS = {
|
|
'id',
|
|
'model_id',
|
|
'method',
|
|
'value',
|
|
'alert',
|
|
'chunk_index',
|
|
'chunk_start_date',
|
|
'chunk_end_date',
|
|
'accurate',
|
|
'created_at',
|
|
}
|
|
|
|
DEFAULT_DRIFT_METHODS = ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
|
|
|
|
|
|
def _drift_input(model_id: int, **overrides) -> dict:
|
|
"""Load the base drift scenario JSON and apply ad-hoc overrides."""
|
|
input_data = load_scenario_input('drift_base.json', model_id=model_id)
|
|
input_data.update(overrides)
|
|
return input_data
|
|
|
|
|
|
def _recent_minute_timestamps(count: int, offset_minutes: int = 6) -> list[str]:
|
|
"""
|
|
Build ``count`` consecutive UTC minute timestamps placed in the recent past.
|
|
|
|
The Drift workflow filters target rows with ``timestamp > NOW() - INTERVAL``,
|
|
so timestamps must be recent for tests to retrieve any data. Snapping to
|
|
minute precision keeps the helper deterministic regardless of clock skew.
|
|
|
|
Args:
|
|
- count (int): How many consecutive minute timestamps to generate.
|
|
- offset_minutes (int): Minutes ago for the EARLIEST generated timestamp.
|
|
|
|
Return:
|
|
list[str]: ISO strings with ``+0000`` offset, one per minute.
|
|
"""
|
|
base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(
|
|
minutes=offset_minutes
|
|
)
|
|
return [
|
|
(base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z') for i in range(count)
|
|
]
|
|
|
|
|
|
def _configure_reference_csv(mlflow_repository_stub, reference_rows: pd.DataFrame) -> None:
|
|
"""
|
|
Wire ``mlflow_repository_stub`` so ``get_reference_data`` returns
|
|
``reference_rows`` by writing them to ``dst_path/evaluation_data.csv``.
|
|
|
|
Args:
|
|
- mlflow_repository_stub: External MLflow repository fixture.
|
|
- reference_rows (pd.DataFrame): Rows to expose as the production reference.
|
|
"""
|
|
|
|
def _download(run_id: str, artifact_path: str, dst_path: str, metadata=None):
|
|
target = Path(dst_path) / 'evaluation_data.csv'
|
|
reference_rows.to_csv(target, index=False)
|
|
|
|
mlflow_repository_stub._client.get_model_version_by_alias.return_value = MagicMock(
|
|
run_id='fake-reference-run'
|
|
)
|
|
mlflow_repository_stub.download_artifacts.side_effect = _download
|
|
|
|
|
|
def _force_reference_unavailable(mlflow_repository_stub) -> None:
|
|
"""Make ``get_reference_data`` return ``None`` by failing alias resolution."""
|
|
mlflow_repository_stub._client.get_model_version_by_alias.side_effect = Exception(
|
|
'no production alias registered'
|
|
)
|
|
|
|
|
|
def _select_drift_rows(postgres_engine, model_id: int) -> list[dict]:
|
|
"""Read every persisted drift row for ``model_id`` ordered by chunk/feature/method."""
|
|
with postgres_engine.connect() as conn:
|
|
rows = (
|
|
conn.execute(
|
|
text(
|
|
'SELECT * FROM sientia_data.drift_metrics '
|
|
'WHERE model_id = :m '
|
|
'ORDER BY chunk_index, feature, method'
|
|
),
|
|
{'m': str(model_id)},
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
return [dict(row) for row in rows]
|
|
|
|
|
|
def _assert_required_columns_populated(rows: list[dict]) -> None:
|
|
"""Validate column presence and NOT NULL constraints on every row."""
|
|
assert rows, 'expected at least one drift row to be persisted'
|
|
seen_columns = set(rows[0].keys())
|
|
for column in EXPECTED_DRIFT_COLUMNS:
|
|
assert column in seen_columns, f'Missing drift column in postgres: {column}'
|
|
for row in rows:
|
|
for column in NON_NULL_DRIFT_COLUMNS:
|
|
assert row[column] is not None, f"Column '{column}' is NULL in {row}"
|
|
assert 'p_value' not in row, 'p_value must not be persisted to drift_metrics'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_drift_happy_path_persists_all_columns_with_reference_data(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_drift: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario D.1.1: Happy path with reference data downloaded from MLflow.
|
|
|
|
Drives the full pipeline against the real ``DriftAnalysis``. Asserts:
|
|
|
|
- One row is persisted per ``(chunk_index, feature, method)`` combination
|
|
plus the multivariate row block, with every column required by
|
|
``sientia_data.drift_metrics`` populated.
|
|
- The three default univariate methods are forwarded to the analyzer.
|
|
- ``model_id`` and ``timestamp`` are stamped by the activity (not by the
|
|
analyzer); ``timestamp`` equals ``max(target_data.timestamp)`` and is
|
|
identical on every persisted row.
|
|
- ``chunk_start_date`` / ``chunk_end_date`` are persisted as ISO text so
|
|
the analyzer's nanosecond-precision boundaries survive the ``text``
|
|
column type.
|
|
- ``accurate=True`` because the reference dataset was available.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 411
|
|
|
|
target_timestamps = _recent_minute_timestamps(count=10)
|
|
insert_target_data_for_drift(
|
|
postgres_engine,
|
|
model_id=model_id,
|
|
timestamps=target_timestamps,
|
|
variables_values={
|
|
'sensor_1': [10.0 + i * 0.1 for i in range(10)],
|
|
'sensor_2': [20.0 + i * 0.5 for i in range(10)],
|
|
},
|
|
)
|
|
|
|
reference_df = pd.DataFrame(
|
|
{
|
|
'timestamp': [
|
|
f'2023-12-31 11:{minute:02d}:00+00:00' for minute in range(10)
|
|
],
|
|
'sensor_1': [9.0 + i * 0.05 for i in range(10)],
|
|
'sensor_2': [18.0 + i * 0.25 for i in range(10)],
|
|
}
|
|
)
|
|
_configure_reference_csv(mlflow_repository_stub, reference_df)
|
|
|
|
input_data = _drift_input(model_id)
|
|
await start_and_await_workflow(
|
|
client, Drift.run, input_data, make_workflow_id('test-drift-happy-path')
|
|
)
|
|
|
|
rows = _select_drift_rows(postgres_engine, model_id)
|
|
|
|
exported_csv_path = '/tmp/test_drift_happy_path_exported.csv'
|
|
pd.DataFrame(rows).to_csv(exported_csv_path, index=False)
|
|
print(
|
|
f'\n[test_drift_happy_path] Exported drift dataframe '
|
|
f'({len(rows)} rows) -> {exported_csv_path}'
|
|
)
|
|
|
|
_assert_required_columns_populated(rows)
|
|
|
|
# The activity drops the target column from the feature list, so only
|
|
# ``sensor_2`` participates in univariate analysis (``sensor_1`` is the
|
|
# configured target). Multivariate produces one row per chunk regardless.
|
|
univariate_rows = [r for r in rows if r['feature'] != 'multivariate']
|
|
multivariate_rows = [r for r in rows if r['feature'] == 'multivariate']
|
|
assert univariate_rows, 'expected univariate drift rows for non-target features'
|
|
assert multivariate_rows, 'expected one multivariate drift row per chunk'
|
|
|
|
# All three default methods must reach the analyzer.
|
|
assert {r['method'] for r in univariate_rows} == set(DEFAULT_DRIFT_METHODS)
|
|
assert all(r['method'] == 'multivariate' for r in multivariate_rows)
|
|
assert {r['feature'] for r in univariate_rows} == {'sensor_2'}
|
|
|
|
# ``timestamp`` is stamped uniformly with ``max(target_data.timestamp)``.
|
|
expected_timestamp = pd.to_datetime(max(target_timestamps), utc=True)
|
|
persisted_timestamps = {pd.to_datetime(r['timestamp'], utc=True) for r in rows}
|
|
assert len(persisted_timestamps) == 1, (
|
|
'timestamp must be uniform across all drift rows '
|
|
f'(got {len(persisted_timestamps)} distinct values)'
|
|
)
|
|
assert pd.Timestamp(persisted_timestamps.pop()) == expected_timestamp, (
|
|
'timestamp must equal max(target_data.timestamp)'
|
|
)
|
|
|
|
# ``model_id`` is stamped by ``calculate_drift`` (not produced by the analyzer).
|
|
assert all(r['model_id'] == str(model_id) for r in rows), (
|
|
'model_id must be stamped on every drift row'
|
|
)
|
|
|
|
# Reference path → accurate=True.
|
|
assert all(r['accurate'] is True for r in rows), (
|
|
'reference path should mark all rows as accurate'
|
|
)
|
|
|
|
# ISO text serialization preserves ordering between start/end of each chunk.
|
|
for row in rows:
|
|
assert 'T' in row['chunk_start_date'], (
|
|
f"chunk_start_date should be ISO text, got {row['chunk_start_date']!r}"
|
|
)
|
|
assert 'T' in row['chunk_end_date'], (
|
|
f"chunk_end_date should be ISO text, got {row['chunk_end_date']!r}"
|
|
)
|
|
assert row['chunk_start_date'] <= row['chunk_end_date'], (
|
|
f'chunk_start_date must precede chunk_end_date '
|
|
f"(start={row['chunk_start_date']}, end={row['chunk_end_date']})"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_drift_uses_30pct_fallback_when_reference_unavailable(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_drift: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
notification_inserts,
|
|
):
|
|
"""
|
|
Scenario D.1.2: ``get_reference_data`` returns ``None`` (production alias
|
|
missing), so ``calculate_drift`` falls back to using the first 30% of
|
|
target rows as reference. Persisted rows must report ``accurate=False``
|
|
and a ``MODEL_METRICS_REFERENCE_DATA_WARNING`` notification must be
|
|
emitted to mongo.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 412
|
|
|
|
target_timestamps = _recent_minute_timestamps(count=10)
|
|
insert_target_data_for_drift(
|
|
postgres_engine,
|
|
model_id=model_id,
|
|
timestamps=target_timestamps,
|
|
variables_values={
|
|
'sensor_1': [10.0 + i * 0.1 for i in range(10)],
|
|
'sensor_2': [20.0 + i * 0.5 for i in range(10)],
|
|
},
|
|
)
|
|
|
|
_force_reference_unavailable(mlflow_repository_stub)
|
|
|
|
input_data = _drift_input(model_id)
|
|
await start_and_await_workflow(
|
|
client, Drift.run, input_data, make_workflow_id('test-drift-fallback')
|
|
)
|
|
|
|
rows = _select_drift_rows(postgres_engine, model_id)
|
|
_assert_required_columns_populated(rows)
|
|
|
|
assert all(r['accurate'] is False for r in rows), (
|
|
'fallback path must mark all rows as inaccurate'
|
|
)
|
|
|
|
fallback_warnings = [
|
|
call
|
|
for call in notification_inserts.call_args_list
|
|
if call.args
|
|
and isinstance(call.args[0], dict)
|
|
and call.args[0].get('notification_id') == 'MODEL_METRICS_REFERENCE_DATA_WARNING'
|
|
]
|
|
assert len(fallback_warnings) >= 1, 'expected reference fallback warning notification'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_drift_empty_target_data_short_circuits_workflow(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_drift: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario D.3.1: When ``load_custom_query`` returns no rows the workflow
|
|
must return early without invoking the analyzer or writing any drift rows.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 431
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
|
|
|
_force_reference_unavailable(mlflow_repository_stub)
|
|
|
|
input_data = _drift_input(model_id)
|
|
await start_and_await_workflow(
|
|
client, Drift.run, input_data, make_workflow_id('test-drift-empty-target')
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
|
|
{'m': str(model_id)},
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_drift_invalid_chunk_period_raises_value_error(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_drift: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario D.4.2: ``calculate_drift`` validates ``chunk_period`` and rejects
|
|
anything other than ``min`` / ``s``. The workflow must surface the
|
|
``ValueError`` and persist nothing.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 442
|
|
|
|
target_timestamps = _recent_minute_timestamps(count=5)
|
|
insert_target_data_for_drift(
|
|
postgres_engine,
|
|
model_id=model_id,
|
|
timestamps=target_timestamps,
|
|
variables_values={
|
|
'sensor_1': [1.0, 2.0, 3.0, 4.0, 5.0],
|
|
'sensor_2': [10.0, 20.0, 30.0, 40.0, 50.0],
|
|
},
|
|
)
|
|
_force_reference_unavailable(mlflow_repository_stub)
|
|
|
|
input_data = _drift_input(model_id, chunk_period='hour')
|
|
|
|
with pytest.raises(Exception) as excinfo:
|
|
await start_and_await_workflow(
|
|
client,
|
|
Drift.run,
|
|
input_data,
|
|
make_workflow_id('test-drift-bad-chunk-period'),
|
|
)
|
|
|
|
# Temporal wraps the activity ValueError in WorkflowFailureError; the
|
|
# message may live on ``.message`` or ``str(exc)`` depending on the SDK
|
|
# error class, so walk the cause chain looking for the guard text.
|
|
cause_descriptions = []
|
|
current: BaseException | None = excinfo.value
|
|
while current is not None:
|
|
cause_descriptions.append(
|
|
getattr(current, 'message', None) or str(current) or repr(current)
|
|
)
|
|
current = current.__cause__
|
|
assert any('Invalid chunk period' in msg for msg in cause_descriptions), (
|
|
f'Expected ValueError about chunk period in chain, got: {cause_descriptions}'
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
|
|
{'m': str(model_id)},
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_drift_chunk_period_seconds_preserves_seconds_in_chunk_start_date(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_drift: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario D.4.3: With ``chunk_period='s'`` the persisted ``chunk_start_date``
|
|
column must preserve second-level precision so consumers can audit the
|
|
actual chunk boundary.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 443
|
|
|
|
base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes=2)
|
|
# Three samples spaced by 30 seconds inside two adjacent minutes.
|
|
target_timestamps = [
|
|
base.strftime('%Y-%m-%d %H:%M:%S%z'),
|
|
(base + timedelta(seconds=30)).strftime('%Y-%m-%d %H:%M:%S%z'),
|
|
(base + timedelta(minutes=1)).strftime('%Y-%m-%d %H:%M:%S%z'),
|
|
]
|
|
|
|
insert_target_data_for_drift(
|
|
postgres_engine,
|
|
model_id=model_id,
|
|
timestamps=target_timestamps,
|
|
variables_values={
|
|
'sensor_1': [1.0, 2.0, 3.0],
|
|
'sensor_2': [10.0, 20.0, 30.0],
|
|
},
|
|
)
|
|
_force_reference_unavailable(mlflow_repository_stub)
|
|
|
|
input_data = _drift_input(model_id, chunk_period='s')
|
|
await start_and_await_workflow(
|
|
client,
|
|
Drift.run,
|
|
input_data,
|
|
make_workflow_id('test-drift-chunk-seconds'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
rows = (
|
|
conn.execute(
|
|
text(
|
|
'SELECT chunk_start_date FROM sientia_data.drift_metrics '
|
|
'WHERE model_id = :m ORDER BY chunk_index'
|
|
),
|
|
{'m': str(model_id)},
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
|
|
assert rows, 'expected at least one drift row to be persisted'
|
|
# At least one chunk must land on second-30, proving the analyzer chunked
|
|
# at sub-minute granularity instead of collapsing everything to minute=0.
|
|
seconds_present = {pd.Timestamp(r['chunk_start_date']).second for r in rows}
|
|
assert 30 in seconds_present, (
|
|
f'expected at least one chunk_start_date with seconds=30, got {seconds_present}'
|
|
)
|