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

SIENTIAPDE-1231

Browse Source 
Enhance README and repository utilities for clarity and functionality

- Updated README.md to improve descriptions and structure, adding detailed sections for features, workflows, and architecture.
- Enhanced MinioRepository with comprehensive docstrings for methods and class attributes, improving usability and documentation.
- Refined MLFlowRepository with clearer method descriptions and improved logging for better observability and maintainability.
This commit is contained in:
vitor-aignosi
2025-10-16 11:01:36 -03:00
parent de47820c4a
commit 644a43093a
3 changed files with 291 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
laborious/utils/repository/minio_repository.py
View File

@@ -1,3 +1,11 @@
"""
MinIO repository utilities.
This module provides a lightweight repository around a MinIO/S3-compatible
object storage using boto3. It supports creating buckets on demand and
storing/loading pandas DataFrames in Parquet format.
"""
from io import BytesIO
from typing import Any
@@ -10,6 +18,21 @@ from sientia_do.observability.logger import Logger
class MinioRepository:
"""
Repository for interacting with a MinIO (S3-compatible) object storage.
This class encapsulates a reusable `boto3` S3 client and convenience
helpers to persist and retrieve pandas DataFrames as Parquet files.
Attributes:
storage_options (dict): Options compatible with pandas s3fs usage.
minio_bucket (str): Default bucket name used for operations.
minio_endpoint_url (str): MinIO endpoint URL.
minio_region_name (str): MinIO region name.
s3_client (Any): Reusable S3 client from `boto3`.
logger (Logger): Observability logger.
notification_handler (NotificationHandler): Notifications handler.
"""
def __init__(
self,
minio_endpoint_url: str,
@@ -20,6 +43,17 @@ class MinioRepository:
logger: Logger,
notification_handler: NotificationHandler,
):
"""Initialize the repository and S3 client.
Args:
minio_endpoint_url (str): MinIO endpoint URL.
minio_access_key (str): Access key (AK).
minio_secret_key (str): Secret key (SK).
minio_region_name (str): Region name for the client.
minio_default_bucket (str): Default bucket name to operate on.
logger (Logger): Logger instance for structured logs.
notification_handler (NotificationHandler): Notification handler.
"""
# MinIO settings shared with pandas s3fs
self.storage_options = {
'key': minio_access_key,
@@ -54,11 +88,14 @@ class MinioRepository:
self.notification_handler = notification_handler
def close(self):
"""Close the underlying S3 client."""
self.s3_client.close()
def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None:
"""
Ensure the MinIO bucket exists; create it if necessary.
"""Ensure the default bucket exists; create it if missing.
Args:
metadata (dict[str, Any]): Metadata used for structured logging.
"""
try:
@@ -71,6 +108,14 @@ class MinioRepository:
def store_dataframe_as_parquet(
self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any]
):
"""Persist a DataFrame as a Parquet object in the default bucket.
Args:
dataframe (DataFrame): DataFrame to persist.
uri (str): Human-friendly URI used for logging context.
object_name (str): Object key (path/key within the bucket).
metadata (dict[str, Any]): Metadata used for structured logging.
"""
self.ensure_bucket_exists(metadata)
self.logger.custom_info(f'Storing dataframe as parquet in {uri}', metadata)
@@ -83,6 +128,15 @@ class MinioRepository:
self.logger.custom_info(f'Dataframe stored as parquet in {uri}', metadata)
def get_parquet_as_dataframe(self, object_key: str, metadata: dict[str, Any]) -> DataFrame:
"""Load a Parquet object from the default bucket into a DataFrame.
Args:
object_key (str): Object key to retrieve from the bucket.
metadata (dict[str, Any]): Metadata used for structured logging.
Returns:
DataFrame: Loaded DataFrame.
"""
self.logger.custom_info(f'Getting parquet as dataframe from {object_key}', metadata)
response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)

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

@@ -1,13 +1,13 @@
"""
MLFlow Repository
MLflow repository utilities
This module contains the MLFlowRepository class,
which is responsible for handling the communication with MLFlow tracking server.
This module provides the `MLFlowRepository` class and helpers to interact with
an MLflow tracking server and model registry. It covers model discovery,
downloading/loading with multiple flavors, cached operations with retention
policies, transformation/prediction interfaces, retraining workflows, and
production model promotion.
It includes methods for model management, caching, retraining, and serving operations
using MLFlow's tracking and model registry capabilities.
The repository provides comprehensive functionality for:
Capabilities:
- Model loading and caching with retention policies
- Data transformation and prediction operations
- Model retraining workflows
@@ -37,6 +37,15 @@ INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'
def force_memory_release(logger: Logger):
"""Attempt to release memory from the Python process.
Executes a garbage collection cycle and calls `malloc_trim(0)` on glibc
where available to return free memory to the OS. This may be a no-op on
non-glibc systems.
Args:
logger (Logger): Logger for observability.
"""
gc.collect()
try:
@@ -48,6 +57,14 @@ def force_memory_release(logger: Logger):
class MLFlowRepository:
def __init__(self, host: str, username: str, password: str, logger: Logger):
"""Initialize MLflow client and base state.
Args:
host (str): MLflow tracking URI.
username (str): MLflow username.
password (str): MLflow password.
logger (Logger): Logger instance.
"""
# set tracking uri
mlflow.set_tracking_uri(host)
@@ -64,15 +81,15 @@ class MLFlowRepository:
"""
def get_model_uri(self, run_id: str, prediction: bool = True):
"""
Get the model URI based on the run_id.
"""Build the artifact URI for a run's model.
Args:
run_id (str): The run_id of the model.
prediction (bool): Whether to get prediction model URI (default: True)
run_id (str): MLflow run identifier.
prediction (bool): If True, return `prediction_model` URI,
otherwise return `data_model` URI.
Returns:
str: The model URI.
str: Artifact URI to the selected model within the run.
"""
run_info = mlflow.get_run(run_id)
if prediction:
@@ -82,15 +99,14 @@ class MLFlowRepository:
return model_uri
def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str:
"""
Get the run_id of a model based on its name and stage.
"""Resolve the run_id for a registered model at a given stage.
Args:
model_name (str): The name of the model.
stage (str): The stage of the model.
model_name (str): Registered model name.
stage (str): Desired stage (e.g., 'Production').
Returns:
str: The run_id of the model.
str: Run ID for the latest version at the given stage.
"""
# Use search_registered_models instead of deprecated get_latest_versions
registered_models = self.client.search_registered_models(
@@ -120,14 +136,14 @@ class MLFlowRepository:
def get_next_run_name(self, model_name: str) -> str:
"""
Generate the next run name for a specific MLFlow model.
Generate the next run name for a specific MLflow model.
This method calculates the next sequential run number for a model
by searching existing runs and incrementing the count. It ensures
unique run names for model training and retraining operations.
Args:
model_name (str): The name of the MLFlow model
model_name (str): The name of the MLflow model
Returns:
str: The next run name in format 'model_name-run_number'
@@ -140,20 +156,20 @@ class MLFlowRepository:
self, experiment_name: str, create_if_not_exists: bool = False
) -> Experiment:
"""
Retrieve MLFlow experiment ID by experiment name.
Retrieve MLflow experiment by name, optionally creating it.
This method searches for an MLFlow experiment by name and
returns its unique identifier. It provides error handling
for non-existent experiments.
Args:
experiment_name (str): Name of the MLFlow experiment
experiment_name (str): Name of the MLflow experiment
Returns:
int: MLFlow experiment ID
Experiment: MLflow experiment object
Raises:
ValueError: If the experiment name is not found
ValueError: If the experiment name is not found and creation is disabled
"""
experiment = mlflow.get_experiment_by_name(experiment_name)
@@ -166,7 +182,14 @@ class MLFlowRepository:
return experiment
def get_model_params(self, run_id: str):
"""Obtém os parâmetros de uma run"""
"""Fetch parameters associated with a given MLflow run.
Args:
run_id (str): Run identifier to inspect.
Returns:
dict: Mapping of parameter names to values.
"""
run_info = mlflow.get_run(run_id)
return run_info.data.params
@@ -176,14 +199,14 @@ class MLFlowRepository:
def dowload_artifacts(self, model_name: str, artifact_path: str = 'data_model') -> str:
"""
Downloads artifacts from a specific MLFlow run.
Download artifacts from the latest production run of a model.
Args:
model_name (str): Name of the model
artifact_path (str): Path to the artifact within the run
model_name (str): Registered model name.
artifact_path (str): Relative path to artifacts within the run.
Returns:
str: Path to the downloaded artifacts
str: Local filesystem path where artifacts are saved.
"""
run_id = self.get_model_run_id(model_name=model_name, stage='Production')
output_dir = f'{ARTIFACTS_PATH}/{model_name}'
@@ -201,7 +224,7 @@ class MLFlowRepository:
def load_predict_model(self, model_name: str, flavor: str = 'sklearn') -> Any:
"""
Downloads a predictive model from the MLflow Model Registry.
Load a predictive model from the MLflow Model Registry.
Args:
model_name (str): The name of the model to download from the registry.
@@ -212,7 +235,7 @@ class MLFlowRepository:
mlflow.pyfunc.PyFuncModel: The loaded predictive model.
Notes:
- The model is fetched from the "production" stage of the MLflow Model Registry.
- The model is fetched from the "Production" stage of the MLflow Model Registry.
- Warnings during the model loading process are suppressed.
"""
model_uri = f'models:/{model_name}/production'
@@ -230,7 +253,7 @@ class MLFlowRepository:
def load_transform_model(self, model_name: str, flavor: str) -> Any:
"""
Downloads the latest production version of a specified transformation model.
Load the latest Production version of a transformation model.
This method retrieves the latest production model run ID for the given
model name, constructs the model URI, and loads the model using MLflow.
@@ -241,11 +264,11 @@ class MLFlowRepository:
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
Any: The loaded model object, as returned by `mlflow.sklearn.load_model`.
Any: The loaded model object, depending on the flavor used.
Raises:
Exception: If the model run ID or URI cannot be retrieved, or if the
model cannot be loaded.
model cannot be loaded.
"""
latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
@@ -266,7 +289,7 @@ class MLFlowRepository:
self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False
) -> tuple[Any, str | None]:
"""
Download model based on type (predict or transform).
Download model based on type ("predict" or "transform").
Args:
model_name (str): Name of the model to download
@@ -275,7 +298,7 @@ class MLFlowRepository:
load_wrapper (bool): Whether to load wrapper
Returns:
tuple[Any, str]: Model object and artifact path if model is compressed
tuple[Any, str | None]: Model object and optional artifact path.
"""
self.logger.info(
@@ -317,17 +340,19 @@ class MLFlowRepository:
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
"""
Detect and parse datetime index from data. index must be a timestamp like column.
This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
If another type or format, must raise an error.
Normalize DataFrame index to the expected timestamp string format.
The index must be timestamp-like. If the index is:
- string: it must match `DATETIME_FORMAT_WITH_TZ`
- datetime or pandas Timestamp: it will be converted to that format
Any other type raises a ValueError.
Args:
data (pd.DataFrame): DataFrame with timestamp index
metadata (dict): Metadata for logging
data (pd.DataFrame): DataFrame with timestamp index.
metadata (dict): Metadata for structured logging.
Returns:
pd.DataFrame: DataFrame with converted datetime index
pd.DataFrame: DataFrame with converted datetime index.
"""
index = data.index
@@ -368,14 +393,14 @@ class MLFlowRepository:
def check_cache_retention(self, cache: dict, retention: int) -> bool:
"""
Check if cache is still valid based on retention time.
Check whether cached model data is still valid.
Args:
cache (dict): Cached model data
retention (int): Retention time in minutes
cache (dict): Cached model data with a 'timestamp' key.
retention (int): Retention time in minutes.
Returns:
bool: True if cache is still valid, False if expired
bool: True if cache is still valid, False if expired.
"""
current_time = datetime.now()
cache_time = cache['timestamp']
@@ -385,17 +410,14 @@ class MLFlowRepository:
def handle_valid_model(self, model_name: str, cache: dict) -> dict:
"""
Handle valid cached model by returning appropriate model configuration.
Return the cached model configuration when retention is valid.
Args:
model_name (str): Name of the model
model_type (str): Type of model ('predict' or 'transform')
compressed (bool): Whether model is compressed
retention_target (str): Retention target ('model' or 'artifact')
cache (dict): Cached model data
model_name (str): Name of the model (for logging/consistency).
cache (dict): Cached model data structure.
Returns:
dict: Model configuration with model and artifact path
dict: Model configuration.
"""
self.logger.debug(f'Model {model_name} is still valid, using cached version')
@@ -406,8 +428,8 @@ class MLFlowRepository:
Clean up outdated cached model and its artifacts.
Args:
model_name (str): Name of the model
model_key (str): Cache key for the model
model_name (str): Name of the model.
model_key (str): Cache key for the model.
Returns:
None
@@ -419,16 +441,16 @@ class MLFlowRepository:
def get_model(self, model_name: str, retention: int, model_type: str, flavor: str) -> Any:
"""
Get model with caching support based on retention policy.
Retrieve a model with caching support based on retention policy.
Args:
model_name (str): Name of the model to retrieve
retention (int): Cache retention time in minutes (0 = no cache)
retention (int): Cache retention time in minutes (0 = no cache).
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
Returns:
Any: Model object
Any: Model object.
"""
# Retention is 0, download a new model
if retention <= 0:
@@ -488,16 +510,16 @@ class MLFlowRepository:
self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str
) -> pd.DataFrame | ndarray:
"""
Get transformed data using cached transform model.
Execute a cached operation using the requested model.
Args:
model_name (str): Name of the transform model
data (pd.DataFrame): Data to transform
retention (int): Cache retention time in minutes
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
model_name (str): Registered model name.
data (pd.DataFrame): Input data.
retention (int): Cache retention in minutes.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
Returns:
pd.DataFrame: Transformed data
pd.DataFrame | ndarray: Operation result.
"""
if operation not in ['transform', 'predict']:
raise ValueError("Invalid operation. Use 'transform' or 'predict'.")
@@ -531,7 +553,7 @@ class MLFlowRepository:
target_name: str | None = None,
) -> dict[str, dict[str, Any]]:
"""
Create a new MLFlow experiment for model retraining.
Prepare models and data for a retraining run.
This method sets up the complete environment for model retraining by:
1. Loading the current production prediction model
@@ -541,18 +563,16 @@ class MLFlowRepository:
5. Setting up the MLFlow experiment context
Args:
model_name (str): Name of the MLFlow model to retrain
data (pd.DataFrame): Training data for model retraining
transform_flavor (str): Flavor for transformation model
predict_flavor (str): Flavor for prediction model
target_name (str): Target name
metadata (dict): Metadata for logging
model_name (str): Name of the MLflow model to retrain.
data (pd.DataFrame): Training data for model retraining.
transform_flavor (str): Flavor for transformation model.
predict_flavor (str): Flavor for prediction model.
target_name (str | None): Optional target column; if None, use model target.
metadata (dict): Metadata for logging.
Returns:
tuple: (prediction_model, data_model, experiment)
- prediction_model: Loaded prediction model for retraining
- data_model: Fitted transformation model
- experiment: MLFlow experiment name
dict[str, dict[str, Any]]: Mapping with prepared `prediction_model` and
`data_model`, including optional artifact paths.
"""
self.logger.custom_info(f'Starting model experiment creation for {model_name}', metadata)
@@ -659,6 +679,14 @@ class MLFlowRepository:
return retrain_data
def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict):
"""Log a model into the active MLflow run.
Args:
model_data (dict): Model holder with keys 'model' and optional 'artifact_path'.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'.
metadata (dict): Metadata for structured logging.
"""
model = model_data['model']
self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata)
@@ -688,7 +716,7 @@ class MLFlowRepository:
predict_flavor: str = 'sklearn',
) -> dict:
"""
Execute the complete model retraining process in MLFlow.
Execute the complete model retraining process in MLflow.
This method performs the actual model retraining by:
1. Starting a new MLFlow run with descriptive metadata
@@ -700,17 +728,15 @@ class MLFlowRepository:
Args:
prediction_model: MLFlow prediction model to retrain
data_model: MLFlow transformation model to retrain
experiment (str): MLFlow experiment name for the retraining
model_name (str): Name of the model being retrained
data (pd.DataFrame): Training data used for retraining
transform_flavor (str): Flavor for transformation model
predict_flavor (str): Flavor for prediction model
metadata (dict): Metadata for logging
experiment (str): MLflow experiment name for the retraining.
model_name (str): Name of the model being retrained.
data (pd.DataFrame): Training data used for retraining.
transform_flavor (str): Flavor for transformation model.
predict_flavor (str): Flavor for prediction model.
metadata (dict): Metadata for logging.
Returns:
tuple: (status_message, experiment_name)
- status_message (str): Success confirmation message
- experiment_name (str): Name of the experiment
dict: Metadata about the created run and experiment.
"""
prediction_model = retrain_data['prediction_model']
@@ -795,16 +821,16 @@ class MLFlowRepository:
self, run_id: str, model_name: str, metadata: dict
) -> dict:
"""
Update production model with a specific MLFlow run.
Promote a specific run's model to Production.
This method promotes a model from a specific MLFlow run to
production stage. It handles model registration, versioning,
and stage transitions with proper error handling.
Args:
run_id (str): MLFlow run ID containing the model to promote
model_name (str): Name of the MLFlow model
metadata (dict): Metadata for logging
run_id (str): MLflow run ID containing the model to promote.
model_name (str): Name of the MLflow model.
metadata (dict): Metadata for logging.
Returns:
dict: Model update metadata containing:
@@ -863,7 +889,7 @@ class MLFlowRepository:
5. Manages model lifecycle based on retention policy (cleanup artifacts if needed)
Parameters:
model_name (str): The name of the MLFlow model to use for transformation.
model_name (str): The name of the MLflow model to use for transformation.
data (pd.DataFrame): The input data to be transformed by the model.
model_config (dict): Model configuration parameters
metadata (dict): Metadata for logging
@@ -934,7 +960,7 @@ class MLFlowRepository:
9. Handles any exceptions and returns structured error information
Parameters:
model_name (str): The name of the MLFlow model to use for prediction.
model_name (str): The name of the MLflow model to use for prediction.
data (pd.DataFrame): The input data to make predictions on.
model_retention (int): Cache retention time in minutes (0 = no caching).
model_config (dict): Model configuration parameters
@@ -1031,15 +1057,13 @@ class MLFlowRepository:
data (pd.DataFrame): Training data for model retraining. Must contain
all features required by both transformation and
prediction models, including target variable.
model_name (str): Name of the MLFlow model to retrain. Must exist
in the MLFlow Model Registry in Production stage.
model_name (str): Name of the MLflow model to retrain. Must exist
in the MLflow Model Registry in Production stage.
model_config (dict): Model configuration parameters
metadata (dict): Metadata for logging
Returns:
tuple: Retraining operation results containing:
- status_message (str): Success confirmation message or error details
- experiment_name (str): MLFlow experiment identifier for tracking
dict: Retraining operation results and experiment details.
Raises:
mlflow.exceptions.MlflowException: If model not found in registry
@@ -1129,7 +1153,7 @@ class MLFlowRepository:
3. Returns comprehensive update metadata
Args:
experiment (str): MLFlow experiment name containing the retraining runs.
experiment (str): MLflow experiment name containing the retraining runs.
Must be a valid experiment that exists in MLFlow.
model_name (str): Name of the MLFlow model to update. Must exist
in the MLFlow Model Registry.