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: Unify artifact code directory handling for all model flavors and standardize joblib pickle path

Implemented a robust mechanism to ensure an artifact's code/ directory is prepended to sys.path before loading for sklearn, pyfunc, and joblib flavors. sklearn and pyfunc now resolve MLflow URIs to local directories first to facilitate this sys.path modification. Standardized the joblib flavor to explicitly read and write model.pkl within the artifact directory. Updated documentation and added comprehensive tests to reflect these changes.
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-28 15:10:41 -03:00
parent f4c110cd8a
commit c1a44a9d7a
6 changed files with 475 additions and 273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -45,6 +45,10 @@ ARTIFACTS_PATH = './tmp/artifacts'
TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl'
PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl'
# Pickle file name the `joblib` flavor writes and looks for first, inside the model artifact's own
# directory (`<run>/artifacts/prediction_model/`, `<run>/artifacts/data_model/`).
MODEL_PICKLE_NAME = 'model.pkl'
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'joblib'."
@@ -358,21 +362,31 @@ class MLFlowRepository(SientiaMonitoring):
"""
Locate the pickle file of a downloaded model artifact.
Primary path of the `joblib` flavor. The target artifact is always a
single `.pkl` at the top level of the artifact directory, so the listing
does not recurse — nested wrapper artifacts must never be deserialized
in place of the model. No other extension is considered and `MLmodel` is
not parsed.
Primary path of the `joblib` flavor. `local_dir` is the directory of the
model artifact itself — `<run>/artifacts/prediction_model` or
`<run>/artifacts/data_model`, the very artifacts the other flavors
resolve — never the run's `artifacts/` root. So the pickle sought is
`<local_dir>/model.pkl`, the name `log_model` writes.
`MODEL_PICKLE_NAME` wins whenever it is present, so write and read agree
by construction. Any other `.pkl` directly inside the artifact is taken
as a fallback, in name order, for artifacts built by hand. The listing
does not recurse: a pyfunc wrapper nested inside the artifact
(`prediction_model/artifacts/stacking_model.pkl`) must never be
deserialized in place of the model. No other extension is considered and
`MLmodel` is not parsed.
Args:
local_dir (str): Local directory holding the downloaded artifact.
local_dir (str): Local directory of the downloaded model artifact.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
Returns:
str: Path to the first top-level `.pkl` file, in name order.
str: Path to `<local_dir>/model.pkl` when it exists, otherwise to
the first other `.pkl` directly inside `local_dir`, in name
order.
Raises:
ValueError: If the top level of `local_dir` has no `.pkl` file.
ValueError: If `local_dir` holds no `.pkl` file directly.
"""
entries = sorted(listdir(local_dir))
pkl_files = [
@@ -383,11 +397,97 @@ class MLFlowRepository(SientiaMonitoring):
if not pkl_files:
raise ValueError(
f'No .pkl file found in {artifact_name} artifact. '
f'Top-level entries in {local_dir}: {entries}'
f'No .pkl file found in {artifact_name} artifact. Entries in {local_dir}: {entries}'
)
return path.join(local_dir, pkl_files[0])
chosen = MODEL_PICKLE_NAME if MODEL_PICKLE_NAME in pkl_files else pkl_files[0]
return path.join(local_dir, chosen)
def prepend_artifact_code_dir(
self,
local_dir: str | None,
flavor: str,
metadata: dict[str, Any] | None = None,
) -> str | None:
"""
Put a downloaded artifact's own `code` directory on `sys.path`.
Every flavor goes through here, not only `joblib`. MLflow's loaders
extend `sys.path` from the **manifest**: `_add_code_from_conf_to_system_path`
acts only when the `MLmodel` flavor config declares a `code` entry. An
artifact that carries `code/` on disk without that entry — anything
whose code was logged with `log_artifacts`, anything re-logged by a
path that dropped the config, anything a `joblib` retrain wrote — loads
with no module path and raises `ModuleNotFoundError` at unpickle time,
with the artifact intact on disk the whole while. Reading the directory
itself depends on no manifest, so it works for all three flavors.
Prepending before the loader runs is deliberate: the loader is what
unpickles, so the path has to be in place first. When the manifest
*does* declare `code`, MLflow inserts the same directory again and the
duplicate entry is harmless.
The entry is never removed. A loaded model may import lazily, long after
the load returned.
Args:
local_dir (str | None): Local directory of the downloaded model
artifact, e.g. the download of `<run>/artifacts/prediction_model`.
flavor (str): Flavor doing the load, for the log line only.
metadata (dict[str, Any] | None): Metadata used for structured
logging.
Returns:
str | None: The directory added, `None` when the artifact carries no
`code` directory or it was already on `sys.path`.
"""
code_dir = self.resolve_artifact_code_dir(local_dir)
if not code_dir:
return None
if code_dir in sys.path:
self.debug(f'{flavor} flavor: {code_dir} already on sys.path', metadata)
return None
sys.path.insert(0, code_dir)
self.debug(f'{flavor} flavor: added {code_dir} to sys.path', metadata)
return code_dir
def _localize_model_uri(
self, model_uri: str, flavor: str, metadata: dict[str, Any] | None = None
) -> str:
"""
Resolve a model URI to a local directory and put its `code` 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.
Args:
model_uri (str): `models:/<name>/production` or a run artifact URI.
flavor (str): Flavor doing the load, for the log line 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.
Raises:
Exception: Any error raised while downloading the artifact.
"""
local_path = mlflow.artifacts.download_artifacts(artifact_uri=model_uri)
self.debug(f'{flavor} flavor: {model_uri} downloaded to {local_path}', metadata)
self.prepend_artifact_code_dir(local_path, flavor, metadata)
return local_path
def _load_model_pickle(
self,
@@ -404,13 +504,14 @@ class MLFlowRepository(SientiaMonitoring):
`._model_impl.python_model` unwrap is applied, because the pickled
object already is the model.
The artifact's top-level `code` directory is prepended to `sys.path`
before deserializing, when it exists, so that classes whose modules ship
inside the artifact resolve at unpickle time. The entry is not removed
afterwards: the deserialized object may import lazily.
The artifact's own `code` directory — `<local_dir>/code` — is prepended
to `sys.path` before deserializing, when it exists, so that classes
whose modules ship inside the artifact resolve at unpickle time. That
step is shared with every other flavor: see `prepend_artifact_code_dir`.
Args:
local_dir (str): Local directory holding the downloaded artifact.
local_dir (str): Local directory of the downloaded model artifact,
e.g. the download of `<run>/artifacts/prediction_model`.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
metadata (dict[str, Any] | None): Metadata used for structured
logging, so the chosen pickle path correlates with the other
@@ -420,15 +521,12 @@ class MLFlowRepository(SientiaMonitoring):
Any: The object deserialized from the artifact's pickle file.
Raises:
ValueError: If the artifact directory has no top-level `.pkl` file.
ValueError: If the artifact directory holds no `.pkl` file directly.
Exception: Any error raised by `joblib.load` itself.
"""
pkl_path = self._find_model_pickle(local_dir, artifact_name)
code_dir = self.resolve_artifact_code_dir(local_dir)
if code_dir and code_dir not in sys.path:
sys.path.insert(0, code_dir)
self.debug(f'joblib flavor: added {code_dir} to sys.path', metadata)
self.prepend_artifact_code_dir(local_dir, 'joblib', metadata)
self.debug(f'joblib flavor: loading {artifact_name} from {pkl_path}', metadata)
@@ -448,8 +546,10 @@ class MLFlowRepository(SientiaMonitoring):
`dst_path`, so MLflow writes to a per-call temporary directory: the
shared `ARTIFACTS_PATH` tree is never written to and
`dowload_artifacts` — which removes its destination before
downloading — is deliberately not reused. The target artifact is always
a `.pkl` and the deserialized object is returned as is.
downloading — is deliberately not reused. `artifact_name` is the same
model artifact the other flavors resolve, so the download yields
`<run>/artifacts/<artifact_name>/` and the pickle read from it is that
directory's `model.pkl`. The deserialized object is returned as is.
Callers that already resolved a run must pass it as `run_id`, so that
the run is resolved exactly once per load. Callers with no run at hand —
@@ -470,7 +570,7 @@ class MLFlowRepository(SientiaMonitoring):
artifact was downloaded to.
Raises:
ValueError: If the downloaded artifact has no top-level `.pkl` file.
ValueError: If the downloaded artifact holds no `.pkl` file directly.
Exception: Any error raised while resolving the Production run,
downloading the artifact, or deserializing the pickle.
"""
@@ -559,10 +659,13 @@ class MLFlowRepository(SientiaMonitoring):
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.
- 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. The artifact's
`code` directory, when present, is put on `sys.path` first.
`joblib.load`, which needs no `MLmodel` manifest.
"""
if flavor not in ('pyfunc', 'sklearn', 'joblib'):
raise ValueError(INVALID_FLAVOR_MESSAGE)
@@ -582,10 +685,12 @@ class MLFlowRepository(SientiaMonitoring):
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(model_uri)
model = mlflow.pyfunc.load_model(local_path)
else:
model = mlflow.sklearn.load_model(model_uri)
model = mlflow.sklearn.load_model(local_path)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
@@ -617,11 +722,14 @@ class MLFlowRepository(SientiaMonitoring):
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.
- 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. The artifact's `code`
directory, when present, is put on `sys.path` first.
which needs no `MLmodel` manifest.
"""
if flavor not in ('sklearn', 'pyfunc', 'joblib'):
@@ -646,10 +754,12 @@ class MLFlowRepository(SientiaMonitoring):
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(model_uri)
model = mlflow.sklearn.load_model(local_path)
else:
model = mlflow.pyfunc.load_model(model_uri)
model = mlflow.pyfunc.load_model(local_path)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
@@ -684,6 +794,10 @@ class MLFlowRepository(SientiaMonitoring):
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`.
"""
self.info(
@@ -706,6 +820,8 @@ class MLFlowRepository(SientiaMonitoring):
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
@@ -1255,7 +1371,7 @@ class MLFlowRepository(SientiaMonitoring):
self.debug('Model uploaded successfully', metadata)
elif flavor == 'joblib':
with tempfile.TemporaryDirectory() as tmp_dir:
pkl_path = path.join(tmp_dir, 'model.pkl')
pkl_path = path.join(tmp_dir, MODEL_PICKLE_NAME)
joblib.dump(model, pkl_path, compress=3)
mlflow.log_artifact(pkl_path, artifact_path=model_type)