7512963e19
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.
143 lines
5.3 KiB
Python
143 lines
5.3 KiB
Python
from temporalio import activity, workflow
|
|
|
|
from laborious.utils.repository.minio_repository import MinioRepository
|
|
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
|
from sientia_do.temporal.activities.postgres import Postgres
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.temporal.constants import now
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from typing import Any
|
|
import traceback
|
|
import pandas as pd
|
|
|
|
DATETIME_FILENAME_FORMAT = "%Y-%m-%d_%H-%M-%S"
|
|
|
|
|
|
class Storage(Postgres):
|
|
"""
|
|
Extensions for Postgres activities with a helper to export query results
|
|
directly to MinIO as Parquet and return the object name.
|
|
"""
|
|
|
|
def __init__(self,
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
dbname: str,
|
|
min_connections: int,
|
|
max_connections: int,
|
|
minio_config: dict[str, Any],
|
|
logger: Logger,
|
|
notification_handler: NotificationHandler):
|
|
super().__init__(host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password,
|
|
dbname=dbname,
|
|
min_connections=min_connections,
|
|
max_connections=max_connections,
|
|
logger=logger,
|
|
notification_handler=notification_handler)
|
|
|
|
if not hasattr(self, 'minio_repository'):
|
|
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'])
|
|
|
|
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'])
|
|
|
|
@activity.defn(name='query_to_minio')
|
|
async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Execute SQL query, write result as Parquet to MinIO, and return object name.
|
|
|
|
Args (input_data):
|
|
metadata (dict): Workflow metadata
|
|
query (str): SQL query
|
|
model_name (str): Model name for object naming
|
|
object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain)
|
|
|
|
Returns:
|
|
dict: { success: bool, object_name: str, uri: str }
|
|
"""
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
object_prefix = input_data.get('object_prefix', 'datasets/retrain')
|
|
|
|
timestamp = now().strftime(DATETIME_FILENAME_FORMAT)
|
|
object_name = f"{object_prefix}_{timestamp}.parquet"
|
|
uri = f"s3://{self.minio_repository.minio_bucket}/{object_name}"
|
|
|
|
try:
|
|
data = await self.load_custom_query(input_data)
|
|
if not data:
|
|
self.error(
|
|
f"query_to_minio failed: No data returned from query", metadata)
|
|
return {"success": False, "message": "No data returned from query"}
|
|
|
|
# Ensure we have a DataFrame
|
|
data = pd.DataFrame(data)
|
|
|
|
# Write parquet to memory and upload via persistent client
|
|
self.minio_repository.store_dataframe_as_parquet(
|
|
dataframe=data,
|
|
uri=uri,
|
|
object_name=object_name,
|
|
metadata=metadata
|
|
)
|
|
|
|
return {"success": True, "object_key": object_name, "uri": uri}
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id="ERROR_LOADING_CUSTOM_QUERY",
|
|
message=f"Error fetching data from query: {e}",
|
|
block="load_custom_query",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace
|
|
)
|
|
|
|
self.error(trace, metadata)
|
|
|
|
return {"success": False, "message": str(e)}
|
|
|
|
def close(self) -> None:
|
|
"""Close Storage resources (MinIO client and Postgres engine)."""
|
|
try:
|
|
if hasattr(self, 's3_client') and self.s3_client is not None:
|
|
try:
|
|
self.s3_client.close()
|
|
finally:
|
|
self.s3_client = None
|
|
finally:
|
|
# Ensure Postgres resources are disposed as well
|
|
try:
|
|
super().close()
|
|
except Exception:
|
|
pass
|
|
|
|
def __del__(self):
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
pass
|