125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
import pytest
|
|
|
|
from ray.tune.utils.object_cache import _ObjectCache
|
|
|
|
|
|
@pytest.mark.parametrize("eager", [False, True])
|
|
def test_no_may_keep_one(eager):
|
|
"""Test object caching.
|
|
|
|
- After init, no objects are cached (as max cached is 0), except when eager caching
|
|
- After increasing max to 2, up to 2 objects are cached
|
|
- Decreasing max objects will evict them on flush
|
|
"""
|
|
cache = _ObjectCache(may_keep_one=eager)
|
|
|
|
# max(A) = 0, so we we only cache when eager caching
|
|
assert cache.cache_object("A", 1) == eager
|
|
assert cache.num_cached_objects == int(eager)
|
|
|
|
# Set max(A) = 2
|
|
cache.increase_max("A", 2)
|
|
|
|
# max(A) = 2, so we cache up to two objects
|
|
if not eager:
|
|
assert cache.cache_object("A", 1)
|
|
|
|
assert cache.cache_object("A", 2)
|
|
assert not cache.cache_object("A", 3)
|
|
|
|
assert cache.num_cached_objects == 2
|
|
|
|
# Nothing has to be evicted
|
|
assert not list(cache.flush_cached_objects())
|
|
|
|
# Set max(A) = 1, so we have one object too much
|
|
cache.decrease_max("A", 1)
|
|
|
|
# First cached object is evicted
|
|
assert list(cache.flush_cached_objects()) == [1]
|
|
assert cache.num_cached_objects == 1
|
|
|
|
# Set max(A) = 0
|
|
cache.decrease_max("A", 1)
|
|
|
|
# Second cached object is evicted if not eager caching
|
|
assert list(cache.flush_cached_objects()) == ([2] if not eager else [])
|
|
assert cache.num_cached_objects == (0 if not eager else 1)
|
|
|
|
|
|
@pytest.mark.parametrize("eager", [False, True])
|
|
def test_multi(eager):
|
|
"""Test caching with multiple objects"""
|
|
cache = _ObjectCache(may_keep_one=eager)
|
|
|
|
# max(A) = 0, so we we only cache when eager caching
|
|
assert cache.cache_object("A", 1) == eager
|
|
assert cache.num_cached_objects == int(eager)
|
|
|
|
# max(B) = 0, so no caching
|
|
assert not cache.cache_object("B", 5)
|
|
assert cache.num_cached_objects == int(eager)
|
|
|
|
# Increase maximums levels
|
|
cache.increase_max("A", 1)
|
|
cache.increase_max("B", 1)
|
|
|
|
# Cache objects (A is already cached if eager)
|
|
assert cache.cache_object("A", 1) != eager
|
|
assert cache.cache_object("B", 5)
|
|
|
|
# No further objects can be cached
|
|
assert not cache.cache_object("A", 2)
|
|
assert not cache.cache_object("B", 6)
|
|
|
|
assert cache.num_cached_objects == 2
|
|
|
|
# Decrease
|
|
cache.decrease_max("A", 1)
|
|
|
|
# Evict A object
|
|
assert list(cache.flush_cached_objects()) == [1]
|
|
|
|
cache.decrease_max("B", 1)
|
|
|
|
# If eager, keep B object, otherwise, evict B
|
|
assert list(cache.flush_cached_objects()) == ([5] if not eager else [])
|
|
assert cache.num_cached_objects == (0 if not eager else 1)
|
|
|
|
|
|
def test_multi_eager_other():
|
|
"""On eager caching, only cache an object if no other object is expected.
|
|
|
|
- Expect up to one cached A object
|
|
- Try to cache object B --> doesn't get cached
|
|
- Remove expectation for A object
|
|
- Try to cache object B --> get's cached
|
|
"""
|
|
cache = _ObjectCache(may_keep_one=True)
|
|
|
|
cache.increase_max("A", 1)
|
|
assert not cache.cache_object("B", 2)
|
|
|
|
cache.decrease_max("A", 1)
|
|
assert cache.cache_object("B", 3)
|
|
|
|
|
|
@pytest.mark.parametrize("eager", [False, True])
|
|
def test_force_all(eager):
|
|
"""Assert that force_all=True will always evict all object."""
|
|
cache = _ObjectCache(may_keep_one=eager)
|
|
|
|
cache.increase_max("A", 2)
|
|
|
|
assert cache.cache_object("A", 1)
|
|
assert cache.cache_object("A", 2)
|
|
|
|
assert list(cache.flush_cached_objects(force_all=True)) == [1, 2]
|
|
assert cache.num_cached_objects == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|