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.
137 lines
4.9 KiB
Python
137 lines
4.9 KiB
Python
"""
|
|
E2E-style tests for MinIO offload using a real MinIO testcontainer.
|
|
"""
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow
|
|
from laborious.activities.activities import Activities
|
|
from laborious.utils.models import minio_dataframe_payload as mdp
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_load_query_with_minio_offload_writes_object_to_bucket(
|
|
postgres_engine,
|
|
minio_container,
|
|
test_activities_real_minio: Activities,
|
|
):
|
|
"""
|
|
With a tiny offload threshold, query results are uploaded as Parquet to MinIO.
|
|
|
|
Uses real MinioRepository against testcontainers MinIO (no MinIO mock).
|
|
"""
|
|
model_id = 501
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
|
)
|
|
insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': model_id,
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
}
|
|
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
|
|
payload = await test_activities_real_minio.load_query_with_minio_offload(
|
|
{
|
|
**metadata,
|
|
'query': (
|
|
'SELECT timestamp, variable, value, created_at '
|
|
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
|
),
|
|
'model_name': 'test_model',
|
|
'datetime_columns': ['timestamp', 'created_at'],
|
|
}
|
|
)
|
|
assert payload.object_key, 'offloaded payload must reference a MinIO object'
|
|
assert payload.data is None or payload.data == {}, (
|
|
'large payloads should not inline tabular dict'
|
|
)
|
|
|
|
df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
|
|
assert len(df) >= 1
|
|
|
|
client = minio_container.get_client()
|
|
listed = list(client.list_objects('test-bucket', recursive=True))
|
|
names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
|
|
assert any(n and 'prediction_datasets' in n for n in names), (
|
|
f'unexpected object listing: {names!r}'
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_predictions_batch_with_minio_offload_path(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker_real_minio: Worker,
|
|
postgres_engine,
|
|
test_activities_real_minio: Activities,
|
|
):
|
|
"""
|
|
Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes.
|
|
"""
|
|
model_id = 502
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
|
)
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
|
)
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
|
)
|
|
insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
|
|
|
|
input_data = {
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': model_id,
|
|
'query': (
|
|
'SELECT timestamp, variable, value, created_at '
|
|
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
|
),
|
|
'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'],
|
|
}
|
|
|
|
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
|
|
await start_and_await_workflow(
|
|
temporal_test_env.client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-batch-minio-offload'),
|
|
)
|
|
|
|
with postgres_engine.connect() as conn:
|
|
count = conn.execute(
|
|
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
|
).scalar()
|
|
assert count == 1
|