diff --git a/README.md b/README.md index f61454e..3ce886b 100644 --- a/README.md +++ b/README.md @@ -1114,8 +1114,15 @@ no fallback between flavors: a failing loader propagates. retrain (`model.pkl` + `code/`, no `MLmodel`, so only this runtime reads it back). Every flavor puts the artifact's own `code/` on `sys.path` before loading, **when that directory -exists**: MLflow's loaders only do it when the `MLmodel` manifest declares `code`, so `sklearn` and -`pyfunc` resolve the URI to a local directory first and load from there. +exists**: MLflow's loaders only do it when the `MLmodel` manifest declares `code`. `download_model` +owns that step for all of them — it asks the registry whether the artifact carries `code/` and, only +then, downloads that one directory. The models themselves are never pre-downloaded: `sklearn` and +`pyfunc` load straight from the model URI and fetch what they need. + +`download_model` does not branch on the flavor: past that shared first step it delegates to +`load_predict_model` / `load_transform_model`, and its only branch left is `load_wrapper`, which is +orthogonal to the flavor. It returns `(model, artifact_dir)`, the directory the retrain path re-logs +`code/` from. Full reference, including how to switch an existing model over: **[docs/model-flavors.md](docs/model-flavors.md)**. diff --git a/docs/model-flavors.md b/docs/model-flavors.md index 74c0d72..9efd64c 100644 --- a/docs/model-flavors.md +++ b/docs/model-flavors.md @@ -22,8 +22,8 @@ reads the artifact's `MLmodel` manifest to guess it. | Flavor | Read | Write (retrain) | How it loads | |---|---|---|---| -| `sklearn` | yes | yes | resolves the URI to a local directory, then `mlflow.sklearn.load_model()` | -| `pyfunc` | yes | yes | resolves the URI to a local directory, then `mlflow.pyfunc.load_model()`; the wrapper path also unwraps `._model_impl.python_model` | +| `sklearn` | yes | yes | `mlflow.sklearn.load_model()` — the loader fetches what it needs | +| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model()`; the wrapper path downloads the artifact instead and unwraps `._model_impl.python_model` | | `joblib` | yes | yes | downloads the same artifact and `joblib.load`s its `/model.pkl` — **no MLflow flavor loader involved** | Anything else raises `ValueError` with `INVALID_FLAVOR_MESSAGE`, **before** any artifact is resolved @@ -71,15 +71,29 @@ which no manifest can misdescribe: - `prepend_artifact_code_dir` inserts `/code` at the front of `sys.path`, **only when that directory exists**, and only once (an entry already present is not duplicated). -- Because the insert must happen *before* the loader unpickles, `sklearn` and `pyfunc` resolve their - URI to a local directory first (`mlflow.artifacts.download_artifacts`) and the loader is handed - that directory. Same bytes, same download count — MLflow would have downloaded it anyway. -- The wrapper path (`download_model(load_wrapper=True)`) already had its local directory and adds - the entry itself; `joblib` adds it inside `_load_model_pickle`. +- `download_artifact_code_dir` is what feeds it on the download paths, and it is deliberately cheap: + the remote listing (`list_artifacts(run_id, '')`) decides first, so a run without `code` + downloads **nothing**, and a run with it downloads only `/code` — a handful of `.py` + files, never the model's weights. +- `download_model` runs that step first, for every flavor and both of its branches, *before* any + model is loaded — the insert has to be in place before the loader unpickles. +- The models themselves are not pre-downloaded: `mlflow.sklearn.load_model` and + `mlflow.pyfunc.load_model` are handed the model URI and fetch what they need. `joblib` downloads + its artifact because it must read the pickle, and adds the entry again inside `_load_model_pickle` + through the same helper, so the flavor works when called directly too. - When the manifest *does* declare `code`, MLflow inserts the same directory a second time. Harmless. The entry is never removed: a loaded model may import lazily, long after the load returned. +The directory `download_artifact_code_dir` downloaded into is also what `download_model` returns +next to the model — `tuple[Any, str | None]`, the second element being the artifact directory the +retrain path re-logs `code/` from, or `None` when the artifact has no `code/`. It falls back to +whatever the loader downloaded on its own, which is a real artifact directory for `joblib` and for +the wrapper branch. `load_predict_model` and `load_transform_model` stay the **single** place where +the flavor is decided: `download_model` has no `joblib` branch, it delegates by `model_type`. Its +only remaining branch is `load_wrapper`, which is derived from `flavor == 'pyfunc'` but is otherwise +orthogonal to the flavor. + > `sys.path` is process-global: two models embedding a package with the same name (`utils/`) on the > same worker resolve to whichever entered first. Mitigation, if it bites: one worker per model. diff --git a/e2e/conftest.py b/e2e/conftest.py index bf4310f..b9ee3d2 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -546,14 +546,7 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts): mock_mlflow.pyfunc = MagicMock() mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model) - # The sklearn and pyfunc paths resolve the URI to a local directory before loading, so that the - # artifact's `code/` can go on `sys.path` whatever the `MLmodel` manifest says. Echoing the URI - # keeps the loader stubs above — which branch on the URI text — working unchanged, and points - # the code-directory lookup at a path that does not exist. mock_mlflow.artifacts = MagicMock() - mock_mlflow.artifacts.download_artifacts = MagicMock( - side_effect=lambda artifact_uri: str(artifact_uri) - ) # Mock other mlflow methods that might be called mock_mlflow.set_tracking_uri = MagicMock() @@ -572,9 +565,23 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts): ), ) - # Only the joblib flavor calls download_artifacts directly on the client, one artifact - # per model type — a single return value would serve `data_model` from `prediction_model`. + # `download_model` asks whether the run's artifact carries a `code` directory before it loads + # anything, and downloads only that directory when it does. Answering "no" here keeps the + # download counts of these tests about the model itself; a test that exercises the `code` step + # overrides `list_artifacts` to answer "yes". + def mock_list_artifacts(_run_id, artifact_path=None): + if artifact_path: + return [] + return [MagicMock(path=artifact_name) for artifact_name in joblib_flavor_artifacts] + + mock_mlflow_client.list_artifacts = MagicMock(side_effect=mock_list_artifacts) + + # The joblib flavor downloads a whole model artifact, one per model type — a single return + # value would serve `data_model` from `prediction_model`. The `code` step asks for the nested + # `/code` instead, and gets that directory back, as MLflow would return it. def mock_download_artifacts(_run_id, artifact_name): + if artifact_name.endswith('/code'): + return joblib_flavor_artifacts[artifact_name.removesuffix('/code')]['code_dir'] return joblib_flavor_artifacts[artifact_name]['path'] mock_mlflow_client.download_artifacts = MagicMock(side_effect=mock_download_artifacts) diff --git a/e2e/scenarios.md b/e2e/scenarios.md index 1654303..28e70c0 100644 --- a/e2e/scenarios.md +++ b/e2e/scenarios.md @@ -1149,9 +1149,34 @@ restores `sys.path` in teardown, because the flavor deliberately never removes w #### Scenario 5.1.3: An MLflow Flavor Is Unaffected **Description**: A model declaring `sklearn` behaves exactly as before +**Input**: +- `client.list_artifacts` answering that the run's artifacts carry no `code` directory + **Assertions**: - The prediction lands with the mocked model's value -- `download_artifacts` NOT called — the joblib path is never entered as a safety net +- `download_artifacts` NOT called — the joblib path is never entered as a safety net, and with no + `code` directory to fetch there is nothing else to download either + +--- + +#### Scenario 5.1.4: Only the Artifact's `code` Directory Is Downloaded +**Description**: A model on an MLflow flavor whose artifact **does** carry `code/`. The flavor +loaders read the model URI themselves and fetch what they need, so the `code` directory — the one +thing MLflow will not put on `sys.path` unless the `MLmodel` manifest declares it — is the whole of +what this workflow downloads through the client + +**Input**: +- `model_config` on the default `sklearn` flavors +- `client.list_artifacts` answering that both artifacts carry `code` + +**Expected Behavior**: +- The `code` directory of each artifact is downloaded before its model is loaded, and goes on + `sys.path` +- Neither model's weights travel through `client.download_artifacts` + +**Assertions**: +- The prediction lands with the mocked model's value +- `download_artifacts` called exactly twice, with `'data_model/code'` and `'prediction_model/code'` --- diff --git a/e2e/test_predictions_batch_joblib_flavor.py b/e2e/test_predictions_batch_joblib_flavor.py index 0993e8f..76a0410 100644 --- a/e2e/test_predictions_batch_joblib_flavor.py +++ b/e2e/test_predictions_batch_joblib_flavor.py @@ -7,6 +7,8 @@ A passing run therefore proves the flavor decompresses the pickle **and** resolv without any `mlflow.*.load_model` being called for that model. """ +from unittest.mock import MagicMock + import pandas as pd import pytest from sqlalchemy import text @@ -208,6 +210,41 @@ async def test_prediction_workflow_with_sklearn_flavor_is_unchanged( mlflow_client.download_artifacts.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.integration +async def test_prediction_workflow_downloads_only_the_artifact_code_directory( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + joblib_flavor_artifacts, + mlflow_client, + mock_mlflow_models, +): + """A flavor model whose artifact carries `code` downloads that directory — and only it. + + The flavor loaders read the model URI themselves, so nothing else is fetched through the + client: the `code` directory is the whole of what this workflow downloads. + """ + mlflow_client.list_artifacts.side_effect = lambda _run_id, artifact_path=None: ( + [MagicMock(path=f'{artifact_path}/code')] if artifact_path else [] + ) + + client = temporal_test_env.client + model_id = 275 + 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-code-only') + ) + + assert_prediction(postgres_engine, model_id, prediction=0.5) + + downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list] + assert sorted(downloaded) == ['data_model/code', 'prediction_model/code'] + + @pytest.mark.asyncio @pytest.mark.integration async def test_prediction_workflow_joblib_flavor_failure_reports_single_cause( diff --git a/e2e/test_retrain_joblib_flavor.py b/e2e/test_retrain_joblib_flavor.py index dabda57..4e6dcdd 100644 --- a/e2e/test_retrain_joblib_flavor.py +++ b/e2e/test_retrain_joblib_flavor.py @@ -117,7 +117,9 @@ async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads( mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written) - reloaded = await repository.load_predict_model(JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib') + reloaded, _local_dir = await repository.load_predict_model( + JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib' + ) features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]}) expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index ea13d07..76fbb4e 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -261,14 +261,21 @@ class MLFlowRepository(SientiaMonitoring): """ Check if an artifact exists in the MLflow Model Registry. + The listing is scoped to the parent of `artifact_path`, so nested paths + such as `prediction_model/code` are answered by listing + `prediction_model` rather than the run's artifact root. A top-level path + has no parent and lists the root, as before. Only the listing travels: + nothing is downloaded to answer the question. + Args: run_id (str): Run identifier to inspect. - artifact_path (str): Path to the artifact to check. + artifact_path (str): Path to the artifact to check, relative to the + run's artifact root, e.g. 'data_model' or 'data_model/code'. Returns: bool: True if the artifact exists, False otherwise. """ - artifacts = self.client.list_artifacts(run_id) + artifacts = self.client.list_artifacts(run_id, path.dirname(artifact_path) or None) self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata) self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata) @@ -456,38 +463,62 @@ class MLFlowRepository(SientiaMonitoring): return code_dir - def _localize_model_uri( - self, model_uri: str, flavor: str, metadata: dict[str, Any] | None = None - ) -> str: + async def download_artifact_code_dir( + self, + model_name: str, + artifact_name: str, + flavor: str, + metadata: dict[str, Any] | None = None, + ) -> str | None: """ - Resolve a model URI to a local directory and put its `code` on `sys.path`. + Download only an artifact's `code` directory and put it on `sys.path`. - The MLflow loaders download the artifact themselves and take a local - path just as readily as a URI, so resolving here costs no extra - download — it only moves the download one step earlier, which is what - makes the artifact's `code` directory visible before the loader - unpickles anything. + Cheap by construction: the remote listing decides first, and nothing is + downloaded when the artifact carries no `code`. When it does, the + download is scoped to `/code` — a handful of `.py` files + — never the artifact's weights. Loading the model itself stays the + loaders' business, and they read straight from the model URI. Args: - model_uri (str): `models://production` or a run artifact URI. - flavor (str): Flavor doing the load, for the log line only. + model_name (str): Registered model name. + artifact_name (str): Model artifact to inspect, i.e. + 'prediction_model' or 'data_model'. + flavor (str): Flavor doing the load, for the log lines only. metadata (dict[str, Any] | None): Metadata used for structured logging. Returns: - str: Local directory the model artifact was downloaded to, to be - handed to the flavor loader in place of the URI. + str | None: Local directory the `code` directory was downloaded + *into* — the artifact directory, i.e. the parent of the + downloaded `code` — or `None` when the run's artifact has no + `code` directory. Raises: - Exception: Any error raised while downloading the artifact. + Exception: Any error raised while resolving the Production run, + listing the artifact or downloading `code`. """ - local_path = mlflow.artifacts.download_artifacts(artifact_uri=model_uri) + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + remote_code_path = f'{artifact_name}/code' - self.debug(f'{flavor} flavor: {model_uri} downloaded to {local_path}', metadata) + if not self.check_artifact_exists(run_id, remote_code_path, metadata): + self.debug( + f'{flavor} flavor: run {run_id} has no {remote_code_path} artifact, ' + 'nothing to download', + metadata, + ) + return None - self.prepend_artifact_code_dir(local_path, flavor, metadata) + local_code_dir = self.client.download_artifacts(run_id, remote_code_path) + artifact_dir = path.dirname(local_code_dir) - return local_path + self.debug( + f'{flavor} flavor: {remote_code_path} of run {run_id} downloaded to {local_code_dir}', + metadata, + ) + + self.prepend_artifact_code_dir(artifact_dir, flavor, metadata) + + return artifact_dir def _load_model_pickle( self, @@ -588,62 +619,9 @@ class MLFlowRepository(SientiaMonitoring): return model, local_dir - async def _load_joblib_artifact( - self, - model_name: str, - artifact_name: str, - operation_type: str, - metadata: dict[str, Any], - run_id: str | None = None, - ) -> tuple[Any, str]: - """ - Load a model artifact through the `joblib` flavor, with observability. - - Single owner of the joblib read path: no MLflow flavor loader is called - here, neither before nor after the pickle load. Metrics are emitted with - the **base** `operation_type` the caller passes — a joblib load is a - load, so no suffixed operation type exists. - - Args: - model_name (str): Registered model name. - artifact_name (str): Artifact to download, e.g. 'prediction_model'. - operation_type (str): Base metric operation type of the calling - method, e.g. 'load_predict_model'. - 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 during the download. - - Returns: - tuple[Any, str]: The deserialized model and the local directory the - artifact was downloaded to. - - Raises: - Exception: Any error raised while resolving the run, downloading the - artifact or deserializing the pickle. It propagates as is: there - is no second attempt. - """ - core_labels = self.get_core_labels(metadata, operation_type=operation_type) - start_time = time.time() - - try: - model, local_dir = await self._download_and_load_pickle( - model_name=model_name, - artifact_name=artifact_name, - metadata=metadata, - run_id=run_id, - ) - 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) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) - - return model, local_dir - async def load_predict_model( self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' - ) -> Any: + ) -> tuple[Any, str | None]: """ Load a predictive model from the MLflow Model Registry. @@ -651,57 +629,57 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): The name of the model to download from the registry. metadata (dict[str, Any]): Metadata used for structured logging. flavor (str): Model flavor ('pyfunc', 'sklearn', 'joblib') - artifact_path (str | None): Path to compressed artifacts if model is compressed Returns: - mlflow.pyfunc.PyFuncModel: The loaded predictive model. + tuple[Any, str | None]: The loaded predictive model and the local + directory its artifact was downloaded to, `None` when nothing + was downloaded here. Notes: - The model is fetched from the "Production" stage of the MLflow Model Registry. - Warnings during the model loading process are suppressed. - - The artifact's `code` directory, when present, is put on - `sys.path` before the load, whatever the flavor. `sklearn` and - `pyfunc` are therefore resolved to a local directory first and the - loader is handed that directory instead of the URI. + - The MLflow flavor loaders read the model URI directly and fetch + only what they need. Nothing is pre-downloaded here, so the local + directory returned alongside the model is `None`; putting the + artifact's `code` directory on `sys.path` belongs to + `download_artifact_code_dir`, which `download_model` runs first. - With `flavor='joblib'` no MLflow flavor loader is called: the `prediction_model` artifact is downloaded and deserialized with - `joblib.load`, which needs no `MLmodel` manifest. + `joblib.load`, which needs no `MLmodel` manifest. That path does + return its local directory, since it downloads by necessity. """ if flavor not in ('pyfunc', 'sklearn', 'joblib'): raise ValueError(INVALID_FLAVOR_MESSAGE) - if flavor == 'joblib': - model, _local_dir = await self._load_joblib_artifact( - model_name=model_name, - artifact_name='prediction_model', - operation_type='load_predict_model', - metadata=metadata, - ) - return model + model_uri = None if flavor == 'joblib' else f'models:/{model_name}/production' - model_uri = f'models:/{model_name}/production' - self.info(f'Loading prediction model {model_name} from {model_uri}') + if model_uri is not None: + 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() try: - local_path = self._localize_model_uri(model_uri, flavor, metadata) - - if flavor == 'pyfunc': - model = mlflow.pyfunc.load_model(local_path) + if flavor == 'joblib': + model, local_dir = await self._download_and_load_pickle( + model_name=model_name, + artifact_name='prediction_model', + metadata=metadata, + ) + elif flavor == 'pyfunc': + model, local_dir = mlflow.pyfunc.load_model(model_uri), None else: - model = mlflow.sklearn.load_model(local_path) + model, local_dir = mlflow.sklearn.load_model(model_uri), None 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) await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) - return model + return model, local_dir async def load_transform_model( self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' - ) -> Any: + ) -> tuple[Any, str | None]: """ Load the latest Production version of a transformation model. @@ -712,24 +690,27 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): The name of the model to download. metadata (dict[str, Any]): Metadata used for structured logging. flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib') - artifact_path (str | None): Path to compressed artifacts if model is compressed Returns: - Any: The loaded model object, depending on the flavor used. + tuple[Any, str | None]: The loaded model object, depending on the + flavor used, and the local directory its artifact was downloaded + to, `None` when nothing was downloaded here. Raises: Exception: If the model run ID or URI cannot be retrieved, or if the model cannot be loaded. Notes: - - The artifact's `code` directory, when present, is put on - `sys.path` before the load, whatever the flavor. `sklearn` and - `pyfunc` are therefore resolved to a local directory first and the - loader is handed that directory instead of the URI. + - The MLflow flavor loaders read the model URI directly and fetch + only what they need. Nothing is pre-downloaded here, so the local + directory returned alongside the model is `None`; putting the + artifact's `code` directory on `sys.path` belongs to + `download_artifact_code_dir`, which `download_model` runs first. - With `flavor='joblib'` no MLflow flavor loader is called and no model URI is built: the `data_model` artifact of the resolved Production run is downloaded and deserialized with `joblib.load`, - which needs no `MLmodel` manifest. + which needs no `MLmodel` manifest. That path does return its local + directory, since it downloads by necessity. """ if flavor not in ('sklearn', 'pyfunc', 'joblib'): @@ -737,36 +718,36 @@ class MLFlowRepository(SientiaMonitoring): latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production') - if flavor == 'joblib': - model, _local_dir = await self._load_joblib_artifact( - model_name=model_name, - artifact_name='data_model', - operation_type='load_transform_model', - metadata=metadata, - run_id=latest_production_id, - ) - return model + model_uri = ( + None + if flavor == 'joblib' + else self.get_model_uri(latest_production_id, prediction=False) + ) - model_uri = self.get_model_uri(latest_production_id, prediction=False) - - self.info(f'Loading data model {model_name} from {model_uri}') + if model_uri is not None: + 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() try: - local_path = self._localize_model_uri(model_uri, flavor, metadata) - - if flavor == 'sklearn': - model = mlflow.sklearn.load_model(local_path) + if flavor == 'joblib': + model, local_dir = await self._download_and_load_pickle( + model_name=model_name, + artifact_name='data_model', + metadata=metadata, + run_id=latest_production_id, + ) + elif flavor == 'sklearn': + model, local_dir = mlflow.sklearn.load_model(model_uri), None else: - model = mlflow.pyfunc.load_model(local_path) + model, local_dir = mlflow.pyfunc.load_model(model_uri), None 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) await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) - return model + return model, local_dir async def download_model( self, @@ -787,17 +768,26 @@ class MLFlowRepository(SientiaMonitoring): load_wrapper (bool): Whether to load wrapper Returns: - tuple[Any, str | None]: Model object and optional artifact path. + tuple[Any, str | None]: Model object and the local directory its + artifact was downloaded to. Notes: - - `flavor='joblib'` returns the local directory the artifact was - downloaded to as the artifact path, because the retrain path needs - it to re-log the artifact's `code` directory. Every other flavor - returns `None` unless the wrapper branch was taken. - - The wrapper branch downloads before it loads, so it puts the - artifact's `code` directory on `sys.path` itself. The two plain - branches delegate, and the code directory is handled inside - `load_predict_model` / `load_transform_model`. + - Single owner of the artifact `code` directory, for every flavor + and both branches: `download_artifact_code_dir` runs first, before + any model is loaded, and is the only place that touches + `sys.path`. It downloads nothing unless the run's artifact really + carries a `code` directory, and even then downloads only that + directory — never the model's weights. + - The directory it downloads into is also what this method returns + alongside the model: it is the artifact directory the retrain path + needs to re-log `code`. It falls back to whatever the loader + downloaded on its own, which is a real artifact directory for + `joblib` and for the wrapper branch, and `None` for the flavor + loaders, which stream from the model URI without downloading. + - Apart from that shared first step this method does not branch on + the flavor. `load_predict_model` and `load_transform_model` own + every flavor, `joblib` included. The only branch left is + `load_wrapper`, which is orthogonal to the flavor. """ self.info( @@ -807,47 +797,36 @@ class MLFlowRepository(SientiaMonitoring): if model_type not in ['predict', 'transform']: raise ValueError("Invalid model_type. Use 'predict' or 'transform'.") - artifact_path = None + target = 'prediction_model' if model_type == 'predict' else 'data_model' + + code_artifact_dir = await self.download_artifact_code_dir( + model_name, target, flavor, metadata + ) if load_wrapper: self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}') - target = 'prediction_model' if model_type == 'predict' else 'data_model' - artifact_path = await self.dowload_artifacts(model_name, metadata, target) self.info( f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' ) - self.prepend_artifact_code_dir(artifact_path, 'pyfunc', metadata) - raw_model = mlflow.pyfunc.load_model(artifact_path) model = raw_model._model_impl.python_model self.debug( f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata ) - elif flavor == 'joblib': - target = 'prediction_model' if model_type == 'predict' else 'data_model' - operation_type = ( - 'load_predict_model' if model_type == 'predict' else 'load_transform_model' - ) - return await self._load_joblib_artifact( - model_name=model_name, - artifact_name=target, - operation_type=operation_type, - metadata=metadata, - ) + return model, code_artifact_dir or artifact_path + + if model_type == 'predict': + model, local_dir = await self.load_predict_model(model_name, metadata, flavor) else: - if model_type == 'predict': - model = await self.load_predict_model(model_name, metadata, flavor) + model, local_dir = await self.load_transform_model(model_name, metadata, flavor) - else: - model = await self.load_transform_model(model_name, metadata, flavor) - - return model, artifact_path + return model, code_artifact_dir or local_dir """ Functions related to data format @@ -896,11 +875,18 @@ class MLFlowRepository(SientiaMonitoring): raise ValueError(f'{message}. Unable to parse given date format: {e}') from e elif index_type == datetime or index_type == pd.Timestamp: - index = data.index - if hasattr(index, 'tz') and index.tz is None: - data.index = index.tz_localize('UTC') # type: ignore[attr-defined] + # Timestamp-like elements do not guarantee a DatetimeIndex: an object dtype + # index holds them just as well, and only DatetimeIndex exposes `tz`, + # `tz_localize` and `strftime`. Converting first keeps the branch typed and + # covers the object dtype case instead of failing on the missing attribute. + datetime_index = ( + index if isinstance(index, pd.DatetimeIndex) else pd.DatetimeIndex(index) + ) - data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined] + if datetime_index.tz is None: + datetime_index = datetime_index.tz_localize('UTC') + + data.index = datetime_index.strftime(DATETIME_FORMAT_WITH_TZ) else: raise ValueError(f'{message}. Got {index_type}.') diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 99f6194..38fab72 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -69,11 +69,6 @@ def mlflow_repository(): @pytest.fixture def mlflow(): with patch('laborious.utils.repository.model_repository.mlflow') as mlflow: - # Every URI load resolves the artifact to a local directory before calling the flavor - # loader, and `os.path` refuses a MagicMock. Echoing the URI back keeps the loader's - # argument readable and points the code-directory lookup at a path that does not exist, - # so no test touches `sys.path` unless it builds a real artifact directory. - mlflow.artifacts.download_artifacts.side_effect = lambda artifact_uri: str(artifact_uri) yield mlflow @@ -253,7 +248,7 @@ def test_check_artifact_exists_true(mlflow_repository): ) assert result is True - mlflow_repository.client.list_artifacts.assert_called_once_with('run_id') + mlflow_repository.client.list_artifacts.assert_called_once_with('run_id', None) def test_check_artifact_exists_false(mlflow_repository): @@ -265,7 +260,19 @@ def test_check_artifact_exists_false(mlflow_repository): ) assert result is False - mlflow_repository.client.list_artifacts.assert_called_once_with('run_id') + mlflow_repository.client.list_artifacts.assert_called_once_with('run_id', None) + + +def test_check_artifact_exists_nested_lists_the_parent(mlflow_repository): + artifact = MagicMock(path='data_model/code') + mlflow_repository.client.list_artifacts.return_value = [artifact] + + result = mlflow_repository.check_artifact_exists( + 'run_id', 'data_model/code', metadata['metadata'] + ) + + assert result is True + mlflow_repository.client.list_artifacts.assert_called_once_with('run_id', 'data_model') @pytest.mark.asyncio @@ -440,9 +447,10 @@ def test_get_experiment_create_error(mlflow, mlflow_repository): @pytest.mark.asyncio async def test_load_predict_model_sklearn(mlflow, mlflow_repository): - result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn') + model, local_dir = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn') - assert result == mlflow.sklearn.load_model.return_value + assert model == mlflow.sklearn.load_model.return_value + assert local_dir is None mlflow.sklearn.load_model.assert_called_once_with('models:/test_model/production') mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) mlflow_repository.emit_metric.assert_called_once_with( @@ -452,8 +460,10 @@ async def test_load_predict_model_sklearn(mlflow, mlflow_repository): @pytest.mark.asyncio async def test_load_predict_model_pyfunc(mlflow, mlflow_repository): - result = await mlflow_repository.load_predict_model('test_model', {}, 'pyfunc') - assert result == mlflow.pyfunc.load_model.return_value + model, local_dir = await mlflow_repository.load_predict_model('test_model', {}, 'pyfunc') + + assert model == mlflow.pyfunc.load_model.return_value + assert local_dir is None mlflow.pyfunc.load_model.assert_called_once_with('models:/test_model/production') mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) mlflow_repository.emit_metric.assert_called_once_with( @@ -473,8 +483,6 @@ async def test_load_predict_model_error(mlflow_repository): mlflow_repository.observe_lag.assert_not_called() -# A real string, because the transform load now resolves its URI to a local directory before -# calling the loader, and that resolution goes through `os.path`. TRANSFORM_MODEL_URI = 'runs:/run-123/data_model' @@ -492,11 +500,12 @@ async def test_load_transform_model_sklearn(mlflow, mlflow_repository): mlflow_repository.get_model_run_id = MagicMock() mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) - result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') + model, local_dir = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') - assert result == mlflow.sklearn.load_model.return_value + assert model == mlflow.sklearn.load_model.return_value + assert local_dir is None mlflow.sklearn.load_model.assert_called_once_with(TRANSFORM_MODEL_URI) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) @@ -510,11 +519,12 @@ async def test_load_transform_model_pyfunc(mlflow, mlflow_repository): mlflow_repository.get_model_run_id = MagicMock() mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) - result = await mlflow_repository.load_transform_model('test_model', {}, 'pyfunc') + model, local_dir = await mlflow_repository.load_transform_model('test_model', {}, 'pyfunc') validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') - assert result == mlflow.pyfunc.load_model.return_value + assert model == mlflow.pyfunc.load_model.return_value + assert local_dir is None mlflow.pyfunc.load_model.assert_called_once_with(TRANSFORM_MODEL_URI) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) @@ -555,6 +565,7 @@ async def test_download_model_invalid_model_type(mlflow_repository): 'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')] ) async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type, tmp_path): + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) result = await mlflow_repository.download_model('test_model', {}, model_type[0], 'pyfunc', True) @@ -573,7 +584,8 @@ async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type @pytest.mark.asyncio async def test_download_model_predict(mlflow_repository): - mlflow_repository.load_predict_model = AsyncMock() + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) + mlflow_repository.load_predict_model = AsyncMock(return_value=('model', 'local_dir')) mlflow_repository.load_transform_model = AsyncMock() result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', False) @@ -581,20 +593,59 @@ async def test_download_model_predict(mlflow_repository): mlflow_repository.load_predict_model.assert_called_once_with('test_model', {}, 'pyfunc') mlflow_repository.load_transform_model.assert_not_called() - assert result == (mlflow_repository.load_predict_model.return_value, None) + assert result == ('model', 'local_dir') @pytest.mark.asyncio async def test_download_model_transform(mlflow_repository): + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.load_predict_model = AsyncMock() - mlflow_repository.load_transform_model = AsyncMock() + mlflow_repository.load_transform_model = AsyncMock(return_value=('model', 'local_dir')) result = await mlflow_repository.download_model('test_model', {}, 'transform', 'pyfunc', False) mlflow_repository.load_predict_model.assert_not_called() mlflow_repository.load_transform_model.assert_called_once_with('test_model', {}, 'pyfunc') - assert result == (mlflow_repository.load_transform_model.return_value, None) + assert result == ('model', 'local_dir') + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'model_type,artifact_name', [('predict', 'prediction_model'), ('transform', 'data_model')] +) +async def test_download_model_prepares_the_code_dir_first( + mlflow_repository, model_type, artifact_name +): + """The `code` step runs before any loader, for the artifact matching the model type.""" + calls = [] + mlflow_repository.download_artifact_code_dir = AsyncMock( + side_effect=lambda *_a, **_k: calls.append('code') or 'code-artifact-dir' + ) + loader = AsyncMock(side_effect=lambda *_a: calls.append('load') or ('model', None)) + mlflow_repository.load_predict_model = loader + mlflow_repository.load_transform_model = loader + + result = await mlflow_repository.download_model('test_model', {}, model_type, 'sklearn', False) + + mlflow_repository.download_artifact_code_dir.assert_called_once_with( + 'test_model', artifact_name, 'sklearn', {} + ) + assert calls == ['code', 'load'] + assert result == ('model', 'code-artifact-dir') + + +@pytest.mark.asyncio +async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlflow_repository): + """With a `code` directory downloaded, the wrapper branch returns it over its own download.""" + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value='code-artifact-dir') + mlflow_repository.dowload_artifacts = AsyncMock(return_value='wrapper-dir') + + _model, artifact_dir = await mlflow_repository.download_model( + 'test_model', {}, 'predict', 'pyfunc', True + ) + + assert artifact_dir == 'code-artifact-dir' def test_detect_and_parse_datetime_index_empty(mlflow_repository): @@ -1803,9 +1854,10 @@ async def test_load_predict_model_joblib_flavor_loads_pickle_directly( tmp_path, 'MLmodel', 'model.pkl' ) - result = await mlflow_repository.load_predict_model('test_model', {}, 'joblib') + model, local_dir = await mlflow_repository.load_predict_model('test_model', {}, 'joblib') - assert result is joblib_mock.load.return_value + assert model is joblib_mock.load.return_value + assert local_dir == str(tmp_path) mlflow_repository.get_model_run_id.assert_called_once_with( model_name='test_model', stage='Production' ) @@ -1827,9 +1879,10 @@ async def test_load_transform_model_joblib_flavor_loads_pickle_directly( tmp_path, 'MLmodel', 'model.pkl' ) - result = await mlflow_repository.load_transform_model('test_model', {}, 'joblib') + model, local_dir = await mlflow_repository.load_transform_model('test_model', {}, 'joblib') - assert result is joblib_mock.load.return_value + assert model is joblib_mock.load.return_value + assert local_dir == str(tmp_path) 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')) mlflow_repository.get_model_uri.assert_not_called() @@ -1903,6 +1956,7 @@ async def test_load_transform_model_flavor_failure_propagates( @patch(JOBLIB_PATH) @pytest.mark.asyncio async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, mlflow_repository): + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.dowload_artifacts = AsyncMock() loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom') mlflow.pyfunc.load_model.side_effect = loader_error @@ -2171,60 +2225,110 @@ def test_prepend_artifact_code_dir_does_not_duplicate( assert sys.path.count(code_dir) == 1 +def stub_remote_code_dir(mlflow_repository, tmp_path, artifact_name, exists=True): + """Point the client at a run whose artifact does — or does not — carry `code`. + + Returns the local `code` directory the download would produce, already on disk. + """ + artifact_dir = tmp_path / artifact_name + artifact_dir.mkdir() + code_dir = build_code_dir(artifact_dir) + + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.client.list_artifacts.return_value = ( + [MagicMock(path=f'{artifact_name}/code')] if exists else [MagicMock(path=artifact_name)] + ) + mlflow_repository.client.download_artifacts.return_value = code_dir + + return code_dir + + +@pytest.mark.asyncio +async def test_download_artifact_code_dir_downloads_only_code( + mlflow_repository, tmp_path, preserved_sys_path +): + """Only `/code` travels — never the artifact's weights.""" + code_dir = stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model') + + result = await mlflow_repository.download_artifact_code_dir( + 'test_model', 'prediction_model', 'sklearn', {} + ) + + mlflow_repository.client.list_artifacts.assert_called_once_with('run-123', 'prediction_model') + mlflow_repository.client.download_artifacts.assert_called_once_with( + 'run-123', 'prediction_model/code' + ) + assert result == str(tmp_path / 'prediction_model') + assert sys.path[0] == code_dir + + +@pytest.mark.asyncio +async def test_download_artifact_code_dir_without_remote_code_downloads_nothing( + mlflow_repository, tmp_path, preserved_sys_path +): + stub_remote_code_dir(mlflow_repository, tmp_path, 'data_model', exists=False) + before = list(sys.path) + + result = await mlflow_repository.download_artifact_code_dir( + 'test_model', 'data_model', 'pyfunc', {} + ) + + assert result is None + mlflow_repository.client.download_artifacts.assert_not_called() + assert sys.path == before + + @pytest.mark.asyncio @pytest.mark.parametrize('flavor', FLAVOR_LOADERS) -async def test_load_predict_model_puts_code_dir_on_sys_path_before_the_loader( +async def test_download_model_puts_code_dir_on_sys_path_before_the_loader( flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path ): """MLflow extends `sys.path` only when `MLmodel` declares `code`; the directory decides here.""" - code_dir = build_code_dir(tmp_path) - local_dir = build_artifact_dir(tmp_path, 'MLmodel', 'model.pkl') - mlflow.artifacts.download_artifacts.side_effect = None - mlflow.artifacts.download_artifacts.return_value = local_dir - sys_path_at_load = [] - getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load) - - await mlflow_repository.load_predict_model('test_model', {}, flavor) - - mlflow.artifacts.download_artifacts.assert_called_once_with( - artifact_uri='models:/test_model/production' - ) - getattr(mlflow, flavor).load_model.assert_called_once_with(local_dir) - assert sys_path_at_load[0] == code_dir - - -@pytest.mark.asyncio -@pytest.mark.parametrize('flavor', FLAVOR_LOADERS) -async def test_load_transform_model_puts_code_dir_on_sys_path_before_the_loader( - flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path -): - code_dir = build_code_dir(tmp_path) - local_dir = build_artifact_dir(tmp_path, 'MLmodel', 'model.pkl') - mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + code_dir = stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model') mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) - mlflow.artifacts.download_artifacts.side_effect = None - mlflow.artifacts.download_artifacts.return_value = local_dir sys_path_at_load = [] getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load) - await mlflow_repository.load_transform_model('test_model', {}, flavor) + _model, artifact_dir = await mlflow_repository.download_model( + 'test_model', {}, 'predict', flavor, False + ) - mlflow.artifacts.download_artifacts.assert_called_once_with(artifact_uri=TRANSFORM_MODEL_URI) - getattr(mlflow, flavor).load_model.assert_called_once_with(local_dir) assert sys_path_at_load[0] == code_dir + assert artifact_dir == str(tmp_path / 'prediction_model') @pytest.mark.asyncio @pytest.mark.parametrize('flavor', FLAVOR_LOADERS) -async def test_load_predict_model_without_code_dir_leaves_sys_path_alone( - flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path +async def test_load_predict_model_loads_from_the_uri_without_downloading( + flavor, mlflow, mlflow_repository, preserved_sys_path ): - mlflow.artifacts.download_artifacts.side_effect = None - mlflow.artifacts.download_artifacts.return_value = build_artifact_dir(tmp_path, 'MLmodel') + """The flavor loaders fetch what they need themselves; nothing is pre-downloaded.""" before = list(sys.path) - await mlflow_repository.load_predict_model('test_model', {}, flavor) + _model, local_dir = await mlflow_repository.load_predict_model('test_model', {}, flavor) + mlflow.artifacts.download_artifacts.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + getattr(mlflow, flavor).load_model.assert_called_once_with('models:/test_model/production') + assert local_dir is None + assert sys.path == before + + +@pytest.mark.asyncio +@pytest.mark.parametrize('flavor', FLAVOR_LOADERS) +async def test_load_transform_model_loads_from_the_uri_without_downloading( + flavor, mlflow, mlflow_repository, preserved_sys_path +): + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) + before = list(sys.path) + + _model, local_dir = await mlflow_repository.load_transform_model('test_model', {}, flavor) + + mlflow.artifacts.download_artifacts.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + getattr(mlflow, flavor).load_model.assert_called_once_with(TRANSFORM_MODEL_URI) + assert local_dir is None assert sys.path == before @@ -2232,8 +2336,8 @@ async def test_load_predict_model_without_code_dir_leaves_sys_path_alone( async def test_download_model_load_wrapper_puts_code_dir_on_sys_path( mlflow, mlflow_repository, tmp_path, preserved_sys_path ): - """The wrapper branch downloads itself, so it is the one that has to add the directory.""" - code_dir = build_code_dir(tmp_path) + """The wrapper branch loads a downloaded artifact, so `code` has to be in place first too.""" + code_dir = stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model') mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) sys_path_at_load = [] mlflow.pyfunc.load_model.side_effect = observe_sys_path(sys_path_at_load) @@ -2323,18 +2427,56 @@ async def test_download_model_joblib_returns_download_dir( @patch(JOBLIB_PATH) @pytest.mark.asyncio -async def test_download_model_sklearn_returns_no_artifact_path( - joblib_mock, mlflow, mlflow_repository +async def test_download_model_sklearn_returns_the_code_artifact_dir( + joblib_mock, mlflow, mlflow_repository, tmp_path, preserved_sys_path ): + """The sklearn loader downloads nothing, so the `code` step is the only source of a directory.""" + stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model') + model, artifact_path = await mlflow_repository.download_model( 'test_model', {}, 'predict', 'sklearn' ) assert model is mlflow.sklearn.load_model.return_value - assert artifact_path is None + assert artifact_path == str(tmp_path / 'prediction_model') joblib_mock.load.assert_not_called() +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_download_model_sklearn_without_code_returns_no_dir( + joblib_mock, mlflow, mlflow_repository, tmp_path, preserved_sys_path +): + stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model', exists=False) + + _model, artifact_path = await mlflow_repository.download_model( + 'test_model', {}, 'predict', 'sklearn' + ) + + assert artifact_path is None + mlflow_repository.client.download_artifacts.assert_not_called() + joblib_mock.load.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('model_type', 'loader'), + [('predict', 'load_predict_model'), ('transform', 'load_transform_model')], +) +async def test_download_model_delegates_every_flavor_to_the_loaders( + mlflow_repository, model_type, loader +): + """`download_model` owns no flavor branch: joblib delegates like the other two.""" + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) + delegate = AsyncMock(return_value=('model', 'local-dir')) + setattr(mlflow_repository, loader, delegate) + + result = await mlflow_repository.download_model('test_model', {}, model_type, 'joblib') + + delegate.assert_called_once_with('test_model', {}, 'joblib') + assert result == ('model', 'local-dir') + + @patch(JOBLIB_PATH) @pytest.mark.asyncio async def test_download_model_transform_joblib_downloads_data_model( @@ -2401,7 +2543,6 @@ def test_module_defines_no_joblib_fallback_routine(): '_find_model_pickle', '_load_model_pickle', '_download_and_load_pickle', - '_load_joblib_artifact', } for handler in ast.walk(ast.parse(source)): @@ -2557,7 +2698,7 @@ async def test_joblib_flavor_real_round_trip_with_code_dir( mlflow_repository.client.download_artifacts.return_value = str(tmp_path) try: - loaded = await mlflow_repository.load_predict_model('test_model', {}, 'joblib') + loaded, _local_dir = await mlflow_repository.load_predict_model('test_model', {}, 'joblib') assert loaded.predict([1.0, 2.0]) == [3.0, 4.0] assert code_dir in sys.path