SIENTIAPDE-1646

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.
This commit is contained in:
vitor-aignosi
2026-05-08 16:39:12 -03:00
parent e6018af23f
commit 10c7e292b9
29 changed files with 843 additions and 1266 deletions

View File

@@ -1,22 +1,31 @@
"""
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:
- Full pipeline persists drift rows with **all** columns expected by the
``predictions_schema.drift`` table (model_id, feature, method, value, drift,
chunk, timestamp, timestamp_end, accurate, plus DB-managed id/created_at/updated_at).
- ``get_reference_data`` happy path (CSV downloaded from MLflow stub) and
fallback path (30% of target data when reference is unavailable).
- ``calculate_drift`` invariants: ``p_value`` dropped, duplicates removed, rows
outside target timestamps filtered, default ``drift_metrics`` propagated.
- Failure paths: ``ModelAnalysis.get_drift_metrics_dataframe`` raises ⇒
workflow completes without writes; empty target data ⇒ workflow short-circuits;
invalid ``chunk_period`` ⇒ activity raises and workflow surfaces the error.
- 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 decimal import Decimal
from pathlib import Path
from unittest.mock import MagicMock
@@ -27,7 +36,6 @@ from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.helpers import (
build_drift_dataframe,
insert_target_data_for_drift,
load_scenario_input,
make_workflow_id,
@@ -36,21 +44,41 @@ from e2e.helpers import (
from laborious.activities.activities import Activities
from laborious.workflows.drift import Drift
# Drift columns that must be present (and non-null where required) on every row.
# Drift columns persisted on every row in ``sientia_data.drift_metrics`` —
# mirrors the production DDL.
EXPECTED_DRIFT_COLUMNS = [
'id',
'model_id',
'feature',
'method',
'value',
'drift',
'chunk',
'alert',
'chunk_index',
'chunk_start_date',
'chunk_end_date',
'accurate',
'timestamp',
'timestamp_end',
'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',
'updated_at',
]
}
DEFAULT_DRIFT_METHODS = ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
def _drift_input(model_id: int, **overrides) -> dict:
@@ -60,42 +88,37 @@ def _drift_input(model_id: int, **overrides) -> dict:
return input_data
def _five_minute_window(offset_minutes: int = 6) -> tuple[list[str], list[str]]:
def _recent_minute_timestamps(count: int, offset_minutes: int = 6) -> list[str]:
"""
Build five consecutive UTC minute timestamps positioned in the recent past.
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. We snap to
minute precision and back off ``offset_minutes`` minutes so all chunks land
well inside the default 60-minute interval defined in ``drift_base.json``.
so timestamps must be recent for tests to retrieve any data. Snapping to
minute precision keeps the helper deterministic regardless of clock skew.
Args:
- offset_minutes: How many minutes ago the most recent chunk should be.
- count (int): How many consecutive minute timestamps to generate.
- offset_minutes (int): Minutes ago for the EARLIEST generated timestamp.
Return:
Tuple ``(target_timestamps, chunk_timestamps)``:
- ``target_timestamps``: ISO strings with ``+0000`` used as ``timestamp``
and ``created_at`` columns when inserting target rows.
- ``chunk_timestamps``: ``YYYY-MM-DD HH:MM`` truncations matching what
``calculate_drift`` filters on for ``chunk_period='min'``.
list[str]: ISO strings with ``+0000`` offset, one per minute.
"""
base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(
minutes=offset_minutes
)
target_timestamps = [
(base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z') for i in range(5)
return [
(base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z') for i in range(count)
]
chunk_timestamps = [
(base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M') for i in range(5)
]
return target_timestamps, chunk_timestamps
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):
@@ -115,22 +138,34 @@ def _force_reference_unavailable(mlflow_repository_stub) -> None:
)
def _assert_drift_columns_complete(rows, *, expected_count: int) -> None:
"""Validate row count and that the canonical drift columns are populated."""
assert len(rows) == expected_count, (
f'Expected {expected_count} drift rows persisted, got {len(rows)}'
)
seen_columns = set(rows[0]._mapping.keys()) if rows else set()
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:
mapping = dict(row._mapping)
for column in EXPECTED_DRIFT_COLUMNS:
if column == 'value':
# ``value`` is nullable in the table; skip null check, only ensure key exists.
continue
assert mapping[column] is not None, f"Column '{column}' is NULL in {mapping}"
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
@@ -141,112 +176,111 @@ async def test_drift_happy_path_persists_all_columns_with_reference_data(
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_stub,
):
"""
Scenario D.1.1: Happy path with reference data downloaded from MLflow.
Validates the full drift pipeline produces one row per (chunk, feature, method)
combination plus the multivariate row block, with **all** columns required by
``predictions_schema.drift`` populated. Uses a deterministic mocked drift
dataframe so the postgres assertions remain stable across runs.
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, chunk_timestamps = _five_minute_window()
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, 11.0, 12.0, 13.0, 14.0],
'sensor_2': [20.0, 21.0, 22.0, 23.0, 24.0],
'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': ['2023-12-31 11:00:00+00:00', '2023-12-31 11:01:00+00:00'],
'sensor_1': [9.5, 9.6],
'sensor_2': [19.5, 19.6],
'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)
features = ['sensor_1', 'sensor_2']
methods = ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
drift_flags = {
('sensor_1', 'wasserstein'): True,
('sensor_2', 'wasserstein'): True,
('sensor_2', 'jensen_shannon'): True,
}
model_analysis_stub.set_drift_dataframe(
lambda univariate, multivariate: build_drift_dataframe(
timestamps=chunk_timestamps,
features=features,
methods=methods,
statistic=0.42,
drift_flags=drift_flags,
multivariate_value=15.5,
multivariate_drift=True,
)
)
input_data = _drift_input(model_id)
await start_and_await_workflow(
client, Drift.run, input_data, make_workflow_id('test-drift-happy-path')
)
with postgres_engine.connect() as conn:
rows = (
conn.execute(
text(
'SELECT * FROM predictions_schema.drift '
'WHERE model_id = :model_id ORDER BY chunk, feature, method'
),
{'model_id': model_id},
)
.mappings()
.all()
)
rows = _select_drift_rows(postgres_engine, model_id)
expected_rows = len(chunk_timestamps) * (len(features) * len(methods) + 1) # univariate + multivariate
assert len(rows) == expected_rows
for column in EXPECTED_DRIFT_COLUMNS:
assert column in rows[0], f'Missing drift column in postgres: {column}'
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}'
)
for row in rows:
for column in EXPECTED_DRIFT_COLUMNS:
if column == 'value':
continue
assert row[column] is not None, f"Column '{column}' is NULL in {dict(row)}"
_assert_required_columns_populated(rows)
multivariate_rows = [r for r in rows if r['feature'] == 'multivariate']
# 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'
assert len(multivariate_rows) == len(chunk_timestamps)
# 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 all(r['drift'] is True for r in multivariate_rows)
assert all(Decimal(str(r['value'])) == Decimal('15.5') for r in multivariate_rows)
assert {r['feature'] for r in univariate_rows} == {'sensor_2'}
assert len(univariate_rows) == len(features) * len(methods) * len(chunk_timestamps)
assert {r['method'] for r in univariate_rows} == set(methods)
assert {r['feature'] for r in univariate_rows} == set(features)
# ``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)'
)
drift_pairs = {(r['feature'], r['method']): r['drift'] for r in univariate_rows}
for (feature, method), expected_drift in drift_flags.items():
assert drift_pairs[(feature, method)] is expected_drift
# ``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'
)
assert all(r['accurate'] is True for r in rows), 'reference path should mark rows as accurate'
# Reference path → accurate=True.
assert all(r['accurate'] is True for r in rows), (
'reference path should mark all rows as accurate'
)
assert all(
r['timestamp_end'].endswith(':59.999999999') and 'T' in r['timestamp_end']
for r in rows
), 'timestamp_end should preserve the high-precision ISO string from ModelAnalysis'
assert all('p_value' not in r for r in rows), 'p_value must be dropped before postgres'
# 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
@@ -257,60 +291,42 @@ async def test_drift_uses_30pct_fallback_when_reference_unavailable(
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_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 warning
notification must be emitted to mongo.
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, chunk_timestamps = _five_minute_window()
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, 10.1, 10.2, 10.3, 10.4],
'sensor_2': [20.0, 20.1, 20.2, 20.3, 20.4],
'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)
model_analysis_stub.set_drift_dataframe(
lambda univariate, multivariate: build_drift_dataframe(
timestamps=chunk_timestamps,
features=['sensor_1'],
methods=['kolmogorov_smirnov'],
statistic=0.7,
include_multivariate=False,
)
)
input_data = _drift_input(model_id)
await start_and_await_workflow(
client, Drift.run, input_data, make_workflow_id('test-drift-fallback')
)
with postgres_engine.connect() as conn:
rows = (
conn.execute(
text('SELECT * FROM predictions_schema.drift WHERE model_id = :m'),
{'m': model_id},
)
.mappings()
.all()
)
rows = _select_drift_rows(postgres_engine, model_id)
_assert_required_columns_populated(rows)
assert len(rows) == len(chunk_timestamps)
assert all(r['accurate'] is False for r in rows), 'fallback path must mark rows as inaccurate'
assert all(r['feature'] == 'sensor_1' for r in rows)
assert all(r['accurate'] is False for r in rows), (
'fallback path must mark all rows as inaccurate'
)
fallback_warnings = [
call
@@ -322,168 +338,6 @@ async def test_drift_uses_30pct_fallback_when_reference_unavailable(
assert len(fallback_warnings) >= 1, 'expected reference fallback warning notification'
@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_drops_p_value_and_dedupes_by_timestamp_method_feature(
temporal_test_env: WorkflowEnvironment,
temporal_worker_drift: Worker,
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_stub,
):
"""
Scenario D.2.1: ModelAnalysis returns duplicate (timestamp, method, feature)
rows and a populated ``p_value`` column. ``calculate_drift`` must deduplicate
keeping the first occurrence and the persisted rows must not contain
``p_value``.
"""
client = temporal_test_env.client
model_id = 421
target_timestamps, chunk_timestamps = _five_minute_window()
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],
},
)
duplicates = [
# Re-emit the first chunk for sensor_1/kolmogorov_smirnov with a different value
# so we can prove the dedup keeps the FIRST row.
{
'timestamp': pd.Timestamp(chunk_timestamps[0]),
'feature': 'sensor_1',
'metric': 'kolmogorov_smirnov',
'statistic': 0.99,
'p_value': 0.02,
'alert': True,
'chunk_index': 0,
'chunk_start_date': pd.Timestamp(chunk_timestamps[0]),
'chunk_end_date': pd.Timestamp(f'{chunk_timestamps[0]}:59.999999999'),
}
]
model_analysis_stub.set_drift_dataframe(
lambda univariate, multivariate: build_drift_dataframe(
timestamps=chunk_timestamps,
features=['sensor_1'],
methods=['kolmogorov_smirnov'],
statistic=0.5,
include_multivariate=False,
extra_rows=duplicates,
)
)
_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-dedupe')
)
with postgres_engine.connect() as conn:
rows = (
conn.execute(
text(
'SELECT chunk, feature, method, value FROM predictions_schema.drift '
'WHERE model_id = :m ORDER BY chunk'
),
{'m': model_id},
)
.mappings()
.all()
)
assert len(rows) == len(chunk_timestamps), 'duplicates must be removed before persistence'
first_chunk_rows = [r for r in rows if r['chunk'] == 0]
assert len(first_chunk_rows) == 1
assert Decimal(str(first_chunk_rows[0]['value'])) == Decimal('0.5'), (
'dedup must keep the FIRST occurrence (statistic=0.5), not the duplicate (statistic=0.99)'
)
columns = set(rows[0].keys())
assert 'p_value' not in columns
@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_drops_rows_outside_target_timestamps(
temporal_test_env: WorkflowEnvironment,
temporal_worker_drift: Worker,
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_stub,
):
"""
Scenario D.2.2: Drift rows whose timestamps are not present in the target
data (e.g. came from the reference distribution) must be discarded so the
saved drift only reflects analysis chunks.
"""
client = temporal_test_env.client
model_id = 422
target_timestamps, chunk_timestamps = _five_minute_window()
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],
},
)
out_of_range_timestamp = pd.Timestamp('2099-12-31 23:59')
extra = [
{
'timestamp': out_of_range_timestamp,
'feature': 'sensor_1',
'metric': 'kolmogorov_smirnov',
'statistic': 0.5,
'p_value': 0.0,
'alert': True,
'chunk_index': 99,
'chunk_start_date': out_of_range_timestamp,
'chunk_end_date': pd.Timestamp('2099-12-31 23:59:59.999999999'),
}
]
model_analysis_stub.set_drift_dataframe(
lambda univariate, multivariate: build_drift_dataframe(
timestamps=chunk_timestamps,
features=['sensor_1'],
methods=['kolmogorov_smirnov'],
include_multivariate=False,
extra_rows=extra,
)
)
_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-tts-filter')
)
with postgres_engine.connect() as conn:
chunks = [
r[0]
for r in conn.execute(
text(
'SELECT chunk FROM predictions_schema.drift WHERE model_id = :m ORDER BY chunk'
),
{'m': model_id},
).all()
]
assert chunks == [0, 1, 2, 3, 4], 'out-of-range timestamps must be filtered out'
@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_empty_target_data_short_circuits_workflow(
@@ -492,33 +346,19 @@ async def test_drift_empty_target_data_short_circuits_workflow(
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_stub,
):
"""
Scenario D.3.1: When ``load_custom_query`` returns no rows the workflow must
return early without invoking ModelAnalysis or writing any drift rows.
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 predictions_schema.laborious_data WHERE model_id = {model_id}'))
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
_force_reference_unavailable(mlflow_repository_stub)
drift_calls = {'count': 0}
def _factory(univariate, multivariate):
drift_calls['count'] += 1
return build_drift_dataframe(
timestamps=['2024-01-01 12:00'],
features=['sensor_1'],
methods=['kolmogorov_smirnov'],
include_multivariate=False,
)
model_analysis_stub.set_drift_dataframe(_factory)
input_data = _drift_input(model_id)
await start_and_await_workflow(
client, Drift.run, input_data, make_workflow_id('test-drift-empty-target')
@@ -526,121 +366,10 @@ async def test_drift_empty_target_data_short_circuits_workflow(
with postgres_engine.connect() as conn:
count = conn.execute(
text('SELECT COUNT(*) FROM predictions_schema.drift WHERE model_id = :m'),
{'m': model_id},
text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
{'m': str(model_id)},
).scalar()
assert count == 0
assert drift_calls['count'] == 0, 'ModelAnalysis must not be invoked when target data is empty'
@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_model_analysis_failure_keeps_workflow_alive_no_writes(
temporal_test_env: WorkflowEnvironment,
temporal_worker_drift: Worker,
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_stub,
notification_inserts,
):
"""
Scenario D.3.2: ``ModelAnalysis.get_drift_metrics_dataframe`` raises. The
activity must catch the error, send an error notification and return ``[]``
so the workflow completes without persisting drift rows.
"""
client = temporal_test_env.client
model_id = 432
target_timestamps, _ = _five_minute_window()
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)
model_analysis_stub.raise_on_get_drift_metrics_dataframe(
RuntimeError('drift analyzer crashed')
)
input_data = _drift_input(model_id)
await start_and_await_workflow(
client, Drift.run, input_data, make_workflow_id('test-drift-analyzer-error')
)
with postgres_engine.connect() as conn:
count = conn.execute(
text('SELECT COUNT(*) FROM predictions_schema.drift WHERE model_id = :m'),
{'m': model_id},
).scalar()
assert count == 0
error_notifications = [
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_GET_DRIFT_METRICS_ERROR'
]
assert len(error_notifications) >= 1
@pytest.mark.asyncio
@pytest.mark.integration
async def test_drift_default_drift_metrics_propagated_to_analyzer(
temporal_test_env: WorkflowEnvironment,
temporal_worker_drift: Worker,
test_activities: Activities,
postgres_engine,
mlflow_repository_stub,
model_analysis_stub,
):
"""
Scenario D.4.1: When ``drift_metrics`` is omitted from input the workflow
must default to ``['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']``
and forward exactly that list to ``ModelAnalysis.detect_univariate_drift``.
"""
client = temporal_test_env.client
model_id = 441
target_timestamps, chunk_timestamps = _five_minute_window()
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)
model_analysis_stub.set_drift_dataframe(
lambda univariate, multivariate: build_drift_dataframe(
timestamps=chunk_timestamps,
features=['sensor_1'],
methods=['kolmogorov_smirnov'],
include_multivariate=False,
)
)
input_data = _drift_input(model_id)
input_data.pop('drift_metrics', None)
await start_and_await_workflow(
client, Drift.run, input_data, make_workflow_id('test-drift-default-methods')
)
instance = model_analysis_stub.last_instance
assert instance is not None, 'ModelAnalysis must have been instantiated'
assert len(instance.detect_univariate_drift_calls) == 1
forwarded_methods = instance.detect_univariate_drift_calls[0]['methods']
assert forwarded_methods == ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
@pytest.mark.asyncio
@@ -655,12 +384,12 @@ async def test_drift_invalid_chunk_period_raises_value_error(
"""
Scenario D.4.2: ``calculate_drift`` validates ``chunk_period`` and rejects
anything other than ``min`` / ``s``. The workflow must surface the
``ValueError`` to the caller and persist nothing.
``ValueError`` and persist nothing.
"""
client = temporal_test_env.client
model_id = 442
target_timestamps, _ = _five_minute_window()
target_timestamps = _recent_minute_timestamps(count=5)
insert_target_data_for_drift(
postgres_engine,
model_id=model_id,
@@ -682,8 +411,9 @@ async def test_drift_invalid_chunk_period_raises_value_error(
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.
# 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:
@@ -697,42 +427,36 @@ async def test_drift_invalid_chunk_period_raises_value_error(
with postgres_engine.connect() as conn:
count = conn.execute(
text('SELECT COUNT(*) FROM predictions_schema.drift WHERE model_id = :m'),
{'m': model_id},
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_timestamp_column(
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,
model_analysis_stub,
):
"""
Scenario D.4.3: With ``chunk_period='s'`` the persisted ``timestamp``
column must preserve second-level precision instead of being flattened to
the start of the minute, and ``chunk_period`` must be propagated to the
analyzer so it actually chunks at second granularity.
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'),
]
chunk_timestamps_full = [
base.strftime('%Y-%m-%d %H:%M:%S'),
(base + timedelta(seconds=30)).strftime('%Y-%m-%d %H:%M:%S'),
(base + timedelta(minutes=1)).strftime('%Y-%m-%d %H:%M:%S'),
]
insert_target_data_for_drift(
postgres_engine,
@@ -745,15 +469,6 @@ async def test_drift_chunk_period_seconds_preserves_seconds_in_timestamp_column(
)
_force_reference_unavailable(mlflow_repository_stub)
model_analysis_stub.set_drift_dataframe(
lambda univariate, multivariate: build_drift_dataframe(
timestamps=chunk_timestamps_full,
features=['sensor_1'],
methods=['kolmogorov_smirnov'],
include_multivariate=False,
)
)
input_data = _drift_input(model_id, chunk_period='s')
await start_and_await_workflow(
client,
@@ -766,21 +481,19 @@ async def test_drift_chunk_period_seconds_preserves_seconds_in_timestamp_column(
rows = (
conn.execute(
text(
'SELECT chunk, timestamp FROM predictions_schema.drift '
'WHERE model_id = :m ORDER BY chunk'
'SELECT chunk_start_date FROM sientia_data.drift_metrics '
'WHERE model_id = :m ORDER BY chunk_index'
),
{'m': model_id},
{'m': str(model_id)},
)
.mappings()
.all()
)
assert [r['chunk'] for r in rows] == [0, 1, 2]
seconds_present = {r['timestamp'].second for r in rows}
assert seconds_present == {0, 30}, (
f'expected seconds 0 and 30 to be preserved, got {seconds_present}'
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}'
)
instance = model_analysis_stub.last_instance
assert instance is not None
assert instance.detect_univariate_drift_calls[0]['chunk_period'] == 's'