vitor-aignosi
67942c45e0
Refactor MinioDataFramePayload usage across activities - Updated instances of MinioDataFramePayload initialization in Gates, MLFlow, and Storage classes to use the new from_dict method for better data reconstruction from dictionaries. - Enhanced the PredictionProcess workflow to utilize the updated payload handling. - Added passthrough fixtures in tests to accommodate the new from_dict method for consistent testing behavior.
259 lines
9.5 KiB
Python
259 lines
9.5 KiB
Python
"""
|
|
MinIO-backed DataFrame payload for Temporal workflows.
|
|
|
|
Data is never stored as a pandas ``DataFrame`` field on the dataclass.
|
|
Instead, the DataFrame is only provided as an input to:
|
|
`from_dataframe` / `from_dataframe_to_dict`.
|
|
|
|
At build time, the DataFrame is evaluated for its serialized size; if it exceeds
|
|
the configured threshold, it is serialized to parquet bytes and uploaded to MinIO.
|
|
Otherwise, it is inlined as a Temporal-friendly ``dict``.
|
|
"""
|
|
|
|
import pickle
|
|
import re
|
|
from collections.abc import Hashable
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from io import BytesIO
|
|
from os import getenv
|
|
from typing import Any, Literal
|
|
|
|
from pandas import DataFrame, read_parquet
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now
|
|
|
|
# Keys that are part of the serialized wire format (not arbitrary metadata).
|
|
_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'})
|
|
|
|
_OBJECT_TIMESTAMP_PATTERN = re.compile(
|
|
r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$'
|
|
)
|
|
|
|
OFFLOAD_THRESHOLD_BYTES = int(
|
|
float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024
|
|
)
|
|
|
|
# Relative prefix used for storing offloaded training datasets in MinIO.
|
|
# It is also the root directory for retention cleanup listing.
|
|
TRAINING_DATASETS_PREFIX = 'training_datasets'
|
|
|
|
OperationKind = Literal['initial', 'transform', 'predict']
|
|
|
|
|
|
def _build_object_key(
|
|
model_name: str, operation: OperationKind, timestamp: str
|
|
) -> tuple[str, str | None]:
|
|
"""
|
|
Build the MinIO object key and the directory prefix used for retention listing.
|
|
|
|
Args:
|
|
model_name: Registered model name used in the pipeline.
|
|
operation: Either initial (pre-transform load) or transform (post-MLFlow transform).
|
|
timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME.
|
|
|
|
Return:
|
|
tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root).
|
|
"""
|
|
# Naming convention:
|
|
# - Directory is always `training_datasets/<model_name>`
|
|
# - Filename follows the retention-parsing pattern
|
|
basename = f'{model_name}-{operation}-{timestamp}.parquet'
|
|
model_dir = model_name.strip().strip('/')
|
|
prefix = f'{TRAINING_DATASETS_PREFIX}/{model_dir}'
|
|
return f'{prefix}/{basename}', prefix
|
|
|
|
|
|
@dataclass
|
|
class MinioDataFramePayload:
|
|
"""
|
|
Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys.
|
|
|
|
Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`.
|
|
Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class.
|
|
"""
|
|
|
|
last_timestamp: str
|
|
status: dict[str, Any] | None = None
|
|
data: dict[Hashable, Any] | None = None
|
|
bucket: str | None = None
|
|
object_key: str | None = None
|
|
object_prefix: str | None = None
|
|
uri: str | None = None
|
|
|
|
@classmethod
|
|
def from_dict(cls, raw: dict[str, Any] | 'MinioDataFramePayload') -> 'MinioDataFramePayload':
|
|
"""
|
|
Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization.
|
|
|
|
Temporal converts dataclass return values into plain dicts when crossing
|
|
workflow/activity boundaries. This method rebuilds the typed instance so
|
|
that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are
|
|
available on the receiving side.
|
|
|
|
If the argument is already a MinioDataFramePayload, it is returned as-is.
|
|
|
|
Args:
|
|
raw: Dict with keys matching the dataclass fields
|
|
(last_timestamp, status, data, bucket, object_key, object_prefix, uri),
|
|
or an existing MinioDataFramePayload instance.
|
|
|
|
Return:
|
|
MinioDataFramePayload: Reconstructed (or original) instance.
|
|
"""
|
|
if isinstance(raw, MinioDataFramePayload):
|
|
return raw
|
|
return cls(
|
|
last_timestamp=raw['last_timestamp'],
|
|
status=raw.get('status'),
|
|
data=raw.get('data'),
|
|
bucket=raw.get('bucket'),
|
|
object_key=raw.get('object_key'),
|
|
object_prefix=raw.get('object_prefix'),
|
|
uri=raw.get('uri'),
|
|
)
|
|
|
|
@staticmethod
|
|
def estimate_size_bytes(df: DataFrame) -> int:
|
|
"""
|
|
Approximate serialized size of the DataFrame as the default-orient dict.
|
|
|
|
Args:
|
|
df: DataFrame whose tabular content size is estimated.
|
|
|
|
Return:
|
|
int: Estimated size in bytes (pickle of dict representation).
|
|
"""
|
|
try:
|
|
return len(pickle.dumps(df.to_dict()))
|
|
except Exception:
|
|
return len(pickle.dumps(df))
|
|
|
|
@staticmethod
|
|
def parse_object_timestamp(object_key: str) -> datetime | None:
|
|
"""
|
|
Parse the timestamp embedded in the object key basename (before .parquet).
|
|
|
|
Args:
|
|
object_key: S3/MinIO object key whose basename follows
|
|
``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``.
|
|
|
|
Return:
|
|
datetime | None: Parsed UTC-naive datetime from the key, or None if not matched.
|
|
"""
|
|
basename = object_key.rsplit('/', 1)[-1]
|
|
match = _OBJECT_TIMESTAMP_PATTERN.search(basename)
|
|
if not match:
|
|
return None
|
|
try:
|
|
return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME)
|
|
except ValueError:
|
|
return None
|
|
|
|
def cleanup_prefix(self) -> str | None:
|
|
"""
|
|
Return True if cleanup is enabled for this payload.
|
|
"""
|
|
if self.object_key is not None and self.data is None:
|
|
return self.object_prefix
|
|
return None
|
|
|
|
def has_data(self) -> bool:
|
|
"""
|
|
Return True if the payload has some data internally or in MinIO.
|
|
"""
|
|
return (self.data is not None and not self.data != {}) or self.object_key is not None
|
|
|
|
@classmethod
|
|
async def from_dataframe(
|
|
cls,
|
|
dataframe: DataFrame | None,
|
|
minio_repo: MinioRepository,
|
|
model_name: str,
|
|
operation: OperationKind,
|
|
status: dict[str, Any] | None = None,
|
|
workflow_metadata: dict | None = None,
|
|
) -> 'MinioDataFramePayload':
|
|
"""
|
|
Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO.
|
|
|
|
The DataFrame is not stored on the returned instance.
|
|
|
|
Args:
|
|
dataframe: Tabular data to evaluate and persist (inline or MinIO).
|
|
metadata: Small metadata dict merged into the payload (e.g. success, message).
|
|
minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`.
|
|
workflow_metadata: Metadata passed to MinIO store for logging/metrics.
|
|
model_name: Registered model name used in the object basename.
|
|
operation: Either ``initial`` (query load) or ``transform`` (post-transform).
|
|
key_prefix: Backward-compatible parameter (currently ignored for object naming).
|
|
size_threshold_bytes: Byte limit before offload. When None, the module-level
|
|
environment-derived default is used.
|
|
|
|
Return:
|
|
MinioDataFramePayload: Instance with data and/or MinIO fields set.
|
|
"""
|
|
|
|
if not dataframe or dataframe.empty:
|
|
return cls(
|
|
data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
|
|
)
|
|
|
|
last_timestamp = max(dataframe['timestamp'].values.tolist())
|
|
|
|
if cls.estimate_size_bytes(dataframe) <= OFFLOAD_THRESHOLD_BYTES:
|
|
return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp)
|
|
|
|
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
|
|
object_key, object_prefix = _build_object_key(model_name, operation, timestamp)
|
|
|
|
# Upload using the relative object key. The upstream repository will
|
|
# prefix it internally under its MinIO namespace.
|
|
parquet_buffer = BytesIO()
|
|
dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
|
|
file_bytes = parquet_buffer.getvalue()
|
|
|
|
upload_result = await minio_repo.upload_file(
|
|
file_bytes=file_bytes,
|
|
relative_key=object_key,
|
|
metadata=workflow_metadata,
|
|
)
|
|
|
|
bucket = minio_repo.bucket
|
|
object_key_full = upload_result.get('minio_object_name', object_key)
|
|
uri = f's3://{bucket}/{object_key_full}' if bucket else None
|
|
|
|
return cls(
|
|
data=None,
|
|
bucket=bucket,
|
|
object_key=object_key_full,
|
|
object_prefix=object_prefix,
|
|
uri=uri,
|
|
last_timestamp=last_timestamp,
|
|
)
|
|
|
|
async def retrieve(
|
|
self, minio_repo: MinioRepository, workflow_metadata: dict[str, Any] | None = None
|
|
) -> DataFrame:
|
|
"""
|
|
Load parquet from MinIO when object_key is set and populate inline data.
|
|
|
|
Args:
|
|
minio_repo: sientia_do MinioRepository (or compatible) with download_file().
|
|
workflow_metadata: Metadata passed to MinIO read for logging/metrics.
|
|
|
|
Return:
|
|
dict[str, Any]: Flat dict with data filled (same keys as to_dict after load).
|
|
"""
|
|
if self.data is not None:
|
|
return DataFrame(self.data)
|
|
|
|
if not self.has_data():
|
|
return DataFrame()
|
|
|
|
file_bytes = await minio_repo.download_file(
|
|
object_name=self.object_key, metadata=workflow_metadata
|
|
)
|
|
df = read_parquet(BytesIO(file_bytes))
|
|
return df
|