diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index ab007d4..e89b685 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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://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) - elif flavor == 'sklearn': - model = mlflow.sklearn.load_model(model_uri) - else: - 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 + return mlflow.pyfunc.load_model(model_uri) + if flavor == 'sklearn': + return mlflow.sklearn.load_model(model_uri) + return mlflow.pytorch.load_model(model_uri) - 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 + 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' @@ -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) - elif flavor == 'pyfunc': - model = mlflow.pyfunc.load_model(model_uri) - else: - 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}', + return mlflow.sklearn.load_model(model_uri) + if flavor == 'pyfunc': + return mlflow.pyfunc.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, + 'data_model', + 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) - - return model + 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, @@ -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) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) + 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 diff --git a/requirements-light.txt b/requirements-light.txt index d5128de..48cb400 100644 --- a/requirements-light.txt +++ b/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 diff --git a/requirements.txt b/requirements.txt index 8f994e9..4b2a2a8 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 5ea4ce7..4a89d69 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/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(