SIENTIAPDE-1273

Enhance MLFlowRepository and Activities classes with new methods and metrics

- Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry.
- Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models.
- Updated Activities class to integrate ModelMetrics for improved metrics handling.
- Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities.
- Updated various workflows to include `transform_table_name` in input data for better data handling.

laborious/activities/activities.py

@@ -11,9 +11,9 @@ with workflow.unsafe.imports_passed_through():
     from laborious.activities.mlflow import MLFlow
     from laborious.activities.opc import OPC
     from laborious.activities.storage import Storage
+    from laborious.activities.model_metrics import ModelMetrics
 
-
-class Activities(Storage, MLFlow, Gates, OPC):
+class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
     """
     Main activities orchestrator for the Laborious system.
 
@@ -108,6 +108,13 @@ class Activities(Storage, MLFlow, Gates, OPC):
             metrics_controller=metrics_controller,
         )
 
+        ModelMetrics.__init__(
+            self,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
+
     async def shutdown(self):
         """
         Gracefully shutdown all activities and clean up resources.
@@ -124,3 +131,4 @@ class Activities(Storage, MLFlow, Gates, OPC):
         MLFlow.close(self)
         Gates.close(self)
         await OPC.close(self)
+        ModelMetrics.close(self)

laborious/activities/mlflow.py

@@ -1,3 +1,4 @@
+from typing import Hashable
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
@@ -418,3 +419,36 @@ class MLFlow(SientiaMonitoring):
             )
             self.error(trace, metadata=metadata)
             raise e
+
+
+    @activity.defn(name='get_reference_data')
+    async def get_reference_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any] | None:
+        """
+        Get reference data from the MLflow Model Registry.
+
+        Args:
+            input_data (dict): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - model_name (str): Name of the MLFlow model to get reference data from
+
+        Returns:
+            dict[Hashable, Any] | None: Reference data from the MLflow Model Registry.
+
+        """
+
+        metadata = input_data['metadata']
+        model_name = input_data['model_name']
+        artifact = "evaluation_data.csv"
+
+        reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
+            model_name=model_name, artifact_path=artifact, metadata=metadata
+        )
+
+        if reference_data is None:
+            self.warning(f'Reference data not found for model {model_name}', metadata)
+            return None
+
+        reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
+        reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
+
+        return reference_data.to_dict()

laborious/activities/model_metrics.py

@@ -0,0 +1,266 @@
+from temporalio import activity, workflow
+
+with workflow.unsafe.imports_passed_through():
+    from typing import Any, Hashable
+
+    from pandas import DataFrame, Index, to_datetime
+
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+    from sientia_do.observability.metrics_controller import MetricsController
+    from sientia_do.observability.logger import Logger
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
+    from sientia_do.notifications.models import NotificationLevel
+    from sientia.ModelAnalysis import ModelAnalysis
+    from laborious import metrics
+    import time
+    from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+    import warnings
+    import traceback
+
+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.
+    """
+
+    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()
+
+    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(f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', metadata)
+
+        self.debug(f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', 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(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
+        
+        return drift_df
+
+
+    @activity.defn(name='calculate_drift')
+    async def calculate_drift(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
+        """
+        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', 'chunk_start_date', 'chunk_end_date'], 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(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
+
+        
+        return drift_df.to_dict()
+

laborious/utils/repository/model_repository.py

@@ -217,6 +217,24 @@ class MLFlowRepository(SientiaMonitoring):
         run_info = mlflow.get_run(run_id)
         return run_info.data.params
 
+    def check_artifact_exists(self, run_id: str,
+    artifact_path: str, metadata: dict[str, Any]) -> bool:
+        """
+        Check if an artifact exists in the MLflow Model Registry.
+
+        Args:
+            run_id (str): Run identifier to inspect.
+            artifact_path (str): Path to the artifact to check.
+
+        Returns:
+            bool: True if the artifact exists, False otherwise.
+        """
+        artifacts = self.client.list_artifacts(run_id)
+
+        self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata)
+        self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata)
+        
+        return any(artifact.path == artifact_path for artifact in artifacts)
     """
     Functions related to download and load models
     """
