Fix create_trial brk when many elements are read in the first attempt (aws#168)

Author: NihalHarish

smdebug/core/index_reader.py

@@ -54,7 +54,8 @@ class ReadIndexFilesCache:
 
     Note: cache_limit is a soft limit.
 
-    In certain cases, the size of the cache can exceed this limit.
+    The size of the cache can exceed this limit if eviction_point is 0.
+    This can happen if start_after_key is None (unset) or is equal to the first element in the cache.
 
     If start_after_key happens to the be first element in sorted(self.lookup_set), then we do not
     evict any element from the cache, but add more elements to cache.
@@ -76,6 +77,9 @@ def has_not_read(self, index_file: str) -> bool:
 
     def _evict_cache(self, start_after_key: str) -> None:
         read_files = sorted(self.lookup_set)
+        start_after_key = "" if start_after_key is None else start_after_key
+        # eviction_point = 0 if start_after_key is None (unset).
+        # This happens if more than self.cache_limit number of files are read on the first read attempt.
         eviction_point = bisect_left(read_files, start_after_key)
         for i in range(eviction_point):
             self.lookup_set.remove(read_files[i])

tests/core/test_utils.py

@@ -98,6 +98,17 @@ def test_index_files_cache():
     )  # Elements in the cache will be file_4, file_5, file_6
 
 
+def test_index_files_cache_insert_many_elements_in_the_first_read():
+    cache = ReadIndexFilesCache()
+    cache.cache_limit = 5
+    elements = ["a", "b", "c", "d", "e", "f", "g", "h"]
+    for e in elements:
+        cache.add(e, None)
+
+    # No files should be evicted because start_after_key has not been set
+    assert len(cache.lookup_set) == len(elements)
+
+
 def test_get_prefix_from_index_file():
     local_index_filepath = "/opt/ml/testing/run_1/index/000000000/000000000000_worker_0.json"
     prefix = IndexFileLocationUtils.get_prefix_from_index_file(local_index_filepath)