Merge branch 'main' into fix/SIENTIAPDE-2072

.gitignore

@@ -55,4 +55,5 @@ mlruns/
 
 relatorio*
 openspec/
-.cursor/
+.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():
@@ -50,7 +57,7 @@ def minio_container():
 def postgres_container():
     """
     Create a PostgreSQL container using testcontainers.
-    
+
     This fixture creates a real PostgreSQL database in a Docker container
     that will be used for all tests in the session.
     """
@@ -64,21 +71,21 @@ def postgres_container():
 def postgres_engine(postgres_container):
     """
     Create SQLAlchemy engine for PostgreSQL test database.
-    
+
     This fixture creates a connection to the PostgreSQL container
     created by the postgres_container fixture.
     """
     engine = create_engine(postgres_container.get_connection_url())
-    
+
     yield engine
-    
+
     engine.dispose()
 
 
 def _create_schema_and_tables(engine):
     """
     Helper function to create schema and tables in the given engine.
-    
+
     Creates predictions_schema with:
     - laborious_data: Input data table for queries
     - predictions: Output predictions table
@@ -87,8 +94,8 @@ 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 = """
             CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
@@ -102,7 +109,7 @@ def _create_schema_and_tables(engine):
             );
         """
         conn.execute(text(create_laborious_data_sql))
-        
+
         # Create predictions table
         create_predictions_sql = """
             CREATE TABLE if not exists predictions_schema.predictions (
@@ -119,7 +126,7 @@ def _create_schema_and_tables(engine):
             );
         """
         conn.execute(text(create_predictions_sql))
-        
+
         # Create transformed_data table
         create_transformed_sql = """
             CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
@@ -139,7 +146,7 @@ def _create_schema_and_tables(engine):
 def setup_postgres_schema_and_tables(postgres_engine):
     """
     Automatically create necessary schema and tables before each test.
-    
+
     This fixture runs automatically (autouse=True) and ensures
     that the predictions_schema and tables exist with the correct structure.
     """
@@ -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
 
 
@@ -186,24 +180,24 @@ def mock_logger():
 def mock_mongo_client():
     """
     Mock MongoDB client to avoid real connections.
-    
+
     This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
     allowing us to use a real NotificationHandler instance without connecting to MongoDB.
     """
     mock_client = MagicMock()
     mock_db = MagicMock()
     mock_collection = MagicMock()
-    
+
     # Configure the mock chain: client[database] -> db[collection] -> collection
     mock_client.__getitem__.return_value = mock_db
     mock_db.__getitem__.return_value = mock_collection
-    
+
     # Mock server_info() to avoid connection attempts
     mock_client.server_info = MagicMock()
-    
+
     # Mock insert_one for notifications
     mock_collection.insert_one = MagicMock()
-    
+
     return mock_client
 
 
@@ -225,7 +219,7 @@ def notification_inserts(mock_mongo_client):
 def notification_handler(mock_logger, mock_mongo_client):
     """
     Create a real NotificationHandler instance with mocked MongoDB client.
-    
+
     This fixture creates a real CoreNotificationHandler instance but mocks
     the underlying MongoDB connection to avoid real database connections.
     """
@@ -258,7 +252,7 @@ def mock_minio_repository():
     parquet_buffer = BytesIO()
     parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
     parquet_bytes = parquet_buffer.getvalue()