@@ -263,7 +281,7 @@ class MLFlowRepository(SientiaMonitoring):
 
 
     async def load_artifact_dataframe(self, model_name: str, artifact_path: str,
-                                         metadata: dict[str, Any]) -> pd.DataFrame:
+                                         metadata: dict[str, Any]) -> pd.DataFrame | None:
         
         """
         Load the dataframe content of an artifact from the MLflow Model Registry.
@@ -277,9 +295,13 @@ class MLFlowRepository(SientiaMonitoring):
             pd.DataFrame: The dataframe content of the artifact.
         """
         run_id = self.get_model_run_id(model_name=model_name, stage='Production')
-        artifact_path = path.join("runs:/", run_id, artifact_path)
         core_labels = self.get_core_labels(metadata, operation_type='load_text')
 
+        if not self.check_artifact_exists(run_id, artifact_path, metadata):
+            return None
+
+        artifact_path = path.join("runs:/", run_id, artifact_path)
+
         start_time = time.time()
         try:
             content = mlflow.artifacts.load_text(artifact_path)
@@ -290,6 +312,8 @@ class MLFlowRepository(SientiaMonitoring):
         await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
         await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
 
+        self.debug(f'Content of {run_id}/{artifact_path}: \n{content}', metadata)
+
         dataframe = pd.read_csv(StringIO(content))
 
         self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata)
@@ -680,6 +704,41 @@ class MLFlowRepository(SientiaMonitoring):
     Functions related to model retraining
     """
 
+    def get_prediction_data(self, prediction_model: Any, retrain_dataset: pd.DataFrame,
+    target_name: str) -> pd.DataFrame:
+        """
+        Get prediction data from prediction model.
+        """
+        input_index = retrain_dataset.index
+
+        prediction_data = prediction_model.predict(retrain_dataset)
+
+        if isinstance(prediction_data, pd.DataFrame):
+
+            prediction_data.columns = pd.Index(['prediction'])
+
+        else:
+            prediction_data = pd.DataFrame(prediction_data, columns=['prediction'])
+
+        prediction_data.index = input_index
+
+        # Merge prediction data with retrain_dataset on index
+        prediction_data = pd.merge(
+            retrain_dataset, prediction_data, left_index=True, right_index=True, how='left')
+
+        # Rename column "target_name" to "target"
+        prediction_data.rename(columns={target_name: 'target'}, inplace=True)
+
+        prediction_data['timestamp'] = prediction_data.index
+
+        prediction_data.reset_index(drop=True, inplace=True)
+
+        prediction_data.sort_values(
+            by='timestamp', ascending=True, inplace=True
+        )
+
+        return prediction_data
+
     async def fit_models(
         self,
         model_name: str,
@@ -689,7 +748,7 @@ class MLFlowRepository(SientiaMonitoring):
         transform_flavor: str = 'sklearn',
         predict_flavor: str = 'sklearn',
         target_name: str | None = None,
-    ) -> dict[str, dict[str, Any]]:
+    ) -> dict[str, Any]:
         """
         Prepare models and data for a retraining run.
 
@@ -799,6 +858,10 @@ class MLFlowRepository(SientiaMonitoring):
 
         prediction_model.fit(retrain_dataset)
 
+        # get prediction data
+        prediction_data = self.get_prediction_data(
+            prediction_model, retrain_dataset, target_name)
+
         self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
 
         retrain_data = {
@@ -807,6 +870,7 @@ class MLFlowRepository(SientiaMonitoring):
                 'artifact_path': prediction_artifact_path,
             },
             'data_model': {'model': data_model, 'artifact_path': data_artifact_path},
+            'prediction_data': prediction_data,
         }
         return retrain_data
 
@@ -885,6 +949,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         prediction_model = retrain_data['prediction_model']
         data_model = retrain_data['data_model']
+        prediction_data = retrain_data['prediction_data']
 
         model_temp_path = path.join(ARTIFACTS_PATH, model_name)
 
@@ -908,10 +973,12 @@ class MLFlowRepository(SientiaMonitoring):
         self.debug(f'Attributes: {retrain_params}', metadata)
 
         data_path = f'{model_temp_path}/retrain_data.csv'
+        prediction_data_path = f'{model_temp_path}/evaluation_data.csv'
 
         makedirs(model_temp_path, exist_ok=True)
 
-        data.to_csv(data_path, index=True)
+        data.to_csv(data_path, index=False)
+        prediction_data.to_csv(prediction_data_path, index=False)
 
         self.info(
             f'Starting model upload for {experiment_name} with run name {current_run_name}',
@@ -945,6 +1012,7 @@ class MLFlowRepository(SientiaMonitoring):
 
                 # log the data raw
                 mlflow.log_artifact(data_path)
+                mlflow.log_artifact(prediction_data_path)
 
         except Exception as e:
             await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)
@@ -1372,59 +1440,3 @@ class MLFlowRepository(SientiaMonitoring):
 
         return metadata_result
 
-    async def get_model_metrics(self, 
-        analyse_data: pd.DataFrame,
-        train_reference_data: pd.DataFrame,
-        test_reference_data: pd.DataFrame,
-        target_name: str,
-        drift_metrics: list[str],
-        performance_metrics: list[str],
-        metadata: dict) -> dict[str, Any]:
-
-        """
-        Calculates drift and preformance metrics for a model.
-
-        Args:
-            analyse_data (pd.DataFrame): The data to analyse.
-            train_reference_data (pd.DataFrame): The train reference data.
-            test_reference_data (pd.DataFrame): The test reference data.
-            target_name (str): The target name.
-            drift_metrics (list[str]): The drift metrics.
-            performance_metrics (list[str]): The performance metrics.
-            metadata (dict): The metadata.
-
-        Returns:
-            dict: The metrics.
-        """
-        columns = train_reference_data.drop(
-            columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore').columns
-
-        config = {
-            "target": target_name,
-            "prediction": "prediction",
-            "timestamp": "timestamp",
-            "columns": columns,
-        }
-
-        model_analysis = ModelAnalysis(config=config)
-
-        core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
-
-        start_time = time.time()
-        try:
-            model_analysis.detect_univariate_drift(
-                reference_df=train_reference_data,
-                analyse_df=analyse_data,
-                features=columns,
-                timestamp_column=config['timestamp'],
-                metrics=drift_metrics,
-                chunk_period="s"
-            )
-        except Exception as e:
-            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)
-
-        return {}

laborious/worker/worker.py

@@ -29,6 +29,7 @@ from temporalio import client, workflow
 from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 from temporalio.worker import PollerBehaviorAutoscaling, Worker
 
+
 with workflow.unsafe.imports_passed_through():
     import asyncio
     import os
@@ -49,6 +50,7 @@ with workflow.unsafe.imports_passed_through():
     )
     from laborious.workflows.minimal_retrain import MinimalRetrain
     from laborious.workflows.predictions_batch import PredictionsBatch
+    from laborious.workflows.drift import Drift
     from laborious.workflows.sub_workflows.format_and_export_prediction import (
         FormatAndExportPrediction,
     )
@@ -156,6 +158,23 @@ async def main():
             workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
             activity_task_poller_behavior=PollerBehaviorAutoscaling(),
         ),
