QTZPOC-13

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.

.gitignore

@@ -56,3 +56,4 @@ mlruns/
 relatorio*
 openspec/
 .cursor/
+.claude/

e2e/conftest.py

@@ -12,28 +12,35 @@ sys.modules.setdefault('sientia', MagicMock())
 sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
 from io import BytesIO
 
+import joblib
 import pandas as pd
 import pytest
 import pytest_asyncio
+from sientia_do.notifications.handlers import CoreNotificationHandler
+from sientia_do.observability.metrics_controller import MetricsController
+from sklearn.compose import ColumnTransformer
+from sklearn.linear_model import LinearRegression
+from sklearn.pipeline import Pipeline
 from sqlalchemy import create_engine, text
-from testcontainers.minio import MinioContainer
-from testcontainers.postgres import PostgresContainer
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
+from testcontainers.minio import MinioContainer
+from testcontainers.postgres import PostgresContainer
 
 from e2e.opc_test_server import OpcE2ETestServer
 from laborious.activities.activities import Activities
 from laborious.workflows.predictions_batch import PredictionsBatch
-from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
 from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
-from sientia_do.notifications.handlers import CoreNotificationHandler
+from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
-from sientia_do.observability.logger import Logger
-from sientia_do.observability.metrics_controller import MetricsController
 
 # Test constants
 TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
 TEST_DATABASE_NAME = 'test_db'
 
+# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised.
+JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model'
+JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact'
+
 
 @pytest_asyncio.fixture(scope='session')
 def minio_container():
@@ -87,7 +94,7 @@ def _create_schema_and_tables(engine):
     # Use begin() to ensure transaction is properly committed
     with engine.begin() as conn:
         # Create predictions_schema
