From a3989b84af54500dc80a9dfea979347ae69ce3fe Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 13 Oct 2025 15:57:23 -0300 Subject: [PATCH] SIENTIAPDE-1231 Refactor memory management in MLFlowRepository to use `force_memory_release` function - Replaced direct calls to `gc.collect()` with `force_memory_release()` for improved memory optimization after model deletion. - This change enhances memory management during model operations, ensuring more efficient resource handling. --- laborious/utils/repository/model_repository.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 904e271..25345ce 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -36,13 +36,16 @@ TRANSFORMED_COMPRESSED_PATH = "artifacts/training_transformer.pkl" PREDICTION_COMPRESSED_PATH = "artifacts/stacking_model.pkl" -def force_memory_release(): +def force_memory_release(logger: Logger, metadata: dict): gc.collect() try: ctypes.CDLL("libc.so.6").malloc_trim(0) - except: - pass + logger.custom_info( + f"Memory released", metadata) + except Exception as e: + logger.custom_info( + f"Memory release failed: {e}", metadata) class MLFlowRepository(): @@ -627,7 +630,7 @@ class MLFlowRepository(): f"Deleting model {model_name} from memory") del model - gc.collect() + force_memory_release() return prediction @@ -656,7 +659,7 @@ class MLFlowRepository(): f"Deleting model {model_name} from memory") del model - gc.collect() + force_memory_release() return prediction