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.

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 sientia_do.observability.logger import Logger
-from sientia_do.observability.metrics_controller import MetricsController
+from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
 
 # 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}")
-    def custom_message(message, _metadata={}):
-        print(f"[LOG] {message}")
+        print(f'[LOG] {message}')
+
+    def custom_message(message, _metadata=None):
+        print(f'[LOG] {message}')
+
     logger = MagicMock()
-    logger.info = MagicMock(
-        side_effect=message
-    )
-    logger.debug = MagicMock(
-        side_effect=message
-    )
-    logger.error = MagicMock(
-        side_effect=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
-    )
+    logger.info = MagicMock(side_effect=message)
+    logger.debug = MagicMock(side_effect=message)
+    logger.error = MagicMock(side_effect=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(
-        return_value=(True, {'response_time': 0.1})
-    )
+    mock_repo.write_data = AsyncMock(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({
-            'feature_1': [0.234] * num_rows,
-            'feature_2': [0.783] * num_rows,
-        })
+        result = pd.DataFrame(
+            {
+                '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(
-            current_stage='Production',
-            version='1',
-            source='runs:/artifacts/test_run_id'
-        )])
-    ))
-    
+        search_model_versions=MagicMock(
+            return_value=[
+                MagicMock(
+                    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:
-        if self._session_bad_on_write:
+    async def _pre_write_callback(self, event, _service) -> None:
+        # 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(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}'))
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
+        )
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
+        )
+        conn.execute(
+            text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
+        )
     insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
 
     input_data = {

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(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, ['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, {
-        'notification_id': 'OPC_WRITE_DATA_ERROR_1',
-        'message': 'OPC server unavailable',
-        'block': 'opc_repository',
-        'level': NotificationLevel.ERROR,
-        'attachment_content': 'OPC server unavailable',
-    })
-    
+    opc_write_data.return_value = (
+        False,
+        {
+            'notification_id': 'OPC_WRITE_DATA_ERROR_1',
+            'message': '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(