gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
QTZPOC-20: Prioritize code/ directory download and sys.path injection before model loading, enabling MLflow flavors to stream models directly from URIs and simplifying joblib artifact handling.
This commit is contained in:
Vitor Pimentel dos Santos
2026-09-01 13:27:32 -03:00
parent c1a44a9d7a
commit 30ff378cb7
8 changed files with 464 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
README.md
View File

@@ -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)**.

28
docs/model-flavors.md
View File

@@ -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(<dir>)` |
| `pyfunc` | yes | yes | resolves the URI to a local directory, then `mlflow.pyfunc.load_model(<dir>)`; the wrapper path also unwraps `._model_impl.python_model` |
| `sklearn` | yes | yes | `mlflow.sklearn.load_model(<uri>)` — the loader fetches what it needs |
| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; 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 `<artifact>/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 `<artifact>/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, '<artifact>')`) decides first, so a run without `code`
downloads **nothing**, and a run with it downloads only `<artifact>/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.

25
e2e/conftest.py
View File

@@ -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
# `<artifact>/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)

27
e2e/scenarios.md
View File

@@ -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'`
---

37
e2e/test_predictions_batch_joblib_flavor.py
View File

@@ -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(

4
e2e/test_retrain_joblib_flavor.py
View File

@@ -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)

294
laborious/utils/repository/model_repository.py
View File

@@ -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 `<artifact_name>/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:/<name>/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'
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,
model_uri = (
None
if flavor == 'joblib'
else self.get_model_uri(latest_production_id, prediction=False)
)
return model
model_uri = self.get_model_uri(latest_production_id, prediction=False)
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,
)
else:
return model, code_artifact_dir or artifact_path
if model_type == 'predict':
model = await self.load_predict_model(model_name, metadata, flavor)
model, local_dir = await self.load_predict_model(model_name, metadata, flavor)
else:
model = await self.load_transform_model(model_name, metadata, flavor)
model, local_dir = 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}.')

275
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -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 `<artifact>/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