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.
187 lines
6.9 KiB
Python
187 lines
6.9 KiB
Python
"""
|
|
Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import datetime
|
|
from decimal import Decimal
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.engine import Engine
|
|
|
|
|
|
async def start_and_await_workflow(
|
|
client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0
|
|
):
|
|
"""
|
|
Start a workflow and wait for its result.
|
|
|
|
Args:
|
|
client: Temporal client from WorkflowEnvironment.
|
|
workflow_run: Workflow run method (e.g. PredictionsBatch.run).
|
|
input_data: Workflow input payload.
|
|
workflow_id: Unique workflow id.
|
|
timeout: Max seconds to wait for completion.
|
|
|
|
Return:
|
|
Workflow result value.
|
|
"""
|
|
handle = await client.start_workflow(
|
|
workflow_run,
|
|
input_data,
|
|
id=workflow_id,
|
|
task_queue='test-queue',
|
|
)
|
|
return await asyncio.wait_for(handle.result(), timeout=timeout)
|
|
|
|
|
|
def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None:
|
|
"""
|
|
Replace laborious_data rows for a model_id with one row per value (sensor_1..n).
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Model id column value.
|
|
values: Per-sensor values; use string 'NULL' for SQL NULL.
|
|
"""
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
|
)
|
|
values_sql = []
|
|
for i, value in enumerate(values):
|
|
values_sql.append(f"""
|
|
({model_id}, 'sensor_{i + 1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
""")
|
|
insert_sql = f"""
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
{', '.join(values_sql)}
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
|
|
|
|
def assert_prediction(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
prediction: float = 0.5,
|
|
prediction_confidence: int | Decimal = 0,
|
|
prediction_status: str = 'Good',
|
|
comments: str | None = None,
|
|
comments_contains: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Assert exactly one prediction row exists for model_id with expected columns.
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Expected model_id.
|
|
prediction: Expected prediction value.
|
|
prediction_confidence: Expected confidence (int or Decimal for numeric column).
|
|
prediction_status: Expected status string.
|
|
comments: Expected exact comments string (optional).
|
|
comments_contains: Substring expected in comments when queued (optional).
|
|
"""
|
|
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
|
|
f'ORDER BY created_at ASC'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1, (
|
|
f'Expected one prediction record, got {len(prediction_rows)}'
|
|
)
|
|
row = prediction_rows[0]
|
|
assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
|
|
assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
|
|
f'Expected prediction={prediction}, got {row[1]}'
|
|
)
|
|
assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
|
|
str(prediction_confidence)
|
|
), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
|
|
assert row[3] == prediction_status, (
|
|
f"Expected prediction_status='{prediction_status}', got {row[3]}"
|
|
)
|
|
if comments is not None:
|
|
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
|
if comments_contains is not None:
|
|
assert comments_contains in row[4], (
|
|
f"Expected comments to contain '{comments_contains}', got {row[4]}"
|
|
)
|
|
|
|
|
|
def assert_continue(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
prediction_confidence: Decimal = Decimal(2),
|
|
comments: str = 'Input data with bad quality',
|
|
) -> None:
|
|
"""Assert one default-style prediction row after CONTINUE gate path."""
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1, 'Expected one prediction record despite warnings'
|
|
row = prediction_rows[0]
|
|
assert row[1] == 0, f'Expected prediction=0, got {row[1]}'
|
|
assert row[2] == prediction_confidence, (
|
|
f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
|
|
)
|
|
assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
|
|
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
|
|
|
|
|
def assert_stop(postgres_engine: Engine, model_id: int) -> None:
|
|
"""Assert no prediction rows for model_id."""
|
|
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
|
)
|
|
count = result_query.scalar()
|
|
assert count == 0, f'Expected no predictions, but found {count} records'
|
|
|
|
|
|
def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None:
|
|
"""
|
|
Assert two prediction rows for model_id both match last_prediction.
|
|
|
|
Rows are compared in created_at order for stability.
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Model id.
|
|
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
|
|
"""
|
|
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status '
|
|
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
|
|
f'ORDER BY created_at ASC'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 2, 'Expected two prediction records'
|
|
assert prediction_rows[0] == last_prediction, (
|
|
f'Expected first row {last_prediction}, got {prediction_rows[0]}'
|
|
)
|
|
assert prediction_rows[1] == last_prediction, (
|
|
f'Expected second row {last_prediction}, got {prediction_rows[1]}'
|
|
)
|
|
|
|
|
|
def make_workflow_id(prefix: str) -> str:
|
|
"""Build a unique workflow id using a prefix and current timestamp."""
|
|
return f'{prefix}-{datetime.now().timestamp()}'
|