5d0d049082
Remove code validation script and refactor imports in activities and workflows - Deleted the `validate.sh` script, which was responsible for running code quality checks. - Cleaned up import statements in `activities.py`, `gates.py`, `mlflow.py`, and `storage.py` by removing unused imports and organizing them. - Refactored initialization methods in `MinioManager` and `MLFlow` classes for improved readability. - Updated various workflows to ensure compatibility with the new structure and removed unnecessary comments. - Enhanced test cases to accommodate changes in the activities and workflows, ensuring proper mocking of dependencies.
286 lines
11 KiB
Python
286 lines
11 KiB
Python
import json
|
|
|
|
from temporalio import activity, workflow
|
|
|
|
from laborious.utils.repository.minio_manager import MinioManager
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
|
import traceback
|
|
from datetime import timedelta
|
|
from io import BytesIO
|
|
from typing import Any
|
|
|
|
import pandas as pd
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_do.temporal.activities.postgres import Postgres
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now
|
|
|
|
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
|
|
|
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
|
|
|
|
|
|
class Storage(Postgres, MinioManager):
|
|
"""
|
|
Extensions for Postgres activities with a helper to export query results
|
|
directly to MinIO as Parquet and return the object name.
|
|
"""
|
|
|
|
minio_repository: MinioRepository | None = None
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
dbname: str,
|
|
min_connections: int,
|
|
max_connections: int,
|
|
retention_hours: int = 24,
|
|
minio_repository: MinioRepository | None = None,
|
|
logger: Logger | None = None,
|
|
notification_handler: NotificationHandler | None = None,
|
|
metrics_controller: MetricsController | None = None,
|
|
):
|
|
self.retention_hours = retention_hours
|
|
Postgres.__init__(
|
|
self,
|
|
host=host,
|
|
port=port,
|
|
user=user,
|
|
password=password,
|
|
dbname=dbname,
|
|
min_connections=min_connections,
|
|
max_connections=max_connections,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
MinioManager.__init__(
|
|
self, minio_repository, logger, notification_handler, metrics_controller
|
|
)
|
|
|
|
@activity.defn(name='load_query_with_minio_offload')
|
|
async def load_query_with_minio_offload(
|
|
self, input_data: dict[str, Any]
|
|
) -> MinioDataFramePayload:
|
|
"""
|
|
Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
|
|
|
|
Args (input_data):
|
|
metadata (dict): Workflow metadata (same as load_custom_query).
|
|
query (str): SQL query.
|
|
datetime_columns (list[str], optional): Datetime column names.
|
|
model_name (str): Model name for object key basename.
|
|
key_prefix (str, optional): Directory prefix inside the bucket.
|
|
size_threshold_bytes (int, optional): Override env offload threshold.
|
|
|
|
Returns:
|
|
dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
|
|
"""
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata: dict = input_data.get('metadata', {})
|
|
model_name = input_data['model_name']
|
|
|
|
rows = await self.load_custom_query(
|
|
input_data,
|
|
)
|
|
if not rows:
|
|
self.error(
|
|
'load_query_with_minio_offload failed: No data returned from query', metadata
|
|
)
|
|
dataframe = None
|
|
else:
|
|
dataframe = pd.DataFrame(rows)
|
|
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe,
|
|
minio_repo=self.minio_repository,
|
|
workflow_metadata=metadata,
|
|
model_name=model_name,
|
|
operation='initial',
|
|
)
|
|
|
|
@activity.defn(name='export_payload_to_postgres')
|
|
async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
|
|
"""
|
|
Export a payload to PostgreSQL.
|
|
"""
|
|
metadata = input_data.get('metadata')
|
|
payload: MinioDataFramePayload = input_data['data']
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
return await self.export_data_to_postgres(
|
|
{
|
|
**input_data,
|
|
'data': data,
|
|
}
|
|
)
|
|
|
|
@activity.defn(name='cleanup_minio_objects_expired')
|
|
async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Delete objects under the given prefixes that are older than the retention window.
|
|
|
|
Args (input_data):
|
|
metadata (dict): Workflow metadata for logging and metrics.
|
|
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
|
|
|
|
Returns:
|
|
dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
|
|
"""
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
prefix = input_data['prefix']
|
|
base = now()
|
|
cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
|
|
hours=self.retention_hours
|
|
)
|
|
|
|
report: dict[str, Any] = {
|
|
'failed': {},
|
|
'deleted': {},
|
|
'failed_count': 0,
|
|
'deleted_count': 0,
|
|
}
|
|
try:
|
|
keys = await self.minio_repository.list_objects(
|
|
prefix=prefix,
|
|
recursive=True,
|
|
metadata=metadata,
|
|
)
|
|
for key in keys:
|
|
try:
|
|
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
|
if ts is None:
|
|
continue
|
|
if ts >= cutoff:
|
|
continue
|
|
await self.minio_repository.delete_file(
|
|
object_name=key,
|
|
metadata=metadata,
|
|
)
|
|
except Exception as e:
|
|
report['failed'][key] = {
|
|
'success': False,
|
|
'message': str(e),
|
|
}
|
|
report['failed_count'] += 1
|
|
continue
|
|
report['deleted'][key] = {
|
|
'success': True,
|
|
'message': 'Deleted',
|
|
}
|
|
report['deleted_count'] += 1
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
|
|
message=f'Error cleaning up MinIO objects: {e}',
|
|
block='cleanup_minio_objects_expired',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata)
|
|
else:
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='CLEANUP_MINIO_OBJECTS_EXPIRED',
|
|
message='MinIO objects cleaned up successfully',
|
|
block='cleanup_minio_objects_expired',
|
|
level=NotificationLevel.INFO,
|
|
attachment_content=json.dumps(report),
|
|
)
|
|
|
|
return report
|
|
|
|
@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 }
|
|
"""
|
|
|
|
if self.minio_repository is None:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
model_name = input_data.get('model_name') or metadata.get('model_name') or 'unknown'
|
|
object_prefix = input_data.get('object_prefix', 'datasets/retrain')
|
|
|
|
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
|
|
# Keep a stable model-level layout for minimal_retrain:
|
|
# training_datasets/<model_name>/<filename>
|
|
# Sanitize object_prefix to avoid extra subdirectories in the relative key.
|
|
safe_prefix = str(object_prefix).strip().strip('/').replace('/', '_')
|
|
filename = f'{safe_prefix}_{timestamp}.parquet'
|
|
relative_key = f'training_datasets/{model_name}/{filename}'
|
|
bucket = getattr(self.minio_repository, 'bucket', 'streamlit-connectors')
|
|
uri = f's3://{bucket}/{relative_key}'
|
|
|
|
try:
|
|
data = await self.load_custom_query(input_data)
|
|
if not data:
|
|
self.error('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)
|
|
|
|
# Convert DataFrame -> parquet bytes, then upload using the new MinIO interface.
|
|
parquet_buffer = BytesIO()
|
|
data.to_parquet(parquet_buffer, engine='pyarrow', index=True)
|
|
file_bytes = parquet_buffer.getvalue()
|
|
|
|
upload_result = await self.minio_repository.upload_file(
|
|
file_bytes=file_bytes,
|
|
relative_key=relative_key,
|
|
metadata=metadata,
|
|
)
|
|
|
|
object_key_full = upload_result.get('minio_object_name', relative_key)
|
|
uri = f's3://{bucket}/{object_key_full}'
|
|
return {'success': True, 'object_key': object_key_full, 'uri': uri}
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='ERROR_STORING_QUERY_TO_MINIO',
|
|
message=f'Error storing query to MinIO: {e}',
|
|
block='query_to_minio',
|
|
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)."""
|
|
Postgres.close(self)
|
|
MinioManager.close(self)
|
|
|
|
def __del__(self):
|
|
self.close()
|