a812acd4e9
Rename PredictionMetadata (prediction_metadata.py) to MongoActivities (mongo.py) and move update_simple_metrics_adaptive_state there from ModelMetrics, so both Mongo-backed activities share one class and repository (PR #59 review, agreed with the class author). write_prediction_metadata is unchanged; the pure hysteresis helpers stay in model_metrics.py. References in activities.py, docstrings and tests updated.
342 lines
13 KiB
Python
342 lines
13 KiB
Python
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
|
import time
|
|
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 import metrics
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
|
from sientia_do.repository.minio_repository_sync import MinioRepository
|
|
from sientia_do.temporal.activities.postgres_sync import Postgres
|
|
from sientia_do.temporal.constants import now
|
|
from sqlalchemy import text
|
|
|
|
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, SientiaMonitoring):
|
|
"""
|
|
Extensions for Postgres activities with a helper to export query results
|
|
directly to MinIO as Parquet and return the object name, and with
|
|
``repeat_last_prediction``, a query specific to this repo's REPEAT path
|
|
(moved here from the shared library — nothing else consumes it).
|
|
"""
|
|
|
|
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,
|
|
)
|
|
|
|
self.minio_repository = minio_repository
|
|
SientiaMonitoring.__init__(
|
|
self,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
@activity.defn(name='load_query_with_minio_offload')
|
|
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 = 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 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')
|
|
def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
|
|
"""
|
|
Resolve a MinIO-aware payload into a DataFrame and persist it into PostgreSQL.
|
|
|
|
This activity accepts the serialized payload produced by previous steps
|
|
(inline dict or MinIO object reference), reconstructs the tabular data,
|
|
and delegates the final write to ``export_data_to_postgres`` using the
|
|
same input contract expected by the Postgres activity mixin.
|
|
|
|
Args:
|
|
- input_data (dict[str, Any]): Activity input containing ``data`` as a
|
|
``MinioDataFramePayload``-compatible dict plus database write options
|
|
(schema/table/on_conflict/metadata and related fields).
|
|
|
|
Return:
|
|
dict: Result dictionary returned by ``export_data_to_postgres``, including
|
|
success status and optional write diagnostics.
|
|
"""
|
|
metadata = input_data.get('metadata')
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = payload.retrieve(self.minio_repository, metadata)
|
|
|
|
return self.export_data_to_postgres(
|
|
{
|
|
**input_data,
|
|
'data': data,
|
|
}
|
|
)
|
|
|
|
def _build_repeat_last_prediction_query(self, schema: str, table_name: str):
|
|
table_identifier = f'"{schema}".{table_name}'
|
|
return text( # noqa: S608
|
|
f"""
|
|
INSERT INTO {table_identifier} (
|
|
model_id,
|
|
prediction,
|
|
timestamp,
|
|
response_time,
|
|
prediction_status,
|
|
prediction_confidence,
|
|
metadata_id,
|
|
created_at
|
|
)
|
|
SELECT
|
|
model_id,
|
|
prediction,
|
|
:last_timestamp,
|
|
response_time,
|
|
prediction_status,
|
|
prediction_confidence,
|
|
:metadata_id,
|
|
NOW()
|
|
FROM {table_identifier}
|
|
WHERE model_id = :model
|
|
ORDER BY timestamp DESC
|
|
LIMIT 1;
|
|
""" # nosec B608
|
|
)
|
|
|
|
@activity.defn(name='repeat_last_prediction')
|
|
def repeat_last_prediction(self, query_items: dict[str, Any]):
|
|
"""
|
|
Insert a new `predictions` row that copies the last one for a model forward in time.
|
|
|
|
REPEAT persists a new row (unlike STOP, which persists nothing), reusing the previous
|
|
row's `prediction`/`response_time`/`prediction_status`/`prediction_confidence` values with
|
|
a fresh `timestamp`/`created_at`. `metadata_id` references the `prediction_metadata`
|
|
document for this cycle, if any (see `MongoActivities.write_prediction_metadata`);
|
|
`None` is a valid value — the column is nullable.
|
|
|
|
Args (query_items):
|
|
metadata (dict): Workflow metadata for logging and metrics.
|
|
schema (str): Postgres schema of the `predictions` table.
|
|
table_name (str): `predictions` table name.
|
|
model (str): `model_id` to repeat the last prediction for.
|
|
last_timestamp: Timestamp to stamp the new row with.
|
|
metadata_id (str | None): `prediction_metadata` document id, or `None`.
|
|
|
|
Raises:
|
|
Exception: Re-raised after logging/metrics if the insert fails.
|
|
"""
|
|
metadata = query_items['metadata']
|
|
schema = query_items['schema']
|
|
table_name = query_items['table_name']
|
|
model = query_items['model']
|
|
last_timestamp = query_items['last_timestamp']
|
|
metadata_id = query_items.get('metadata_id')
|
|
|
|
repeat_query = self._build_repeat_last_prediction_query(schema, table_name)
|
|
self.info(f'Repeating last prediction for model {model}', metadata)
|
|
self.debug(
|
|
f'Query parameters: schema={schema}, model={model}, last_timestamp={last_timestamp}',
|
|
metadata,
|
|
)
|
|
core_labels = {
|
|
**self.get_core_labels(metadata, 'execute_sql'),
|
|
'conflict_mode': 'error',
|
|
}
|
|
|
|
with self.session_factory.begin() as session:
|
|
try:
|
|
start_time = time.time()
|
|
session.execute(
|
|
repeat_query,
|
|
{
|
|
'last_timestamp': last_timestamp,
|
|
'model': model,
|
|
'metadata_id': metadata_id,
|
|
},
|
|
)
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id='ERROR_REPEATING_LAST_PREDICTION',
|
|
message=f'Error repeating last prediction: {e}',
|
|
block='repeat_last_prediction',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata)
|
|
self.emit_metric_sync(
|
|
metric_object=metrics.POSTGRES_WRITE_ERROR_COUNT,
|
|
tags=core_labels,
|
|
)
|
|
raise
|
|
else:
|
|
self.observe_lag_sync(start_time, metrics.POSTGRES_WRITE_LAG, core_labels)
|
|
self.emit_metric_sync(
|
|
metric_object=metrics.POSTGRES_WRITE_COUNT,
|
|
tags=core_labels,
|
|
)
|
|
|
|
@activity.defn(name='cleanup_minio_objects_expired')
|
|
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(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 = 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
|
|
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()
|
|
self.send_notification(
|
|
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:
|
|
"""
|
|
Shutdown Storage resources in deterministic order.
|
|
|
|
The method first closes Postgres resources via ``Postgres.close`` (engine,
|
|
sessions, and monitoring hooks), then closes the optional MinIO repository
|
|
and clears the local reference to avoid accidental reuse after shutdown.
|
|
"""
|
|
if hasattr(self, 'engine'):
|
|
Postgres.close(self)
|
|
if self.minio_repository is not None:
|
|
try:
|
|
self.minio_repository.close()
|
|
finally:
|
|
self.minio_repository = None
|