gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1231

Browse Source 
Enhance MLFlow and MLFlowRepository with improved data handling and logging

- Refactored MLFlow class to sort data by 'created_at' and drop duplicates for better input preparation.
- Updated MLFlowRepository methods to include detailed logging for artifact downloads and model predictions.
- Introduced LzmaPayloadCodec for efficient payload compression in the worker, optimizing data handling for large payloads.
- Enhanced timestamp handling in treated data to ensure compatibility with model expectations.
This commit is contained in:
vitor-aignosi
2025-10-06 13:32:36 -03:00
parent 9b71ad7556
commit 36af84f056
3 changed files with 123 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
laborious/utils/repository/model_repository.py
View File

@@ -235,6 +235,9 @@ class MLFlowRepository():
if not path.exists(output_dir):
makedirs(output_dir)
self.logger.info(
f"Downloading artifacts from {run_id} to {output_dir}")
return self.client.download_artifacts(
run_id,
artifact_path,
@@ -347,18 +350,13 @@ class MLFlowRepository():
Raises:
ValueError: If model cannot be loaded with any compression method
"""
code_path = path.join(
artifact_path, "code")
if type == "transformer":
code_path = path.join(
artifact_path, "transformer_pyfunc", "code")
pickle_path = path.join(
artifact_path, "transformer_pyfunc", "artifacts", "training_transformer.pkl")
pickle_file = "training_transformer.pkl" if type == "transformer" else "stacking_model.pkl"
elif type == "prediction":
code_path = path.join(
artifact_path, "stacking_model", "code")
pickle_path = path.join(
artifact_path, "stacking_model", "artifacts", "stacking_model.pkl")
pickle_path = path.join(
artifact_path, "artifacts", pickle_file)
if code_path not in sys_path:
sys_path.insert(0, code_path)
@@ -468,6 +466,9 @@ class MLFlowRepository():
f"{message}")
elif index_type == datetime or index_type == pd.Timestamp:
if data.index.tz is None:
data.index = data.index.tz_localize('UTC')
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
else:
raise ValueError(
@@ -690,6 +691,9 @@ class MLFlowRepository():
self.logger.custom_debug(
f"Fitting transformation model with training data (shape: {data.shape})", metadata)
data.to_csv(
f"tmp/retrain_data_{model_name}.csv", index=True)
data_model = data_model.fit(data)
self.logger.custom_debug(
@@ -697,6 +701,20 @@ class MLFlowRepository():
treated_data = data_model.predict(data)
# Stores current index as timestamp, Courier model expects a timestamp column
# with specific format
treated_data['timestamp'] = treated_data.index
# Parses timestamp column to datetime format to align with data
treated_data = self.detect_and_parse_datetime_index(
treated_data, metadata)
self.logger.custom_debug(
f"Treated data index: {treated_data.index}", metadata)
treated_data.to_csv(
f"tmp/retrain_treated_data_{model_name}.csv", index=True)
self.logger.custom_debug(
f"Transformed data shape: {treated_data.shape}", metadata)
@@ -708,12 +726,24 @@ class MLFlowRepository():
self.logger.custom_debug(
f"Using provided target variable: {target_name}", metadata)
# Check if treated_data contains target variable
if target_name not in treated_data.columns:
self.logger.custom_debug(
f"Target variable {target_name} not found in treated data, aligning data with treated data indexes", metadata)
# Aligns data with treated data indexes to get target variable
aligned_data = data.loc[treated_data.index]
else:
# Uses target variable from treated data
self.logger.custom_debug(
f"Target variable {target_name} found in treated data, using it", metadata)
aligned_data = treated_data
if fit_config.get('y_type', 'series').lower() == 'series':
y = data[target_name]
y = aligned_data[target_name]
self.logger.custom_debug(
f"Extracting target as series, shape: {y.shape}", metadata)
else:
y = data[[target_name]]
y = aligned_data[[target_name]]
self.logger.custom_debug(
f"Extracting target as dataframe, shape: {y.shape}", metadata)
@@ -726,22 +756,25 @@ class MLFlowRepository():
self.logger.custom_debug(
f"Combined data shape for prediction model fit: {treated_data.shape}", metadata)
treated_data.to_csv(
f"tmp/retrain_treated_data_combined_{model_name}.csv", index=True)
prediction_model = prediction_model.fit(treated_data)
self.logger.custom_debug(
"Prediction model fitted with combined data", metadata)
else:
# Keep data separated
fit_order = fit_config.get('split_fit_first', 'x').lower()
y.to_csv(
f"tmp/retrain_y_{model_name}.csv", index=True)
self.logger.custom_debug(
f"Fitting prediction model with separated data, order: {fit_order}", metadata)
f"Fitting prediction model with separated data, first arg: {fit_order}", metadata)
if fit_order == 'x':
prediction_model = prediction_model.fit(treated_data, y)
self.logger.custom_debug(
"Prediction model fitted with X, y order", metadata)
else:
prediction_model = prediction_model.fit(y, treated_data)
self.logger.custom_debug(
"Prediction model fitted with y, X order", metadata)
experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment)
@@ -1012,30 +1045,32 @@ class MLFlowRepository():
data.to_csv(
f"tmp/treated_data_{model_name}.csv", index=True)
data = self.get_cached_predict(
predict_data = self.get_cached_predict(
model_name, data, model_retention, flavor)
end_time = datetime.now()
if isinstance(data, pd.DataFrame):
if isinstance(predict_data, pd.DataFrame):
self.logger.custom_debug(
f"Data received from model prediction: {data.head(5).to_csv()}", metadata)
data.to_csv(
predict_data.to_csv(
f"tmp/predicted_data_{model_name}.csv", index=True)
data.columns = ['prediction']
else:
data = pd.DataFrame(data, columns=['prediction'])
data.to_csv(
predict_data = pd.DataFrame(
predict_data, columns=['prediction'])
predict_data.to_csv(
f"tmp/predicted_data_{model_name}.csv", index=True)
data.index = input_index
data['response_time'] = (end_time - start_time).total_seconds()
predict_data.index = input_index
predict_data['response_time'] = (
end_time - start_time).total_seconds()
return {
'success': True,
'content': data.to_dict()
'content': predict_data.to_dict()
}
except Exception as e: