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: Implement explicit joblib model flavor and enhance model import workflow with extensive E2E tests.

This commit introduces a first-class joblib model flavor for direct loading and saving of Python pickle artifacts, including custom code within the artifact's code/ directory, bypassing MLflow's flavor loaders. The model import workflow's status database configuration is now streamlined to default to the POSTGRES_* server, with explicit IMPORT_STATUS_DB_* overrides.

Comprehensive end-to-end tests have been added for the model import workflow, covering success paths, gate rejections, digest enforcement, provisioning failures, status record integrity, cleanup, and retry/recovery scenarios using a real-clock environment. New E2
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-26 15:45:33 -03:00
parent dac4f1718f
commit e95e65b72b
32 changed files with 4247 additions and 1040 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -16,11 +16,11 @@ Capabilities:
import ctypes
import gc
import inspect
import sys
import tempfile
import threading
import time
import traceback
from collections.abc import Callable
from datetime import datetime, timedelta
from io import StringIO
from os import environ, listdir, makedirs, path
@@ -45,7 +45,7 @@ ARTIFACTS_PATH = './tmp/artifacts'
TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl'
PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl'
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch' or 'joblib'."
def force_memory_release(logger: Logger):
@@ -354,23 +354,19 @@ class MLFlowRepository(SientiaMonitoring):
self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata)
return dataframe
def _find_model_pickle(
self, local_dir: str, artifact_name: str, flavor_error: Exception
) -> str:
def _find_model_pickle(self, local_dir: str, artifact_name: str) -> str:
"""
Locate the pickle file of a downloaded model artifact.
Fallback path only: this is reached after the MLflow flavor loader has
already failed. 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. 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.
Args:
local_dir (str): Local directory holding the downloaded artifact.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader,
reported in the message so the original cause is not lost.
Returns:
str: Path to the first top-level `.pkl` file, in name order.
@@ -388,7 +384,6 @@ class MLFlowRepository(SientiaMonitoring):
if not pkl_files:
raise ValueError(
f'No .pkl file found in {artifact_name} artifact. '
f'Original error: {flavor_error}. '
f'Top-level entries in {local_dir}: {entries}'
)
@@ -398,24 +393,28 @@ class MLFlowRepository(SientiaMonitoring):
self,
local_dir: str,
artifact_name: str,
flavor_error: Exception,
metadata: dict[str, Any] | None = None,
) -> Any:
"""
Deserialize a downloaded model artifact with `joblib`.
Fallback path only. `joblib.load` reads the compressed `.pkl` produced
by `joblib.dump`, which is what the MLflow flavor loader fails to read.
The deserialized object is returned as is: no `._model_impl.python_model`
unwrap is applied, because the pickled object already is the model.
Primary path of the `joblib` flavor. `joblib.load` reads the compressed
`.pkl` produced by `joblib.dump`, which is what the MLflow flavor
loaders do not read. The deserialized object is returned as is: no
`._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.
Args:
local_dir (str): Local directory holding the downloaded artifact.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader.
metadata (dict[str, Any] | None): Metadata used for structured
logging, so the chosen pickle path correlates with the other
fallback log lines.
joblib log lines.
Returns:
Any: The object deserialized from the artifact's pickle file.
@@ -424,9 +423,14 @@ class MLFlowRepository(SientiaMonitoring):
ValueError: If the artifact directory has no top-level `.pkl` file.
Exception: Any error raised by `joblib.load` itself.
"""
pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error)
pkl_path = self._find_model_pickle(local_dir, artifact_name)
self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}', metadata)
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.debug(f'joblib flavor: loading {artifact_name} from {pkl_path}', metadata)
return joblib.load(pkl_path)
@@ -434,37 +438,36 @@ class MLFlowRepository(SientiaMonitoring):
self,
model_name: str,
artifact_name: str,
flavor_error: Exception,
metadata: dict[str, Any],
run_id: str | None = None,
) -> Any:
) -> tuple[Any, str]:
"""
Download a model artifact and deserialize it with `joblib`.
Fallback path only. The download runs without `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.
Primary path of the `joblib` flavor. The download runs without
`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.
Callers that already resolved a run to build the flavor URI must pass it
as `run_id`: resolving the Production run again here would open a window
where a promotion between the two lookups makes the fallback serve a
different version than the one the flavor attempt tried and logged. Only
callers whose flavor attempt used a symbolic URI — `models:/<name>/production`,
resolved by MLflow itself — leave it as `None`.
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 —
`load_predict_model`, whose flavor URI is symbolic — leave it as `None`.
The local directory is returned alongside the model because the retrain
path needs it to re-log the artifact's `code` directory.
Args:
model_name (str): Registered model name.
artifact_name (str): Artifact to download, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader.
metadata (dict[str, Any]): Metadata used for structured logging.
run_id (str | None): Run already resolved by the caller. When `None`,
the Production run is resolved here.
Returns:
Any: The object deserialized from the artifact's pickle file.
tuple[Any, str]: The deserialized object and the local directory the
artifact was downloaded to.
Raises:
ValueError: If the downloaded artifact has no top-level `.pkl` file.
@@ -475,86 +478,68 @@ class MLFlowRepository(SientiaMonitoring):
run_id = self.get_model_run_id(model_name=model_name, stage='Production')
self.info(
f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}',
f'joblib flavor: downloading {artifact_name} of {model_name} from run {run_id}',
metadata,
)
local_dir = self.client.download_artifacts(run_id, artifact_name)
return self._load_model_pickle(local_dir, artifact_name, flavor_error, metadata)
model = self._load_model_pickle(local_dir, artifact_name, metadata)
async def _load_with_joblib_fallback(
return model, local_dir
async def _load_joblib_artifact(
self,
primary_loader: Callable[[], Any],
fallback_loader: Callable[[Exception], Any],
failure_description: str,
model_name: str,
artifact_name: str,
operation_type: str,
metadata: dict[str, Any],
fallback_labels: dict[str, Any],
core_labels: dict[str, Any] | None = None,
) -> Any:
run_id: str | None = None,
) -> tuple[Any, str]:
"""
Run a model load with the joblib fallback and its observability.
Load a model artifact through the `joblib` flavor, with observability.
Single owner of the fallback flow shared by every load path: warn on the
flavor failure without emitting an error metric, try the fallback, and
emit lag/count under the fallback labels — or the error count chained to
the flavor error when the fallback also fails. Fixes to the flow belong
here, not in the call sites.
Both loaders are called with no observability concern of their own:
`fallback_loader` receives the flavor error and may be sync or async,
which is what lets the download-based paths and the already-downloaded
wrapper path share this method.
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:
primary_loader (Callable[[], Any]): Flavor load attempt. Any
exception it raises triggers the fallback.
fallback_loader (Callable[[Exception], Any]): joblib fallback,
receiving the flavor error. Its result is awaited when
awaitable.
failure_description (str): What failed, e.g. 'Flavor loader sklearn
failed for prediction model x'. The fallback and original error
are appended to it.
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.
fallback_labels (dict[str, Any]): Metric labels for the fallback
path, already carrying its `operation_type`.
core_labels (dict[str, Any] | None): Metric labels for a successful
flavor load. `None` emits no metric on that path.
run_id (str | None): Run already resolved by the caller. When `None`,
the Production run is resolved during the download.
Returns:
Any: The model loaded by the flavor loader or by the fallback.
tuple[Any, str]: The deserialized model and the local directory the
artifact was downloaded to.
Raises:
Exception: The fallback error, chained from the flavor error, when
both attempts fail.
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 = primary_loader()
except Exception as flavor_error:
self.warning(
f'{failure_description}, trying joblib fallback. Original error: {flavor_error}',
metadata,
model, local_dir = await self._download_and_load_pickle(
model_name=model_name,
artifact_name=artifact_name,
metadata=metadata,
run_id=run_id,
)
try:
model = fallback_loader(flavor_error)
if inspect.isawaitable(model):
model = await model
except Exception as fallback_error:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
)
raise fallback_error from flavor_error
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, fallback_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
return model
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
if core_labels is not None:
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
return model, local_dir
async def load_predict_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -565,7 +550,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
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', 'pytorch')
flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
@@ -574,37 +559,42 @@ 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.
- 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.
"""
if flavor not in ('pyfunc', 'sklearn', 'pytorch'):
if flavor not in ('pyfunc', 'sklearn', 'pytorch', '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 = f'models:/{model_name}/production'
self.info(f'Loading prediction model {model_name} from {model_uri}')
core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
fallback_labels = self.get_core_labels(
metadata, operation_type='load_predict_model_joblib_fallback'
)
def load_flavor() -> Any:
start_time = time.time()
try:
if flavor == 'pyfunc':
return mlflow.pyfunc.load_model(model_uri)
if flavor == 'sklearn':
return mlflow.sklearn.load_model(model_uri)
return mlflow.pytorch.load_model(model_uri)
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
return await self._load_with_joblib_fallback(
primary_loader=load_flavor,
fallback_loader=lambda flavor_error: self._download_and_load_pickle(
model_name, 'prediction_model', flavor_error, metadata
),
failure_description=(
f'Flavor loader {flavor} failed for prediction model {model_name}'
),
metadata=metadata,
fallback_labels=fallback_labels,
core_labels=core_labels,
)
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
async def load_transform_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -618,7 +608,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
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', 'pytorch')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
@@ -629,45 +619,48 @@ class MLFlowRepository(SientiaMonitoring):
model cannot be loaded.
Notes:
- The Production run is resolved once and reused by the joblib
fallback, so a promotion between both steps cannot make the
fallback serve a different version than the flavor attempt.
- 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.
"""
if flavor not in ('sklearn', 'pyfunc', 'pytorch'):
if flavor not in ('sklearn', 'pyfunc', 'pytorch', 'joblib'):
raise ValueError(INVALID_FLAVOR_MESSAGE)
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 = self.get_model_uri(latest_production_id, prediction=False)
self.info(f'Loading data model {model_name} from {model_uri}')
core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
fallback_labels = self.get_core_labels(
metadata, operation_type='load_transform_model_joblib_fallback'
)
def load_flavor() -> Any:
start_time = time.time()
try:
if flavor == 'sklearn':
return mlflow.sklearn.load_model(model_uri)
if flavor == 'pyfunc':
return mlflow.pyfunc.load_model(model_uri)
return mlflow.pytorch.load_model(model_uri)
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
return await self._load_with_joblib_fallback(
primary_loader=load_flavor,
fallback_loader=lambda flavor_error: self._download_and_load_pickle(
model_name,
'data_model',
flavor_error,
metadata,
run_id=latest_production_id,
),
failure_description=f'Flavor loader {flavor} failed for data model {model_name}',
metadata=metadata,
fallback_labels=fallback_labels,
core_labels=core_labels,
)
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
async def download_model(
self,
@@ -684,11 +677,17 @@ class MLFlowRepository(SientiaMonitoring):
model_name (str): Name of the model to download
metadata (dict[str, Any]): Metadata used for structured logging.
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib')
load_wrapper (bool): Whether to load wrapper
Returns:
tuple[Any, str | None]: Model object and optional artifact path.
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.
"""
self.info(
@@ -711,27 +710,24 @@ class MLFlowRepository(SientiaMonitoring):
f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
)
fallback_labels = self.get_core_labels(
metadata, operation_type='download_model_wrapper_joblib_fallback'
)
def load_wrapper_flavor() -> Any:
raw_model = mlflow.pyfunc.load_model(artifact_path)
return raw_model._model_impl.python_model
model = await self._load_with_joblib_fallback(
primary_loader=load_wrapper_flavor,
fallback_loader=lambda flavor_error: self._load_model_pickle(
artifact_path, target, flavor_error, metadata
),
failure_description=f'Wrapper loader failed for {model_type} model {model_name}',
metadata=metadata,
fallback_labels=fallback_labels,
)
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:
if model_type == 'predict':
model = await self.load_predict_model(model_name, metadata, flavor)
@@ -1165,6 +1161,30 @@ class MLFlowRepository(SientiaMonitoring):
}
return retrain_data
def resolve_artifact_code_dir(self, artifact_path: str | None) -> str | None:
"""Resolve the top-level `code` directory of a downloaded artifact.
Sibling of `resolve_code_paths`: both look at the same
`<artifact_path>/code` root, but this one returns the directory itself,
which is what `sys.path` and `mlflow.log_artifacts` take, while
`resolve_code_paths` returns the entries inside it, which is what
`mlflow.pyfunc`'s `code_path=` takes.
Args:
artifact_path (str | None): Local directory of the downloaded model
artifact.
Returns:
str | None: `<artifact_path>/code` when it is a directory, `None`
otherwise — including when `artifact_path` itself is `None`.
"""
if not artifact_path:
return None
code_dir = path.join(artifact_path, 'code')
return code_dir if path.isdir(code_dir) else None
def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]:
"""Resolve the code dependencies to forward when re-logging a pyfunc model.
@@ -1187,11 +1207,12 @@ class MLFlowRepository(SientiaMonitoring):
self.warning('No artifact path available, logging model without code paths', metadata)
return []
code_dir = path.join(artifact_path, 'code')
code_dir = self.resolve_artifact_code_dir(artifact_path)
if not path.isdir(code_dir):
if not code_dir:
self.warning(
f'No code directory found at {code_dir}, logging model without code paths',
f'No code directory found at {path.join(artifact_path, "code")}, '
'logging model without code paths',
metadata,
)
return []
@@ -1207,9 +1228,16 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_data (dict): Model holder with keys 'model' and optional 'artifact_path'.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib').
model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'.
metadata (dict): Metadata for structured logging.
Notes:
- `flavor='joblib'` writes a `joblib`-compressed `model.pkl` under
`model_type` and re-logs the source artifact's `code` directory,
so the written model stays deserializable. No `mlflow.*.log_model`
runs and no `MLmodel` manifest is written, which means only this
runtime's `joblib` flavor can read that artifact back.
"""
model = model_data['model']
@@ -1231,6 +1259,19 @@ class MLFlowRepository(SientiaMonitoring):
self.debug('Model uploaded successfully', metadata)
elif flavor == 'pytorch':
mlflow.pytorch.log_model(model, model_type)
elif flavor == 'joblib':
with tempfile.TemporaryDirectory() as tmp_dir:
pkl_path = path.join(tmp_dir, 'model.pkl')
joblib.dump(model, pkl_path, compress=3)
mlflow.log_artifact(pkl_path, artifact_path=model_type)
code_dir = self.resolve_artifact_code_dir(model_data.get('artifact_path'))
if code_dir:
self.debug(f'Re-logging code directory {code_dir}', metadata)
mlflow.log_artifacts(code_dir, artifact_path=f'{model_type}/code')
else:
self.debug(f'No code directory to re-log for {model_type}', metadata)
else:
raise ValueError(INVALID_FLAVOR_MESSAGE)