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.
This commit is contained in:
vitor-aignosi
2026-03-31 15:05:40 -03:00
parent d9ec7fc496
commit a3775279e1

View File

@@ -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)