vitor-aignosi
c12f2996c2
QTZPOC-13: Introduce joblib fallback for MLflow model loading Enhances model loading resilience by attempting to deserialize models directly with joblib if MLflow's native flavor loaders fail. This mechanism handles cases where .pkl files are present in the artifact. Adds joblib as a core dependency and includes new end-to-end and unit tests to validate fallback behavior, covering success, no-fallback, and double-failure scenarios. Also updates .gitignore for .claude/ files and incorporates minor refactorings and linting adjustments across E2E tests and utility functions.
303 lines
10 KiB
Python
303 lines
10 KiB
Python
"""
|
|
End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import make_workflow_id, start_and_await_workflow
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_1_1_happy_path_complete_success(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms."""
|
|
client = temporal_test_env.client
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123'))
|
|
insert_sql = """
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
(123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
|
(123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
|
(123, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 123,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 123,
|
|
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 123',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'input_filters': {
|
|
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_transform_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_predict_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
'opc_output_config': {},
|
|
'pi_web_api_output_config': {},
|
|
'save_transform': True,
|
|
'prediction_store_policy': 'lts:1',
|
|
'model_config': {
|
|
'retention_minutes': 0,
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'sklearn',
|
|
},
|
|
'datetime_columns': ['timestamp', 'created_at'],
|
|
}
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-predictions-batch'),
|
|
)
|
|
|
|
schema_name = 'predictions_schema'
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments '
|
|
f'FROM {schema_name}.predictions WHERE model_id = 123'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1
|
|
row = prediction_rows[0]
|
|
assert row[0] == 123
|
|
assert row[1] == 0.5
|
|
assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}'
|
|
assert row[3] is not None
|
|
assert row[4] == 'Good'
|
|
assert row[5] == ''
|
|
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123'
|
|
)
|
|
)
|
|
transformed_rows = result_query.fetchall()
|
|
assert len(transformed_rows) == 2
|
|
assert transformed_rows[0][0] == 123
|
|
assert transformed_rows[0][1] == 'feature_1'
|
|
assert float(transformed_rows[0][2]) == 0.234
|
|
assert transformed_rows[1][0] == 123
|
|
assert transformed_rows[1][1] == 'feature_2'
|
|
assert float(transformed_rows[1][2]) == 0.783
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_2_1_sql_query_execution_error(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Invalid SQL: workflow may complete with early exit; no prediction rows."""
|
|
client = temporal_test_env.client
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 128,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 128,
|
|
'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'input_filters': {
|
|
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_transform_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_predict_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
'opc_output_config': {},
|
|
'pi_web_api_output_config': {},
|
|
'save_transform': True,
|
|
'prediction_store_policy': 'lts:1',
|
|
'model_config': {
|
|
'retention_minutes': 0,
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'sklearn',
|
|
},
|
|
}
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-sql-error'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_2_2_missing_required_parameters(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Missing query: workflow does not produce predictions and is terminated explicitly."""
|
|
client = temporal_test_env.client
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 129,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 129,
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
}
|
|
|
|
handle = await client.start_workflow(
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
id=make_workflow_id('test-missing-param'),
|
|
task_queue='test-queue',
|
|
)
|
|
|
|
# Let Temporal process a few workflow tasks; for this case, result() can hang.
|
|
await asyncio.sleep(2.0)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129')
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
await handle.terminate('expected failure path in e2e test (missing required parameters)')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_1_2_3_invalid_datetime_column_specification(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""Invalid datetime column: no predictions persisted; workflow terminated after validation."""
|
|
client = temporal_test_env.client
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130'))
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
"""
|
|
)
|
|
)
|
|
|
|
input_data = {
|
|
'metadata': {
|
|
'metadata': {
|
|
'model_id': 130,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
},
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 130,
|
|
'query': 'SELECT timestamp, variable, value FROM predictions_schema.laborious_data WHERE model_id = 130',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'input_filters': {
|
|
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_transform_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_predict_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
'opc_output_config': {},
|
|
'pi_web_api_output_config': {},
|
|
'save_transform': True,
|
|
'prediction_store_policy': 'lts:1',
|
|
'model_config': {
|
|
'retention_minutes': 0,
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'sklearn',
|
|
},
|
|
'datetime_columns': ['nonexistent_column'],
|
|
}
|
|
|
|
handle = await client.start_workflow(
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
id=make_workflow_id('test-invalid-datetime-col'),
|
|
task_queue='test-queue',
|
|
)
|
|
|
|
# Let Temporal process and surface the failure path internally.
|
|
await asyncio.sleep(2.0)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130')
|
|
).scalar()
|
|
assert count == 0
|
|
|
|
await handle.terminate('expected failure path in e2e test (invalid datetime column)')
|