-    
+
     # sientia_do MinioRepository API
     mock_repo.bucket = 'test-bucket'
     mock_repo.upload_file = AsyncMock(
@@ -273,7 +267,7 @@ def mock_minio_repository():
     mock_repo.list_objects = AsyncMock(return_value=[])
     mock_repo.delete_file = AsyncMock()
     mock_repo.close = MagicMock()
-    
+
     return mock_repo
 
 
@@ -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,44 +325,54 @@ 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
     mock_pyfunc_model = MagicMock()
     mock_pyfunc_model._model_impl = MagicMock()
     mock_pyfunc_model._model_impl.python_model = mock_transform_model
-    
+
     return {
         'transform_model': mock_transform_model,
         'predict_model': mock_predict_model,
@@ -374,50 +380,130 @@ 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)
-    
+
     # Mock pytorch.load_model
     mock_mlflow.pytorch = MagicMock()
     mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
-    
+
     # 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(
@@ -203,17 +212,17 @@ async def test_scenario_3_1_2_export_with_opc_only(
 ):
     """
     Scenario 3.1.2: Export with OPC only
-    
+
     Description:
     Export to PostgreSQL and OPC server only (no PI Web API).
-    
+
     Expected Behavior:
     - Normal formatting
     - PostgreSQL export executed
     - OPC export executed
     - PI Web API activity skipped
     - Metrics written with OPC metrics
-    
+
     Assertions:
     - PI Web API activity NOT called
     - OPC activity called
@@ -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(
@@ -291,17 +299,17 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
 ):
     """
     Scenario 3.1.3: Export with PI Web API only
-    
+
     Description:
     Export to PostgreSQL and PI Web API only (no OPC).
-    
+
     Expected Behavior:
     - Normal formatting
     - PostgreSQL export executed
     - PI Web API export executed
     - OPC activity skipped
     - Metrics written without OPC metrics
-    
+
     Assertions:
     - OPC activity NOT called
     - PI Web API activity called
@@ -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(
@@ -375,16 +382,16 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
 ):
     """
     Scenario 3.1.4: Export Without Optional Outputs
-    
+
     Description:
     Export only to PostgreSQL (no OPC or PI Web API).
-    
+
     Expected Behavior:
     - Normal formatting
     - Only PostgreSQL export executed
     - OPC and PI Web API activities skipped
     - Metrics written without OPC metrics
-    
+
     Assertions:
     - PI Web API activity NOT called
     - OPC activity NOT called
@@ -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(
@@ -423,15 +429,15 @@ async def test_scenario_3_1_5_export_without_transformed_data(
 ):
     """
     Scenario 3.1.5: Export Without Transformed Data
-    
+
     Description:
     Only prediction exported, no transform table.
-    
+
     Expected Behavior:
     - Only prediction formatted and exported
     - Transform export skipped
     - Single PostgreSQL write
-    
+
     Assertions:
     - format_transformed_data NOT called
     - One PostgreSQL export
@@ -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,14 +632,17 @@ 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'] = {
         '1': {
@@ -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

@@ -16,15 +16,18 @@ Capabilities:
 
 import ctypes
 import gc
+import inspect
 import threading
 import time
 import traceback
+from collections.abc import Callable
 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 +354,208 @@ 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,
+        metadata: dict[str, Any] | None = None,
+    ) -> 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.
+            metadata (dict[str, Any] | None): Metadata used for structured
+                logging, so the chosen pickle path correlates with the other
+                fallback log lines.
+
+        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}', metadata)
+
+        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],
+        run_id: str | None = None,
+    ) -> 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.
+
+        Callers that already resolved a run to build the flavor URI must pass it
+        as `run_id`: resolving the Production run again here would open a window
+        where a promotion between the two lookups makes the fallback serve a
+        different version than the one the flavor attempt tried and logged. Only
+        callers whose flavor attempt used a symbolic URI — `models:/<name>/production`,
+        resolved by MLflow itself — leave it as `None`.
+
+        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.
+            run_id (str | None): Run already resolved by the caller. When `None`,
+                the Production run is resolved here.
+
+        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.
+        """
+        if run_id is None:
+            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, metadata)
+
+    async def _load_with_joblib_fallback(
+        self,
+        primary_loader: Callable[[], Any],
+        fallback_loader: Callable[[Exception], Any],
+        failure_description: str,
+        metadata: dict[str, Any],
+        fallback_labels: dict[str, Any],
+        core_labels: dict[str, Any] | None = None,
+    ) -> Any:
+        """
+        Run a model load with the joblib fallback and its observability.
+
+        Single owner of the fallback flow shared by every load path: warn on the
+        flavor failure without emitting an error metric, try the fallback, and
+        emit lag/count under the fallback labels — or the error count chained to
+        the flavor error when the fallback also fails. Fixes to the flow belong
+        here, not in the call sites.
+
+        Both loaders are called with no observability concern of their own:
+        `fallback_loader` receives the flavor error and may be sync or async,
+        which is what lets the download-based paths and the already-downloaded
+        wrapper path share this method.
+
+        Args:
+            primary_loader (Callable[[], Any]): Flavor load attempt. Any
+                exception it raises triggers the fallback.
+            fallback_loader (Callable[[Exception], Any]): joblib fallback,
+                receiving the flavor error. Its result is awaited when
+                awaitable.
+            failure_description (str): What failed, e.g. 'Flavor loader sklearn
+                failed for prediction model x'. The fallback and original error
+                are appended to it.
+            metadata (dict[str, Any]): Metadata used for structured logging.
+            fallback_labels (dict[str, Any]): Metric labels for the fallback
+                path, already carrying its `operation_type`.
+            core_labels (dict[str, Any] | None): Metric labels for a successful
+                flavor load. `None` emits no metric on that path.
+
+        Returns:
+            Any: The model loaded by the flavor loader or by the fallback.
+
+        Raises:
+            Exception: The fallback error, chained from the flavor error, when
+                both attempts fail.
+        """
+        start_time = time.time()
+        try:
+            model = primary_loader()
+        except Exception as flavor_error:
+            self.warning(
+                f'{failure_description}, trying joblib fallback. Original error: {flavor_error}',
+                metadata,
+            )
+            try:
+                model = fallback_loader(flavor_error)
+                if inspect.isawaitable(model):
+                    model = await model
+            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
+
+        if core_labels is not None:
+            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)
+
+        return model
+
     async def load_predict_model(
         self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
     ) -> Any:
@@ -370,27 +575,36 @@ 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')
-        start_time = time.time()
+        fallback_labels = self.get_core_labels(
-        try:
+            metadata, operation_type='load_predict_model_joblib_fallback'
-            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)
-        except Exception as e:
-            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
-            raise e
 
-        await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
+        def load_flavor() -> Any:
-        await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
+            if flavor == 'pyfunc':
-        return model
+                return mlflow.pyfunc.load_model(model_uri)
+            if flavor == 'sklearn':
+                return mlflow.sklearn.load_model(model_uri)
+            return mlflow.pytorch.load_model(model_uri)
+
+        return await self._load_with_joblib_fallback(
+            primary_loader=load_flavor,
+            fallback_loader=lambda flavor_error: self._download_and_load_pickle(
+                model_name, 'prediction_model', flavor_error, metadata
+            ),
+            failure_description=(
+                f'Flavor loader {flavor} failed for prediction model {model_name}'
+            ),
+            metadata=metadata,
+            fallback_labels=fallback_labels,
+            core_labels=core_labels,
+        )
 
     async def load_transform_model(
         self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -413,32 +627,47 @@ class MLFlowRepository(SientiaMonitoring):
         Raises:
             Exception: If the model run ID or URI cannot be retrieved, or if the
                 model cannot be loaded.
+
+        Notes:
+            - The Production run is resolved once and reused by the joblib
+              fallback, so a promotion between both steps cannot make the
+              fallback serve a different version than the flavor attempt.
         """
 
