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/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'],