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.
77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
"""
|
|
Direct E2E execution of child workflows (smaller surface than PredictionsBatch).
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
|
|
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.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_format_and_export_prediction_default_path_e2e(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
postgres_engine,
|
|
):
|
|
"""
|
|
Run FormatAndExportPrediction with path_flag set (format_default_prediction path).
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 401
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
|
)
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': model_id,
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test-schedule',
|
|
'workflow_name': 'subworkflow.format_and_export_prediction',
|
|
}
|
|
}
|
|
input_data = {
|
|
'metadata': metadata,
|
|
'path_flag': 'CONTINUE',
|
|
'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'},
|
|
'prediction_confidence': 2,
|
|
'timestamp': '2024-01-01 12:00:00+00:00',
|
|
'model_id': model_id,
|
|
'model_name': 'test_model',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'comment': 'e2e child workflow default path',
|
|
'opc_output_config': {},
|
|
'pi_web_api_output_config': {},
|
|
'prediction_store_policy': 'lts:1',
|
|
}
|
|
|
|
await start_and_await_workflow(
|
|
client,
|
|
FormatAndExportPrediction.run,
|
|
input_data,
|
|
make_workflow_id('e2e-format-export-child'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
row = conn.execute(
|
|
text(
|
|
f'SELECT prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
|
|
)
|
|
).fetchone()
|
|
assert row is not None
|
|
assert row[0] == 0
|
|
assert row[1] == Decimal(2)
|
|
assert row[2] == 'Bad'
|
|
assert row[3] == 'e2e child workflow default path'
|