+        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')
-        start_time = time.time()
+        fallback_labels = self.get_core_labels(
-        try:
+            metadata, operation_type='load_transform_model_joblib_fallback'
+        )
+
+        def load_flavor() -> Any:
             if flavor == 'sklearn':
-                model = mlflow.sklearn.load_model(model_uri)
+                return mlflow.sklearn.load_model(model_uri)
-            elif flavor == 'pyfunc':
+            if flavor == 'pyfunc':
-                model = mlflow.pyfunc.load_model(model_uri)
+                return mlflow.pyfunc.load_model(model_uri)
-            elif flavor == 'pytorch':
+            return mlflow.pytorch.load_model(model_uri)
-                model = mlflow.pytorch.load_model(model_uri)
-            else:
-                raise ValueError(INVALID_FLAVOR_MESSAGE)
-        except Exception as e:
-            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
-            raise e
 
-        await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
+        return await self._load_with_joblib_fallback(
-        await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
+            primary_loader=load_flavor,
-
+            fallback_loader=lambda flavor_error: self._download_and_load_pickle(
-        return model
+                model_name,
+                'data_model',
+                flavor_error,
+                metadata,
+                run_id=latest_production_id,
+            ),
+            failure_description=f'Flavor loader {flavor} failed for data model {model_name}',
+            metadata=metadata,
+            fallback_labels=fallback_labels,
+            core_labels=core_labels,
+        )
 
     async def download_model(
         self,
@@ -482,8 +711,23 @@ 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'
+            )
+
+            def load_wrapper_flavor() -> Any:
+                raw_model = mlflow.pyfunc.load_model(artifact_path)
+                return raw_model._model_impl.python_model
+
+            model = await self._load_with_joblib_fallback(
+                primary_loader=load_wrapper_flavor,
+                fallback_loader=lambda flavor_error: self._load_model_pickle(
+                    artifact_path, target, flavor_error, metadata
+                ),
+                failure_description=f'Wrapper loader failed for {model_type} model {model_name}',
+                metadata=metadata,
+                fallback_labels=fallback_labels,
+            )
 
             self.debug(
                 f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata
@@ -921,6 +1165,43 @@ class MLFlowRepository(SientiaMonitoring):
         }
         return retrain_data
 
