a3da800cab
Refactor monitoring and metrics integration across various components - Removed coverage options from `pyproject.toml`. - Updated prediction metrics in `README.md` to replace `pipeline_name` with `workflow_name`. - Upgraded `sientia-dataops-library` dependency version in `requirements-light.txt` and `requirements.txt`. - Enhanced metrics handling in `laborious` activities, including `Activities`, `Gates`, `MLFlow`, and `OPC`, to utilize a new `MetricsController`. - Refactored metric emission methods to improve clarity and consistency across the codebase. - Updated tests to reflect changes in metrics handling and ensure proper functionality.
139 lines
5.1 KiB
Python
139 lines
5.1 KiB
Python
from temporalio import activity, workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
|
import traceback
|
|
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.temporal.activities.postgres import Postgres
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now
|
|
|
|
from laborious.utils.repository.minio_repository import MinioRepository
|
|
|
|
|
|
class Storage(Postgres):
|
|
"""
|
|
Extensions for Postgres activities with a helper to export query results
|
|
directly to MinIO as Parquet and return the object name.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
user: str,
|
|
password: str,
|
|
dbname: str,
|
|
min_connections: int,
|
|
max_connections: int,
|
|
minio_config: dict[str, Any],
|
|
logger: Logger,
|
|
notification_handler: NotificationHandler,
|
|
metrics_controller: MetricsController,
|
|
):
|
|
super().__init__(
|
|
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,
|
|
)
|
|
|
|
if not hasattr(self, 'minio_repository'):
|
|
self.minio_repository: MinioRepository | None = None
|
|
|
|
if self.minio_repository is None:
|
|
self.minio_repository = MinioRepository(
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
minio_endpoint_url=minio_config['endpoint_url'],
|
|
minio_access_key=minio_config['access_key'],
|
|
minio_secret_key=minio_config['secret_key'],
|
|
minio_region_name=minio_config['region_name'],
|
|
minio_default_bucket=minio_config['default_bucket'],
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
@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', {})
|
|
object_prefix = input_data.get('object_prefix', 'datasets/retrain')
|
|
|
|
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
|
|
object_name = f'{object_prefix}_{timestamp}.parquet'
|
|
uri = f's3://{self.minio_repository.minio_bucket}/{object_name}'
|
|
|
|
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)
|
|
|
|
# Write parquet to memory and upload via persistent client
|
|
await self.minio_repository.store_dataframe_as_parquet(
|
|
dataframe=data, uri=uri, object_name=object_name, metadata=metadata
|
|
)
|
|
|
|
return {'success': True, 'object_key': object_name, '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)."""
|
|
try:
|
|
if hasattr(self, 'minio_repository') and self.minio_repository is not None:
|
|
try:
|
|
self.minio_repository.close()
|
|
finally:
|
|
self.minio_repository = None
|
|
finally:
|
|
# Ensure Postgres resources are disposed as well
|
|
try:
|
|
super().close()
|
|
except Exception:
|
|
self.logger.error('Error closing Postgres resources')
|
|
|
|
def __del__(self):
|
|
self.close()
|