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

QTZPOC-13

Browse Source 
QTZPOC-13: Centralize MLflow joblib fallback logic, ensuring consistent model version resolution by reusing the initial run ID, improving logging, and updating dependencies.
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-20 10:23:24 -03:00
parent 03ac4b31e5
commit 291d16efa5
4 changed files with 179 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -16,9 +16,11 @@ Capabilities:
import ctypes
import gc
import inspect
import threading
import time
import traceback
from collections.abc import Callable
from datetime import datetime, timedelta
from io import StringIO
from os import environ, listdir, makedirs, path
@@ -393,7 +395,11 @@ class MLFlowRepository(SientiaMonitoring):
return path.join(local_dir, pkl_files[0])
def _load_model_pickle(
self, local_dir: str, artifact_name: str, flavor_error: Exception
self,
local_dir: str,
artifact_name: str,
flavor_error: Exception,
metadata: dict[str, Any] | None = None,
) -> Any:
"""
Deserialize a downloaded model artifact with `joblib`.
@@ -407,6 +413,9 @@ class MLFlowRepository(SientiaMonitoring):
local_dir (str): Local directory holding the downloaded artifact.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader.
metadata (dict[str, Any] | None): Metadata used for structured
logging, so the chosen pickle path correlates with the other
fallback log lines.
Returns:
Any: The object deserialized from the artifact's pickle file.
@@ -417,7 +426,7 @@ class MLFlowRepository(SientiaMonitoring):
"""
pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error)
self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}')
self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}', metadata)
return joblib.load(pkl_path)
@@ -427,6 +436,7 @@ class MLFlowRepository(SientiaMonitoring):
artifact_name: str,
flavor_error: Exception,
metadata: dict[str, Any],
run_id: str | None = None,
) -> Any:
"""
Download a model artifact and deserialize it with `joblib`.
@@ -438,11 +448,20 @@ class MLFlowRepository(SientiaMonitoring):
artifact is always a `.pkl` and the deserialized object is returned as
is.
Callers that already resolved a run to build the flavor URI must pass it
as `run_id`: resolving the Production run again here would open a window
where a promotion between the two lookups makes the fallback serve a
different version than the one the flavor attempt tried and logged. Only
callers whose flavor attempt used a symbolic URI — `models:/<name>/production`,
resolved by MLflow itself — leave it as `None`.
Args:
model_name (str): Registered model name.
artifact_name (str): Artifact to download, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader.
metadata (dict[str, Any]): Metadata used for structured logging.
run_id (str | None): Run already resolved by the caller. When `None`,
the Production run is resolved here.
Returns:
Any: The object deserialized from the artifact's pickle file.
@@ -452,7 +471,8 @@ class MLFlowRepository(SientiaMonitoring):
Exception: Any error raised while resolving the Production run,
downloading the artifact, or deserializing the pickle.
"""
run_id = self.get_model_run_id(model_name=model_name, stage='Production')
if run_id is None:
run_id = self.get_model_run_id(model_name=model_name, stage='Production')
self.info(
f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}',
@@ -461,7 +481,80 @@ class MLFlowRepository(SientiaMonitoring):
local_dir = self.client.download_artifacts(run_id, artifact_name)
return self._load_model_pickle(local_dir, artifact_name, flavor_error)
return self._load_model_pickle(local_dir, artifact_name, flavor_error, metadata)
async def _load_with_joblib_fallback(
self,
primary_loader: Callable[[], Any],
fallback_loader: Callable[[Exception], Any],
failure_description: str,
metadata: dict[str, Any],
fallback_labels: dict[str, Any],
core_labels: dict[str, Any] | None = None,
) -> Any:
"""
Run a model load with the joblib fallback and its observability.
Single owner of the fallback flow shared by every load path: warn on the
flavor failure without emitting an error metric, try the fallback, and
emit lag/count under the fallback labels — or the error count chained to
the flavor error when the fallback also fails. Fixes to the flow belong
here, not in the call sites.
Both loaders are called with no observability concern of their own:
`fallback_loader` receives the flavor error and may be sync or async,
which is what lets the download-based paths and the already-downloaded
wrapper path share this method.
Args:
primary_loader (Callable[[], Any]): Flavor load attempt. Any
exception it raises triggers the fallback.
fallback_loader (Callable[[Exception], Any]): joblib fallback,
receiving the flavor error. Its result is awaited when
awaitable.
failure_description (str): What failed, e.g. 'Flavor loader sklearn
failed for prediction model x'. The fallback and original error
are appended to it.
metadata (dict[str, Any]): Metadata used for structured logging.
fallback_labels (dict[str, Any]): Metric labels for the fallback
path, already carrying its `operation_type`.
core_labels (dict[str, Any] | None): Metric labels for a successful
flavor load. `None` emits no metric on that path.
Returns:
Any: The model loaded by the flavor loader or by the fallback.
Raises:
Exception: The fallback error, chained from the flavor error, when
both attempts fail.
"""
start_time = time.time()
try:
model = primary_loader()
except Exception as flavor_error:
self.warning(
f'{failure_description}, trying joblib fallback. Original error: {flavor_error}',
metadata,
)
try:
model = fallback_loader(flavor_error)
if inspect.isawaitable(model):
model = await model
except Exception as fallback_error:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
)
raise fallback_error from flavor_error
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
return model
if core_labels is not None:
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
async def load_predict_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -492,37 +585,26 @@ class MLFlowRepository(SientiaMonitoring):
fallback_labels = self.get_core_labels(
metadata, operation_type='load_predict_model_joblib_fallback'
)
start_time = time.time()
try:
def load_flavor() -> Any:
if flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
except Exception as flavor_error:
self.warning(
f'Flavor loader {flavor} failed for prediction model {model_name}, '
f'trying joblib fallback. Original error: {flavor_error}',
metadata,
)
try:
model = await self._download_and_load_pickle(
model_name, 'prediction_model', flavor_error, metadata
)
except Exception as fallback_error:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
)
raise fallback_error from flavor_error
return mlflow.pyfunc.load_model(model_uri)
if flavor == 'sklearn':
return mlflow.sklearn.load_model(model_uri)
return mlflow.pytorch.load_model(model_uri)
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
return model
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
return await self._load_with_joblib_fallback(
primary_loader=load_flavor,
fallback_loader=lambda flavor_error: self._download_and_load_pickle(
model_name, 'prediction_model', flavor_error, metadata
),
failure_description=(
f'Flavor loader {flavor} failed for prediction model {model_name}'
),
metadata=metadata,
fallback_labels=fallback_labels,
core_labels=core_labels,
)
async def load_transform_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -545,6 +627,11 @@ class MLFlowRepository(SientiaMonitoring):
Raises:
Exception: If the model run ID or URI cannot be retrieved, or if the
model cannot be loaded.
Notes:
- The Production run is resolved once and reused by the joblib
fallback, so a promotion between both steps cannot make the
fallback serve a different version than the flavor attempt.
"""
if flavor not in ('sklearn', 'pyfunc', 'pytorch'):
@@ -559,38 +646,28 @@ class MLFlowRepository(SientiaMonitoring):
fallback_labels = self.get_core_labels(
metadata, operation_type='load_transform_model_joblib_fallback'
)
start_time = time.time()
try:
def load_flavor() -> Any:
if flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
except Exception as flavor_error:
self.warning(
f'Flavor loader {flavor} failed for data model {model_name}, '
f'trying joblib fallback. Original error: {flavor_error}',
return mlflow.sklearn.load_model(model_uri)
if flavor == 'pyfunc':
return mlflow.pyfunc.load_model(model_uri)
return mlflow.pytorch.load_model(model_uri)
return await self._load_with_joblib_fallback(
primary_loader=load_flavor,
fallback_loader=lambda flavor_error: self._download_and_load_pickle(
model_name,
'data_model',
flavor_error,
metadata,
)
try:
model = await self._download_and_load_pickle(
model_name, 'data_model', flavor_error, metadata
)
except Exception as fallback_error:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
)
raise fallback_error from flavor_error
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
return model
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
run_id=latest_production_id,
),
failure_description=f'Flavor loader {flavor} failed for data model {model_name}',
metadata=metadata,
fallback_labels=fallback_labels,
core_labels=core_labels,
)
async def download_model(
self,
@@ -637,26 +714,20 @@ class MLFlowRepository(SientiaMonitoring):
fallback_labels = self.get_core_labels(
metadata, operation_type='download_model_wrapper_joblib_fallback'
)
start_time = time.time()
try:
raw_model = mlflow.pyfunc.load_model(artifact_path)
model = raw_model._model_impl.python_model
except Exception as flavor_error:
self.warning(
f'Wrapper loader failed for {model_type} model {model_name}, '
f'trying joblib fallback. Original error: {flavor_error}',
metadata,
)
try:
model = self._load_model_pickle(artifact_path, target, flavor_error)
except Exception as fallback_error:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
)
raise fallback_error from flavor_error
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
def load_wrapper_flavor() -> Any:
raw_model = mlflow.pyfunc.load_model(artifact_path)
return raw_model._model_impl.python_model
model = await self._load_with_joblib_fallback(
primary_loader=load_wrapper_flavor,
fallback_loader=lambda flavor_error: self._load_model_pickle(
artifact_path, target, flavor_error, metadata
),
failure_description=f'Wrapper loader failed for {model_type} model {model_name}',
metadata=metadata,
fallback_labels=fallback_labels,
)
self.debug(
f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata