237 lines
8.7 KiB
Python
237 lines
8.7 KiB
Python
import gc
|
|
import time
|
|
from typing import Any, Dict, Optional, Union
|
|
from enum import Enum
|
|
|
|
from ..helpers.metaclasses import Singleton
|
|
|
|
|
|
class CacheCleanupMethod(Enum):
|
|
ROUND_ROBIN = 'round_robin'
|
|
LEAST_USED = 'least_used'
|
|
MOST_OFTEN_USED = 'most_often_used'
|
|
|
|
|
|
class ComfyCache(metaclass=Singleton):
|
|
def __init__(self) -> None:
|
|
self.data: Dict[str, Any] = {}
|
|
self.metadata: Dict[str, Dict[str, Union[int, float, CacheCleanupMethod, Optional[int]]]] = {}
|
|
|
|
def __del__(self) -> None:
|
|
self.data.clear()
|
|
self.metadata.clear()
|
|
del self.data
|
|
del self.metadata
|
|
gc.collect()
|
|
|
|
@classmethod
|
|
def cache_pop(cls, property: str) -> Any:
|
|
method = cls.get_cachemethod(property)
|
|
if method == CacheCleanupMethod.ROUND_ROBIN:
|
|
return cls.cache_pop_roundrobin(property)
|
|
elif method == CacheCleanupMethod.LEAST_USED:
|
|
return cls.cache_pop_least_used(property)
|
|
elif method == CacheCleanupMethod.MOST_OFTEN_USED:
|
|
return cls.cache_pop_most_often_used(property)
|
|
else:
|
|
raise ValueError("Unsupported cache cleanup method")
|
|
|
|
@classmethod
|
|
def cache_pop_roundrobin(cls, property: str) -> Any:
|
|
instance = cls()
|
|
if property not in instance.metadata:
|
|
return None
|
|
oldest_property = min(instance.metadata[property].items(), key=lambda item: item[1]['timestamp'])[0]
|
|
value = instance.data[property].pop(oldest_property)
|
|
instance.metadata[property].pop(oldest_property)
|
|
if not instance.data[property]:
|
|
del instance.data[property]
|
|
del instance.metadata[property]
|
|
return value
|
|
|
|
@classmethod
|
|
def cache_pop_least_used(cls, property: str) -> Any:
|
|
instance = cls()
|
|
if property not in instance.metadata:
|
|
return None
|
|
least_used_property = min(instance.metadata[property].items(), key=lambda item: item[1]['usage'])[0]
|
|
value = instance.data[property].pop(least_used_property)
|
|
instance.metadata[property].pop(least_used_property)
|
|
if not instance.data[property]:
|
|
del instance.data[property]
|
|
del instance.metadata[property]
|
|
return value
|
|
|
|
@classmethod
|
|
def cache_pop_most_often_used(cls, property: str) -> Any:
|
|
instance = cls()
|
|
if property not in instance.metadata:
|
|
return None
|
|
most_often_used_property = max(instance.metadata[property].items(), key=lambda item: item[1]['usage'])[0]
|
|
value = instance.data[property].pop(most_often_used_property)
|
|
instance.metadata[property].pop(most_often_used_property)
|
|
if not instance.data[property]:
|
|
del instance.data[property]
|
|
del instance.metadata[property]
|
|
return value
|
|
|
|
@classmethod
|
|
def get(cls, property: Optional[str] = None) -> Any:
|
|
instance = cls()
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
if property is not None:
|
|
keys = property.split(".")
|
|
for key in keys:
|
|
if key not in cache:
|
|
return None
|
|
cache = cache[key]
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if key not in metadata[property]:
|
|
metadata[property][key] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][key]['timestamp'] = time.time()
|
|
metadata[property][key]['usage'] += 1
|
|
return cache
|
|
|
|
return cache
|
|
|
|
@classmethod
|
|
def set(cls, property: Union[str, None], value: Any) -> None:
|
|
instance = cls()
|
|
if property is not None:
|
|
keys = property.split(".")
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
for key in keys[:-1]:
|
|
if key not in cache:
|
|
cache[key] = {}
|
|
cache = cache[key]
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if key not in metadata[property]:
|
|
metadata[property][key] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][key]['timestamp'] = time.time()
|
|
metadata[property][key]['usage'] += 1
|
|
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if keys[-1] not in metadata[property]:
|
|
metadata[property][keys[-1]] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][keys[-1]]['timestamp'] = time.time()
|
|
metadata[property][keys[-1]]['usage'] += 1
|
|
|
|
max_size = metadata[property].get(keys[-1], {}).get('max_size')
|
|
if max_size is not None and len(cache) >= max_size:
|
|
cls.cache_pop(property)
|
|
|
|
cache[keys[-1]] = value
|
|
else:
|
|
instance.data = value
|
|
instance.metadata = {}
|
|
|
|
return
|
|
|
|
@classmethod
|
|
def set_cachemethod(cls, property: str, method: CacheCleanupMethod) -> None:
|
|
instance = cls()
|
|
keys = property.split(".")
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
for key in keys[:-1]:
|
|
if key not in cache:
|
|
cache[key] = {}
|
|
cache = cache[key]
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if key not in metadata[property]:
|
|
metadata[property][key] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][key]['timestamp'] = time.time()
|
|
metadata[property][key]['usage'] += 1
|
|
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if keys[-1] not in metadata[property]:
|
|
metadata[property][keys[-1]] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][keys[-1]]['cachemethod'] = method
|
|
return
|
|
|
|
@classmethod
|
|
def get_cachemethod(cls, property: str) -> CacheCleanupMethod:
|
|
instance = cls()
|
|
keys = property.split(".")
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
for key in keys:
|
|
if key not in cache:
|
|
return CacheCleanupMethod.ROUND_ROBIN # Default method if not set
|
|
cache = cache[key]
|
|
|
|
if property in metadata and keys[-1] in metadata[property]:
|
|
return metadata[property][keys[-1]]['cachemethod']
|
|
return CacheCleanupMethod.ROUND_ROBIN
|
|
|
|
@classmethod
|
|
def set_max_size(cls, property: str, max_size: int) -> None:
|
|
instance = cls()
|
|
keys = property.split(".")
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
for key in keys[:-1]:
|
|
if key not in cache:
|
|
cache[key] = {}
|
|
cache = cache[key]
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if key not in metadata[property]:
|
|
metadata[property][key] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][key]['timestamp'] = time.time()
|
|
metadata[property][key]['usage'] += 1
|
|
|
|
if property not in metadata:
|
|
metadata[property] = {}
|
|
if keys[-1] not in metadata[property]:
|
|
metadata[property][keys[-1]] = {'timestamp': time.time(), 'usage': 0, 'cachemethod': CacheCleanupMethod.ROUND_ROBIN, 'max_size': None}
|
|
metadata[property][keys[-1]]['max_size'] = max_size
|
|
return
|
|
|
|
@classmethod
|
|
def get_max_size(cls, property: str) -> Optional[int]:
|
|
instance = cls()
|
|
keys = property.split(".")
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
for key in keys:
|
|
if key not in cache:
|
|
return None
|
|
cache = cache[key]
|
|
|
|
if property in metadata and keys[-1] in metadata[property]:
|
|
return metadata[property][keys[-1]].get('max_size')
|
|
return None
|
|
|
|
@classmethod
|
|
def flush(cls, property: str) -> None:
|
|
instance = cls()
|
|
keys = property.split(".")
|
|
cache = instance.data
|
|
metadata = instance.metadata
|
|
|
|
for key in keys[:-1]:
|
|
if key not in cache:
|
|
return
|
|
cache = cache[key]
|
|
metadata = metadata.get(property, {})
|
|
metadata = metadata.get(key, {})
|
|
|
|
if keys[-1] in cache:
|
|
del cache[keys[-1]]
|
|
if keys[-1] in metadata:
|
|
del metadata[keys[-1]]
|