431 lines
14 KiB
Python
431 lines
14 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
|
|
import pandas as pd
|
|
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
|
|
from sientia_model.wrappers.sientia_model import SientiaModel
|
|
|
|
# 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',
|
|
]
|
|
|
|
# Matches ``retrain_model`` return ``message`` when ``success`` is True (also written to ``log_retrain.status``).
|
|
RETRAIN_ACTIVITY_SUCCESS_MESSAGE = 'Model retrained successfully.'
|
|
|
|
|
|
class _FakeSientiaModelForMinimalRetrain(SientiaModel):
|
|
"""
|
|
Fake SientiaModel that uses the real SientiaModel lifecycle to surface
|
|
index-alignment issues during ``retrain()``.
|
|
|
|
It intentionally performs strict alignment inside ``_retrain_model``:
|
|
``y.loc[x.index]``.
|
|
"""
|
|
|
|
def __init__(self, *, target: str = 'sensor_1'):
|
|
super().__init__(
|
|
model_type='FakeMinimalRetrain',
|
|
model_version='0.0.0',
|
|
model=object(),
|
|
transformer=object(),
|
|
)
|
|
self.target = target
|
|
self.model_is_fitted = True
|
|
self.force_retrain_error = False
|
|
|
|
def store_model( # type: ignore[override]
|
|
self,
|
|
name: str,
|
|
signature=None,
|
|
pip_requirements=None,
|
|
code_path=None,
|
|
) -> None:
|
|
# No-op: E2E tests validate workflow persistence, not real MLflow artifacts.
|
|
return None
|
|
|
|
def _predict(self, data: pd.DataFrame):
|
|
pred = pd.DataFrame({'prediction': [0.5] * len(data)}, index=data.index)
|
|
return pred, {}
|
|
|
|
def _transform(self, data: pd.DataFrame):
|
|
out = data.drop(columns=[self.target], errors='ignore').copy()
|
|
out.index = data.index
|
|
return out, {}
|
|
|
|
def _train_transformer(self, train_data: pd.DataFrame, val_data: pd.DataFrame) -> None:
|
|
return None
|
|
|
|
def _train_model(
|
|
self,
|
|
x: pd.DataFrame,
|
|
y: pd.DataFrame,
|
|
x_val: pd.DataFrame | None = None,
|
|
y_val: pd.DataFrame | None = None,
|
|
) -> None:
|
|
return None
|
|
|
|
def _retrain_transformer(self, data: pd.DataFrame) -> None:
|
|
return None
|
|
|
|
def _retrain_model(self, x: pd.DataFrame, y: pd.DataFrame | None) -> None:
|
|
if self.force_retrain_error:
|
|
raise RuntimeError('training did not converge')
|
|
if y is None:
|
|
return
|
|
# Strict alignment on purpose to reproduce the production failure mode.
|
|
_ = y.loc[x.index]
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_repository_stub():
|
|
"""
|
|
Override the shared E2E fixture: return a real fake ``SientiaModel`` wrapper
|
|
instead of a MagicMock wrapper.
|
|
"""
|
|
repo = MagicMock()
|
|
repo._client = MagicMock()
|
|
|
|
wrapper = _FakeSientiaModelForMinimalRetrain(target='sensor_1')
|
|
repo.get_cached_model = MagicMock(return_value=wrapper)
|
|
return repo
|
|
|
|
|
|
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``).
|
|
- ``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
|
|
|
|
@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'),
|
|
)
|
|
|
|
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]
|
|
for column in EXPECTED_RETRAIN_REPORT_COLUMNS:
|
|
assert column in row, f'Missing log_retrain column: {column}'
|
|
|
|
assert row['status'] == RETRAIN_ACTIVITY_SUCCESS_MESSAGE, (
|
|
"Expected retrain_model to return success (experiment_response['success'] is True). "
|
|
'Persisted log_retrain.status is the activity message; when success is False the run '
|
|
'never reaches mlflow.log_artifact — diagnose the retrain failure from status below, '
|
|
'not from a skipped artifact upload. '
|
|
f"Got status={row['status']!r}, version={row.get('version')!r}, "
|
|
f"mlflow_run_id={row.get('mlflow_run_id')!r}."
|
|
)
|
|
|
|
assert log_artifact_mock.called, (
|
|
'After a successful retrain, retrain_model must call mlflow.log_artifact for the '
|
|
'input CSV inside start_run.'
|
|
)
|
|
|
|
# ``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['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.force_retrain_error = True
|
|
|
|
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
|