10c7e292b9
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.
347 lines
12 KiB
Python
347 lines
12 KiB
Python
"""
|
|
End-to-end tests for the MinimalRetrain workflow.
|
|
|
|
The MLflow registry is fully stubbed because no real artifacts exist in a
|
|
test container; we only validate that the workflow:
|
|
|
|
- Loads training data via ``load_query_with_minio_offload``.
|
|
- Calls ``retrain_model`` with a payload pointing at MinIO.
|
|
- Calls ``update_production_model`` only when retrain succeeds.
|
|
- Persists ``sientia_data.log_retrain`` rows with all required columns;
|
|
success rows carry the new ``version`` / ``mlflow_run_id`` /
|
|
``mlflow_experiment_id`` while failure rows leave them ``NULL``.
|
|
|
|
The production DDL drops the legacy ``id`` / ``created_at`` columns and
|
|
moves ``mlflow_experiment_id`` to ``int8`` and ``model_id`` to ``text``.
|
|
The stubs used here therefore emit ``experiment_id`` as an integer to fit
|
|
the new column type.
|
|
"""
|
|
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timedelta, timezone
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
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.minimal_retrain import MinimalRetrain
|
|
|
|
# Columns defined by the production DDL for ``sientia_data.log_retrain``.
|
|
# The legacy ``retrain_reports`` table had ``id`` and ``created_at``; the new
|
|
# DDL drops both. ``mlflow_experiment_id`` is ``int8`` and ``model_id`` is
|
|
# ``text``.
|
|
EXPECTED_RETRAIN_REPORT_COLUMNS = [
|
|
'mlflow_experiment_id',
|
|
'mlflow_run_id',
|
|
'model_id',
|
|
'model_name',
|
|
'status',
|
|
'timestamp',
|
|
'version',
|
|
]
|
|
|
|
|
|
def _retrain_input(model_id: int, **overrides) -> dict:
|
|
"""Load and override the minimal-retrain base scenario."""
|
|
payload = load_scenario_input('minimal_retrain_base.json', model_id=model_id)
|
|
payload.update(overrides)
|
|
return payload
|
|
|
|
|
|
def _seed_retrain_training_rows(postgres_engine, model_id: int) -> None:
|
|
"""
|
|
Insert training rows in long format that pivot cleanly into
|
|
``index=timestamp`` / ``columns={sensor_1, sensor_2}`` for ``retrain_model``.
|
|
"""
|
|
target_timestamps = [
|
|
(datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes=10 - i))
|
|
.strftime('%Y-%m-%d %H:%M:%S%z')
|
|
for i in range(5)
|
|
]
|
|
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],
|
|
},
|
|
)
|
|
|
|
|
|
def _configure_retrain_happy_path(mlflow_repository_stub) -> None:
|
|
"""
|
|
Wire ``mlflow_repository_stub`` so retrain + update_production succeed.
|
|
|
|
Mocks (in order of consumption):
|
|
|
|
- ``_client.get_model_version_by_alias``: returns ``mv`` with a stable
|
|
``run_id`` (used as ``source_run_id``).
|
|
- ``get_cached_model``: returns a wrapper exposing inert ``retrain`` and
|
|
``store_model`` methods.
|
|
- ``start_run``: returns a context manager yielding a ``run_info`` with
|
|
run/experiment ids.
|
|
- ``log_params``: inert.
|
|
- ``_client.search_model_versions``: returns one registry entry whose
|
|
``version`` is promoted by ``update_production_model``.
|
|
- ``promote_to_alias``: inert success.
|
|
"""
|
|
mv_src = MagicMock()
|
|
mv_src.run_id = 'source-run-id'
|
|
|
|
new_version = MagicMock()
|
|
new_version.version = '7'
|
|
new_version.run_id = 'retrain-run-id'
|
|
|
|
mlflow_repository_stub._client.get_model_version_by_alias.return_value = mv_src
|
|
|
|
cached_wrapper = MagicMock()
|
|
cached_wrapper.retrain = MagicMock(return_value=None)
|
|
cached_wrapper.store_model = MagicMock(return_value=None)
|
|
mlflow_repository_stub.get_cached_model.return_value = cached_wrapper
|
|
|
|
@contextmanager
|
|
def fake_start_run(**kwargs):
|
|
run_info = MagicMock()
|
|
run_info.run_id = 'retrain-run-id'
|
|
# ``mlflow_experiment_id`` is ``int8`` in the new DDL, so we feed an
|
|
# integer-compatible id from the stubbed run info.
|
|
run_info.experiment_id = 4242
|
|
yield run_info
|
|
|
|
mlflow_repository_stub.start_run.side_effect = fake_start_run
|
|
mlflow_repository_stub.log_params = MagicMock(return_value=None)
|
|
mlflow_repository_stub._client.search_model_versions.return_value = [new_version]
|
|
mlflow_repository_stub.promote_to_alias = MagicMock(return_value=None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_minimal_retrain_happy_path_writes_success_report(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_minimal_retrain: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario MR.1.1: Retrain succeeds. ``sientia_data.log_retrain`` must
|
|
contain a success row with version/mlflow_run_id/mlflow_experiment_id
|
|
populated and the registry must have been told to promote the new version
|
|
to the configured alias.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 711
|
|
|
|
_seed_retrain_training_rows(postgres_engine, model_id)
|
|
_configure_retrain_happy_path(mlflow_repository_stub)
|
|
|
|
input_data = _retrain_input(model_id)
|
|
|
|
with patch('laborious.activities.mlflow.mlflow.log_artifact') as log_artifact_mock:
|
|
await start_and_await_workflow(
|
|
client,
|
|
MinimalRetrain.run,
|
|
input_data,
|
|
make_workflow_id('test-retrain-happy'),
|
|
)
|
|
|
|
assert log_artifact_mock.called, 'retrain_model should log the input CSV artifact'
|
|
|
|
with postgres_engine.connect() as conn:
|
|
rows = (
|
|
conn.execute(
|
|
text(
|
|
'SELECT * FROM sientia_data.log_retrain '
|
|
'WHERE model_id = :m'
|
|
),
|
|
{'m': str(model_id)},
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
|
|
assert len(rows) == 1
|
|
for column in EXPECTED_RETRAIN_REPORT_COLUMNS:
|
|
assert column in rows[0], f'Missing log_retrain column: {column}'
|
|
|
|
row = rows[0]
|
|
# ``model_id`` is now ``text``; compare against the stringified id.
|
|
assert row['model_id'] == str(model_id)
|
|
assert row['model_name'] == 'test_model'
|
|
assert row['status'] == 'Model retrained successfully.'
|
|
assert row['version'] == '7'
|
|
assert row['mlflow_run_id'] == 'retrain-run-id'
|
|
# ``mlflow_experiment_id`` is now ``int8``; assert the integer value
|
|
# provided by the stubbed run info.
|
|
assert row['mlflow_experiment_id'] == 4242
|
|
assert row['timestamp'] is not None
|
|
|
|
mlflow_repository_stub.promote_to_alias.assert_called_once()
|
|
promote_kwargs = mlflow_repository_stub.promote_to_alias.call_args.kwargs
|
|
assert promote_kwargs['model_name'] == 'test_model'
|
|
assert promote_kwargs['version'] == '7'
|
|
assert promote_kwargs['alias'] == 'production'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_minimal_retrain_failure_writes_report_without_version_columns(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_minimal_retrain: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario MR.2.1: ``wrapper.retrain`` raises. The activity must catch the
|
|
error, return ``success=False`` so ``update_production_model`` is skipped,
|
|
and ``format_retrain_report`` must produce a row with the error message
|
|
and NULL version columns.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 721
|
|
|
|
_seed_retrain_training_rows(postgres_engine, model_id)
|
|
_configure_retrain_happy_path(mlflow_repository_stub)
|
|
mlflow_repository_stub.get_cached_model.return_value.retrain.side_effect = RuntimeError(
|
|
'training did not converge'
|
|
)
|
|
|
|
input_data = _retrain_input(model_id)
|
|
|
|
with patch('laborious.activities.mlflow.mlflow.log_artifact'):
|
|
await start_and_await_workflow(
|
|
client,
|
|
MinimalRetrain.run,
|
|
input_data,
|
|
make_workflow_id('test-retrain-failure'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
rows = (
|
|
conn.execute(
|
|
text(
|
|
'SELECT * FROM sientia_data.log_retrain '
|
|
'WHERE model_id = :m'
|
|
),
|
|
{'m': str(model_id)},
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
|
|
assert len(rows) == 1
|
|
row = rows[0]
|
|
assert row['model_id'] == str(model_id)
|
|
assert row['model_name'] == 'test_model'
|
|
assert 'training did not converge' in row['status']
|
|
assert row['version'] is None
|
|
assert row['mlflow_run_id'] is None
|
|
assert row['mlflow_experiment_id'] is None
|
|
|
|
mlflow_repository_stub.promote_to_alias.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_minimal_retrain_missing_target_writes_failure_report(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_minimal_retrain: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario MR.2.2: ``model_config`` does not declare ``target``. The retrain
|
|
activity must short-circuit before any MLflow call and the report row must
|
|
carry the explicit guard message.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 722
|
|
|
|
_seed_retrain_training_rows(postgres_engine, model_id)
|
|
_configure_retrain_happy_path(mlflow_repository_stub)
|
|
|
|
input_data = _retrain_input(model_id, model_config={})
|
|
|
|
with patch('laborious.activities.mlflow.mlflow.log_artifact'):
|
|
await start_and_await_workflow(
|
|
client,
|
|
MinimalRetrain.run,
|
|
input_data,
|
|
make_workflow_id('test-retrain-missing-target'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
row = (
|
|
conn.execute(
|
|
text(
|
|
'SELECT * FROM sientia_data.log_retrain '
|
|
'WHERE model_id = :m'
|
|
),
|
|
{'m': str(model_id)},
|
|
)
|
|
.mappings()
|
|
.first()
|
|
)
|
|
|
|
assert row is not None
|
|
assert 'target' in row['status'].lower(), (
|
|
f"expected target-missing message, got status={row['status']!r}"
|
|
)
|
|
assert row['version'] is None
|
|
mlflow_repository_stub.get_cached_model.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_minimal_retrain_no_training_data_does_not_persist_report(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_minimal_retrain: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
Scenario MR.3.1: When the training query returns no rows the workflow must
|
|
not persist any report row. The workflow currently raises plain
|
|
``ValueError`` which Temporal treats as a workflow-task failure (causing
|
|
indefinite retries until the test environment times out), so the assertion
|
|
here is constrained to the persistence side-effect. See ``e2e/CODE_ISSUES.md``
|
|
issue MR-1 for the recommended ``ApplicationError`` fix.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 731
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
|
|
|
_configure_retrain_happy_path(mlflow_repository_stub)
|
|
|
|
input_data = _retrain_input(model_id)
|
|
|
|
with pytest.raises(Exception), patch('laborious.activities.mlflow.mlflow.log_artifact'):
|
|
await start_and_await_workflow(
|
|
client,
|
|
MinimalRetrain.run,
|
|
input_data,
|
|
make_workflow_id('test-retrain-no-data'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM sientia_data.log_retrain WHERE model_id = :m'),
|
|
{'m': str(model_id)},
|
|
).scalar()
|
|
assert count == 0
|