wehub-resource-sync
174 lines
5.4 KiB
Python
174 lines
5.4 KiB
Python
from collections import Counter, defaultdict
|
|
from typing import Dict, Generator, List, Optional, TypeVar
|
|
|
|
# Grouping key - must be hashable
|
|
T = TypeVar("T")
|
|
# Objects to cache
|
|
U = TypeVar("U")
|
|
|
|
|
|
class _ObjectCache:
|
|
"""Cache up to some maximum count given a grouping key.
|
|
|
|
This object cache can e.g. be used to cache Ray Tune trainable actors
|
|
given their resource requirements (reuse_actors=True).
|
|
|
|
If the max number of cached objects for a grouping key is reached,
|
|
no more objects for this group will be cached.
|
|
|
|
However, if `may_keep_one=True`, one object (globally across all grouping
|
|
keys) may be cached, even if the max number of objects is 0. This is to
|
|
allow to cache an object if the max number of objects of this key
|
|
will increase shortly after (as is the case e.g. in the Ray Tune control
|
|
loop).
|
|
|
|
Args:
|
|
may_keep_one: If True, one object (globally) may be cached if no desired
|
|
maximum objects are defined.
|
|
|
|
"""
|
|
|
|
def __init__(self, may_keep_one: bool = True):
|
|
self._num_cached_objects: int = 0
|
|
self._cached_objects: Dict[T, List[U]] = defaultdict(list)
|
|
self._max_num_objects: Counter[T] = Counter()
|
|
|
|
self._may_keep_one = may_keep_one
|
|
|
|
@property
|
|
def num_cached_objects(self):
|
|
return self._num_cached_objects
|
|
|
|
@property
|
|
def total_max_objects(self):
|
|
# Counter.total() is only available for python 3.10+
|
|
return sum(self._max_num_objects.values())
|
|
|
|
def increase_max(self, key: T, by: int = 1) -> None:
|
|
"""Increase number of max objects for this key.
|
|
|
|
Args:
|
|
key: Group key.
|
|
by: Decrease by this amount.
|
|
"""
|
|
self._max_num_objects[key] += by
|
|
|
|
def decrease_max(self, key: T, by: int = 1) -> None:
|
|
"""Decrease number of max objects for this key.
|
|
|
|
Args:
|
|
key: Group key.
|
|
by: Decrease by this amount.
|
|
"""
|
|
self._max_num_objects[key] -= by
|
|
|
|
def has_cached_object(self, key: T) -> bool:
|
|
"""Return True if at least one cached object exists for this key.
|
|
|
|
Args:
|
|
key: Group key.
|
|
|
|
Returns:
|
|
True if at least one cached object exists for this key.
|
|
"""
|
|
return bool(self._cached_objects[key])
|
|
|
|
def cache_object(self, key: T, obj: U) -> bool:
|
|
"""Cache object for a given key.
|
|
|
|
This will put the object into a cache, assuming the number
|
|
of cached objects for this key is less than the number of
|
|
max objects for this key.
|
|
|
|
An exception is made if `max_keep_one=True` and no other
|
|
objects are cached globally. In that case, the object can
|
|
still be cached.
|
|
|
|
Args:
|
|
key: Group key.
|
|
obj: Object to cache.
|
|
|
|
Returns:
|
|
True if the object has been cached. False otherwise.
|
|
|
|
"""
|
|
# If we have more objects cached already than we desire
|
|
if len(self._cached_objects[key]) >= self._max_num_objects[key]:
|
|
# If may_keep_one is False, never cache
|
|
if not self._may_keep_one:
|
|
return False
|
|
|
|
# If we have more than one other cached object, don't cache
|
|
if self._num_cached_objects > 0:
|
|
return False
|
|
|
|
# If any other objects are expected to be cached, don't cache
|
|
if any(v for v in self._max_num_objects.values()):
|
|
return False
|
|
|
|
# Otherwise, cache (for now).
|
|
|
|
self._cached_objects[key].append(obj)
|
|
self._num_cached_objects += 1
|
|
return True
|
|
|
|
def pop_cached_object(self, key: T) -> Optional[U]:
|
|
"""Get one cached object for a key.
|
|
|
|
This will remove the object from the cache.
|
|
|
|
Args:
|
|
key: Group key.
|
|
|
|
Returns:
|
|
Cached object.
|
|
"""
|
|
if not self.has_cached_object(key):
|
|
return None
|
|
|
|
self._num_cached_objects -= 1
|
|
return self._cached_objects[key].pop(0)
|
|
|
|
def flush_cached_objects(self, force_all: bool = False) -> Generator[U, None, None]:
|
|
"""Return a generator over cached objects evicted from the cache.
|
|
|
|
This method yields all cached objects that should be evicted from the
|
|
cache for cleanup by the caller.
|
|
|
|
If the number of max objects is lower than the number of
|
|
cached objects for a given key, objects are evicted until
|
|
the numbers are equal.
|
|
|
|
If `max_keep_one=True` (and ``force_all=False``), one cached object
|
|
may be retained.
|
|
|
|
Objects are evicted FIFO.
|
|
|
|
If ``force_all=True``, all objects are evicted.
|
|
|
|
Args:
|
|
force_all: If True, all objects are flushed. This takes precedence
|
|
over ``keep_one``.
|
|
|
|
Yields:
|
|
U: Evicted objects to be cleaned up by caller.
|
|
|
|
"""
|
|
# If force_all=True, don't keep one.
|
|
keep_one = self._may_keep_one and not force_all
|
|
|
|
for key, objs in self._cached_objects.items():
|
|
max_cached = self._max_num_objects[key] if not force_all else 0
|
|
|
|
if (
|
|
self._num_cached_objects == 1
|
|
and keep_one
|
|
# Only keep this object if we don't expect a different one
|
|
and not any(v for v in self._max_num_objects.values())
|
|
):
|
|
break
|
|
|
|
while len(objs) > max_cached:
|
|
self._num_cached_objects -= 1
|
|
yield objs.pop(0)
|