QTZPOC-13

QTZPOC-13: Centralize MLflow joblib fallback logic, ensuring consistent model version resolution by reusing the initial run ID, improving logging, and updating dependencies.

laborious/utils/repository/model_repository.py

@@ -16,9 +16,11 @@ 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, listdir, makedirs, path
@@ -393,7 +395,11 @@ class MLFlowRepository(SientiaMonitoring):
         return path.join(local_dir, pkl_files[0])
 
     def _load_model_pickle(
-        self, local_dir: str, artifact_name: str, flavor_error: Exception
+        self,
+        local_dir: str,
+        artifact_name: str,
+        flavor_error: Exception,
+        metadata: dict[str, Any] | None = None,
     ) -> Any:
         """
         Deserialize a downloaded model artifact with `joblib`.
@@ -407,6 +413,9 @@ class MLFlowRepository(SientiaMonitoring):
             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.
@@ -417,7 +426,7 @@ class MLFlowRepository(SientiaMonitoring):
         """
         pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error)
 
-        self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}')
+        self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}', metadata)
 
         return joblib.load(pkl_path)
 
@@ -427,6 +436,7 @@ class MLFlowRepository(SientiaMonitoring):
         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`.
@@ -438,11 +448,20 @@ class MLFlowRepository(SientiaMonitoring):
         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.
@@ -452,7 +471,8 @@ class MLFlowRepository(SientiaMonitoring):
             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')
+        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}',
@@ -461,7 +481,80 @@ class MLFlowRepository(SientiaMonitoring):
 
         local_dir = self.client.download_artifacts(run_id, artifact_name)
 
-        return self._load_model_pickle(local_dir, artifact_name, flavor_error)
+        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'
@@ -492,37 +585,26 @@ class MLFlowRepository(SientiaMonitoring):
         fallback_labels = self.get_core_labels(
             metadata, operation_type='load_predict_model_joblib_fallback'
         )
-        start_time = time.time()
+
-        try:
+        def load_flavor() -> Any:
             if flavor == 'pyfunc':
-                model = mlflow.pyfunc.load_model(model_uri)
+                return mlflow.pyfunc.load_model(model_uri)
-            elif flavor == 'sklearn':
+            if flavor == 'sklearn':
-                model = mlflow.sklearn.load_model(model_uri)
+                return mlflow.sklearn.load_model(model_uri)
-            else:
+            return mlflow.pytorch.load_model(model_uri)
-                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)
+        return await self._load_with_joblib_fallback(
-            await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
+            primary_loader=load_flavor,
-            return model
+            fallback_loader=lambda flavor_error: self._download_and_load_pickle(
-
+                model_name, 'prediction_model', flavor_error, metadata
-        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)
+            failure_description=(
-        return model
+                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'
@@ -545,6 +627,11 @@ 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'):
@@ -559,38 +646,28 @@ class MLFlowRepository(SientiaMonitoring):
         fallback_labels = self.get_core_labels(
             metadata, operation_type='load_transform_model_joblib_fallback'
         )
-        start_time = time.time()
+
-        try:
+        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)
-            else:
+            return mlflow.pytorch.load_model(model_uri)
-                model = mlflow.pytorch.load_model(model_uri)
+
-        except Exception as flavor_error:
+        return await self._load_with_joblib_fallback(
-            self.warning(
+            primary_loader=load_flavor,
-                f'Flavor loader {flavor} failed for data model {model_name}, '
+            fallback_loader=lambda flavor_error: self._download_and_load_pickle(
-                f'trying joblib fallback. Original error: {flavor_error}',
+                model_name,
+                'data_model',
+                flavor_error,
                 metadata,
-            )
+                run_id=latest_production_id,
-            try:
+            ),
-                model = await self._download_and_load_pickle(
+            failure_description=f'Flavor loader {flavor} failed for data model {model_name}',
-                    model_name, 'data_model', flavor_error, metadata
+            metadata=metadata,
-                )
+            fallback_labels=fallback_labels,
-            except Exception as fallback_error:
+            core_labels=core_labels,
-                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)
-
-        return model
 
     async def download_model(
         self,
@@ -637,26 +714,20 @@ class MLFlowRepository(SientiaMonitoring):
             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)
+            def load_wrapper_flavor() -> Any:
-                await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
+                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

requirements-light.txt

@@ -15,5 +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,5 +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

@@ -1903,6 +1903,28 @@ async def test_load_transform_model_fallback_recovers_load(
     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(