+        Worker(
+            temporal_client,
+            task_queue='drift-queue',
+            workflows=[Drift],
+            activities=[
+                activities.load_custom_query,
+                activities.get_reference_data,
+                activities.calculate_drift,
+                activities.export_data_to_postgres,
+            ],
+            max_concurrent_workflow_tasks=50,
+            max_concurrent_activities=50,
+            max_concurrent_local_activities=50,
+            max_cached_workflows=2,
+            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
+            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
+        ),
         Worker(
             temporal_client,
             task_queue='predictions_batch-queue',
@@ -168,6 +187,7 @@ async def main():
                 activities.input_gate,
                 activities.mlflow_response_gate,
                 activities.mlflow_content_gate,
+                activities.format_transformed_data,
                 activities.format_prediction,
                 activities.format_default_prediction,
                 activities.get_last_timestamp,

laborious/workflows/drift.py

@@ -0,0 +1,99 @@
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
+    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+
+
+@workflow.defn(name='drift')
+class Drift:
+    @workflow.run
+    async def run(self, input_data: dict[str, Any]):
+        """
+        Execute the drift workflow.
+
+        This method orchestrates the complete drift process by:
+        1. Loading data using the provided custom SQL query
+        2. Preparing prediction configuration and filters
+        3. Delegating to the PredictionProcess workflow for ML operations
+        """
+
+        metadata = {
+            'metadata': {
+                'schedule_name': input_data['schedule_name'],
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id'],
+                'workflow_name': 'drift',
+            }
+        }
+
+        gathering_query = f"""
+            SELECT *
+            FROM {input_data['schema']}.{input_data['source_table_name']}
+            WHERE
+                model_id = {input_data['model_id']} AND
+                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
+            ORDER BY timestamp ASC
+        """
+
+        target_data_handler = workflow.start_local_activity_method(
+            Activities.load_custom_query,
+            {
+                **metadata,
+                'query': gathering_query,
+                'datetime_columns': ['timestamp', 'created_at'],
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        reference_data_handler = workflow.start_local_activity_method(
+            Activities.get_reference_data,
+            {
+                **metadata,
+                'model_name': input_data['model_name']
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        target_data = await target_data_handler
+        reference_data = await reference_data_handler
+
+        if not target_data:
+            return
+
+        drift_data = await workflow.execute_local_activity_method(
+            Activities.calculate_drift,
+            {
+                **metadata,
+                'target_data': target_data,
+                'reference_data': reference_data,
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id'],
+                'target_name': input_data['target_name'],
+                'drift_metrics': input_data['drift_metrics'],
+                'chunk_period': input_data.get('chunk_period', 'min'),
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        if drift_data:
+            await workflow.execute_activity_method(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': drift_data,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['target_table_name'],
+                    'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+                },
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=300),
+            )

laborious/workflows/predictions_batch.py

@@ -98,6 +98,7 @@ class PredictionsBatch:
             'data': data,
             'schema': input_data['schema'],
             'table_name': input_data['table_name'],
+            'transform_table_name': input_data['transform_table_name'],
             'model_id': input_data['model_id'],
             'model_name': input_data['model_name'],
             'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -95,12 +95,13 @@ class FormatAndExportPrediction:
                     {
                         **metadata,
                         'data': transformed_data,
+                        'model_id': input_data['model_id'],
                     },
                     retry_policy=retry_policy,
                     start_to_close_timeout=timedelta(seconds=60),
                 )
 
-                write_transformed_handler = workflow.execute_activity_method(
+                write_transformed_handler = workflow.start_activity_method(
                     Activities.export_data_to_postgres,
                     {
                         **metadata,

laborious/workflows/sub_workflows/prediction_process.py

@@ -209,6 +209,7 @@ class PredictionProcess:
                 'opc_output_config': input_data['opc_output_config'],
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
+                'transform_table_name': input_data['transform_table_name'],
                 'comment': comment,
                 'prediction_store_policy': input_data['prediction_store_policy'],
             },
@@ -250,6 +251,7 @@ class PredictionProcess:
 
         schema = input_data['schema']
         table_name = input_data['table_name']
+        transform_table_name = input_data['transform_table_name']
         model_id = input_data['model_id']
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
@@ -289,6 +291,7 @@ class PredictionProcess:
                     'model_config': model_config,
                     'schema': schema,
                     'table_name': table_name,
+                    'transform_table_name': transform_table_name,
                     'comment': comment,
                     'opc_output_config': input_data['opc_output_config'],
                     'prediction_store_policy': input_data['prediction_store_policy'],