-        conn.execute(text("CREATE SCHEMA IF NOT EXISTS predictions_schema"))
+        conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema'))
 
         # Create laborious_data table (input data from sensors)
         create_laborious_data_sql = """
@@ -150,35 +157,22 @@ def setup_postgres_schema_and_tables(postgres_engine):
 @pytest_asyncio.fixture
 def mock_logger():
     """Mock logger for testing."""
+
     def message(message):
-        print(f"[LOG] {message}")
+        print(f'[LOG] {message}')
-    def custom_message(message, _metadata={}):
+
-        print(f"[LOG] {message}")
+    def custom_message(message, _metadata=None):
+        print(f'[LOG] {message}')
+
     logger = MagicMock()
-    logger.info = MagicMock(
+    logger.info = MagicMock(side_effect=message)
-        side_effect=message
+    logger.debug = MagicMock(side_effect=message)
-    )
+    logger.error = MagicMock(side_effect=message)
-    logger.debug = MagicMock(
+    logger.warning = MagicMock(side_effect=message)
-        side_effect=message
+    logger.custom_info = MagicMock(side_effect=custom_message)
-    )
+    logger.custom_debug = MagicMock(side_effect=custom_message)
-    logger.error = MagicMock(
+    logger.custom_error = MagicMock(side_effect=custom_message)
-        side_effect=message
+    logger.custom_warning = MagicMock(side_effect=custom_message)
-    )
-    logger.warning = MagicMock(
-        side_effect=message
-    )
-    logger.custom_info = MagicMock(
-        side_effect=custom_message
-    )
-    logger.custom_debug = MagicMock(
-        side_effect=custom_message
-    )
-    logger.custom_error = MagicMock(
-        side_effect=custom_message
-    )
-    logger.custom_warning = MagicMock(
-        side_effect=custom_message
-    )
     return logger
 
 
@@ -295,13 +289,12 @@ def mock_pi_web_api_repository():
     mock_repo.close = MagicMock()
     return mock_repo
 
+
 @pytest_asyncio.fixture
 def mock_opc_repository():
     """Mock OPC repository for OPC operations."""
     mock_repo = MagicMock()
-    mock_repo.write_data = AsyncMock(
+    mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1}))
-        return_value=(True, {'response_time': 0.1})
-    )
     mock_repo.disconnect = AsyncMock()
     return mock_repo
 
@@ -318,10 +311,13 @@ async def opc_e2e_server():
     finally:
         await server.stop()
 
+
 @pytest_asyncio.fixture
 def patch_create_engine(postgres_engine):
     """Patch create_engine to return test postgres_engine."""
-    with patch('sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine):
+    with patch(
+        'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine
+    ):
         yield
 
 
@@ -329,37 +325,47 @@ def patch_create_engine(postgres_engine):
 def patch_minio_repository(mock_minio_repository):
     """Patch MinioRepository to return mock."""
     # Patch where Activities resolves the symbol (import binds the original class).
-    with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository):
+    with patch(
+        'laborious.activities.activities.MinioRepository', return_value=mock_minio_repository
+    ):
         yield
 
+
 @pytest_asyncio.fixture
 def patch_pi_web_api_repository(mock_pi_web_api_repository):
     """Patch MLflowRepository to return mock."""
     with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
         yield
 
+
 @pytest_asyncio.fixture
 def mock_mlflow_models():
     """Create mock models for MLflow load_model methods."""
     # Mock transform model - returns DataFrame with same index as input
     mock_transform_model = MagicMock()
+
     def mock_transform_predict(data):
         num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
         print(data.to_csv())
         print(data.index)
-        result = pd.DataFrame({
+        result = pd.DataFrame(
-            'feature_1': [0.234] * num_rows,
+            {
-            'feature_2': [0.783] * num_rows,
+                'feature_1': [0.234] * num_rows,
-        })
+                'feature_2': [0.783] * num_rows,
+            }
+        )
         result.index = data.index
         return result
+
     mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
 
     # Mock predict model - returns array/list of predictions
     mock_predict_model = MagicMock()
+
     def mock_predict_predict(data):
         num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
         return [0.5] * num_rows
+
     mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
 
     # Mock PyFuncModel for compressed models
@@ -374,26 +380,93 @@ def mock_mlflow_models():
     }
 
 
+def _build_joblib_fallback_estimator():
+    """
+    Build a real sklearn estimator for the joblib fallback artifact.
+
+    `request_predict` appends a string `timestamp` column to the transformed
+    data, so the estimator selects the transform output columns explicitly and
+    drops the rest.
+    """
+    features = ['feature_1', 'feature_2']
+    estimator = Pipeline(
+        steps=[
+            (
+                'select',
+                ColumnTransformer([('features', 'passthrough', features)], remainder='drop'),
+            ),
+            ('regressor', LinearRegression()),
+        ]
+    )
+    train = pd.DataFrame(
+        {
+            'feature_1': [0.0, 1.0, 2.0, 3.0],
+            'feature_2': [0.0, 1.0, 0.0, 1.0],
+            'timestamp': ['2024-01-01 12:00:00'] * 4,
+        }
+    )
+    estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3)
+    return estimator
+
+
+@pytest.fixture
+def joblib_fallback_artifact(tmp_path):
+    """
+    Real artifact directory holding a joblib-compressed sklearn estimator.
+
+    Mirrors what the joblib fallback finds after downloading `prediction_model`:
+    a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an
+    `MLmodel` placeholder.
+
+    Return:
+        dict with 'path' (the artifact directory) and 'estimator' (the object
+        that was serialized, so tests can compute the expected prediction).
+    """
+    artifact_dir = tmp_path / 'prediction_model'
+    artifact_dir.mkdir()
+
+    estimator = _build_joblib_fallback_estimator()
+    joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3)
+    (artifact_dir / 'MLmodel').write_text('flavors:\n  sklearn: {}\n')
+
+    return {'path': str(artifact_dir), 'estimator': estimator}
+
+
+@pytest.fixture
+def joblib_fallback_artifact_without_pickle(tmp_path):
+    """Artifact directory with no top-level `.pkl`, so the joblib fallback also fails."""
+    artifact_dir = tmp_path / 'prediction_model_without_pickle'
+    artifact_dir.mkdir()
+    (artifact_dir / 'MLmodel').write_text('flavors:\n  sklearn: {}\n')
+    return str(artifact_dir)
+
+
 @pytest_asyncio.fixture
-def patch_mlflow(mock_mlflow_models):
+def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
     """Patch mlflow module in repository with load_model mocks."""
     mock_mlflow = MagicMock()
 
     # Mock sklearn.load_model
     def mock_sklearn_load_model(model_uri):
+        if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
+            raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
         if 'data_model' in model_uri or 'transform' in model_uri.lower():
             return mock_mlflow_models['transform_model']
         return mock_mlflow_models['predict_model']
+
     mock_mlflow.sklearn = MagicMock()
     mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
 
     # Mock pyfunc.load_model
     def mock_pyfunc_load_model(model_uri):
+        if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
+            raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
         if 'artifacts' in model_uri or 'tmp' in model_uri:
             return mock_mlflow_models['pyfunc_model']
         if 'data_model' in model_uri or 'transform' in model_uri.lower():
             return mock_mlflow_models['transform_model']
         return mock_mlflow_models['predict_model']
+
     mock_mlflow.pyfunc = MagicMock()
     mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
 
@@ -403,21 +476,34 @@ def patch_mlflow(mock_mlflow_models):
 
     # Mock other mlflow methods that might be called
     mock_mlflow.set_tracking_uri = MagicMock()
-    mock_mlflow.get_run = MagicMock(return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id')))
+    mock_mlflow.get_run = MagicMock(
+        return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))
+    )
     mock_mlflow.tracking = MagicMock()
-    mock_mlflow.tracking.MlflowClient = MagicMock(return_value=MagicMock(
+    mock_mlflow_client = MagicMock(
         search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
-        search_model_versions=MagicMock(return_value=[MagicMock(
+        search_model_versions=MagicMock(
-            current_stage='Production',
+            return_value=[
-            version='1',
+                MagicMock(
-            source='runs:/artifacts/test_run_id'
+                    current_stage='Production', version='1', source='runs:/artifacts/test_run_id'
-        )])
+                )
-    ))
+            ]
+        ),
+    )
+    # Only the joblib fallback calls download_artifacts directly on the client.
+    mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path'])
+    mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)
 
     with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
         yield mock_mlflow
 
 
+@pytest.fixture
+def mlflow_client(patch_mlflow):
+    """The MlflowClient mock the repository holds as `self.client`."""
+    return patch_mlflow.tracking.MlflowClient.return_value
+
+
 @pytest_asyncio.fixture(scope='function')
 async def test_activities(
     postgres_engine,
@@ -430,7 +516,7 @@ async def test_activities(
     patch_minio_repository,
     patch_mlflow,
     patch_pi_web_api_repository,
-    mock_opc_repository
+    mock_opc_repository,
 ):
     """
     Create Activities instance with test dependencies.

e2e/helpers.py

