SIENTIAPDE-1222

Refactor datetime index handling in MLFlow and MLFlowRepository

- Moved the detect_and_parse_datetime_index method from MLFlow to MLFlowRepository for better organization and reusability.
- Updated the method to include enhanced logging and error handling for invalid datetime formats.
- Adjusted the transform method in MLFlowRepository to utilize the new datetime index parsing logic.
- Added unit tests for both valid and invalid datetime index cases to ensure robustness.

laborious/utils/repository/model_repository.py

@@ -16,17 +16,57 @@ import pandas as pd
 import mlflow
 from os import makedirs, path, remove
 from sientia.ModelServing import ModelServing
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+from sientia_do.observability.logger import Logger
 
 
 class MLFlowRepository():
-    def __init__(self, host, username, password, logger):
+    def __init__(self, host, username, password, logger: Logger):
 
         self.model_serving = ModelServing(tracking_uri=host,
                                           username=username, password=password,
                                           logger=logger)
         self.logger = logger
 
-    def transform(self, model_name: str, data: pd.DataFrame, model_config: dict) -> dict:
+    def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
+        """
+        Detect and parse datetime index from data. index must be a timestamp like column.
+        This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
+        If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
+        If another type or format, must raise an error.
+        """
+        index = data.index
+
+        # Get type of first element of index
+        index_type = type(index[0])
+
+        self.logger.custom_info(f"Index type: {index_type}", metadata)
+
+        message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}"
+
+        # Check if all in index are of the same type
+        if not all(isinstance(i, index_type) for i in index):
+            raise ValueError(
+                f"{message}")
+
+        # Check type and converts to DATETIME_FORMAT_WITH_TZ
+        if index_type == str:
+            # Validate format of string and return error if not valid
+            try:
+                pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
+            except ValueError:
+                raise ValueError(
+                    f"{message}")
+
+        elif index_type == datetime or index_type == pd.Timestamp:
+            data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
+        else:
+            raise ValueError(
+                f"{message}")
+
+        return data
+
+    def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
         """
         Transform data using a model.
 
@@ -40,6 +80,9 @@ class MLFlowRepository():
         """
 
         try:
+            self.logger.custom_debug(
+                f"Data received for model transformation: {data.to_csv()}", metadata)
+
             model_retention = model_config.get('retention_minutes', 0)
             flavor = model_config.get('transform_flavor', 'sklearn')
             compressed = model_config.get('is_compressed', False)
@@ -47,12 +90,20 @@ class MLFlowRepository():
             transform_keyword = model_config.get(
                 'transform_function_keyword', 'predict')
 
+            transformed_data = self.model_serving.get_cached_transform(
+                model_name, data, model_retention, flavor,
+                compressed, retention_target, transform_keyword
+            )
+
+            self.logger.custom_debug(
+                f"Data received from model transformation: {transformed_data.to_csv()}", metadata)
+
+            transformed_data = self.detect_and_parse_datetime_index(
+                transformed_data, metadata)
+
             return {
                 'success': True,
-                'content': self.model_serving.get_cached_transform(
-                    model_name, data, model_retention, flavor,
-                    compressed, retention_target, transform_keyword
-                ).to_dict()
+                'content': transformed_data.to_dict()
             }
 
         except Exception as e:
@@ -64,7 +115,7 @@ class MLFlowRepository():
                 }
             }
 
-    def predict(self, model_name: str, data: pd.DataFrame, model_config: dict) -> dict:
+    def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
         """
         Predict data using a model.
 
@@ -85,8 +136,8 @@ class MLFlowRepository():
             input_index = data.index
             start_time = datetime.now()
 
-            self.logger.debug(
-                f"Data received for model prediction: {data.to_string()}")
+            self.logger.custom_debug(
+                f"Data received for model prediction: {data.to_csv()}", metadata)
             data = self.model_serving.get_cached_predict(
                 model_name, data, model_retention, flavor,
                 compressed, retention_target
@@ -94,8 +145,8 @@ class MLFlowRepository():
 
             end_time = datetime.now()
             data = pd.DataFrame(data, columns=['prediction'])
-            self.logger.debug(
-                f"Data received from model prediction: {data.to_string()}")
+            self.logger.custom_debug(
+                f"Data received from model prediction: {data.to_csv()}", metadata)
             data.index = input_index
             data['response_time'] = (end_time - start_time).total_seconds()