gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-13

Browse Source 
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.
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-11 11:36:22 -03:00
parent 6a8c41b328
commit c12f2996c2
17 changed files with 1063 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
e2e/helpers.py
View File

@@ -11,7 +11,9 @@ 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):
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.
@@ -44,7 +46,9 @@ def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]
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}'))
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"""
@@ -79,7 +83,6 @@ def assert_prediction(
comments: Expected exact comments string (optional).
comments_contains: Substring expected in comments when queued (optional).
"""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
@@ -90,7 +93,9 @@ def assert_prediction(
)
)
prediction_rows = result_query.fetchall()
assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
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)), (
@@ -99,7 +104,9 @@ def assert_prediction(
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]}"
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:
@@ -135,7 +142,6 @@ def assert_continue(
def assert_stop(postgres_engine: Engine, model_id: int) -> None:
"""Assert no prediction rows for model_id."""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(
@@ -156,7 +162,6 @@ def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple
model_id: Model id.
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
"""
import pytest
with postgres_engine.connect() as conn:
result_query = conn.execute(