@@ -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(

e2e/opc_test_server.py

@@ -147,8 +147,11 @@ class OpcE2ETestServer:
         value = await self._confidence_node.read_value()
         return float(value)
 
-    async def _pre_write_callback(self, _event, _service) -> None:
+    async def _pre_write_callback(self, event, _service) -> None:
-        if self._session_bad_on_write:
+        # Only client writes are faulted. asyncua's own clock task writes CurrentTime
+        # through this same PreWrite dispatch, and raising there kills its time_task;
+        # the stored exception then resurfaces as a teardown failure in server.stop().
+        if self._session_bad_on_write and getattr(event, 'is_external', False):
             raise ServiceError(ua.StatusCodes.BadSessionIdInvalid)
 
 

e2e/test_child_workflows_e2e.py

@@ -26,7 +26,9 @@ async def test_format_and_export_prediction_default_path_e2e(
     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}'))
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
+        )
 
     metadata = {
         'metadata': {

e2e/test_minio_offload.py

@@ -29,7 +29,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
     """
     model_id = 501
     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}')
+        )
     insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
 
     metadata = {
@@ -53,7 +55,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
             }
         )
     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'
+    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
@@ -61,7 +65,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
     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}'
+    assert any(n and 'prediction_datasets' in n for n in names), (
+        f'unexpected object listing: {names!r}'
+    )
 
 
 @pytest.mark.asyncio
@@ -77,9 +83,15 @@ async def test_predictions_batch_with_minio_offload_path(
     """
     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(
-        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+            text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
-        conn.execute(text(f'DELETE FROM predictions_schema.transformed_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 = {

e2e/test_opc_real_server.py

@@ -11,7 +11,12 @@ import pytest
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
 
-from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
+from e2e.helpers import (
+    assert_prediction,
+    insert_sample_data,
+    make_workflow_id,
+    start_and_await_workflow,
+)
 from e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config
 from e2e.test_predictions_batch_format_export import get_base_input_data
 from laborious.activities.activities import Activities
@@ -20,16 +25,24 @@ from laborious.utils.repository.opc_repository import OpcRepository
 from laborious.workflows.predictions_batch import PredictionsBatch
 
 
-async def _slow_reconnect_under_lock(repo: OpcRepository, hold_seconds: float = 0.75) -> None:
+async def _reconnect_under_held_lock(
+    repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event
+) -> None:
     """
-    Hold the connection lock briefly so concurrent writes see reconnect_in_progress.
+    Hold the connection lock until released, then reconnect.
+
+    The lock must stay held for the whole workflow, and how long the workflow takes to
+    reach the OPC write is not bounded — a fixed sleep here lets `_reconnect_locked`
+    finish first, which sets `_session_ready` again and the write then succeeds.
 
     Args:
         repo (OpcRepository): Connected repository.
-        hold_seconds (float): Time to keep the lock before reconnecting.
+        acquired (asyncio.Event): Set once the lock is held, before any write can race.
+        release (asyncio.Event): Awaited while holding the lock; set by the caller.
     """
     async with repo._connection_lock:
-        await asyncio.sleep(hold_seconds)
+        acquired.set()
+        await release.wait()
         await repo._reconnect_locked()
 
 
@@ -170,7 +183,10 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
 
     repo = test_activities_real_opc.opc_repository['1']
     repo._session_ready.clear()
-    reconnect_task = asyncio.create_task(_slow_reconnect_under_lock(repo))
+    acquired = asyncio.Event()
+    release = asyncio.Event()
+    reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release))
+    await acquired.wait()
 
     input_data = get_base_input_data(model_id)
     input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
@@ -184,6 +200,7 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
             make_workflow_id('test-opc-real-reconnect-block'),
         )
     finally:
+        release.set()
         await reconnect_task
 
     assert_prediction(

e2e/test_predictions_batch_format_export.py

@@ -4,7 +4,7 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
 
 from decimal import Decimal
 from typing import Any, cast
-from unittest.mock import ANY, AsyncMock, call
+from unittest.mock import AsyncMock, call
 
 import pytest
 from sientia_do.notifications.models import NotificationLevel
@@ -12,7 +12,12 @@ from sqlalchemy import text
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
 
-from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
+from e2e.helpers import (
+    assert_prediction,
+    insert_sample_data,
+    make_workflow_id,
+    start_and_await_workflow,
+)
 from laborious.activities.activities import Activities
 from laborious.workflows.predictions_batch import PredictionsBatch
 
@@ -46,7 +51,8 @@ base_input_data = {
     'datetime_columns': ['timestamp', 'created_at'],
 }
 
-base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}"
+base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
+
 
 def get_base_input_data(model_id):
     return {
@@ -79,8 +85,12 @@ async def test_scenario_3_1_1_default_prediction_export(
     model_id = 311
 
     with postgres_engine.begin() as conn:
-        conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}"))
+        conn.execute(
-        conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
+            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, ['NULL', 78.2])
 
     input_data = get_base_input_data(model_id)
@@ -176,7 +186,9 @@ async def test_scenario_3_1_1_default_prediction_export(
 
     with postgres_engine.connect() as conn:
         tf_count = conn.execute(
-            text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
+            text(
+                f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
+            )
         ).scalar()
         assert tf_count == 0, 'transform export must be skipped when path_flag is set'
 
@@ -190,9 +202,6 @@ async def test_scenario_3_1_1_default_prediction_export(
     )
 
 
-
-
-
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_3_1_2_export_with_opc_only(
@@ -280,7 +289,6 @@ async def test_scenario_3_1_2_export_with_opc_only(
     assert_prediction(postgres_engine, model_id)
 
 
-
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_3_1_3_export_with_pi_web_api_only(
@@ -364,7 +372,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
     assert_prediction(postgres_engine, model_id)
 
 
-
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_3_1_4_export_without_optional_outputs(
@@ -412,7 +419,6 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
     assert_prediction(postgres_engine, model_id)
 
 
-
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_3_1_5_export_without_transformed_data(
@@ -443,7 +449,9 @@ async def test_scenario_3_1_5_export_without_transformed_data(
 
     insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
     with postgres_engine.begin() as conn:
-        conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
+        )
 
     input_data = get_base_input_data(model_id)
     input_data['save_transform'] = False  # Don't save transformed data
@@ -533,15 +541,16 @@ async def test_scenario_3_1_5_export_without_transformed_data(
 
     with postgres_engine.connect() as conn:
         result_query = conn.execute(
-            text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
+            text(
+                f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
+            )
         )
         count = result_query.scalar()
-        assert count == 0, f"Expected transform table to be empty, but found {count} records"
+        assert count == 0, f'Expected transform table to be empty, but found {count} records'
 
     assert_prediction(postgres_engine, model_id)
 
 
-
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_3_2_1_pi_web_api_write_error(
@@ -565,7 +574,8 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
     insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
 
     test_activities.pi_web_api_client.write_value.side_effect = Exception(
-        "PI Web API service unavailable")
+        'PI Web API service unavailable'
+    )
 
     input_data = get_base_input_data(model_id)
     input_data['pi_web_api_output_config'] = {
@@ -593,14 +603,14 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
     )
 
     assert_prediction(
-        postgres_engine, model_id,
+        postgres_engine,
+        model_id,
         prediction_confidence=13,
         comments='PI Web API service unavailable',
     )
     assert notification_inserts.call_count >= 1
 
 
-
 @pytest.mark.asyncio
 @pytest.mark.integration
 async def test_scenario_3_2_2_opc_write_error(
@@ -622,13 +632,16 @@ async def test_scenario_3_2_2_opc_write_error(
     insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
 
     opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
-    opc_write_data.return_value = (False, {
+    opc_write_data.return_value = (
-        'notification_id': 'OPC_WRITE_DATA_ERROR_1',
+        False,
-        'message': 'OPC server unavailable',
+        {
-        'block': 'opc_repository',
+            'notification_id': 'OPC_WRITE_DATA_ERROR_1',
-        'level': NotificationLevel.ERROR,
+            'message': 'OPC server unavailable',
-        'attachment_content': 'OPC server unavailable',
+            'block': 'opc_repository',
-    })
+            'level': NotificationLevel.ERROR,
+            'attachment_content': 'OPC server unavailable',
+        },
+    )
 
     input_data = get_base_input_data(model_id)
     input_data['opc_output_config'] = {
@@ -656,7 +669,8 @@ async def test_scenario_3_2_2_opc_write_error(
     )
 
     assert_prediction(
-        postgres_engine, model_id,
+        postgres_engine,
+        model_id,
         prediction_confidence=12,
         comments='Some data could not be written to OPC servers',
     )
@@ -779,8 +793,8 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
     )
 
     assert_prediction(
-        postgres_engine, model_id,
+        postgres_engine,
+        model_id,
         prediction_confidence=13,
         comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.",
     )
-    

e2e/test_predictions_batch_joblib_fallback.py

@@ -0,0 +1,186 @@
+"""
+End-to-end tests for the joblib fallback on the MLflow model load paths.
+
+`joblib.load` is NOT mocked here: the artifact fixture holds a real sklearn
+estimator serialized with `joblib.dump(..., compress=3)`, so a successful run
+proves the fallback really deserializes a compressed pickle and serves it.
+"""
+
+import pandas as pd
+import pytest
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.conftest import JOBLIB_FALLBACK_FLAVOR_ERROR, JOBLIB_FALLBACK_MODEL_NAME
+from e2e.helpers import (
+    assert_prediction,
+    assert_stop,
+    insert_sample_data,
+    make_workflow_id,
+    start_and_await_workflow,
+)
+from laborious.activities.activities import Activities
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
+
+
+def build_input_data(model_id: int, model_name: str) -> dict:
+    """Workflow input for a plain transform + predict run of a given model."""
+    return {
+        'schedule_name': 'test-schedule',
+        'model_name': model_name,
+        'model_id': model_id,
+        'query': base_query.format(model_id=model_id),
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+        'input_filters': {
+            'SPECIFIC_VARIABLES_NULL_VALUES': {
+                'POLICY': 'CONTINUE',
+                'CONFIG': {'variables': ['sensor_1']},
+            },
+        },
+        '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'],
+    }
+
+
+def expected_fallback_prediction(estimator) -> float:
+    """Prediction the fixture estimator produces for the mocked transform output."""
+    features = pd.DataFrame(
+        {
+            'feature_1': [0.234],
+            'feature_2': [0.783],
+            'timestamp': ['2024-01-01 12:00:00'],
+        }
+    )
+    return float(estimator.predict(features)[0])
+
+
+def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
+    """Assert exactly one Good prediction row whose value matches `expected`."""
+    with postgres_engine.connect() as conn:
+        rows = conn.execute(
+            text(
+                'SELECT prediction, prediction_status '
+                'FROM predictions_schema.predictions '
+                f'WHERE model_id = {model_id}'
+            )
+        ).fetchall()
+
+    assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
+    assert abs(float(rows[0][0]) - expected) < tolerance, (
+        f'Expected prediction close to {expected}, got {rows[0][0]}'
+    )
+    assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"
+
+
+def logged_messages(mock_logger, level: str) -> str:
+    """Concatenate every message logged at a given custom level."""
+    calls = getattr(mock_logger, f'custom_{level}').call_args_list
+    return ' '.join(str(logged.args[0]) for logged in calls)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_prediction_workflow_recovered_by_joblib_fallback(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    joblib_fallback_artifact,
+    mlflow_client,
+    mock_logger,
+):
+    """Flavor loader fails, joblib fallback loads the real compressed pickle, prediction lands."""
+    client = temporal_test_env.client
+    model_id = 271
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback')
+    )
+
+    assert_single_prediction_close(
+        postgres_engine,
+        model_id,
+        expected_fallback_prediction(joblib_fallback_artifact['estimator']),
+    )
+
+    assert mlflow_client.download_artifacts.call_count == 1
+    assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
+
+    assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged_messages(mock_logger, 'warning')
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_prediction_workflow_without_fallback_is_unchanged(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mlflow_client,
+    mock_mlflow_models,
+):
+    """A model whose flavor loader works keeps the current result and downloads nothing."""
+    client = temporal_test_env.client
+    model_id = 272
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = build_input_data(model_id, 'test_model')
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-no-fallback')
+    )
+
+    assert_prediction(postgres_engine, model_id, prediction=0.5)
+    mock_mlflow_models['predict_model'].predict.assert_called()
+    mlflow_client.download_artifacts.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_prediction_workflow_reports_both_causes_when_fallback_fails(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    joblib_fallback_artifact_without_pickle,
+    mlflow_client,
+    mock_logger,
+):
+    """Both attempts fail: workflow stops and the logged traceback carries both causes."""
+    mlflow_client.download_artifacts.return_value = joblib_fallback_artifact_without_pickle
+
+    client = temporal_test_env.client
+    model_id = 273
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback-failed')
+    )
+
+    assert_stop(postgres_engine, model_id)
+
+    logged = logged_messages(mock_logger, 'debug')
+    assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged
+    assert 'No .pkl file found in prediction_model artifact' in logged

e2e/test_predictions_batch_main_workflow.py

@@ -4,12 +4,11 @@ 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
 
-import pytest
-
 from e2e.helpers import make_workflow_id, start_and_await_workflow
 from laborious.activities.activities import Activities
 from laborious.workflows.predictions_batch import PredictionsBatch

e2e/test_predictions_batch_prediction_process.py

@@ -3,7 +3,7 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
 """
 
 from decimal import Decimal
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 import numpy as np
 import pandas as pd
@@ -69,7 +69,9 @@ def get_base_input_data(model_id):
 
 def insert_sample_prediction(postgres_engine, model_id):
     with postgres_engine.begin() as conn:
-        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
+        )
         insert_sql = f"""
             INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
             VALUES
@@ -180,7 +182,9 @@ async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
     model_id = 214
     insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
     with postgres_engine.begin() as conn:
-        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
+        )
     input_data = get_base_input_data(model_id)
     input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
     await start_and_await_workflow(
@@ -370,7 +374,9 @@ async def test_scenario_2_4_1_input_empty_data_stop(
     client = temporal_test_env.client
     model_id = 241
     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}')
+        )
     input_data = get_base_input_data(model_id)
     input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
     await start_and_await_workflow(

laborious/utils/repository/model_repository.py

@@ -21,10 +21,11 @@ import time
 import traceback
 from datetime import datetime, timedelta
 from io import StringIO
-from os import environ, makedirs, path
+from os import environ, listdir, makedirs, path
 from shutil import rmtree
 from typing import Any, Literal, overload
 
+import joblib
 import mlflow
 import pandas as pd
 from mlflow.entities import Experiment
@@ -351,6 +352,117 @@ class MLFlowRepository(SientiaMonitoring):
         self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata)
         return dataframe
 
+    def _find_model_pickle(
+        self, local_dir: str, artifact_name: str, flavor_error: Exception
+    ) -> str:
+        """
+        Locate the pickle file of a downloaded model artifact.
+
+        Fallback path only: this is reached after the MLflow flavor loader has
+        already failed. The target artifact is always a single `.pkl` at the top
+        level of the artifact directory, so the listing does not recurse —
+        nested wrapper artifacts must never be deserialized in place of the
+        model. No other extension is considered and `MLmodel` is not parsed.
+
+        Args:
+            local_dir (str): Local directory holding the downloaded artifact.
+            artifact_name (str): Artifact name, e.g. 'prediction_model'.
+            flavor_error (Exception): Error raised by the MLflow flavor loader,
+                reported in the message so the original cause is not lost.
+
+        Returns:
+            str: Path to the first top-level `.pkl` file, in name order.
+
+        Raises:
+            ValueError: If the top level of `local_dir` has no `.pkl` file.
+        """
+        entries = sorted(listdir(local_dir))
+        pkl_files = [
+            entry
+            for entry in entries
+            if entry.endswith('.pkl') and path.isfile(path.join(local_dir, entry))
+        ]
+
+        if not pkl_files:
+            raise ValueError(
+                f'No .pkl file found in {artifact_name} artifact. '
+                f'Original error: {flavor_error}. '
+                f'Top-level entries in {local_dir}: {entries}'
+            )
+
+        return path.join(local_dir, pkl_files[0])
+
+    def _load_model_pickle(
+        self, local_dir: str, artifact_name: str, flavor_error: Exception
+    ) -> Any:
+        """
+        Deserialize a downloaded model artifact with `joblib`.
+
+        Fallback path only. `joblib.load` reads the compressed `.pkl` produced
+        by `joblib.dump`, which is what the MLflow flavor loader fails to read.
+        The deserialized object is returned as is: no `._model_impl.python_model`
+        unwrap is applied, because the pickled object already is the model.
+
+        Args:
+            local_dir (str): Local directory holding the downloaded artifact.
+            artifact_name (str): Artifact name, e.g. 'prediction_model'.
+            flavor_error (Exception): Error raised by the MLflow flavor loader.
+
+        Returns:
+            Any: The object deserialized from the artifact's pickle file.
+
+        Raises:
+            ValueError: If the artifact directory has no top-level `.pkl` file.
+            Exception: Any error raised by `joblib.load` itself.
+        """
+        pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error)
+
+        self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}')
+
+        return joblib.load(pkl_path)
+
+    async def _download_and_load_pickle(
+        self,
+        model_name: str,
+        artifact_name: str,
+        flavor_error: Exception,
+        metadata: dict[str, Any],
+    ) -> Any:
+        """
+        Download a model artifact and deserialize it with `joblib`.
+
+        Fallback path only. The download runs without `dst_path`, so MLflow
+        writes to a per-call temporary directory: the shared `ARTIFACTS_PATH`
+        tree is never written to and `dowload_artifacts` — which removes its
+        destination before downloading — is deliberately not reused. The target
+        artifact is always a `.pkl` and the deserialized object is returned as
+        is.
+
+        Args:
+            model_name (str): Registered model name.
+            artifact_name (str): Artifact to download, e.g. 'prediction_model'.
+            flavor_error (Exception): Error raised by the MLflow flavor loader.
+            metadata (dict[str, Any]): Metadata used for structured logging.
+
+        Returns:
+            Any: The object deserialized from the artifact's pickle file.
+
+        Raises:
+            ValueError: If the downloaded artifact has no top-level `.pkl` file.
+            Exception: Any error raised while resolving the Production run,
+                downloading the artifact, or deserializing the pickle.
+        """
+        run_id = self.get_model_run_id(model_name=model_name, stage='Production')
+
+        self.info(
+            f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}',
+            metadata,
+        )
+
+        local_dir = self.client.download_artifacts(run_id, artifact_name)
+
+        return self._load_model_pickle(local_dir, artifact_name, flavor_error)
+
     async def load_predict_model(
         self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
     ) -> Any:
@@ -370,23 +482,43 @@ class MLFlowRepository(SientiaMonitoring):
             - The model is fetched from the "Production" stage of the MLflow Model Registry.
             - Warnings during the model loading process are suppressed.
         """
+        if flavor not in ('pyfunc', 'sklearn', 'pytorch'):
+            raise ValueError(INVALID_FLAVOR_MESSAGE)
+
         model_uri = f'models:/{model_name}/production'
         self.info(f'Loading prediction model {model_name} from {model_uri}')
 
         core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
+        fallback_labels = self.get_core_labels(
+            metadata, operation_type='load_predict_model_joblib_fallback'
+        )
         start_time = time.time()
         try:
             if flavor == 'pyfunc':
                 model = mlflow.pyfunc.load_model(model_uri)
             elif flavor == 'sklearn':
                 model = mlflow.sklearn.load_model(model_uri)
-            elif flavor == 'pytorch':
-                model = mlflow.pytorch.load_model(model_uri)
             else:
-                raise ValueError(INVALID_FLAVOR_MESSAGE)
+                model = mlflow.pytorch.load_model(model_uri)
-        except Exception as e:
+        except Exception as flavor_error:
-            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
+            self.warning(
-            raise e
+                f'Flavor loader {flavor} failed for prediction model {model_name}, '
+                f'trying joblib fallback. Original error: {flavor_error}',
+                metadata,
+            )
+            try:
+                model = await self._download_and_load_pickle(
+                    model_name, 'prediction_model', flavor_error, metadata
+                )
+            except Exception as fallback_error:
+                await self.emit_metric(
+                    metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
+                )
+                raise fallback_error from flavor_error
+
+            await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
+            await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
+            return model
 
         await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
         await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
@@ -415,25 +547,45 @@ class MLFlowRepository(SientiaMonitoring):
                 model cannot be loaded.
         """
 
+        if flavor not in ('sklearn', 'pyfunc', 'pytorch'):
+            raise ValueError(INVALID_FLAVOR_MESSAGE)
+
         latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
         model_uri = self.get_model_uri(latest_production_id, prediction=False)
 
         self.info(f'Loading data model {model_name} from {model_uri}')
 
         core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
+        fallback_labels = self.get_core_labels(
+            metadata, operation_type='load_transform_model_joblib_fallback'
+        )
         start_time = time.time()
         try:
             if flavor == 'sklearn':
                 model = mlflow.sklearn.load_model(model_uri)
             elif flavor == 'pyfunc':
                 model = mlflow.pyfunc.load_model(model_uri)
-            elif flavor == 'pytorch':
-                model = mlflow.pytorch.load_model(model_uri)
             else:
-                raise ValueError(INVALID_FLAVOR_MESSAGE)
+                model = mlflow.pytorch.load_model(model_uri)
-        except Exception as e:
+        except Exception as flavor_error:
-            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
+            self.warning(
-            raise e
+                f'Flavor loader {flavor} failed for data model {model_name}, '
+                f'trying joblib fallback. Original error: {flavor_error}',
+                metadata,
+            )
+            try:
+                model = await self._download_and_load_pickle(
+                    model_name, 'data_model', flavor_error, metadata
+                )
+            except Exception as fallback_error:
+                await self.emit_metric(
+                    metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
+                )
+                raise fallback_error from flavor_error
+
+            await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
+            await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
+            return model
 
         await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
         await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
@@ -482,8 +634,29 @@ class MLFlowRepository(SientiaMonitoring):
                 f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
             )
 
-            raw_model = mlflow.pyfunc.load_model(artifact_path)
+            fallback_labels = self.get_core_labels(
-            model = raw_model._model_impl.python_model
+                metadata, operation_type='download_model_wrapper_joblib_fallback'
+            )
+            start_time = time.time()
+            try:
+                raw_model = mlflow.pyfunc.load_model(artifact_path)
+                model = raw_model._model_impl.python_model
+            except Exception as flavor_error:
+                self.warning(
+                    f'Wrapper loader failed for {model_type} model {model_name}, '
+                    f'trying joblib fallback. Original error: {flavor_error}',
+                    metadata,
+                )
+                try:
+                    model = self._load_model_pickle(artifact_path, target, flavor_error)
+                except Exception as fallback_error:
+                    await self.emit_metric(
+                        metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
+                    )
+                    raise fallback_error from flavor_error
+
+                await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
+                await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
 
             self.debug(
                 f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata

pyproject.toml

@@ -60,6 +60,11 @@ ignore = [
     "S105",  # hardcoded passwords ok in tests
     "S106",  # hardcoded passwords ok in tests
 ]
+"e2e/conftest.py" = [
+    # Imports below the sientia/ModelAnalysis stub are load-order critical: the stub
+    # must be in sys.modules before Activities is imported, so they cannot move up.
+    "E402",
+]
 
 [tool.ruff.lint.mccabe]
 max-complexity = 15

requirements-light.txt

@@ -16,3 +16,4 @@ shap
 pycurl
 scipy<1.14.0
 scikit-learn==1.5.2
+joblib

requirements-local.txt

@@ -15,4 +15,6 @@ hyperopt
 shap
 pycurl
 scipy<1.14.0
-scikit-learn==1.5.2
+scikit-learn==1.4.2
+pycaret==3.3.2
+joblib

requirements.txt

@@ -16,3 +16,4 @@ shap
 pycurl
 scipy<1.14.0
 scikit-learn==1.5.2
+joblib

tests/laborious/utils/repository/test_model_repository.py

@@ -1,13 +1,20 @@
 from datetime import UTC, datetime
 from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
 
+import joblib as joblib_lib
 import mlflow as mlflow_lib
 import numpy as np
 import pytest
 from pandas import DataFrame, Timestamp
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+from sklearn.linear_model import LinearRegression
 
 from laborious import metrics
-from laborious.utils.repository.model_repository import MLFlowRepository, force_memory_release
+from laborious.utils.repository.model_repository import (
+    INVALID_FLAVOR_MESSAGE,
+    MLFlowRepository,
+    force_memory_release,
+)
 
 
 @patch('laborious.utils.repository.model_repository.ctypes')
@@ -608,32 +615,46 @@ def test_detect_and_parse_datetime_index_empty(mlflow_repository):
     assert response.equals(input_data)
 
 
+# Each case is (index, fragments expected in the error message).
+#
+# The index is passed explicitly instead of through a dict-of-dict: building a
+# DataFrame from nested dicts makes pandas sort the keys to infer the index, which a
+# mixed-type index cannot do.
+#
+# Fragments, not whole messages: the parse branch appends pandas' own multi-line
+# error, and the type branch appends the fully qualified class name, so neither is
+# stable enough to compare literally.
 invalid_cases = [
     (
-        {'value': {'2024-01-01 12:00:00': 1, 2024: 2}},
+        ['2024-01-01 12:00:00', 2024],
-        "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Elements are <class 'str'>, <class 'int'>.",
+        ["Elements are <class 'str'>,<class 'int'>"],
     ),
     (
-        {'value': {'2024-01-01': 1, '2024-01-02': 2}},
+        ['2024-01-01', '2024-01-02'],
-        'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Unable to parse given date format',
+        ['Unable to parse given date format'],
     ),
     (
-        {'value': {Any(): 1, Any(): 2}},
+        [Any(), Any()],
-        'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Got <class Any>.',
+        ['Got <class', "Any'>."],
     ),
 ]
 
 
 @pytest.mark.parametrize('data', invalid_cases)
 def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
-    input_data = DataFrame(data[0])
+    input_data = DataFrame({'value': [1, 2]}, index=data[0])
 
-    message = data[1]
+    expected_fragments = data[1]
 
     with pytest.raises(ValueError) as e:
         mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
 
-        assert str(e) == message
+    message = str(e.value)
+
+    assert 'Index must be all timestamp like column' in message
+    assert f'string in format {DATETIME_FORMAT_WITH_TZ}' in message
+    for fragment in expected_fragments:
+        assert fragment in message, f'Expected {fragment!r} in {message!r}'
 
 
 valid_cases = [
@@ -1713,3 +1734,357 @@ async def test_predict_success_ndarray(mlflow_repository):
     assert 'prediction' in content.columns
     assert 'response_time' in content.columns
     assert content.index.tolist() == data.index.tolist()
+
+
+"""
+Tests for the joblib fallback on the model load paths
+"""
+
+JOBLIB_PATH = 'laborious.utils.repository.model_repository.joblib'
+
+
+def build_artifact_dir(tmp_path, *filenames):
+    """Create a fake downloaded artifact directory with the given top-level files."""
+    for filename in filenames:
+        (tmp_path / filename).write_bytes(b'pickle-bytes')
+    return str(tmp_path)
+
+
+def emitted_tags(emit_metric_mock, metric_object):
+    """Collect the `tags` kwarg of every emit_metric call for a given metric."""
+    return [
+        emitted.kwargs['tags']
+        for emitted in emit_metric_mock.call_args_list
+        if emitted.kwargs['metric_object'] is metric_object
+    ]
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_predict_model_flavor_success_skips_fallback(
+    joblib_mock, mlflow, mlflow_repository
+):
+    result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
+
+    assert result == mlflow.sklearn.load_model.return_value
+    joblib_mock.load.assert_not_called()
+    mlflow_repository.client.download_artifacts.assert_not_called()
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_transform_model_flavor_success_skips_fallback(
+    joblib_mock, mlflow, mlflow_repository
+):
+    mlflow_repository.get_model_run_id = MagicMock()
+    mlflow_repository.get_model_uri = MagicMock()
+
+    result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
+
+    assert result == mlflow.sklearn.load_model.return_value
+    joblib_mock.load.assert_not_called()
+    mlflow_repository.client.download_artifacts.assert_not_called()
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_download_model_wrapper_success_skips_fallback(
+    joblib_mock, mlflow, mlflow_repository
+):
+    mlflow_repository.dowload_artifacts = AsyncMock()
+
+    result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
+
+    assert result == (
+        mlflow.pyfunc.load_model.return_value._model_impl.python_model,
+        mlflow_repository.dowload_artifacts.return_value,
+    )
+    joblib_mock.load.assert_not_called()
+    mlflow_repository.client.download_artifacts.assert_not_called()
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_predict_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository):
+    with pytest.raises(ValueError) as exc_info:
+        await mlflow_repository.load_predict_model('test_model', {}, 'invalid')
+
+    assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
+    joblib_mock.load.assert_not_called()
+    mlflow_repository.client.download_artifacts.assert_not_called()
+    mlflow_repository.emit_metric.assert_not_called()
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_transform_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository):
+    with pytest.raises(ValueError) as exc_info:
+        await mlflow_repository.load_transform_model('test_model', {}, 'invalid')
+
+    assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
+    joblib_mock.load.assert_not_called()
+    mlflow_repository.client.download_artifacts.assert_not_called()
+    mlflow_repository.emit_metric.assert_not_called()
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_predict_model_fallback_recovers_load(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'MLmodel', 'model.pkl'
+    )
+
+    result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
+
+    assert result is joblib_mock.load.return_value
+    mlflow_repository.get_model_run_id.assert_called_once_with(
+        model_name='test_model', stage='Production'
+    )
+    mlflow_repository.client.download_artifacts.assert_called_once_with(
+        'run-123', 'prediction_model'
+    )
+    joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_transform_model_fallback_recovers_load(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.get_model_uri = MagicMock()
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'MLmodel', 'model.pkl'
+    )
+
+    result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
+
+    assert result is joblib_mock.load.return_value
+    mlflow_repository.get_model_run_id.assert_called_with(
+        model_name='test_model', stage='Production'
+    )
+    mlflow_repository.client.download_artifacts.assert_called_once_with('run-123', 'data_model')
+    joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_download_model_wrapper_fallback_reuses_downloaded_dir(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    artifact_dir = build_artifact_dir(tmp_path, 'MLmodel', 'python_model.pkl')
+    mlflow_repository.dowload_artifacts = AsyncMock(return_value=artifact_dir)
+    mlflow.pyfunc.load_model.side_effect = Exception('flavor boom')
+
+    model, returned_path = await mlflow_repository.download_model(
+        'test_model', {}, 'predict', 'pyfunc', True
+    )
+
+    mlflow_repository.dowload_artifacts.assert_called_once_with(
+        'test_model', {}, 'prediction_model'
+    )
+    mlflow_repository.client.download_artifacts.assert_not_called()
+    joblib_mock.load.assert_called_once_with(str(tmp_path / 'python_model.pkl'))
+
+    assert model is joblib_mock.load.return_value
+    assert returned_path == artifact_dir
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_predict_model_both_attempts_fail_chains_causes(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    flavor_error = Exception('flavor boom')
+    mlflow.sklearn.load_model.side_effect = flavor_error
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'model.pkl'
+    )
+    joblib_mock.load.side_effect = ValueError('joblib boom')
+
+    with pytest.raises(ValueError) as exc_info:
+        await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
+
+    assert str(exc_info.value) == 'joblib boom'
+    assert exc_info.value.__cause__ is flavor_error
+
+    error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
+    assert len(error_tags) == 1
+    assert error_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback'
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_transform_model_both_attempts_fail_emits_error_metric(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.get_model_uri = MagicMock()
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'model.pkl'
+    )
+    joblib_mock.load.side_effect = ValueError('joblib boom')
+
+    with pytest.raises(ValueError):
+        await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
+
+    error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
+    assert len(error_tags) == 1
+    assert error_tags[0]['operation_type'] == 'load_transform_model_joblib_fallback'
+    assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == []
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_load_predict_model_fallback_success_metrics(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'model.pkl'
+    )
+    mlflow_repository.warning = MagicMock()
+
+    await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
+
+    read_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT)
+    assert len(read_tags) == 1
+    assert read_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback'
+    assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) == []
+
+    mlflow_repository.observe_lag.assert_called_once()
+    lag_call = mlflow_repository.observe_lag.call_args
+    assert lag_call.args[1] is metrics.MODEL_READ_LAG
+    assert lag_call.args[2]['operation_type'] == 'load_predict_model_joblib_fallback'
+
+    mlflow_repository.warning.assert_called_once()
+    assert 'flavor boom' in mlflow_repository.warning.call_args.args[0]
+
+
+@pytest.mark.parametrize(
+    ('filenames', 'expected'),
+    [
+        (('MLmodel', 'model.pkl'), 'model.pkl'),
+        (('MLmodel', 'python_model.pkl'), 'python_model.pkl'),
+        (('b_model.pkl', 'a_model.pkl'), 'a_model.pkl'),
+    ],
+)
+def test_find_model_pickle_selection(mlflow_repository, tmp_path, filenames, expected):
+    local_dir = build_artifact_dir(tmp_path, *filenames)
+
+    result = mlflow_repository._find_model_pickle(
+        local_dir, 'prediction_model', Exception('flavor boom')
+    )
+
+    assert result == str(tmp_path / expected)
+
+
+def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path):
+    nested = tmp_path / 'artifacts'
+    nested.mkdir()
+    (nested / 'stacking_model.pkl').write_bytes(b'pickle-bytes')
+
+    with pytest.raises(ValueError) as exc_info:
+        mlflow_repository._find_model_pickle(
+            str(tmp_path), 'prediction_model', Exception('flavor boom')
+        )
+
+    assert 'No .pkl file found' in str(exc_info.value)
+
+
+def test_find_model_pickle_empty_directory(mlflow_repository, tmp_path):
+    with pytest.raises(ValueError) as exc_info:
+        mlflow_repository._find_model_pickle(str(tmp_path), 'data_model', Exception('flavor boom'))
+
+    assert 'data_model' in str(exc_info.value)
+
+
+def test_find_model_pickle_error_message_carries_context(mlflow_repository, tmp_path):
+    local_dir = build_artifact_dir(tmp_path, 'MLmodel')
+
+    with pytest.raises(ValueError) as exc_info:
+        mlflow_repository._find_model_pickle(
+            local_dir, 'prediction_model', Exception('flavor boom')
+        )
+
+    message = str(exc_info.value)
+    assert 'prediction_model' in message
+    assert 'flavor boom' in message
+    assert 'MLmodel' in message
+
+
+def test_find_model_pickle_ignores_directory_named_pkl(mlflow_repository, tmp_path):
+    (tmp_path / 'algo.pkl').mkdir()
+
+    with pytest.raises(ValueError) as exc_info:
+        mlflow_repository._find_model_pickle(
+            str(tmp_path), 'prediction_model', Exception('flavor boom')
+        )
+
+    assert 'algo.pkl' in str(exc_info.value)
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_get_model_caches_model_loaded_by_fallback(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'model.pkl'
+    )
+
+    first = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn')
+    second = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn')
+
+    assert first is joblib_mock.load.return_value
+    assert second is first
+    assert mlflow.sklearn.load_model.call_count == 1
+    assert joblib_mock.load.call_count == 1
+
+
+@patch(JOBLIB_PATH)
+@pytest.mark.asyncio
+async def test_predict_reports_both_causes_when_fallback_fails(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
+    mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
+        tmp_path, 'model.pkl'
+    )
+    joblib_mock.load.side_effect = ValueError('joblib boom')
+
+    data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
+
+    output = await mlflow_repository.predict(
+        'test_model',
+        data,
+        {'retention_minutes': 0, 'predict_flavor': 'sklearn'},
+        metadata['metadata'],
+    )
+
+    assert output['success'] is False
+    assert 'joblib boom' in output['content']['traceback']
+    assert 'flavor boom' in output['content']['traceback']
+
+
+def test_load_model_pickle_real_joblib_round_trip(mlflow_repository, tmp_path):
+    estimator = LinearRegression().fit(np.array([[0.0], [1.0], [2.0]]), np.array([0.0, 1.0, 2.0]))
+    joblib_lib.dump(estimator, tmp_path / 'model.pkl', compress=3)
+
+    loaded = mlflow_repository._load_model_pickle(
+        str(tmp_path), 'prediction_model', Exception('flavor boom')
+    )
+
+    assert loaded is not estimator
+    assert loaded.predict(np.array([[3.0]]))[0] == pytest.approx(3.0)