SIENTIAPDE-1712
Implement MinIO Offload and Retention Features - Added configuration options for MinIO retention hours and offload threshold in README. - Introduced MinIO payload offloading for large DataFrame-derived payloads, storing them as parquet files. - Updated activities to utilize MinIO for data loading and cleanup, including new methods for offloading and retention management. - Refactored existing activities to integrate MinIO functionality, ensuring compatibility with previous workflows. - Removed the legacy MinioRepository class, consolidating MinIO operations under a new manager structure. - Updated requirements to use the latest version of the sientia-dataops-library.
This commit is contained in:
vitor-aignosi
@@ -1,11 +1,15 @@
|
||||
from re import M
|
||||
from temporalio import activity, workflow
|
||||
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from pandas import DataFrame, to_datetime
|
||||
from io import BytesIO
|
||||
from pandas import DataFrame, read_parquet, to_datetime
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
@@ -19,11 +23,12 @@ with workflow.unsafe.imports_passed_through():
|
||||
)
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
|
||||
from laborious.utils.repository.minio_repository import MinioRepository
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
class MLFlow(SientiaMonitoring):
|
||||
class MLFlow(MinioManager):
|
||||
"""
|
||||
MLFlow integration activities for model inference operations.
|
||||
|
||||
@@ -47,11 +52,11 @@ class MLFlow(SientiaMonitoring):
|
||||
mlflow_host: str,
|
||||
mlflow_port: int,
|
||||
mlflow_username: str,
|
||||
minio_config: dict[str, Any],
|
||||
mlflow_password: str,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize MLFlow activities with server configuration.
|
||||
@@ -67,7 +72,7 @@ class MLFlow(SientiaMonitoring):
|
||||
Raises:
|
||||
Exception: If MLFlowRepository initialization fails
|
||||
"""
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
MinioManager.__init__(self, minio_repository, logger, notification_handler, metrics_controller)
|
||||
self.mlflow_host = mlflow_host
|
||||
self.mlflow_port = mlflow_port
|
||||
self.mlflow_username = mlflow_username
|
||||
@@ -82,32 +87,17 @@ class MLFlow(SientiaMonitoring):
|
||||
metrics_controller,
|
||||
)
|
||||
|
||||
if not hasattr(self, 'minio_repository'):
|
||||
self.minio_repository: MinioRepository | None = None
|
||||
|
||||
if self.minio_repository is None:
|
||||
self.minio_repository = MinioRepository(
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
minio_endpoint_url=minio_config['endpoint_url'],
|
||||
minio_access_key=minio_config['access_key'],
|
||||
minio_secret_key=minio_config['secret_key'],
|
||||
minio_region_name=minio_config['region_name'],
|
||||
minio_default_bucket=minio_config['default_bucket'],
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the MLFlow activity and clean up resources.
|
||||
"""
|
||||
SientiaMonitoring.shutdown(self)
|
||||
MinioManager.close(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
@activity.defn(name='request_transform')
|
||||
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Transform input data using MLFlow models.
|
||||
|
||||
@@ -138,7 +128,10 @@ class MLFlow(SientiaMonitoring):
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Transforming data...', metadata)
|
||||
data = DataFrame(input_data['data'])
|
||||
|
||||
payload: MinioDataFramePayload = input_data['data']
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
model_name = input_data['model_name']
|
||||
model_config = input_data.get('model_config', {})
|
||||
|
||||
@@ -176,10 +169,29 @@ class MLFlow(SientiaMonitoring):
|
||||
|
||||
self.info('Data transformed successfully', metadata)
|
||||
|
||||
return response_data
|
||||
if not response_data.get('success', False):
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
operation='transform',
|
||||
status=response_data,
|
||||
workflow_metadata=metadata,
|
||||
)
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=response_data['content'],
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
operation='transform',
|
||||
workflow_metadata=metadata,
|
||||
status={
|
||||
'success': True,
|
||||
},
|
||||
)
|
||||
|
||||
@activity.defn(name='request_predict')
|
||||
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Execute predictions using MLFlow models.
|
||||
|
||||
@@ -210,7 +222,10 @@ class MLFlow(SientiaMonitoring):
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Predicting data...', metadata)
|
||||
data = DataFrame(input_data['data'])
|
||||
|
||||
payload: MinioDataFramePayload = input_data['data']
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
model_name = input_data['model_name']
|
||||
model_config = input_data.get('model_config', {})
|
||||
|
||||
@@ -236,7 +251,28 @@ class MLFlow(SientiaMonitoring):
|
||||
|
||||
self.info('Data predicted successfully', metadata)
|
||||
|
||||
return response_data
|
||||
|
||||
if not response_data.get('success', False):
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
operation='predict',
|
||||
status=response_data,
|
||||
workflow_metadata=metadata,
|
||||
)
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=response_data['content'],
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
operation='predict',
|
||||
workflow_metadata=metadata,
|
||||
status={
|
||||
'success': True,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@activity.defn(name='retrain_model')
|
||||
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -275,14 +311,24 @@ class MLFlow(SientiaMonitoring):
|
||||
raise ValueError('Minio repository not initialized')
|
||||
|
||||
metadata = input_data['metadata']
|
||||
object_key = input_data['object_key']
|
||||
|
||||
self.info(f'Loading retrain data from Key: {object_key}', metadata)
|
||||
|
||||
try:
|
||||
data = await self.minio_repository.get_parquet_as_dataframe(
|
||||
object_key=object_key, metadata=metadata
|
||||
)
|
||||
if 'data' in input_data:
|
||||
# New path: payload-based retrain input (inline or MinIO offloaded).
|
||||
data = await MinioDataFramePayload.dataframe_from_wire(
|
||||
input_data['data'],
|
||||
self.minio_repository,
|
||||
metadata,
|
||||
)
|
||||
else:
|
||||
# Backward compatibility: legacy query_to_minio contract.
|
||||
object_key = input_data['object_key']
|
||||
self.info(f'Loading retrain data from Key: {object_key}', metadata)
|
||||
file_bytes = await self.minio_repository.download_file(
|
||||
object_name=object_key,
|
||||
metadata=metadata,
|
||||
)
|
||||
data = read_parquet(BytesIO(file_bytes))
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
|
||||
Reference in New Issue
Block a user