Code import - branch release/SIENTIAPDE-1646
This commit is contained in:
430
e2e/test_minimal_retrain.py
Normal file
430
e2e/test_minimal_retrain.py
Normal file
@@ -0,0 +1,430 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user