From 331df2dd18d6c5c4726c806be4a62a93b0efb553 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 13 Oct 2025 15:31:04 -0300 Subject: [PATCH] SIENTIAPDE-1231 Enhance MLFlowRepository memory management by adding garbage collection and logging for model deletion - Introduced garbage collection after model deletion to optimize memory usage. - Added logging to inform when a model is deleted from memory, improving traceability during predictions. --- laborious/utils/repository/model_repository.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 33611bb..11944cc 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -27,7 +27,7 @@ import gzip import pickle from numpy import ndarray from typing import Any - +import gc from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ ARTIFACTS_PATH = "./tmp/artifacts" @@ -613,8 +613,12 @@ class MLFlowRepository(): prediction = model.predict(data) if retention == 0: + self.logger.info( + f"Deleting model {model_name} from memory") del model + gc.collect() + return prediction def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int, @@ -638,8 +642,12 @@ class MLFlowRepository(): prediction = model.predict(data) if retention == 0: + self.logger.info( + f"Deleting model {model_name} from memory") del model + gc.collect() + return prediction """