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.

laborious/utils/repository/model_repository.py

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