chore: import upstream snapshot with attribution

python/ray/tune/tests/test_util_object_cache.py

@@ -0,0 +1,124 @@
+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__]))