From a3775279e12e5f3c774e8338280eb0a8fa21bdcf Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 31 Mar 2026 15:05:40 -0300 Subject: [PATCH] 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. --- .../utils/repository/model_repository.py | 70 ++++++------------- 1 file changed, 21 insertions(+), 49 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index fb404b7..337eda2 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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.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) + self.debug(f"Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}", metadata) 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( - 'Data received from model transformation:', transformed_data, metadata + self.debug( + 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)