From aed0282987582262d04f633f823accbbf09e0374 Mon Sep 17 00:00:00 2001 From: Akshith Ambekar Date: Mon, 29 Sep 2025 18:51:17 -0400 Subject: [PATCH] lru cache for faster custom params usage --- tueri/transformers_helpers.py | 218 ++++++++++++++++++++++++++------- tueri_api/app/app.py | 4 - tueri_api/app/scanner.py | 78 +++++++++--- tueri_api/app/scanner_cache.py | 109 +++++++++++++++++ 4 files changed, 342 insertions(+), 67 deletions(-) create mode 100644 tueri_api/app/scanner_cache.py diff --git a/tueri/transformers_helpers.py b/tueri/transformers_helpers.py index 1cf58632..f7c07cad 100644 --- a/tueri/transformers_helpers.py +++ b/tueri/transformers_helpers.py @@ -1,6 +1,8 @@ from __future__ import annotations +import hashlib import importlib +import json from functools import lru_cache from typing import Literal, get_args @@ -11,6 +13,54 @@ LOGGER = get_logger() +def _create_model_cache_key(model: Model, use_onnx: bool = False) -> str: + """ + Create a cache key for a model based on its configuration. + + Args: + model (Model): The model to create a cache key for. + use_onnx (bool): Whether ONNX is being used. + + Returns: + str: A unique cache key for the model configuration. + """ + # Convert any object to a JSON-serializable format + def make_serializable(obj): + if obj is None: + return None + elif isinstance(obj, (str, int, float, bool)): + return obj + elif isinstance(obj, dict): + return {k: make_serializable(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return [make_serializable(item) for item in obj] + else: + # Convert any complex object to string representation + return str(obj) + + model_config = { + "path": model.path, + "subfolder": model.subfolder, + "revision": model.revision, + "onnx_path": model.onnx_path if use_onnx else None, + "onnx_revision": model.onnx_revision if use_onnx else None, + "onnx_subfolder": model.onnx_subfolder if use_onnx else None, + "onnx_filename": model.onnx_filename if use_onnx else None, + "kwargs": make_serializable(model.kwargs), + "pipeline_kwargs": make_serializable(model.pipeline_kwargs), + "tokenizer_kwargs": make_serializable(model.tokenizer_kwargs), + "use_onnx": use_onnx, + "device": str(device()), + } + + # Create a deterministic hash of the configuration + config_str = json.dumps(model_config, sort_keys=True) + return hashlib.md5(config_str.encode()).hexdigest() + + +# Cache for tokenizers +_tokenizer_cache = {} + def get_tokenizer(model: Model): """ This function loads a tokenizer given a model identifier and caches it. @@ -19,10 +69,25 @@ def get_tokenizer(model: Model): Args: model (Model): The model to load the tokenizer for. """ + cache_key = _create_model_cache_key(model, use_onnx=False) + + if cache_key in _tokenizer_cache: + LOGGER.debug("Using cached tokenizer", model=model) + return _tokenizer_cache[cache_key] + transformers = lazy_load_dep("transformers") tokenizer = transformers.AutoTokenizer.from_pretrained( model.path, revision=model.revision, **model.tokenizer_kwargs ) + + # Cache the tokenizer (limit cache size to prevent memory issues) + if len(_tokenizer_cache) >= 32: + # Remove oldest entry (simple FIFO eviction) + oldest_key = next(iter(_tokenizer_cache)) + del _tokenizer_cache[oldest_key] + + _tokenizer_cache[cache_key] = tokenizer + LOGGER.debug("Cached new tokenizer", model=model) return tokenizer @@ -40,9 +105,18 @@ def is_onnx_supported() -> bool: return is_supported +# Cache for ONNX models +_onnx_model_cache = {} + def _ort_model_for_sequence_classification( model: Model, ): + cache_key = _create_model_cache_key(model, use_onnx=True) + + if cache_key in _onnx_model_cache: + LOGGER.debug("Using cached ONNX classification model", model=model) + return _onnx_model_cache[cache_key] + provider = "CPUExecutionProvider" package_name = "optimum[onnxruntime]" if device().type == "cuda": @@ -60,11 +134,22 @@ def _ort_model_for_sequence_classification( provider=provider, **model.kwargs, ) - LOGGER.debug("Initialized classification ONNX model", model=model, device=device()) + + # Cache the model (limit cache size to prevent memory issues) + if len(_onnx_model_cache) >= 16: + # Remove oldest entry (simple FIFO eviction) + oldest_key = next(iter(_onnx_model_cache)) + del _onnx_model_cache[oldest_key] + + _onnx_model_cache[cache_key] = tf_model + LOGGER.debug("Cached new ONNX classification model", model=model, device=device()) return tf_model +# Cache for PyTorch classification models +_pytorch_classification_cache = {} + def get_tokenizer_and_model_for_classification( model: Model, use_onnx: bool = False, @@ -78,27 +163,41 @@ def get_tokenizer_and_model_for_classification( use_onnx (bool): Whether to use the ONNX version of the model. Defaults to False. """ tf_tokenizer = get_tokenizer(model) - transformers = lazy_load_dep("transformers") if use_onnx and is_onnx_supported() is False: LOGGER.warning("ONNX is not supported on this machine. Using PyTorch instead of ONNX.") use_onnx = False if use_onnx is False: - tf_model = transformers.AutoModelForSequenceClassification.from_pretrained( - model.path, - subfolder=model.subfolder, - revision=model.revision, - torch_dtype="auto", - low_cpu_mem_usage=False, - **model.kwargs, - ) - # Handle meta device properly - use to_empty() when moving from meta device - if hasattr(tf_model, 'device') and str(tf_model.device) == 'meta': - tf_model = tf_model.to_empty(device=device()) + cache_key = _create_model_cache_key(model, use_onnx=False) + + if cache_key in _pytorch_classification_cache: + LOGGER.debug("Using cached PyTorch classification model", model=model) + tf_model = _pytorch_classification_cache[cache_key] else: - tf_model = tf_model.to(device()) - LOGGER.debug("Initialized classification model", model=model, device=device()) + transformers = lazy_load_dep("transformers") + tf_model = transformers.AutoModelForSequenceClassification.from_pretrained( + model.path, + subfolder=model.subfolder, + revision=model.revision, + torch_dtype="auto", + low_cpu_mem_usage=False, + **model.kwargs, + ) + # Handle meta device properly - use to_empty() when moving from meta device + if hasattr(tf_model, 'device') and str(tf_model.device) == 'meta': + tf_model = tf_model.to_empty(device=device()) + else: + tf_model = tf_model.to(device()) + + # Cache the model (limit cache size to prevent memory issues) + if len(_pytorch_classification_cache) >= 16: + # Remove oldest entry (simple FIFO eviction) + oldest_key = next(iter(_pytorch_classification_cache)) + del _pytorch_classification_cache[oldest_key] + + _pytorch_classification_cache[cache_key] = tf_model + LOGGER.debug("Cached new PyTorch classification model", model=model, device=device()) return tf_tokenizer, tf_model @@ -107,6 +206,11 @@ def get_tokenizer_and_model_for_classification( return tf_tokenizer, tf_model +# Cache for PyTorch NER models +_pytorch_ner_cache = {} +# Cache for ONNX NER models +_onnx_ner_cache = {} + def get_tokenizer_and_model_for_ner( model: Model, use_onnx: bool = False, @@ -120,45 +224,73 @@ def get_tokenizer_and_model_for_ner( use_onnx (bool): Whether to use the ONNX version of the model. Defaults to False. """ tf_tokenizer = get_tokenizer(model) - transformers = lazy_load_dep("transformers") if use_onnx and is_onnx_supported() is False: LOGGER.warning("ONNX is not supported on this machine. Using PyTorch instead of ONNX.") use_onnx = False if use_onnx is False: - tf_model = transformers.AutoModelForTokenClassification.from_pretrained( - model.path, - subfolder=model.subfolder, - revision=model.revision, - torch_dtype="auto", - low_cpu_mem_usage=False, - **model.kwargs, - ) - # Handle meta device properly - use to_empty() when moving from meta device - if hasattr(tf_model, 'device') and str(tf_model.device) == 'meta': - tf_model = tf_model.to_empty(device=device()) + cache_key = _create_model_cache_key(model, use_onnx=False) + + if cache_key in _pytorch_ner_cache: + LOGGER.debug("Using cached PyTorch NER model", model=model) + tf_model = _pytorch_ner_cache[cache_key] else: - tf_model = tf_model.to(device()) - LOGGER.debug("Initialized NER model", model=model, device=device()) + transformers = lazy_load_dep("transformers") + tf_model = transformers.AutoModelForTokenClassification.from_pretrained( + model.path, + subfolder=model.subfolder, + revision=model.revision, + torch_dtype="auto", + low_cpu_mem_usage=False, + **model.kwargs, + ) + # Handle meta device properly - use to_empty() when moving from meta device + if hasattr(tf_model, 'device') and str(tf_model.device) == 'meta': + tf_model = tf_model.to_empty(device=device()) + else: + tf_model = tf_model.to(device()) + + # Cache the model (limit cache size to prevent memory issues) + if len(_pytorch_ner_cache) >= 16: + # Remove oldest entry (simple FIFO eviction) + oldest_key = next(iter(_pytorch_ner_cache)) + del _pytorch_ner_cache[oldest_key] + + _pytorch_ner_cache[cache_key] = tf_model + LOGGER.debug("Cached new PyTorch NER model", model=model, device=device()) return tf_tokenizer, tf_model - optimum_onnxruntime = lazy_load_dep( - "optimum.onnxruntime", - ("optimum[onnxruntime]" if device().type != "cuda" else "optimum[onnxruntime-gpu]"), - ) + cache_key = _create_model_cache_key(model, use_onnx=True) - tf_model = optimum_onnxruntime.ORTModelForTokenClassification.from_pretrained( - model.onnx_path, - export=False, - subfolder=model.onnx_subfolder, - provider=("CUDAExecutionProvider" if device().type == "cuda" else "CPUExecutionProvider"), - revision=model.onnx_revision, - file_name=model.onnx_filename, - **model.kwargs, - ) - LOGGER.debug("Initialized NER ONNX model", model=model, device=device()) + if cache_key in _onnx_ner_cache: + LOGGER.debug("Using cached ONNX NER model", model=model) + tf_model = _onnx_ner_cache[cache_key] + else: + optimum_onnxruntime = lazy_load_dep( + "optimum.onnxruntime", + ("optimum[onnxruntime]" if device().type != "cuda" else "optimum[onnxruntime-gpu]"), + ) + + tf_model = optimum_onnxruntime.ORTModelForTokenClassification.from_pretrained( + model.onnx_path, + export=False, + subfolder=model.onnx_subfolder, + provider=("CUDAExecutionProvider" if device().type == "cuda" else "CPUExecutionProvider"), + revision=model.onnx_revision, + file_name=model.onnx_filename, + **model.kwargs, + ) + + # Cache the model (limit cache size to prevent memory issues) + if len(_onnx_ner_cache) >= 16: + # Remove oldest entry (simple FIFO eviction) + oldest_key = next(iter(_onnx_ner_cache)) + del _onnx_ner_cache[oldest_key] + + _onnx_ner_cache[cache_key] = tf_model + LOGGER.debug("Cached new ONNX NER model", model=model, device=device()) return tf_tokenizer, tf_model diff --git a/tueri_api/app/app.py b/tueri_api/app/app.py index 467a2016..e5458a2e 100644 --- a/tueri_api/app/app.py +++ b/tueri_api/app/app.py @@ -129,11 +129,9 @@ def _get_input_scanners_function(config: Config, vault: Vault) -> Callable: def get_cached_scanners(runtime_params: dict = None) -> List[InputScanner]: nonlocal scanners - # If runtime parameters are provided, create new scanner instances if runtime_params: LOGGER.debug("Creating input scanners with overridden parameters", params=runtime_params) return get_input_scanners([], vault, runtime_params) - # Otherwise use cached scanners if not scanners and config.app.lazy_load: LOGGER.debug("Lazy loading input scanners from MongoDB") scanners = get_input_scanners([], vault) @@ -151,11 +149,9 @@ def _get_output_scanners_function(config: Config, vault: Vault) -> Callable: def get_cached_scanners(runtime_params: dict = None) -> List[OutputScanner]: nonlocal scanners - # If runtime parameters are provided, create new scanner instances if runtime_params: LOGGER.debug("Creating output scanners with overridden parameters", params=runtime_params) return get_output_scanners([], vault, runtime_params) - # Otherwise use cached scanners if not scanners and config.app.lazy_load: LOGGER.debug("Lazy loading output scanners from MongoDB") scanners = get_output_scanners([], vault) diff --git a/tueri_api/app/scanner.py b/tueri_api/app/scanner.py index cbeb4266..45dcc0be 100644 --- a/tueri_api/app/scanner.py +++ b/tueri_api/app/scanner.py @@ -24,6 +24,7 @@ from tueri.vault import Vault from .config import ScannerConfig +from .scanner_cache import get_scanner_cache_manager from .util import get_resource_utilization torch.set_num_threads(1) @@ -63,41 +64,78 @@ def _fetch_scanners_from_mongo(scanner_type: str) -> List[ScannerConfig]: return scanners def get_input_scanners(scanners: List[ScannerConfig], vault: Vault, runtime_params: Optional[Dict[str, Dict]] = None) -> List[InputScanner]: - """Load input scanners from MongoDB.""" + """Load input scanners from MongoDB""" input_scanners_config = _fetch_scanners_from_mongo("input") loaded_input_scanners: List[InputScanner] = [] - for scanner in input_scanners_config: - scanner_params = scanner.params.copy() if scanner.params else {} - if runtime_params and scanner.type in runtime_params: - scanner_params.update(runtime_params[scanner.type]) - # LOGGER.debug("overriding parameters", scanner=scanner.type, overrides=runtime_params[scanner.type]) - loaded_input_scanners.append( - _get_input_scanner( + + # use caching when runtime parameters are provided + if runtime_params: + scanner_cache = get_scanner_cache_manager() + for scanner in input_scanners_config: + scanner_params = scanner.params.copy() if scanner.params else {} + if scanner.type in runtime_params: + scanner_params.update(runtime_params[scanner.type]) + LOGGER.debug("overriding parameters", scanner=scanner.type, overrides=runtime_params[scanner.type]) + + def scanner_factory(scanner_type: str, params: dict) -> InputScanner: + return _get_input_scanner(scanner_type, params, vault=vault) + + cached_scanner = scanner_cache.get_input_scanner( scanner.type, scanner_params, - vault=vault, + scanner_factory + ) + loaded_input_scanners.append(cached_scanner) + + else: + for scanner in input_scanners_config: + scanner_params = scanner.params.copy() if scanner.params else {} + loaded_input_scanners.append( + _get_input_scanner( + scanner.type, + scanner_params, + vault=vault, + ) ) - ) return loaded_input_scanners def get_output_scanners(scanners: List[ScannerConfig], vault: Vault, runtime_params: Optional[Dict[str, Dict]] = None) -> List[OutputScanner]: - """Load output scanners from MongoDB.""" + """Load output scanners from MongoDB""" output_scanners_config = _fetch_scanners_from_mongo("output") loaded_output_scanners: List[OutputScanner] = [] - for scanner in output_scanners_config: - scanner_params = scanner.params.copy() if scanner.params else {} - if runtime_params and scanner.type in runtime_params: - scanner_params.update(runtime_params[scanner.type]) - # LOGGER.debug("overriding parameters", scanner=scanner.type, overrides=runtime_params[scanner.type]) - loaded_output_scanners.append( - _get_output_scanner( + + # use caching when runtime parameters are provided + if runtime_params: + scanner_cache = get_scanner_cache_manager() + + for scanner in output_scanners_config: + scanner_params = scanner.params.copy() if scanner.params else {} + if scanner.type in runtime_params: + scanner_params.update(runtime_params[scanner.type]) + LOGGER.debug("overriding parameters", scanner=scanner.type, overrides=runtime_params[scanner.type]) + + def scanner_factory(scanner_type: str, params: dict) -> OutputScanner: + return _get_output_scanner(scanner_type, params, vault=vault) + + cached_scanner = scanner_cache.get_output_scanner( scanner.type, scanner_params, - vault=vault, + scanner_factory + ) + loaded_output_scanners.append(cached_scanner) + + else: + for scanner in output_scanners_config: + scanner_params = scanner.params.copy() if scanner.params else {} + loaded_output_scanners.append( + _get_output_scanner( + scanner.type, + scanner_params, + vault=vault, + ) ) - ) return loaded_output_scanners diff --git a/tueri_api/app/scanner_cache.py b/tueri_api/app/scanner_cache.py new file mode 100644 index 00000000..c1bf7568 --- /dev/null +++ b/tueri_api/app/scanner_cache.py @@ -0,0 +1,109 @@ +import hashlib +import json +from collections import OrderedDict +from typing import Dict, List, Optional, Any +import structlog + +from tueri.input_scanners.base import Scanner as InputScanner +from tueri.output_scanners.base import Scanner as OutputScanner +from tueri.vault import Vault +from .config import ScannerConfig + +LOGGER = structlog.getLogger(__name__) + +class ScannerCacheManager: + """ + Manages caching of scanner instances using LRU (Least Recently Used) eviction. + """ + + def __init__(self, max_cache_size: int = 64): + """Initialize the scanner cache manager with LRU eviction.""" + self._input_cache: OrderedDict[str, InputScanner] = OrderedDict() + self._output_cache: OrderedDict[str, OutputScanner] = OrderedDict() + self._max_cache_size = max_cache_size + + def _create_scanner_cache_key(self, scanner_type: str, params: Dict[str, Any]) -> str: + """Create a cache key for a scanner based on its type and parameters.""" + cache_data = { + "scanner_type": scanner_type, + "params": params + } + config_str = json.dumps(cache_data, sort_keys=True) + return hashlib.md5(config_str.encode()).hexdigest() + + def _evict_lru_if_needed(self, cache: OrderedDict) -> None: + """Remove least recently used entry if cache is at capacity.""" + if len(cache) >= self._max_cache_size: + # Remove least recently used (first item in OrderedDict) + evicted_key, _ = cache.popitem(last=False) + LOGGER.debug("Evicted LRU scanner from cache", cache_key=evicted_key) + + def get_input_scanner(self, scanner_type: str, params: Dict[str, Any], scanner_factory_func: callable) -> InputScanner: + """Get an input scanner from cache or create a new one using LRU eviction.""" + cache_key = self._create_scanner_cache_key(scanner_type, params) + + # Check if we have a cached scanner + if cache_key in self._input_cache: + # Move to end (mark as most recently used) + self._input_cache.move_to_end(cache_key) + LOGGER.debug("Using cached input scanner", scanner_type=scanner_type, cache_key=cache_key) + return self._input_cache[cache_key] + + # Create new scanner + scanner = scanner_factory_func(scanner_type, params) + + # Evict LRU if cache is full + self._evict_lru_if_needed(self._input_cache) + + # Cache the new scanner (automatically becomes most recently used) + self._input_cache[cache_key] = scanner + LOGGER.debug("Cached new input scanner", scanner_type=scanner_type, cache_key=cache_key) + + return scanner + + def get_output_scanner(self, scanner_type: str, params: Dict[str, Any], scanner_factory_func: callable) -> OutputScanner: + """Get an output scanner from cache or create a new one using LRU eviction.""" + cache_key = self._create_scanner_cache_key(scanner_type, params) + + # Check if we have a cached scanner + if cache_key in self._output_cache: + # Move to end (mark as most recently used) + self._output_cache.move_to_end(cache_key) + LOGGER.debug("Using cached output scanner", scanner_type=scanner_type, cache_key=cache_key) + return self._output_cache[cache_key] + + # Create new scanner + scanner = scanner_factory_func(scanner_type, params) + + # Evict LRU if cache is full + self._evict_lru_if_needed(self._output_cache) + + # Cache the new scanner (automatically becomes most recently used) + self._output_cache[cache_key] = scanner + LOGGER.debug("Cached new output scanner", scanner_type=scanner_type, cache_key=cache_key) + + return scanner + + def get_cache_stats(self) -> Dict[str, Any]: + """Get statistics about the cache usage.""" + return { + "input_cache_size": len(self._input_cache), + "output_cache_size": len(self._output_cache), + "max_cache_size": self._max_cache_size, + "eviction_policy": "LRU" + } + + def clear_cache(self) -> None: + """Clear all cached scanners from both input and output caches.""" + self._input_cache.clear() + self._output_cache.clear() + LOGGER.info("Cleared all scanner caches") + + +# Global scanner cache manager instance with LRU eviction +_scanner_cache_manager = ScannerCacheManager() + + +def get_scanner_cache_manager() -> ScannerCacheManager: + """Get the global scanner cache manager instance.""" + return _scanner_cache_manager