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

302
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'
self.info(f'Loading prediction model {model_name} from {model_uri}')
if model_uri is not None:
self.info(f'Loading prediction model {model_name} from {model_uri}')
core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
start_time = time.time()
try:
local_path = self._localize_model_uri(model_uri, flavor, metadata)
if flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(local_path)
if flavor == 'joblib':
model, local_dir = await self._download_and_load_pickle(
model_name=model_name,
artifact_name='prediction_model',
metadata=metadata,
)
elif flavor == 'pyfunc':
model, local_dir = mlflow.pyfunc.load_model(model_uri), None
else:
model = mlflow.sklearn.load_model(local_path)
model, local_dir = mlflow.sklearn.load_model(model_uri), None
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
return model, local_dir
async def load_transform_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
) -> Any:
) -> tuple[Any, str | None]:
"""
Load the latest Production version of a transformation model.
@@ -712,24 +690,27 @@ class MLFlowRepository(SientiaMonitoring):
model_name (str): The name of the model to download.
metadata (dict[str, Any]): Metadata used for structured logging.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
Any: The loaded model object, depending on the flavor used.
tuple[Any, str | None]: The loaded model object, depending on the
flavor used, and the local directory its artifact was downloaded
to, `None` when nothing was downloaded here.
Raises:
Exception: If the model run ID or URI cannot be retrieved, or if the
model cannot be loaded.
Notes:
- The artifact's `code` directory, when present, is put on
`sys.path` before the load, whatever the flavor. `sklearn` and
`pyfunc` are therefore resolved to a local directory first and the
loader is handed that directory instead of the URI.
- The MLflow flavor loaders read the model URI directly and fetch
only what they need. Nothing is pre-downloaded here, so the local
directory returned alongside the model is `None`; putting the
artifact's `code` directory on `sys.path` belongs to
`download_artifact_code_dir`, which `download_model` runs first.
- With `flavor='joblib'` no MLflow flavor loader is called and no
model URI is built: the `data_model` artifact of the resolved
Production run is downloaded and deserialized with `joblib.load`,
which needs no `MLmodel` manifest.
which needs no `MLmodel` manifest. That path does return its local
directory, since it downloads by necessity.
"""
if flavor not in ('sklearn', 'pyfunc', 'joblib'):
@@ -737,36 +718,36 @@ class MLFlowRepository(SientiaMonitoring):
latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
if flavor == 'joblib':
model, _local_dir = await self._load_joblib_artifact(
model_name=model_name,
artifact_name='data_model',
operation_type='load_transform_model',
metadata=metadata,
run_id=latest_production_id,
)
return model
model_uri = (
None
if flavor == 'joblib'
else self.get_model_uri(latest_production_id, prediction=False)
)
model_uri = self.get_model_uri(latest_production_id, prediction=False)
self.info(f'Loading data model {model_name} from {model_uri}')
if model_uri is not None:
self.info(f'Loading data model {model_name} from {model_uri}')
core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
start_time = time.time()
try:
local_path = self._localize_model_uri(model_uri, flavor, metadata)
if flavor == 'sklearn':
model = mlflow.sklearn.load_model(local_path)
if flavor == 'joblib':
model, local_dir = await self._download_and_load_pickle(
model_name=model_name,
artifact_name='data_model',
metadata=metadata,
run_id=latest_production_id,
)
elif flavor == 'sklearn':
model, local_dir = mlflow.sklearn.load_model(model_uri), None
else:
model = mlflow.pyfunc.load_model(local_path)
model, local_dir = mlflow.pyfunc.load_model(model_uri), None
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
return model, local_dir
async def download_model(
self,
@@ -787,17 +768,26 @@ class MLFlowRepository(SientiaMonitoring):
load_wrapper (bool): Whether to load wrapper
Returns:
tuple[Any, str | None]: Model object and optional artifact path.
tuple[Any, str | None]: Model object and the local directory its
artifact was downloaded to.
Notes:
- `flavor='joblib'` returns the local directory the artifact was
downloaded to as the artifact path, because the retrain path needs
it to re-log the artifact's `code` directory. Every other flavor
returns `None` unless the wrapper branch was taken.
- The wrapper branch downloads before it loads, so it puts the
artifact's `code` directory on `sys.path` itself. The two plain
branches delegate, and the code directory is handled inside
`load_predict_model` / `load_transform_model`.
- Single owner of the artifact `code` directory, for every flavor
and both branches: `download_artifact_code_dir` runs first, before
any model is loaded, and is the only place that touches
`sys.path`. It downloads nothing unless the run's artifact really
carries a `code` directory, and even then downloads only that
directory — never the model's weights.
- The directory it downloads into is also what this method returns
alongside the model: it is the artifact directory the retrain path
needs to re-log `code`. It falls back to whatever the loader
downloaded on its own, which is a real artifact directory for
`joblib` and for the wrapper branch, and `None` for the flavor
loaders, which stream from the model URI without downloading.
- Apart from that shared first step this method does not branch on
the flavor. `load_predict_model` and `load_transform_model` own
every flavor, `joblib` included. The only branch left is
`load_wrapper`, which is orthogonal to the flavor.
"""
self.info(
@@ -807,47 +797,36 @@ class MLFlowRepository(SientiaMonitoring):
if model_type not in ['predict', 'transform']:
raise ValueError("Invalid model_type. Use 'predict' or 'transform'.")
artifact_path = None
target = 'prediction_model' if model_type == 'predict' else 'data_model'
code_artifact_dir = await self.download_artifact_code_dir(
model_name, target, flavor, metadata
)
if load_wrapper:
self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}')
target = 'prediction_model' if model_type == 'predict' else 'data_model'
artifact_path = await self.dowload_artifacts(model_name, metadata, target)
self.info(
f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
)
self.prepend_artifact_code_dir(artifact_path, 'pyfunc', metadata)
raw_model = mlflow.pyfunc.load_model(artifact_path)
model = raw_model._model_impl.python_model
self.debug(
f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata
)
elif flavor == 'joblib':
target = 'prediction_model' if model_type == 'predict' else 'data_model'
operation_type = (
'load_predict_model' if model_type == 'predict' else 'load_transform_model'
)
return await self._load_joblib_artifact(
model_name=model_name,
artifact_name=target,
operation_type=operation_type,
metadata=metadata,
)
return model, code_artifact_dir or artifact_path
if model_type == 'predict':
model, local_dir = await self.load_predict_model(model_name, metadata, flavor)
else:
if model_type == 'predict':
model = await self.load_predict_model(model_name, metadata, flavor)
model, local_dir = await self.load_transform_model(model_name, metadata, flavor)
else:
model = await self.load_transform_model(model_name, metadata, flavor)
return model, artifact_path
return model, code_artifact_dir or local_dir
"""
Functions related to data format
@@ -896,11 +875,18 @@ class MLFlowRepository(SientiaMonitoring):
raise ValueError(f'{message}. Unable to parse given date format: {e}') from e
elif index_type == datetime or index_type == pd.Timestamp:
index = data.index
if hasattr(index, 'tz') and index.tz is None:
data.index = index.tz_localize('UTC') # type: ignore[attr-defined]
# Timestamp-like elements do not guarantee a DatetimeIndex: an object dtype
# index holds them just as well, and only DatetimeIndex exposes `tz`,
# `tz_localize` and `strftime`. Converting first keeps the branch typed and
# covers the object dtype case instead of failing on the missing attribute.
datetime_index = (
index if isinstance(index, pd.DatetimeIndex) else pd.DatetimeIndex(index)
)
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined]
if datetime_index.tz is None:
datetime_index = datetime_index.tz_localize('UTC')
data.index = datetime_index.strftime(DATETIME_FORMAT_WITH_TZ)
else:
raise ValueError(f'{message}. Got {index_type}.')