+    def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]:
+        """Resolve the code dependencies to forward when re-logging a pyfunc model.
+
+        MLflow flattens whatever was given as `code_paths` into `<artifact>/code/`,
+        so the layout differs per model: some ship a `utils` package, others a
+        single self-contained module (`<model>_model.py`). Re-logging must forward
+        the entries that are actually on disk — assuming one fixed layout makes
+        `_validate_and_copy_code_paths` raise `FileNotFoundError` for the other.
+
+        Args:
+            artifact_path (str | None): Local directory of the downloaded model
+                artifact, e.g. `./tmp/artifacts/<model>/prediction_model`.
+            metadata (dict): Metadata for structured logging.
+
+        Returns:
+            list[str]: Existing paths inside `<artifact_path>/code`, empty when
+                the model carries no code directory.
+        """
+        if not artifact_path:
+            self.warning('No artifact path available, logging model without code paths', metadata)
+            return []
+
+        code_dir = path.join(artifact_path, 'code')
+
+        if not path.isdir(code_dir):
+            self.warning(
+                f'No code directory found at {code_dir}, logging model without code paths',
+                metadata,
+            )
+            return []
+
+        return [
+            path.join(code_dir, entry)
+            for entry in sorted(listdir(code_dir))
+            if entry != '__pycache__'
+        ]
+
     async def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict):
         """Log a model into the active MLflow run.
 
@@ -941,7 +1222,7 @@ class MLFlowRepository(SientiaMonitoring):
             if flavor == 'sklearn':
                 mlflow.sklearn.log_model(model, model_type)
             elif flavor == 'pyfunc':
-                code_path = [path.join(model_data['artifact_path'], 'code', 'utils')]
+                code_path = self.resolve_code_paths(model_data.get('artifact_path'), metadata)
 
                 self.debug(f'Code path: {code_path}', 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

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

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

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 = [
@@ -1028,21 +1049,52 @@ async def test_log_model_sklearn(mlflow, mlflow_repository):
     )
 
 
-@patch('laborious.utils.repository.model_repository.path')
+def test_resolve_code_paths_package_layout(tmp_path, mlflow_repository):
+    code_dir = tmp_path / 'code'
+    (code_dir / 'utils').mkdir(parents=True)
+    (code_dir / '__pycache__').mkdir()
+
+    assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [
+        str(code_dir / 'utils')
+    ]
+
+
+def test_resolve_code_paths_single_module_layout(tmp_path, mlflow_repository):
+    code_dir = tmp_path / 'code'
+    code_dir.mkdir()
+    (code_dir / 'loop_valve_classifier_model.py').touch()
+
+    assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [
+        str(code_dir / 'loop_valve_classifier_model.py')
+    ]
+
+
+def test_resolve_code_paths_without_code_dir(tmp_path, mlflow_repository):
+    assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == []
+
+
+def test_resolve_code_paths_without_artifact_path(mlflow_repository):
+    assert mlflow_repository.resolve_code_paths(None, metadata['metadata']) == []
+
+
 @pytest.mark.asyncio
-async def test_log_model_pyfunc(path, mlflow, mlflow_repository):
+async def test_log_model_pyfunc(tmp_path, mlflow, mlflow_repository):
+    code_dir = tmp_path / 'code'
+    code_dir.mkdir()
+    (code_dir / 'loop_valve_classifier_model.py').touch()
+
     model_mock = MagicMock()
-    model_data = {'model': model_mock, 'artifact_path': 'artifact_path'}
+    model_data = {'model': model_mock, 'artifact_path': str(tmp_path)}
     await mlflow_repository.log_model(
         model_data, 'pyfunc', 'prediction_model', metadata['metadata']
     )
 
     mlflow.pyfunc.log_model.assert_not_called()
 
-    path.join.assert_called_once_with('artifact_path', 'code', 'utils')
-
     model_mock.store_model.assert_called_once_with(
-        artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False
+        artifact_path='prediction_model',
+        code_path=[str(code_dir / 'loop_valve_classifier_model.py')],
+        to_disk=False,
     )
 
     mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY)
@@ -1713,3 +1765,403 @@ 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_load_transform_model_fallback_reuses_resolved_run(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    """A promotion between both steps must not change the run served."""
+    mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
+    mlflow_repository.get_model_run_id = MagicMock(side_effect=['run-A', 'run-B'])
+    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_once_with(
+        model_name='test_model', stage='Production'
+    )
+    mlflow_repository.client.download_artifacts.assert_called_once_with('run-A', 'data_model')
+
+
+@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_download_model_wrapper_both_attempts_fail_chains_causes(
+    joblib_mock, mlflow, mlflow_repository, tmp_path
+):
+    flavor_error = Exception('flavor boom')
+    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 = flavor_error
+    joblib_mock.load.side_effect = ValueError('joblib boom')
+
+    with pytest.raises(ValueError) as exc_info:
+        await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
+
+    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'] == 'download_model_wrapper_joblib_fallback'
+    assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == []
+    mlflow_repository.observe_lag.assert_not_called()
+
+
+@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)