e6018af23f
Update README, requirements, and E2E tests for improved configuration and functionality - Enhanced the README with updated model configuration examples, including the addition of an alias for production. - Removed the `requirements-light.txt` file and updated `requirements-local.txt` and `requirements.txt` to replace `asyncua` with `opcua`. - Refactored E2E test scenarios to utilize scenario input files for better maintainability and clarity. - Improved test coverage for MinIO offload functionality and added new helper functions for loading scenario inputs. - Updated `values.yaml` to reflect new global configurations and environment variables for the laborious worker.
336 lines
11 KiB
Python
336 lines
11 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 ``retrain_reports`` rows with all required columns; success rows
|
|
carry the new ``version`` / ``mlflow_run_id`` / ``mlflow_experiment_id``
|
|
while failure rows leave them ``NULL``.
|
|
"""
|
|
|
|
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
|
|
|
|
EXPECTED_RETRAIN_REPORT_COLUMNS = [
|
|
'id',
|
|
'model_id',
|
|
'model_name',
|
|
'timestamp',
|
|
'status',
|
|
'version',
|
|
'mlflow_run_id',
|
|
'mlflow_experiment_id',
|
|
'created_at',
|
|
]
|
|
|
|
|
|
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'
|
|
run_info.experiment_id = 'experiment-id'
|
|
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. ``retrain_reports`` 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 predictions_schema.retrain_reports '
|
|
'WHERE model_id = :m'
|
|
),
|
|
{'m': model_id},
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
|
|
assert len(rows) == 1
|
|
for column in EXPECTED_RETRAIN_REPORT_COLUMNS:
|
|
assert column in rows[0], f'Missing retrain report column: {column}'
|
|
|
|
row = rows[0]
|
|
assert row['model_id'] == 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'
|
|
assert row['mlflow_experiment_id'] == 'experiment-id'
|
|
assert row['timestamp'] is not None
|
|
assert row['created_at'] 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 predictions_schema.retrain_reports '
|
|
'WHERE model_id = :m'
|
|
),
|
|
{'m': model_id},
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
|
|
assert len(rows) == 1
|
|
row = rows[0]
|
|
assert row['model_id'] == 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 predictions_schema.retrain_reports '
|
|
'WHERE model_id = :m'
|
|
),
|
|
{'m': 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 predictions_schema.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 predictions_schema.retrain_reports WHERE model_id = :m'),
|
|
{'m': model_id},
|
|
).scalar()
|
|
assert count == 0
|