vitor-aignosi
192834bde3
QTZPOC-20: Refactor model import workflow, introduce atomic model ID generation, and clarify MinIO expiry. Fuses bundle download, decryption, MLflow experiment creation, and artifact upload into a single activity to prevent cross-worker issues. Implements an atomic counter for MongoDB model IDs, replacing the max(id) + 1 scan to avoid collisions during concurrent imports. Adds comprehensive documentation for MinIO object expiry mechanisms and refines lifecycle rule application. Improves error handling with clearer responsibility classification and centralizes MLflow repository injection.
216 lines
7.9 KiB
Python
216 lines
7.9 KiB
Python
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 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 now
|
|
|
|
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
|
|
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
|
from laborious.utils.time import to_naive
|
|
|
|
_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(MISSING_MINIO_REPOSITORY)
|
|
|
|
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',
|
|
logger=self.logger,
|
|
)
|
|
|
|
@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.from_dict(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.
|
|
|
|
See also: `ModelImport.ensure_import_bundle_retention`
|
|
(laborious/activities/model_import.py) — a newer, unrelated lifecycle-policy expiry
|
|
mechanism for the import bundle prefix. The two are independent by design for now; see
|
|
docs/minio-object-expiry.md.
|
|
|
|
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(MISSING_MINIO_REPOSITORY)
|
|
|
|
metadata = input_data.get('metadata', {})
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
prefix = payload.cleanup_prefix()
|
|
base = now()
|
|
cutoff = to_naive(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:
|
|
# Cleanup success is expected in normal flow; avoid noisy INFO notifications
|
|
# that do not impact behavior and can flood observability in test runs.
|
|
self.info('MinIO objects cleaned up successfully', metadata)
|
|
|
|
return report
|
|
|
|
def close(self) -> None:
|
|
"""Close Storage resources (MinIO client and Postgres engine)."""
|
|
if hasattr(self, 'engine'):
|
|
Postgres.close(self)
|
|
MinioManager.close(self)
|
|
|
|
def __del__(self):
|
|
self.close()
|