SIENTIAPDE-1314
Enhance thread safety in MLFlowRepository model caching - Introduced a reentrant lock to synchronize access to the model cache, ensuring thread safety during cache checks and updates. - Updated the cache management logic to acquire the lock when checking for existing models and when updating the cache after downloading a new model. - Reduced the maximum cached workflows in the worker configuration for improved resource management.
This commit is contained in:
@@ -16,6 +16,7 @@ Capabilities:
|
||||
|
||||
import ctypes
|
||||
import gc
|
||||
import threading
|
||||
import traceback
|
||||
from datetime import datetime, timedelta
|
||||
from os import environ, makedirs, path
|
||||
@@ -74,6 +75,7 @@ class MLFlowRepository:
|
||||
self.client = mlflow.tracking.MlflowClient()
|
||||
|
||||
self.model_cache: dict[str, Any] = {}
|
||||
self._cache_lock = threading.RLock()
|
||||
self.logger = logger
|
||||
|
||||
"""
|
||||
@@ -466,28 +468,31 @@ class MLFlowRepository:
|
||||
|
||||
model_key = f'{model_name}_{model_type}'
|
||||
|
||||
if model_key in self.model_cache:
|
||||
cache = self.model_cache[model_key]
|
||||
# Acquire lock to check cache
|
||||
with self._cache_lock:
|
||||
if model_key in self.model_cache:
|
||||
cache = self.model_cache[model_key]
|
||||
|
||||
# Check if config has changed or is outdated
|
||||
if self.check_cache_retention(cache, retention):
|
||||
return self.handle_valid_model(model_name=model_name, cache=cache)
|
||||
# Check if config has changed or is outdated
|
||||
if self.check_cache_retention(cache, retention):
|
||||
return self.handle_valid_model(model_name=model_name, cache=cache)
|
||||
else:
|
||||
# Model is outdated, delete old model files
|
||||
self.handle_outdated_model(model_name=model_name, model_key=model_key)
|
||||
else:
|
||||
# Model is outdated, delete old model files
|
||||
self.handle_outdated_model(model_name=model_name, model_key=model_key)
|
||||
else:
|
||||
self.logger.debug(
|
||||
f'Model {model_name} is not in {model_type} cache, downloading a new one'
|
||||
)
|
||||
self.logger.debug(
|
||||
f'Model {model_name} is not in {model_type} cache, downloading a new one'
|
||||
)
|
||||
|
||||
# Donwload new model
|
||||
# Donwload new model (without lock to avoid blocking other threads)
|
||||
model, _artifact_path = self.download_model(
|
||||
model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False
|
||||
)
|
||||
|
||||
cache = {'target': model, 'timestamp': datetime.now()}
|
||||
|
||||
self.model_cache[model_key] = cache
|
||||
# Update cache with lock
|
||||
with self._cache_lock:
|
||||
cache = {'target': model, 'timestamp': datetime.now()}
|
||||
self.model_cache[model_key] = cache
|
||||
|
||||
return model
|
||||
|
||||
|
||||
Reference in New Issue
Block a user