SIENTIAPDE-1712

SIENTIAPDE-1712 Refactor logging in MLFlowRepository class by removing the debug_dataframe method and replacing it with direct debug statements for improved clarity. This change enhances the logging of DataFrame content during model transformation, prediction, and retraining processes, ensuring better observability without excessive log output.

laborious/utils/repository/model_repository.py

@@ -36,7 +36,6 @@ from sientia_do.observability.sientia_monitoring import SientiaMonitoring
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 from laborious import metrics
-from laborious.utils.dataframe_debug import build_dataframe_debug_message
 
 ARTIFACTS_PATH = './tmp/artifacts'
 TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl'
@@ -65,8 +64,6 @@ def force_memory_release(logger: Logger):
 
 
 class MLFlowRepository(SientiaMonitoring):
-    _MAX_DEBUG_DATAFRAME_ROWS = 100
-
     def __init__(
         self,
         host: str,
@@ -97,24 +94,6 @@ class MLFlowRepository(SientiaMonitoring):
         self._cache_lock = threading.RLock()
         self.logger = logger
 
-    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 object expected to expose shape and to_csv
-         - metadata (dict[str, Any]): Metadata for contextual logging
-        """
-        self.debug(
-            build_dataframe_debug_message(
-                message=message,
-                data=data,
-                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
-            ),
-            metadata,
-        )
-
     """
     Functions related to get model registry parameters
     """
@@ -459,26 +438,17 @@ class MLFlowRepository(SientiaMonitoring):
             raise ValueError("Invalid model_type. Use 'predict' or 'transform'.")
 
         artifact_path = None
+        if model_type == 'predict':
+            model = await self.load_predict_model(model_name, metadata, flavor)
 
+        else:
+            model = await self.load_transform_model(model_name, metadata, flavor)
         if load_wrapper:
             self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}')
 
-            target = 'prediction_model' if model_type == 'predict' else 'data_model'
+            model = model._model_impl.python_model
 
-            artifact_path = await self.dowload_artifacts(model_name, metadata, target)
+            self.debug(f"Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}", metadata)
-
-            self.info(
-                f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
-            )
-
-            raw_model = mlflow.pyfunc.load_model(artifact_path)
-            model = raw_model._model_impl.python_model
-        else:
-            if model_type == 'predict':
-                model = await self.load_predict_model(model_name, metadata, flavor)
-
-            else:
-                model = await self.load_transform_model(model_name, metadata, flavor)
 
         return model, artifact_path
 
@@ -1155,7 +1125,7 @@ class MLFlowRepository(SientiaMonitoring):
 
     async def transform(
         self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
-    ) -> dict[str, Any]:
+    ):
         """
         Transform data using a cached transformation model.
 
@@ -1189,7 +1159,7 @@ class MLFlowRepository(SientiaMonitoring):
                       and returned in the response structure rather than propagated.
         """
 
-        self._debug_dataframe('Data received for model transformation:', data, metadata)
+        self.debug(f'Data received for model transformation: {data.to_csv()}', metadata)
 
         # data.to_csv(
         #     f"tmp/data_{model_name}.csv", index=True)
@@ -1207,8 +1177,9 @@ class MLFlowRepository(SientiaMonitoring):
                 metadata=metadata,
             )
 
-            self._debug_dataframe(
+            self.debug(
-                'Data received from model transformation:', transformed_data, metadata
+                f'Data received from model transformation: {transformed_data.head(5).to_csv()}',
+                metadata,
             )
 
             # transformed_data.to_csv(
@@ -1216,7 +1187,7 @@ class MLFlowRepository(SientiaMonitoring):
 
             transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata)
 
-            return {'success': True, 'content': transformed_data}
+            return {'success': True, 'content': transformed_data.to_dict()}
 
         except Exception as e:
             return {
@@ -1274,7 +1245,9 @@ class MLFlowRepository(SientiaMonitoring):
             input_index = data.index
             start_time = datetime.now()
 
-            self._debug_dataframe('Data received for model prediction:', data, metadata)
+            self.debug(
+                f'Data received for model prediction: {data.to_dict(orient="records")}', metadata
+            )
 
             # data.to_csv(
             #     f"tmp/treated_data_{model_name}.csv", index=True)
@@ -1291,17 +1264,16 @@ class MLFlowRepository(SientiaMonitoring):
             end_time = datetime.now()
 
             if isinstance(predict_data, pd.DataFrame):
-                self._debug_dataframe('Data received from model prediction:', predict_data, metadata)
+                self.debug(
+                    f'Data received from model prediction: {predict_data.to_dict(orient="records")}',
+                    metadata,
+                )
 
                 # predict_data.to_csv(
                 #     f"tmp/predicted_data_{model_name}.csv", index=True)
                 predict_data.columns = pd.Index(['prediction'])
 
             else:
-                self.debug(
-                    f'Data received from model prediction (not a DataFrame): {predict_data}',
-                    metadata,
-                )
                 predict_data = pd.DataFrame(predict_data, columns=['prediction'])
                 # predict_data.to_csv(
                 #     f"tmp/predicted_data_{model_name}.csv", index=True)
@@ -1309,7 +1281,7 @@ class MLFlowRepository(SientiaMonitoring):
             predict_data.index = input_index
             predict_data['response_time'] = (end_time - start_time).total_seconds()
 
-            return {'success': True, 'content': predict_data}
+            return {'success': True, 'content': predict_data.to_dict()}
 
         except Exception as e:
             return {
@@ -1363,7 +1335,7 @@ class MLFlowRepository(SientiaMonitoring):
         """
 
         self.info(f'Starting model retraining workflow for {model_name}', metadata)
-        self._debug_dataframe('Data received for model retraining:', data, metadata)
+        self.debug(f'Data received for model retraining: {data.to_csv()}', metadata)
 
         target_name = model_config.get('target', None)