sientia-dataops-scouter_temporal/scouter/activities/gates.py

from collections.abc import Hashable

from sientia_do.observability.metrics_controller import MetricsController
from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import traceback
    from typing import Any

    from pandas import DataFrame
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring

    from scouter import metrics
    from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter

quality_gate_filters = {
    'NULL_VALUES_FILTER': null_values_filter,
    'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter,
}


class Gates(SientiaMonitoring):
    """
    Data quality gates and filtering operations.

    This class implements data quality validation and filtering for industrial
    time-series data. It provides:
    - Configurable data quality filters
    - Data aggregation functions for time-series data
    - Comprehensive error handling and notification
    - Metrics collection for quality monitoring

    The class supports multiple aggregation strategies and quality filters to
    ensure data integrity and enable flexible data processing workflows.
    """

    def __init__(
        self,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        """
        Initialize the Gates class with logging and notification services.

        Args:
            logger (Logger): Logger instance for operation logging
            notification_handler (NotificationHandler): Handler for system notifications
            metrics_controller (MetricsController): Metrics controller instance
        """
        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)

    def close(self):
        """
        Close the Gates class.
        """
        SientiaMonitoring.shutdown(self)

    async def apply_aggregation(
        self, values: DataFrame, aggr_function: str, metadata: dict[str, Any]
    ) -> float | None | str:
        """
        Apply aggregation function to a group of time-series data.

        This method applies the specified aggregation function to a group of
        data points. It handles edge cases and provides comprehensive error
        reporting for invalid aggregation functions.

        Args:
            values (DataFrame): Group of data points to aggregate (pre-sorted by timestamp)
            aggr_function (str): Aggregation function to apply.
                Supported functions: 'lts' (latest), 'avg' (average), 'mdn' (median),
                'max' (maximum), 'min' (minimum)
            metadata (dict[str, Any]): Workflow metadata for error reporting

        Returns:
            float | None | str: Aggregated value, None if no valid data, or 'continue' for errors

        Raises:
            NotificationError: If invalid aggregation function is specified
        """
        # Fast path for single value
        if len(values) == 1:
            return values['value'].iloc[0]

        if aggr_function == 'lts':
            return values['value'].iloc[-1]

        # Remove NaN values without inplace operation
        clean_values = values['value'].dropna()

        if clean_values.empty:
            return None

        # Use dictionary lookup for aggregation functions (faster than if-elif chain)
        aggregation_map = {
            'avg': lambda x: x.mean(),
            'mdn': lambda x: x.median(),
            'max': lambda x: x.max(),
            'min': lambda x: x.min(),
        }

        if aggr_function in aggregation_map:
            return aggregation_map[aggr_function](clean_values)
        else:
            await self.send_notification_async(
                metadata=metadata,
                notification_id='AGGREGATION_ISSUES',
                message=f'Invalid aggregation function: {aggr_function}',
                block='aggregate_data',
                level=NotificationLevel.ERROR,
                attachment_content=traceback.format_exc(),
            )
            return 'continue'

    @activity.defn(name='aggregate_data')
    async def aggregate_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
        """
        Aggregate time-series data by tag and name using specified functions.

        This activity processes time-series data by grouping it by tag and name,
        then applying the configured aggregation functions. It handles data
        validation and provides comprehensive error reporting.

        Args:
            input_data (dict[str, Any]): Activity input parameters.
                Required fields:
                - data (dict[str, Any]): Time-series data to aggregate
                - model_tags (dict[str, Any]): Tag configuration with aggregation functions

        Returns:
            dict[Hashable, Any]: Aggregated data organized by tag and name

        Raises:
            Exception: If aggregation operation fails
        """

        metadata = input_data['metadata']

        try:
            # Convert input data to DataFrame
            df = DataFrame(input_data['data'])

            self.info(f'Aggregating time series data for {len(df)} rows', metadata=metadata)

            # Sort once by timestamp for all data (more efficient than sorting each group)
            df = df.sort_values(['tag', 'name', 'timestamp'])

            # Group by tag and name
            # sort=False since we already sorted
            grouped = df.groupby(['tag', 'name'], sort=False)

            # Prepare aggregation functions mapping
            model_tags = input_data['model_tags']

            # Process groups efficiently
            results = []
            for (tag, name), group in grouped:
                # Get the aggregation function from model_tags
                aggr_function = model_tags.get(name, {}).get('aggr_func', 'lts')

                # Get the latest timestamp (last row since data is sorted)
                latest_timestamp = group['timestamp'].iloc[-1]

                aggr_value = await self.apply_aggregation(group, aggr_function, metadata)

                if aggr_value == 'continue':
                    continue

                # Store the result directly in list for better performance
                results.append(
                    {
                        'tag': tag,
                        'name': name,
                        'value': aggr_value,
                        'timestamp': latest_timestamp,
                        'aggregation_function': aggr_function,
                    }
                )

            self.info(f'Aggregated data has {len(results)} rows', metadata=metadata)

            # Convert to DataFrame only once at the end if we have results
            if results:
                result_df = DataFrame(results)

                self.debug(f'Final aggregated data:\n{result_df.to_string()}', metadata=metadata)

                return result_df.to_dict()
            else:
                # Return empty DataFrame dict structure
                return DataFrame().to_dict()

        except Exception as e:
            trace = traceback.format_exc()

            await self.send_notification_async(
                metadata=metadata,
                notification_id='AGGREGATION_ISSUES',
                message=f'Error aggregating data: {e}',
                block='aggregate_data',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            self.error(trace, metadata=metadata)
            raise e

    @activity.defn(name='data_quality_gate')
    async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
        """
        Apply data quality filters to incoming data.

        This activity applies configurable quality filters to validate incoming
        data. It supports multiple filter types and provides comprehensive
        error reporting for quality issues.

        Args:
            input_data (dict[str, Any]): Activity input parameters.
                Required fields:
                - data (dict[str, Any]): Data to validate
                - filters (dict[str, str]): Filter configuration
                - model_tags (dict[str, Any]): Tag-specific validation rules

        Returns:
            dict[Hashable, Any]: Filtered data that passes quality validation

        Raises:
            Exception: If quality validation fails
        """

        metadata = input_data['metadata']

        filters = input_data['filters']
        data = DataFrame(input_data['data'])
        model_tags = input_data['model_tags']

        self.info(f'Applying quality gate to data to {len(data)} rows', metadata=metadata)

        tags = list(model_tags.keys())

        data = data[data['name'].isin(tags)]

        for filter_name, config in filters.items():
            policy = config['policy']
            if filter_name not in quality_gate_filters:
                self.warning(f'Filter {filter_name} not found', metadata=metadata)
                continue

            try:
                filtered_data = quality_gate_filters[filter_name](data, model_tags)

            except Exception as e:
                trace = traceback.format_exc()
                await self.send_notification_async(
                    metadata=metadata,
                    notification_id='DATA_QUALITY_GATE_ISSUES',
                    message=f'Error applying filter {filter_name}: {e}',
                    block='data_quality_gate',
                    level=NotificationLevel.ERROR,
                    attachment_content=trace,
                )

                self.error(trace, metadata=metadata)

            else:
                if filtered_data.empty:
                    continue

                message = f'{len(filtered_data)} rows has quality issues: {filter_name}: {policy}'
                attachment = filtered_data.to_string()

                await self.send_notification_async(
                    metadata=metadata,
                    notification_id=f'DATA_QUALITY_GATE_ISSUES__{filter_name}',
                    message=message,
                    block='data_quality_gate',
                    level=NotificationLevel.WARNING,
                    attachment_content=attachment,
                )

                if policy == 'DISCARD':
                    data = data[~data.index.isin(filtered_data.index)]

        self.info(f'Data quality gate applied, final data has {len(data)} rows', metadata=metadata)

        return data.to_dict()

    @activity.defn(name='write_metrics')
    async def write_metrics(self, input_data: dict[str, Any]) -> None:
        """
        Write metrics to the database.
        input_data:
            metadata: dict[str, Any]
        """
        metadata = input_data['metadata']
        tag_values = DataFrame(input_data['tag_values'])

        self.info(f'Writing metrics for {metadata["model_name"]}', metadata=metadata)

        metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels(
            pod_id=self.pod_id,
            model_name=metadata['model_name'],
            workflow_name=metadata['workflow_name'],
        ).inc()

        # Register metrics
        for _, row in tag_values.iterrows():
            metrics.TAG_CHANGES_MONITOR.labels(
                pod_id=self.pod_id,
                model_name=metadata['model_name'],
                workflow_name=metadata['workflow_name'],
                tag_name=row['variable'],
            ).set(row['value'])

        self.info(f'Metrics written for {metadata["model_name"]}', metadata=metadata)