sientia-dataops-laborious_temporal/laborious/activities/model_metrics.py

from temporalio import activity, workflow

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

    import numpy as np
    from pandas import DataFrame, Index, to_datetime
    from sientia.ModelAnalysis import ModelAnalysis
    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.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ

    from laborious import metrics
    from laborious.utils.dataframe_debug import build_dataframe_debug_message

warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0')
warnings.filterwarnings(
    'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
)


class ModelMetrics(SientiaMonitoring):
    """
    Metrics activities for the Laborious system.

    This class provides activities for writing metrics to the Prometheus monitoring system.
    """

    _MAX_DEBUG_DATAFRAME_ROWS = 100

    def __init__(
        self,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)

    def close(self) -> None:
        """
        Close the model metrics activity and clean up resources.
        """
        SientiaMonitoring.shutdown(self)

    def __del__(self):
        self.close()

    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
        """
        Log dataframe content only when row count is below the configured threshold

        Args:
         - message (str): Base log message to identify the dataframe in logs
         - data (Any): Dataframe-like payload to be logged
         - metadata (dict[str, Any]): Workflow metadata for contextual logging
        """
        self.debug(
            build_dataframe_debug_message(
                message=message,
                data=data,
                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
            ),
            metadata,
        )

    async def get_drift_metrics(
        self,
        reference_data: DataFrame,
        target_data: DataFrame,
        target_name: str,
        reference_columns: Index,
        drift_metrics: list[str],
        chunk_period: str,
        metadata: dict[str, Any],
    ) -> DataFrame:
        """
        Calculate univariate drift metrics for a model.
        Args:
            model_analysis (ModelAnalysis): Model analysis object
            reference_data (DataFrame): Reference data
            target_data (DataFrame): Target data
            reference_columns (list[str]): Reference columns
            drift_metrics (list[str]): Drift metrics
            metadata (dict[str, Any]): Workflow execution metadata
        """

        config = {
            'target': target_name,
            'prediction': 'prediction',
            'timestamp': 'timestamp',
            'features': reference_columns,
        }

        model_analysis = ModelAnalysis(config=config)

        self._debug_dataframe(
            f'Reference data: Size {reference_data.shape}', reference_data, metadata
        )

        self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata)

        core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
        start_time = time.time()
        try:
            univariate_drift = model_analysis.detect_univariate_drift(
                reference_df=reference_data,
                analysis_df=target_data,
                features=reference_columns,
                timestamp_col=config['timestamp'],
                methods=drift_metrics,
                chunk_period=chunk_period,
            )
        except Exception as e:
            self.error(f'Error detecting univariate drift: {e}', metadata)
            await self.emit_metric(
                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
            )
            raise e
        await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
        await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)

        core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
        start_time = time.time()
        try:
            multivariate_drift = model_analysis.detect_multivariate_drift(
                reference_df=reference_data,
                analysis_df=target_data,
                features=reference_columns,
                timestamp_col=config['timestamp'],
                chunk_period=chunk_period,
            )
        except Exception as e:
            self.error(f'Error detecting multivariate drift: {e}', metadata)
            await self.emit_metric(
                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
            )
            raise e
        await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
        await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)

        start_time = time.time()
        core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe')
        try:
            drift_df = model_analysis.get_drift_metrics_dataframe(
                univariate_drift=univariate_drift,
                multivariate_drift=multivariate_drift,
            )
        except Exception as e:
            self.error(f'Error getting drift metrics: {e}', metadata)
            await self.emit_metric(
                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
            )
            raise e
        await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
        await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)

        self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)

        return drift_df

    @activity.defn(name='calculate_drift')
    async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
        """
        Calculate drift metrics for a model.

        Args:
            input_data (dict[str, Any]): Input data containing:
                - metadata (dict): Workflow execution metadata
                - model_name (str): Name of the MLFlow model to calculate drift for
                - reference_data (pd.DataFrame): Reference data for the model
                - target_data (pd.DataFrame): Target data for calculating drift
                - target_name (str): Name of the target column
                - drift_metrics (list[str]): List of drift metrics to calculate
        """
        metadata = input_data['metadata']
        model_name = input_data['model_name']
        model_id = input_data['model_id']
        reference_raw_data = input_data['reference_data']
        target_data = DataFrame(input_data['target_data'])
        target_name = input_data['target_name']
        drift_metrics = input_data['drift_metrics']
        chunk_period = input_data['chunk_period']

        if chunk_period not in ['min', 's']:
            self.error(f'Invalid chunk period: {chunk_period}', metadata)
            raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"')

        self.info(f'Calculating drift for model {model_name}', metadata)

        target_data = target_data.pivot(index='timestamp', columns='variable', values='value')
        target_data['timestamp'] = target_data.index
        target_data['timestamp'] = to_datetime(target_data['timestamp'])
        target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT)
        target_data = target_data.reset_index(drop=True)
        target_data.dropna(inplace=True)

        if reference_raw_data is not None:
            self.info('Using reference data', metadata)
            reference_data = DataFrame(reference_raw_data)
            accurate = True
        else:
            # Get 30% first rows of target_data
            self.warning('Using 30% first rows of target data as reference data', metadata)
            target_data.sort_values(by='timestamp', ascending=True, inplace=True)
            reference_data = target_data.head(int(len(target_data) * 0.3))
            accurate = False

            await self.send_notification_async(
                metadata=metadata,
                notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
                message='Using 30% first rows of target data as reference data',
                block='model_metrics',
                level=NotificationLevel.WARNING,
                attachment_content=reference_data.to_csv(),
            )

        reference_columns = reference_data.drop(
            columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
        ).columns

        try:
            drift_df = await self.get_drift_metrics(
                reference_data=reference_data,
                target_data=target_data,
                target_name=target_name,
                reference_columns=reference_columns,
                drift_metrics=drift_metrics,
                chunk_period=chunk_period,
                metadata=metadata,
            )
        except Exception as e:
            self.error(f'Error getting drift metrics: {e}', metadata)
            await self.send_notification_async(
                metadata=metadata,
                notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
                message=f'Error getting drift metrics: {e}',
                block='model_metrics',
                level=NotificationLevel.ERROR,
                attachment_content=traceback.format_exc(),
            )
            return []

        if drift_df.empty:
            self.warning('No drift metrics found', metadata)
            return []

        # Drop unnecessary columns
        drift_df.drop(columns=['p_value'], inplace=True)

        # Extract timestamps only until minutes
        if chunk_period == 'min':
            target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
        else:
            target_timestamps = target_data['timestamp']

        # Drop rows where timestamp is not in target data, to avoid save drift from reference
        drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]

        if drift_df.empty:
            self.warning(
                'No drift metrics found after dropping rows where timestamp is not in target data',
                metadata,
            )
            return []

        # Rename columns to match database columns
        drift_df.rename(
            columns={
                'metric': 'method',
                'statistic': 'value',
            },
            inplace=True,
        )

        # Drop duplicates
        drift_df.drop_duplicates(
            subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
        )

        drift_df['model_id'] = model_id
        drift_df['accurate'] = accurate

        drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
        drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
        drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)

        self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)

        return drift_df.to_dict(orient='records')

    @activity.defn(name='calculate_simple_metrics')
    async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
        """
        Calculate simple metrics for a model. Metrics available are:
            - rmse
            - mse
            - mae
            - r2
            - accuracy
            - precision
            - recall
            - f1
        Args:
            input_data (dict[str, Any]): Input data containing:
                - metadata (dict): Workflow execution metadata
                - model_id (str): ID of the MLFlow model
                - target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns
                - metrics (list[str]): List of metrics to calculate
        Returns:
            dict[Hashable, Any]: Dictionary containing the calculated metrics
        """

        metadata = input_data['metadata']
        model_id = input_data['model_id']
        target_data = DataFrame(input_data['target_data'])
        metrics = input_data['metrics']
        interval_minutes = input_data['interval_minutes']

        data_size = target_data.shape[0]

        output_data = []

        diff = target_data['target'] - target_data['prediction']
        diff_squared = diff**2

        self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)

        for metric in metrics:
            if metric == 'rmse':
                output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
            elif metric == 'mse':
                output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
            elif metric == 'mae':
                output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
            elif metric == 'r2':
                y_true = target_data['target']
                y_mean = np.mean(y_true)

                ss_res = np.sum(diff_squared)
                ss_tot = np.sum((y_true - y_mean) ** 2)

                # Evita divisão por zero
                if ss_tot == 0:
                    r2_score = 0.0
                else:
                    r2_score = 1 - (ss_res / ss_tot)

                output_data.append({'metric': 'r2', 'value': r2_score})

        data = DataFrame(output_data)
        data['model_id'] = model_id
        data['timestamp'] = target_data['timestamp'].max()
        data['data_size'] = data_size
        data['interval_minutes'] = interval_minutes

        self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata)

        return data.to_dict(orient='records')