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

SIENTIAPDE-1231

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2025-10-13 15:31:04 -03:00
parent fdaf48ddc9
commit 331df2dd18
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -27,7 +27,7 @@ import gzip
import pickle import pickle
from numpy import ndarray from numpy import ndarray
from typing import Any from typing import Any
import gc
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
ARTIFACTS_PATH = "./tmp/artifacts" ARTIFACTS_PATH = "./tmp/artifacts"
@@ -613,8 +613,12 @@ class MLFlowRepository():
prediction = model.predict(data) prediction = model.predict(data)
if retention == 0: if retention == 0:
self.logger.info(
f"Deleting model {model_name} from memory")
del model del model
gc.collect()
return prediction return prediction
def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int, def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int,
@@ -638,8 +642,12 @@ class MLFlowRepository():
prediction = model.predict(data) prediction = model.predict(data)
if retention == 0: if retention == 0:
self.logger.info(
f"Deleting model {model_name} from memory")
del model del model
gc.collect()
return prediction return prediction
""" """