Add laborious
This commit is contained in:
vitor-aignosi
0
laborious/utils/models/__init__.py
Normal file
0
laborious/utils/models/__init__.py
Normal file
348
laborious/utils/models/minio_dataframe_payload.py
Normal file
348
laborious/utils/models/minio_dataframe_payload.py
Normal file
@@ -0,0 +1,348 @@
|
||||
"""
|
||||
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.observability.logger import Logger
|
||||
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 prediction datasets in MinIO.
|
||||
# It is also the root directory for retention cleanup listing.
|
||||
PREDICTION_DATASETS_PREFIX = 'prediction_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 `prediction_datasets/<model_name>`
|
||||
# - Filename follows the retention-parsing pattern
|
||||
basename = f'{model_name}-{operation}-{timestamp}.parquet'
|
||||
model_dir = model_name.strip().strip('/')
|
||||
prefix = f'{PREDICTION_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
|
||||
|
||||
@staticmethod
|
||||
def _debug(
|
||||
logger: Logger | None,
|
||||
message: str,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Emit debug logs only when logger is provided
|
||||
|
||||
Args:
|
||||
- logger (Logger | None): Logger instance used for debug messages
|
||||
- message (str): Message to be logged
|
||||
- metadata (dict[str, Any] | None): Optional workflow metadata context
|
||||
"""
|
||||
if logger is None:
|
||||
return
|
||||
logger.custom_debug(message, metadata)
|
||||
|
||||
@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,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
logger: Logger | None = None,
|
||||
) -> 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:
|
||||
size = len(pickle.dumps(df.to_dict()))
|
||||
except Exception:
|
||||
size = len(pickle.dumps(df))
|
||||
|
||||
MinioDataFramePayload._debug(
|
||||
logger,
|
||||
f'DataFrame size: {size} bytes',
|
||||
metadata,
|
||||
)
|
||||
return size
|
||||
|
||||
@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 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,
|
||||
last_timestamp: str | None = None,
|
||||
logger: Logger | 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 dataframe is None or dataframe.empty:
|
||||
cls._debug(
|
||||
logger,
|
||||
'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload',
|
||||
workflow_metadata,
|
||||
)
|
||||
return cls(
|
||||
data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
|
||||
)
|
||||
|
||||
if last_timestamp is None:
|
||||
last_timestamp = max(dataframe['timestamp'].values.tolist())
|
||||
|
||||
dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger)
|
||||
cls._debug(
|
||||
logger,
|
||||
(
|
||||
f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes '
|
||||
f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)'
|
||||
),
|
||||
workflow_metadata,
|
||||
)
|
||||
|
||||
if dataframe_size <= OFFLOAD_THRESHOLD_BYTES:
|
||||
cls._debug(
|
||||
logger,
|
||||
'MinioDataFramePayload.from_dataframe using inline payload',
|
||||
workflow_metadata,
|
||||
)
|
||||
return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status)
|
||||
|
||||
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
|
||||
object_key, object_prefix = _build_object_key(model_name, operation, timestamp)
|
||||
cls._debug(
|
||||
logger,
|
||||
(
|
||||
'MinioDataFramePayload.from_dataframe offloading payload to MinIO '
|
||||
f'with key {object_key}'
|
||||
),
|
||||
workflow_metadata,
|
||||
)
|
||||
|
||||
# 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
|
||||
cls._debug(
|
||||
logger,
|
||||
f'MinioDataFramePayload.from_dataframe upload completed: {uri}',
|
||||
workflow_metadata,
|
||||
)
|
||||
|
||||
return cls(
|
||||
data=None,
|
||||
bucket=bucket,
|
||||
object_key=object_key_full,
|
||||
object_prefix=object_prefix,
|
||||
uri=uri,
|
||||
last_timestamp=last_timestamp,
|
||||
status=status,
|
||||
)
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
minio_repo: MinioRepository,
|
||||
workflow_metadata: dict[str, Any] | None = None,
|
||||
logger: Logger | 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:
|
||||
self._debug(
|
||||
logger,
|
||||
'MinioDataFramePayload.retrieve using inline payload data',
|
||||
workflow_metadata,
|
||||
)
|
||||
return DataFrame(self.data)
|
||||
|
||||
if not self.has_data():
|
||||
self._debug(
|
||||
logger,
|
||||
'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe',
|
||||
workflow_metadata,
|
||||
)
|
||||
return DataFrame()
|
||||
|
||||
self._debug(
|
||||
logger,
|
||||
f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}',
|
||||
workflow_metadata,
|
||||
)
|
||||
file_bytes = await minio_repo.download_file(
|
||||
object_name=self.object_key, metadata=workflow_metadata
|
||||
)
|
||||
df = read_parquet(BytesIO(file_bytes))
|
||||
self._debug(
|
||||
logger,
|
||||
f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}',
|
||||
workflow_metadata,
|
||||
)
|
||||
return df
|
||||
Reference in New Issue
Block a user