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 MinIO integration and update environment configurations

- Added MinIO configuration parameters to .env.example and values.yaml for improved storage management.
- Updated requirements.txt to include necessary libraries for MinIO support.
- Refactored Activities class to utilize Storage for MinIO interactions.
- Enhanced MLFlow class to integrate MinIO for data retrieval during model retraining.
- Introduced build_minio_config function to streamline MinIO configuration setup.
- Updated minimal_retrain workflow to support data storage in MinIO.
This commit is contained in:
vitor-aignosi
2025-10-07 16:56:16 -03:00
parent 16a3aa7022
commit 7512963e19
11 changed files with 454 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
laborious/utils/connectors_config.py
View File

@@ -127,3 +127,26 @@ def build_mongodb_config() -> Dict[str, Any]:
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
}
def build_minio_config() -> Dict[str, Any]:
"""
Build MinIO (S3-compatible) configuration from environment variables.
Environment Variables:
MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000)
MINIO_ACCESS_KEY: Access key (default: minioadmin)
MINIO_SECRET_KEY: Secret key (default: minioadmin)
MINIO_REGION: Region name for S3 client (default: us-east-1)
MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious)
Returns:
dict: MinIO configuration dictionary
"""
return {
'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'),
'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious')
}

115
laborious/utils/repository/minio_repository.py Normal file
View File

@@ -0,0 +1,115 @@
from io import BytesIO
import traceback
import boto3
from botocore.config import Config
from pandas import DataFrame, read_parquet
from sientia_do.observability.logger import Logger
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from typing import Any
from botocore.exceptions import ClientError
class MinioRepository():
def __init__(self,
minio_endpoint_url: str,
minio_access_key: str,
minio_secret_key: str,
minio_region_name: str,
minio_default_bucket: str,
logger: Logger,
notification_handler: NotificationHandler):
# MinIO settings shared with pandas s3fs
self.storage_options = {
'key': minio_access_key,
'secret': minio_secret_key,
'client_kwargs': {'endpoint_url': minio_endpoint_url}
}
self.minio_bucket = minio_default_bucket
self.minio_endpoint_url = minio_endpoint_url
self.minio_region_name = minio_region_name
logger.info(
f"Connecting to MinIO at {self.minio_endpoint_url}, default bucket: {self.minio_bucket}")
# Reusable MinIO client
self.s3_client = boto3.client(
's3',
endpoint_url=self.minio_endpoint_url,
aws_access_key_id=self.storage_options['key'],
aws_secret_access_key=self.storage_options['secret'],
region_name=self.minio_region_name,
config=Config(
signature_version='s3v4',
s3={'addressing_style': 'path'},
retries={'max_attempts': 5, 'mode': 'standard'},
connect_timeout=5,
read_timeout=120,
),
)
self._bucket_checked = False
self.logger = logger
self.notification_handler = notification_handler
def ensure_bucket_exists(self, metadata: dict[str, Any]) -> bool:
"""
Ensure the MinIO bucket exists; create it if necessary.
"""
if self._bucket_checked:
return True
try:
self.logger.custom_info(
f"Checking if bucket '{self.minio_bucket}' exists", metadata)
self.s3_client.head_bucket(Bucket=self.minio_bucket)
self._bucket_checked = True
return True
except ClientError:
try:
self.logger.custom_info(
f"Creating bucket '{self.minio_bucket}'", metadata)
self.s3_client.create_bucket(Bucket=self.minio_bucket)
self._bucket_checked = True
return True
except ClientError as ce:
trace = traceback.format_exc()
self.notification_handler.send_notification(
metadata=metadata,
notification_id="ERROR_CREATING_MINIO_BUCKET",
message=f"Failed to ensure bucket '{self.minio_bucket}': {ce}",
block="ensure_bucket_exists",
level=NotificationLevel.ERROR,
attachment_content=str(ce)
)
self.logger.custom_error(trace, metadata)
return False
def store_dataframe_as_parquet(self, dataframe: DataFrame, uri: str,
object_name: str, metadata: dict[str, Any]):
self.ensure_bucket_exists(metadata)
self.logger.custom_info(
f"Storing dataframe as parquet in {uri}", metadata)
buffer = BytesIO()
dataframe.to_parquet(buffer, engine='pyarrow', index=True)
buffer.seek(0)
self.s3_client.put_object(
Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue())
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:
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)
# Read the content into a BytesIO buffer to support seek operations
buffer = BytesIO(response['Body'].read())
return read_parquet(buffer)

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

@@ -244,7 +244,7 @@ class MLFlowRepository():
output_dir
)
def load_predict_model(self, model_name: str, flavor: str = 'pyfunc',
def load_predict_model(self, model_name: str, flavor: str = 'sklearn',
artifact_path: str | None = None) -> Any:
"""
Downloads a predictive model from the MLflow Model Registry.
@@ -386,7 +386,7 @@ class MLFlowRepository():
f"Could not load model from {pickle_path} - unknown or corrupted format")
def download_model(self, model_name: str, model_type: str, flavor: str,
download_artifacts: bool = False) -> Any:
download_artifacts: bool = False) -> tuple[Any, str]:
"""
Download model based on type (predict or transform).
@@ -397,7 +397,7 @@ class MLFlowRepository():
download_artifacts (bool): Whether to download artifacts
Returns:
Any: Model object
tuple[Any, str]: Model object and artifact path if model is compressed
"""
self.logger.info(
@@ -422,7 +422,7 @@ class MLFlowRepository():
model = self.load_transform_model(
model_name, flavor, artifact_path)
return model
return model, artifact_path
"""
Functions related to data format
@@ -633,9 +633,9 @@ class MLFlowRepository():
"""
def create_model_experiment(self, model_name: str, data: pd.DataFrame,
transform_flavor: str = 'sklearn', predict_flavor: str = 'pyfunc',
transform_config: dict = {}, predict_config: dict = {},
fit_config: dict = {}, target_name: str = None,
metadata: dict = {}) -> tuple:
is_compressed: bool = False, metadata: dict = {}) -> tuple:
"""
Create a new MLFlow experiment for model retraining.
@@ -649,10 +649,11 @@ class MLFlowRepository():
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
transform_config (dict): Configuration for transformation model
predict_config (dict): Configuration for prediction model
fit_config (dict): Fit configuration
target_name (str): Target name
is_compressed (bool): Whether model is compressed
metadata (dict): Metadata for logging
Returns:
@@ -664,7 +665,7 @@ class MLFlowRepository():
self.logger.custom_info(
f"Starting model experiment creation for {model_name}", metadata)
self.logger.custom_debug(
f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata)
f"Model configuration - transform_config: {transform_config}, predict_config: {predict_config}, fit_config: {fit_config}, target_name: {target_name}", metadata)
latest_production_id = self.get_model_run_id(
model_name, stage="Production"
@@ -675,17 +676,24 @@ class MLFlowRepository():
self.logger.custom_info(
f"Loading transformation model for {model_name}", metadata)
data_model = self.download_model(
transform_flavor = transform_config.get('flavor', 'sklearn')
transformed_is_compressed = transform_config.get(
'is_compressed', False)
data_model, data_artifact_path = self.download_model(
model_name=model_name, model_type="transform", flavor=transform_flavor,
download_artifacts=(transform_flavor == 'pyfunc')
download_artifacts=transformed_is_compressed
)
self.logger.custom_info(
f"Loading prediction model for {model_name}", metadata)
prediction_model = self.download_model(
predict_flavor = predict_config.get('flavor', 'pyfunc')
predicted_is_compressed = predict_config.get('is_compressed', False)
prediction_model, prediction_artifact_path = self.download_model(
model_name=model_name, model_type="predict", flavor=predict_flavor,
download_artifacts=(predict_flavor == 'pyfunc')
download_artifacts=predicted_is_compressed
)
self.logger.custom_debug(
@@ -784,14 +792,36 @@ class MLFlowRepository():
self.logger.custom_info(
f"Model experiment creation completed successfully for {model_name}", metadata)
return prediction_model, data_model, experiment
retrain_data = {
'prediction_model': prediction_model,
'prediction_artifact_path': prediction_artifact_path,
'data_model': data_model,
'data_artifact_path': data_artifact_path,
'experiment': experiment
}
return retrain_data
def log_model(self, model: Any, artifact_local_path: str, flavor: str, model_type: str):
if artifact_local_path:
mlflow.log_artifact(artifact_local_path, artifact_path="")
else:
if flavor == 'sklearn':
mlflow.sklearn.log_model(model, model_type)
elif flavor == 'pyfunc':
mlflow.pyfunc.log_model(model, model_type)
elif flavor == 'pytorch':
mlflow.pytorch.log_model(model, model_type)
else:
raise ValueError(
"Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.")
def perform_model_retrain(self,
prediction_model,
data_model,
experiment: str,
model_name: str,
data: pd.DataFrame,
retrain_data: dict,
transform_config: dict = {},
predict_config: dict = {},
metadata: dict = {}):
"""
Execute the complete model retraining process in MLFlow.
@@ -809,6 +839,7 @@ class MLFlowRepository():
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
is_compressed (bool): Whether model is compressed
metadata (dict): Metadata for logging
Returns:
@@ -816,9 +847,19 @@ class MLFlowRepository():
- status_message (str): Success confirmation message
- experiment_name (str): Name of the experiment
"""
prediction_model = retrain_data['prediction_model']
data_model = retrain_data['data_model']
experiment = retrain_data['experiment']
prediction_artifact_path = retrain_data['prediction_artifact_path']
data_artifact_path = retrain_data['data_artifact_path']
self.logger.custom_info(
f"Starting model retraining process for {model_name} in experiment {experiment}", metadata)
transform_flavor = transform_config.get('flavor', 'sklearn')
predict_flavor = predict_config.get('flavor', 'sklearn')
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
experiment_description = f"Retrain model {model_name} with new data"
@@ -845,7 +886,8 @@ class MLFlowRepository():
mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(data_model, "data_model")
self.log_model(data_model, data_artifact_path,
transform_flavor, "data_model")
makedirs("tmp/retrain_data", exist_ok=True)
@@ -856,7 +898,8 @@ class MLFlowRepository():
mlflow.log_artifact(file_path)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(prediction_model, "prediction_model")
self.log_model(prediction_model, prediction_artifact_path,
predict_flavor, "prediction_model")
mlflow.log_param("retrain", True)
# clear temp file
@@ -1044,7 +1087,7 @@ class MLFlowRepository():
"""
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('predict_flavor', 'pyfunc')
flavor = model_config.get('predict_flavor', 'sklearn')
try:
@@ -1147,8 +1190,8 @@ class MLFlowRepository():
target_name = model_config.get('target', None)
transform_flavor = model_config.get('transform_flavor', 'sklearn')
predict_flavor = model_config.get('predict_flavor', 'pyfunc')
transform_config = model_config.get('transform_config', {})
predict_config = model_config.get('predict_config', {})
fit_config = {
'split_fit_data': model_config.get('split_fit_data', False),
@@ -1157,21 +1200,19 @@ class MLFlowRepository():
}
self.logger.custom_debug(
f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata)
f"Model configuration - transform_config: {transform_config}, predict_config: {predict_config}, fit_config: {fit_config}, target_name: {target_name}", metadata)
try:
self.logger.custom_info(
"Creating model experiment environment", metadata)
prediction_model, data_model, experiment = self.create_model_experiment(
model_name=model_name, data=data, transform_flavor=transform_flavor,
predict_flavor=predict_flavor, fit_config=fit_config, target_name=target_name,
metadata=metadata)
retrain_data = self.create_model_experiment(
model_name=model_name, data=data, transform_config=transform_config, predict_config=predict_config, fit_config=fit_config, target_name=target_name, metadata=metadata)
self.logger.custom_info(
f"Model experiment created successfully: {experiment}", metadata)
f"Model experiment created successfully: {retrain_data}", metadata)
self.logger.custom_info("Saving model retrain", metadata)
experiment = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data, metadata)
model_name, data, retrain_data, transform_config, predict_config, metadata)
self.logger.custom_info(
f"Model retraining completed successfully for experiment: {experiment}", metadata)